first implementation of aircraft example

components/aircraft.py

+#!/usr/bin/env python3
+
+import numpy as np
+import xarray as xr
+import pandas as pd
+import sys
+from argparse import ArgumentParser, FileType
+from pathlib import Path
+
+this_dir = Path(__file__).parent
+
+parser = ArgumentParser()
+parser.add_argument("--descriptorfile", type=FileType('r'), default=this_dir.parent/"descriptorfiles"/"icon_1mom.py", help="")
+parser.add_argument("--track", type=str, help="File containing time-lon-lat data") # TODO Height?
+parser.add_argument("--frequency", type=float, nargs="+", default=94)
+parser.add_argument("--dt", type=str, default=None, help="Timestep") #TODO output time? for now I stay strict to model timestep and assume mean aircraft location ...
+parser.add_argument("--source_component", type=str, default="atmo_output", help="component name of the source component")
+parser.add_argument("--source_grid", type=str, default="icon_atmos_grid", help="grid name of the source grid")
+args = parser.parse_args()
+
+
+from yac import YAC, UnstructuredGrid, Location, Field, TimeUnit, \
+    InterpolationStack, NNNReductionType, Reduction, Action
+import pyPamtra
+
+
+yac = YAC()
+
+comp = yac.def_comp("pamtra")
+
+track = xr.open_dataset(args.track)
+lats = np.deg2rad(track.lat.values)
+lons = np.deg2rad(track.lon.values) # TODO knowing the model timesteps I can in principle calculate the mean locations in advance and form a reduced grid, maybe pass as an argument?
+                                    # Also, knowing the model domain I can exclude the track points outside the domain in advance?? probably not worth the complication
+# alts = track.alt.values # need to think about this
+times = track.time.values
+ncoords = len(times)
+
+grid = UnstructuredGrid("pamtra-grid", [3]*ncoords,
+                        [*lons, *(lons+0.1), *lons],
+                        [*lats, *lats, *(lats+0.1)],
+                        np.array(list(zip(range(ncoords), range(ncoords, 2*ncoords), range(2*ncoords, 3*ncoords)))).ravel())
+points = grid.def_points(Location.CELL,
+                         lons, lats)
+
+yac.sync_def()
+
+# print(yac.get_field_names(*source), file=sys.stderr)
+
+interp = InterpolationStack()
+interp.add_nnn(NNNReductionType.AVG, 1, 1.0)
+
+source = (args.source_component, args.source_grid)
+
+dt = yac.get_field_timestep(*source, "temp")
+
+fields = {}
+for name in ["temp", "pres", "qv", "qi", "qc", "qs", "qr", "qg", "z_ifc", "u", "v", "w"]:
+    nlevel = yac.get_field_collection_size(*source, name)
+    fields[name] = Field.create(name, comp, points, nlevel,
+                                dt, TimeUnit.ISO_FORMAT)
+    yac.def_couple(*source, name,
+                   "pamtra", "pamtra-grid", name,
+                   dt, TimeUnit.ISO_FORMAT, Reduction.TIME_NONE,
+                   interp)
+
+yac.enddef()
+
+with open(args.descriptorfile, "r") as dfile:
+    exec(dfile.read())  # adds a variable "descriptorFile"
+
+
+def call_pamtra(temp, pres, qv, qi, qc, qs, qr, qg, z_ifc, u, v, w,
+                datetime):
+    pamtra = pyPamtra.pyPamtra()
+    for hydro in descriptorFile:
+        pamtra.df.addHydrometeor(hydro)
+    pamtra.nmlSet["passive"] = False
+    pamtra.nmlSet["active"] = True
+    pamtra.nmlSet["radar_mode"] = "simple"
+    rh = 100*pyPamtra.meteoSI.q2rh(qv, temp, pres)
+    pamtra.createProfile(hgt_lev=z_ifc.T[None, :, ::-1],
+                         press=pres.T[None, :, ::-1],
+                         temp=temp.T[None, :, ::-1],
+                         relhum=rh.T[None, :, ::-1],
+                         hydro_q=np.stack([qc.T, qi.T, qr.T, qs.T, qg.T], axis=-1)[:, ::-1, :],
+                         )
+    pamtra.runPamtra(args.frequency)
+
+    return pamtra # return the entire pamtra object and let the caller decide what to take as results
+
+# prepare output
+Ntimesperoutput = 6 # TODO should be user config...
+pam_time_idx = xr.IndexVariable(
+        dims='time_idx',
+        data=np.arange(0, Ntimesperoutput),
+        attrs={'name':'time_idx',
+               'long_name':'index of the time step within the output file'})
+
+pam_hgt_idx = xr.IndexVariable(
+        dims='hgt_idx',
+        data=np.arange(0, nlevel-1),
+        attrs={'name':'hgt_idx',
+               'long_name':'height indx above surface',
+               'direction':'bottom-up',
+               'comment':'look at variable height to get the height profile above each location'})
+
+pam_time = xr.DataArray(
+        dims=['time_idx'],
+        coords={'time_idx':pam_time_idx},
+        data=np.arange(0, Ntimesperoutput),
+        attrs={'name':'time',
+               'long_name':'unixtime nanoseconds since 1970',
+               'units':'nanoseconds'})
+
+pam_lat = xr.DataArray(
+        dims=['time_idx'],
+        coords={'time_idx':pam_time_idx},
+        data=np.empty((Ntimesperoutput,)),
+        attrs={'name':'lat',
+               'long_name':'location latitude',
+               'unit':'degrees North'})
+
+pam_lon = xr.DataArray(
+        dims=['time_idx'],
+        coords={'time_idx':pam_time_idx},
+        data=np.empty((Ntimesperoutput,)),
+        attrs={'name':'lon',
+               'long_name':'location longitude',
+               'unit':'degrees East'})
+
+pam_hgt = xr.DataArray(
+        dims=['hgt_idx', 'time_idx'],
+        coords={'hgt_idx':pam_hgt_idx,
+                'time_idx':pam_time_idx},
+        data=np.empty((nlevel-1, Ntimesperoutput)),
+        attrs={'name':'hgt',
+               'long_name':'height above surface',
+               'units':'meters'})
+
+pam_Ze = xr.DataArray(
+        dims=['hgt_idx', 'time_idx'],
+        coords={'hgt_idx':pam_hgt_idx,
+                'time_idx':pam_time_idx},
+        data=np.empty((nlevel-1, Ntimesperoutput)),
+        attrs={'name':'Ze',
+               'long_name':'unattenuated reflectivity',
+               'units':'dBZ'})
+
+variables = {'height':pam_hgt,
+             'time':pam_time,
+             'lat':pam_lat,
+             'lon':pam_lon,
+             'Ze':pam_Ze}
+
+coordinates = {'time_idx':pam_time_idx,
+               'hgt_idx':pam_hgt_idx}
+
+import os
+import socket
+from datetime import datetime as datetimelib ### ugly but datetime is reused later and can be uglier
+global_attributes = {'created_by':os.environ['USER'],
+                     'host_machine':socket.gethostname(),
+                     'created_on':str(datetimelib.now()),
+                     'comment':'This is a test output of a aircraft'}
+
+dataset = xr.Dataset(data_vars=variables,
+                     coords=coordinates,
+                     attrs=global_attributes)
+
+itime = 0
+while True:
+    datetime = fields["temp"].datetime
+    data = {}
+
+    # TODO if datetime < first time or  > last time do nothing?
+    timestep = pd.to_datetime(datetime)
+    point_idx = list(track.time.values).index(track.sel(time=timestep, method='nearest').time) # TODO if point outside domain do nothing
+    #print("point index {}".format(type(point_idx)), file=sys.stderr)
+    pam_time[itime] = timestep
+
+    for v, f in fields.items():
+        varfield, info = f.get()
+        datav = varfield[..., point_idx] # we want this to be (Nhgt, something)
+        if datav.ndim == 1:
+            data[v] = datav[:, np.newaxis]
+        else:
+            data[v] = datav
+
+    pamtra = call_pamtra(**data, datetime=datetime)
+    pam_Ze[:, itime] = pamtra.r['Ze'][0,0,:,0,0,0] # this extracts only 1 frequency... need to think about it
+    pam_hgt[:, itime] = pamtra.r['radar_hgt'][0, 0, :]
+
+
+    itime += 1
+    if itime == Ntimesperoutput:
+        dataset.to_netcdf('aircraft_{}.nc'.format(pd.to_datetime(pam_time[0].values).strftime('%Y%m%d-%H%M%S')))
+        itime = 0
+
+    if info == Action.GET_FOR_RESTART:
+        break
+