implemented basic output with multiple timesteps per file. At the moment it is...

implemented basic output with multiple timesteps per file. At the moment it is hardcoded to print 1 file every 6 steps but it should be user configurable

components/groundstation.py

@@ -84,15 +84,116 @@ def call_pamtra(temp, pres, qv, qi, qc, qs, qr, qg, z_ifc, u, v, w,
                          hydro_q=np.stack([qc.T, qi.T, qr.T, qs.T, qg.T], axis=-1)[:, ::-1, :],
                          )
     pamtra.runPamtra(args.frequency)
-    pamtra.writeResultsToNetCDF(f"output_{datetime}.nc", xarrayCompatibleOutput=True)
-
-
+    #pamtra.writeResultsToNetCDF(f"output_{datetime}.nc", xarrayCompatibleOutput=True)
+    return pamtra # return the entire pamtra object and let the caller decide what to take as results
+
+# prepare output
+import xarray as xr
+import pandas as pd
+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_station_idx = xr.IndexVariable(
+        dims='sta_idx',
+        data=np.arange(0, ncoords),
+        attrs={'name':'sta_idx',
+               'long_name':'index of the station',
+               'comment':'look at variables sta_coords for location'})
+
+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 station'})
+
+#pam_sta_name.xr.DataArray(dims='station') # TODO perhaps one can provide station names as input, or instead of giving coordinates can also ask for certain locations to be read from a catalog
+pam_time = xr.DataArray(
+        dims=['time_idx'],
+        coords={'time_idx':pam_time_idx},
+        data=np.arange(0, Ntimesperoutput),
+        attrs={'name':'time',
+               'long_name':'unixtime milliseconds since 1970',
+               'units':'milliseconds'})
+
+pam_lat = xr.DataArray(
+        dims=['sta_idx'],
+        coords={'sta_idx':pam_station_idx},
+        data=np.rad2deg(coords[:, 1]),
+        attrs={'name':'lat',
+               'long_name':'station latitude',
+               'unit':'degrees North'})
+
+pam_lon = xr.DataArray(
+        dims=['sta_idx'],
+        coords={'sta_idx':pam_station_idx},
+        data=np.rad2deg(coords[:, 0]),
+        attrs={'name':'lon',
+               'long_name':'station longitude',
+               'unit':'degrees East'})
+
+pam_hgt = xr.DataArray(
+        dims=['sta_idx', 'hgt_idx'],
+        coords={'sta_idx':pam_station_idx,
+                'hgt_idx':pam_hgt_idx},
+        data=np.empty((ncoords, nlevel-1)),
+        attrs={'name':'hgt',
+               'long_name':'height above surface',
+               'units':'meters'})
+
+pam_Ze = xr.DataArray(
+        dims=['sta_idx', 'hgt_idx', 'time_idx'],
+        coords={'sta_idx':pam_station_idx,
+                'hgt_idx':pam_hgt_idx,
+                'time_idx':pam_time_idx},
+        data=np.empty((ncoords, nlevel-1, Ntimesperoutput)),
+        attrs={'name':'Ze',
+               'long_name':'unattenuated reflecitvity',
+               'units':'dBZ'})
+
+variables = {'height':pam_hgt,
+             'time':pam_time,
+             'lat':pam_lat,
+             'lon':pam_lon,
+             'Ze':pam_Ze}
+
+coordinates = {'time_idx':pam_time_idx,
+               'station_idx':pam_station_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'}
+
+dataset = xr.Dataset(data_vars=variables,
+                     coords=coordinates,
+                     attrs=global_attributes)
+
+itime = 0
 while True:
     datetime = fields["temp"].datetime
     data = {}
     for v, f in fields.items():
         data[v], info = f.get()
-        print(v, np.linalg.norm(data[v]), file=sys.stderr)
-    call_pamtra(**data, datetime=datetime)
+    pamtra = call_pamtra(**data, datetime=datetime)
+    pam_time[itime] = pd.to_datetime(datetime) 
+    pam_Ze[:, :, itime] = pamtra.r['Ze'][:,0,:,0,0,0] # this extracts only 1 frequency... need to think about it
+
+    itime += 1
+    if itime == Ntimesperoutput:
+        pam_hgt[:, :] = pamtra.r['radar_hgt'][:, 0, :]
+        dataset.to_netcdf('pamtraout_{}.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
+