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

Fix histograms attributes

parent 87b21281
......@@ -102,14 +102,18 @@ def main(args):
hists.append(h)
hist = xr.combine_by_coords(hists)
hist = hist.rename({var: 'hist_' + var})
# Get a DataArray with the correct name
hist = hist[var].rename('hist_' + var)
hist = hist.assign_coords({bins_name: bins.edges[:-1]})
hist[bins_name].attrs['right_edge'] = VAR_RANGE[var][1]
hist[bins_name].attrs['VAR'] = var
hist.attrs['VAR'] = var
hists_var.append(hist)
m_next("Merging results")
hist = xr.merge(hists_var)
hist.attrs['VARS'] = VARS
hist.attrs.pop('VAR')
hist = hist.expand_dims({d: [args[d]] for d in INFILE_PARAMS})
......
......@@ -107,7 +107,8 @@ def main(args):
hists.append(h)
hist = xr.combine_by_coords(hists)
hist = hist.rename({var: 'hist_' + var})
# Get a DataArray with the correct name
hist = hist[var].rename('hist_' + var)
hist = hist.assign_coords({bins_name: bins.edges[:-1]})
hist[bins_name].attrs['right_edge'] = VAR_RANGE[var][1]
hist[bins_name].attrs['VAR'] = var
......@@ -117,6 +118,7 @@ def main(args):
m_next("Merging results")
hist = xr.merge(hists_var)
hist.attrs['VARS'] = VARS
hist.attrs.pop('VAR')
hist = hist.expand_dims({d: [args[d]] for d in INFILE_PARAMS})
......
......@@ -117,12 +117,17 @@ def variables(name_in='hist', allowconcat=True):
def decorated(ds, concat=allowconcat, variables=None, **kwargs):
if not allowconcat and concat:
raise ValueError("Cannot concatenate with this function.")
single_out = False # Return a single DataArray
# DataArray with no variables indication
if variables is None and isinstance(ds, xr.DataArray):
if 'VAR' not in ds.attrs:
return func(ds, **kwargs)
variables = ds.attrs['VAR']
# If string or array, we apply the function to a single variable
if isinstance(variables, str):
return func(ds[variables], **kwargs)
if isinstance(ds, xr.DataArray):
return func(ds, **kwargs)
variables = [variables]
single_out = True
# List of available variables
if variables is None:
......@@ -132,15 +137,18 @@ def variables(name_in='hist', allowconcat=True):
out_vars = []
for var, var_ext in zip(variables, variables_ext):
da = ds[var_ext]
da.attrs['VAR'] = var
da = ds if isinstance(ds, xr.DataArray) else ds[var_ext]
da.attrs['VAR'] = var # normally already set
out = func(da, **kwargs)
if concat:
out.attrs['VAR'] = var
if concat and not single_out:
out = out.expand_dims(var=[var])
else:
out = out.rename(var_name(out.name, var))
out_vars.append(out)
if single_out:
return out_vars[0]
if concat:
return xr.concat(out_vars, 'var')
out = xr.merge(out_vars)
......@@ -210,7 +218,7 @@ def get_median(pdf):
def get_total(hist):
"""Return total amount of quantity counted in histograms."""
bins_name = var_name('bins', hist.VAR)
return (hist * hist[bins_name]).sum(bins_name)
return (hist * hist[bins_name]).sum(bins_name).rename('total')
def resample_8D_multiple_years(ds, days=8):
......
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