Commit 58a6e407 authored by BorjaEst's avatar BorjaEst
Browse files

Edit source to allow output split by time periods

parent 6410a0f5
......@@ -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):
......
......@@ -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')
......@@ -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")
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment