refactor dataset writing

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