diff --git a/toargridding/grids.py b/toargridding/grids.py index e114e9579e5b5c8ee77b6334c6bef5e3e2321c5c..ae8abd93aecd2a91c115886a62538a909981ae0e 100644 --- a/toargridding/grids.py +++ b/toargridding/grids.py @@ -14,7 +14,7 @@ from toargridding.toar_rest_client import ( STATION_LON, COORDS, ) -from toargridding.metadata import GLOBAL, LATITUDE, LONGITUDE +from toargridding.metadata import GLOBAL, LATITUDE, LONGITUDE, get_variable_attributes GridType = Enum("GridType", ["regular"]) @@ -116,18 +116,36 @@ class RegularGrid(GridDefinition): gridded_ds = self.get_empty_grid(time, metadata) for cell_statistic, grouped_timeseries in cell_statistics.items(): - values = np.empty((self.lat.size, self.lon.size, time.size)) - index = self._as_xy_index[grouped_timeseries.index] - values[index.T[0], index.T[1]] = grouped_timeseries.values.reshape( - -1, time.size + gridded_ds = self.assign_statistic_to_grid( + gridded_ds, cell_statistic, grouped_timeseries, time ) - gridded_ds = gridded_ds.assign( # maybe better all in one ? (memory) - {cell_statistic: (["latitude", "longitude", "time"], values)} - ) + return gridded_ds + + def assign_statistic_to_grid( + self, gridded_ds: xr.Dataset, cell_statistic, grouped_timeseries, time + ): + values = self.create_gridded_statistics(time, grouped_timeseries) + metadata = get_variable_attributes(cell_statistic) + + print(cell_statistic) + + gridded_ds = gridded_ds.assign( # maybe better all in one ? (memory) + {cell_statistic: (["latitude", "longitude", "time"], values)} + ) return gridded_ds + def create_gridded_statistics(self, time, grouped_timeseries): + values = np.empty((self.lat.size, self.lon.size, time.size)) + + index = self._as_xy_index[grouped_timeseries.index] + values[index.T[0], index.T[1]] = grouped_timeseries.values.reshape( + -1, time.size + ) + + return values + def as_cell_index(self, coords): id_x = self.coord_to_index( coords[STATION_LAT], self.lat.min(), self.lat_resolution