Commit f23ea5cc authored by Clément Haëck's avatar Clément Haëck
Browse files

Add preprocess_finder kwargs

to open_mf_kw
so that it automatically upgrade a preprocess function
to the finder kind.

Also have preprocess_finder_(args|kwargs) key to pass
to preprocess function.
Watch out! preprocess_finder_args includes relative=True keyword !
parent 7ba77577
......@@ -50,6 +50,12 @@ def create_data(module, pregex, get_root, ARGS_DIR,
def get_data(args=None, **kwargs):
finder = get_finder(args, **kwargs)
if 'preprocess_finder' in open_mf_kw_def:
pp_args = open_mf_kw_def.pop('preprocess_finder_args', [])
pp_kw = open_mf_kw_def.pop('preprocess_finder_kwargs', {})
open_mf_kw_def['preprocess'] = finder.get_func_process_filename(
open_mf_kw_def.pop('preprocess_finder'), *pp_args, **pp_kw)
ds = xr.open_mfdataset(finder.get_files(), **open_mf_kw_def)
return ds
......
......@@ -17,19 +17,13 @@ def get_root(args, **kwargs):
return dirname
auto_attr = lib.data.create_data(__name__, pregex, get_root, ARGS_DIR)
def add_dims(ds, filename, finder):
matches = finder.get_matches(filename)
for dim in ['data', 'zone']:
ds = ds.assign_coords(**{dim: matches.get_matches(dim).get_match()})
ds = ds.expand_dims(dim)
return ds
def get_data(args=None, **kwargs):
def preprocess(ds, filename, finder):
matches = finder.get_matches(filename)
for dim in ['data', 'zone']:
ds = ds.assign_coords(**{dim: matches.get_matches(dim).get_match()})
ds = ds.expand_dims(dim)
return ds
finder = auto_attr['get_finder'](args, **kwargs)
f_pp = finder.get_func_process_filename(preprocess)
ds = xr.open_mfdataset(finder.get_files(), parallel=True,
preprocess=f_pp)
return ds
lib.data.create_data(__name__, pregex, get_root, ARGS_DIR,
open_mf_kw=dict(preprocess_finder=add_dims))
......@@ -17,4 +17,11 @@ def get_root(args=None, **kwargs):
return root
lib.data.create_data(__name__, pregex, get_root, ARGS_DIR)
def preprocess(ds):
ds = ds.rename(longitude='lon', latitude='lat')
ds = ds.assign_coords(lon=ds.lon-360)
return ds
lib.data.create_data(__name__, pregex, get_root, ARGS_DIR,
open_mf_kw=dict(preprocess=preprocess))
......@@ -22,5 +22,5 @@ def preprocess(ds):
return ds.expand_dims('time')
auto_attr = lib.data.create_data(__name__, pregex, get_root, ARGS_DIR,
open_mf_kw=dict(preprocess=preprocess))
lib.data.create_data(__name__, pregex, get_root, ARGS_DIR,
open_mf_kw=dict(preprocess=preprocess))
......@@ -18,21 +18,16 @@ def get_root(args=None, **kwargs):
return root
auto_attr = lib.data.create_data(
def k2c(ds, variable='SST'):
ds[variable] -= 273.15
ds[variable].attrs['units'] = 'celsius'
return ds
lib.data.create_data(
__name__, pregex, get_root, ARGS_DIR,
open_mf_kw=dict(
drop_variables=['analysis_uncertainty',
'mask', 'sea_ice_fraction'],
preprocess=k2c
))
def get_data(args=None, **kwargs):
ds = auto_attr['get_data'](args, **kwargs)
ds = k2c(ds, 'analysed_sst')
return ds
def k2c(ds, variable='SST'):
ds[variable] -= 273.15
ds[variable].attrs['units'] = 'celsius'
return ds
......@@ -19,19 +19,13 @@ def get_root(args=None, **kwargs):
return root
auto_attr = lib.data.create_data(__name__, pregex, get_root, ARGS_DIR)
def add_dims(ds, filename, finder):
matches = finder.get_matches(filename)
for dim in ['variable', 'zone']:
ds = ds.assign_coords(**{dim: matches.get_matches(dim).get_match()})
ds = ds.expand_dims(dim)
return ds
def get_data(args=None, **kwargs):
def preprocess(ds, filename, finder):
matches = finder.get_matches(filename)
for dim in ['variable', 'zone']:
ds = ds.assign_coords(**{dim: matches.get_matches(dim).get_match()})
ds = ds.expand_dims(dim)
return ds
finder = auto_attr['get_finder'](args, **kwargs)
f_pp = finder.get_func_process_filename(preprocess)
ds = xr.open_mfdataset(finder.get_files(), parallel=True,
preprocess=f_pp)
return ds
lib.data.create_data(__name__, pregex, get_root, ARGS_DIR,
open_mf_kw=dict(preprocess_finder=add_dims))
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