Edit source to allow output split by time periods

o3skim/sources.py

@@ -19,12 +19,12 @@ class Source:
             logging.info("Load model '%s'", name)
             self._models[name] = Model(variables)
 
-    def skim(self):
+    def skim(self, groupby=None):
         for name, model in self._models.items():
             dirname = self._name + "_" + name
             os.makedirs(dirname, exist_ok=True)
             logger.info("Skim data from '%s'", dirname)
-            model.skim(dirname)
+            model.skim(dirname, groupby)
 
 
 class Model:
@@ -38,13 +38,13 @@ class Model:
             logger.debug("Load 'vmro3_zm' data")
             self.__get_vmro3_zm(**variables)
 
-    def skim(self, dirname):
+    def skim(self, dirname, groupby=None):
         if hasattr(self, '_tco3_zm'):
             logger.debug("Skim 'tco3_zm' data")
-            utils.to_netcdf(dirname, "tco3_zm", self._tco3_zm)
+            utils.to_netcdf(dirname, "tco3_zm", self._tco3_zm, groupby)
         if hasattr(self, '_vmro3_zm'):
             logger.debug("Skim 'vmro3_zm' data")
-            utils.to_netcdf(dirname, "vmro3_zm", self._vmro3_zm)
+            utils.to_netcdf(dirname, "vmro3_zm", self._vmro3_zm, groupby)
 
     @utils.return_on_failure("Error when loading 'tco3_zm'")
     def __get_tco3_zm(self, tco3_zm, **kwarg):

o3skim/utils.py

@@ -50,10 +50,25 @@ def create_empty_netCDF(fname):
     root_grp.close()
 
 
-def to_netcdf(dirname, name, dataset):
+def to_netcdf(dirname, name, dataset, groupby=None):
     """Creates or appends data to named netcdf files"""
-    years, dsx = zip(*dataset.groupby("time.year"))
-    fnames = [dirname + "/" + name + "_%s.nc" % y for y in years]
+    def split_by_year(dataset):
+        """Splits a dataset by year"""
+        years, dsx = zip(*dataset.groupby("time.year"))
+        fnames = [dirname + "/" + name + "_%s.nc" % y for y in years]
+        return fnames, dsx
+
+    def no_split(dataset):
+        """Does not split a dataset"""
+        dsx = (dataset,)
+        fnames = [dirname + "/" + name  + ".nc"]
+        return fnames, dsx
+
+    split_by = {
+        "year":    split_by_year
+    }
+    fnames, dsx = split_by.get(groupby, no_split)(dataset)
+    
     logging.info("Save dataset into: %s", fnames)
     [create_empty_netCDF(fn) for fn in fnames if not os.path.isfile(fn)]
     xr.save_mfdataset(dsx, fnames, mode='a')

tests/mockup_data.py

@@ -39,4 +39,4 @@ def dataset(name, coordinades):
 def netcdf(dirname, name, coordinades, **kwarg):
     """Creates or appends data to a mock netcdf file"""
     ds = dataset(name, coordinades)
-    utils.to_netcdf(dirname, name, ds)
+    utils.to_netcdf(dirname, name, ds, groupby="year")