diff --git a/components/groundstation.py b/components/groundstation.py index 22f6dc66d900bf9049b4b8e77cc7e9937404434c..a61db0d12f7b5ff8c3ff5386b43a850efb8026f7 100644 --- a/components/groundstation.py +++ b/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 +