.

toargridding/grids.py

@@ -13,8 +13,9 @@ from toargridding.toar_rest_client import (
 )
 from toargridding.metadata import (
     Variables,
+    Variable,
+    Coordinates,
     Coordinate,
-    get_variable_attributes,
     get_global_attributes,
 )
 
@@ -55,19 +56,23 @@ class RegularGrid(GridDefinition):
         # TODO make sure only sensible resolutions
 
         self.lat = Coordinate.from_resolution(
-            Variables.latitude, lat_resolution, min=-90, max=90
+            Coordinates.latitude, lat_resolution, min=-90, max=90
         )
         self.lon = Coordinate.from_resolution(
-            Variables.longitude, lon_resolution, min=-180, max=180
+            Coordinates.longitude, lon_resolution, min=-180, max=180
         )
-        self.spatial_shape = (self.lon.size, self.lat.size)
+        spatial_shape = (self.lon.size, self.lat.size)
+        spatial_size = self.lon.size * self.lat.size
+        self.dims = [
+            Coordinates.latitude.name,
+            Coordinates.longitude.name,
+            Coordinates.time.name,
+        ]
 
         self._as_xy_index = np.dstack(
             np.meshgrid(range(self.lat.size), range(self.lon.size))
         ).reshape(-1, 2)
-        self._as_i_index = (
-            np.arange(self.lon.size * self.lat.size).reshape(self.spatial_shape).T
-        )
+        self._as_i_index = np.arange(spatial_size).reshape(spatial_shape).T
 
     def as_xarray(self, timeseries_per_statistic, metadata):
         datasets = dict()
@@ -79,7 +84,6 @@ class RegularGrid(GridDefinition):
             datasets[statistic] = self.create_dataset(cell_statistics, metadata)
 
         return datasets
-        print(cls)
 
     def clean_coords(self, coords: pd.DataFrame):
         valid_coords = coords.notna().all(axis=1)
@@ -103,13 +107,13 @@ class RegularGrid(GridDefinition):
 
     def get_cell_statistics(self, groups):
         stats = {
-            Variables.mean.name: groups.mean(),
-            Variables.std.name: groups.std(),
-            Variables.n.name: groups.count(),
+            Variables.mean: groups.mean(),
+            Variables.std: groups.std(),
+            Variables.n: groups.count(),
         }
-        idx = stats[Variables.n.name] == 0
-        stats[Variables.mean.name][idx] = self.fill_value
-        stats[Variables.std.name][idx] = self.fill_value
+        idx = stats[Variables.n] == 0
+        stats[Variables.mean][idx] = self.fill_value
+        stats[Variables.std][idx] = self.fill_value
 
         return stats
 
@@ -118,34 +122,18 @@ class RegularGrid(GridDefinition):
         time = pd.DatetimeIndex(random_df.columns)
 
         gridded_ds = self.get_empty_grid(time, metadata)
-        for cell_statistic, aggregated_timeseries in cell_statistics.items():
-            gridded_ds = self.assign_statistic_to_grid(
-                gridded_ds, cell_statistic, aggregated_timeseries, time
-            )
+        for variable, aggregated_data in cell_statistics.items():
+            data_array_dict = self.get_data_array_dict(time, aggregated_data, variable)
+            gridded_ds = gridded_ds.assign(data_array_dict)
 
         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)
-        print(values.shape)
-        gridded_ds = gridded_ds.assign(  # maybe better all in one ? (memory)
-            {
-                cell_statistic: (
-                    [
-                        Variables.latitude.name,
-                        Variables.longitude.name,
-                        Variables.time.name,
-                    ],
-                    values,
-                )
-            }
-        )
-
-        return gridded_ds
+    def get_data_array_dict(self, time, aggregated_data, variable):
+        gridded_data = self.create_gridded_data(time, aggregated_data)
+        gridded_variable = Variable.from_data(gridded_data, variable)
+        return {variable.name: gridded_variable.as_data_array(self.dims)}
 
-    def create_gridded_statistics(self, time, grouped_timeseries):
+    def create_gridded_data(self, time, grouped_timeseries):
         values = np.empty((self.lat.size, self.lon.size, time.size))
         values[...] = self.fill_value