less verbose
This commit is contained in:
Noe Brucy
+57
-22
@@ -131,7 +131,7 @@ class StudyProcessor(Aggregator, HDF5Container):
|
||||
super(StudyProcessor, self)._save_data(name_full, data, description, unit)
|
||||
self.save.get_node(name_full)._v_attrs.nums = self.nums
|
||||
|
||||
def _time_series(self, getter, arg=None):
|
||||
def time_series(self, getter, arg=None):
|
||||
series = {}
|
||||
for run in self.runs:
|
||||
series[run] = []
|
||||
@@ -140,17 +140,47 @@ class StudyProcessor(Aggregator, HDF5Container):
|
||||
series[run] = np.array(series[run], dtype=float)
|
||||
return series
|
||||
|
||||
def _compare(self, getter, use_num=True):
|
||||
prop = []
|
||||
for i, run in enumerate(self.runs):
|
||||
def compare(self, getter, use_num=True, select=None):
|
||||
|
||||
if select is None:
|
||||
runs = self.runs
|
||||
else:
|
||||
runs, _ = self.selector.select(**select)
|
||||
|
||||
prop = {}
|
||||
for i, run in enumerate(runs):
|
||||
if use_num:
|
||||
num = self.nums[run][0]
|
||||
prop.append(getter(run, num))
|
||||
prop[run] = getter(run, num)
|
||||
else:
|
||||
prop.append(getter(run))
|
||||
return np.array(prop)
|
||||
prop[run] = getter(run)
|
||||
return np.array(list(prop.keys()))
|
||||
|
||||
def _time_avg(self, name, start=None, end=None, span=None, group="/series"):
|
||||
def time_avg(self, name, start=None, end=None, span=None, unit_time=U.Myr, group="/series", select=None):
|
||||
"""Do the time average and quantiles of a time series
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
name of the array to average
|
||||
start : float, optional
|
||||
The average is taken between start and end or start + span, by default None
|
||||
end : float, optional
|
||||
The average is taken between start and end or end - span, by default None
|
||||
span : _type_, optional
|
||||
length of the averaging period (overrrided if both start and end are set), by default None
|
||||
unit_time : _type_, optional
|
||||
Time unit to use, by default U.Myr
|
||||
group : str, optional
|
||||
group of the data to average, by default "/series"
|
||||
select : dict, optional
|
||||
arguments to selector, by default None
|
||||
|
||||
Returns
|
||||
-------
|
||||
dict
|
||||
time average and quantiles
|
||||
"""
|
||||
serie0 = self.save.get_node(group + "/" + name + "/" + self.runs[0]).read()
|
||||
if len(serie0.shape) > 1:
|
||||
shape = [len(self.runs)] + list(serie0.shape[1:])
|
||||
@@ -166,11 +196,16 @@ class StudyProcessor(Aggregator, HDF5Container):
|
||||
q16 = np.zeros(shape)
|
||||
q84 = np.zeros(shape)
|
||||
|
||||
for i, run in enumerate(self.runs):
|
||||
serie = self.save.get_node(group + "/" + name + "/" + run).read()
|
||||
time = self.save.get_node(group + "/time/" + run).read()
|
||||
if select is None:
|
||||
runs = self.runs
|
||||
else:
|
||||
runs, _ = self.selector.select(**select)
|
||||
|
||||
for i, run in enumerate(runs):
|
||||
serie = self.get_value(group + "/" + name + "/" + run)
|
||||
time = self.get_value(group + "/time/" + run, unit=unit_time)
|
||||
if len(serie.shape) <= 1:
|
||||
mask = abs(serie) != np.inf
|
||||
mask = np.isfinite(serie)
|
||||
|
||||
if not ((start, end, span) == (None, None, None)):
|
||||
start_r, end_r = start, end
|
||||
@@ -202,7 +237,7 @@ class StudyProcessor(Aggregator, HDF5Container):
|
||||
v_max[i],
|
||||
) = np.percentile(serie, [0, 2.5, 16, 50, 84, 97.5, 100], axis=0)
|
||||
return {
|
||||
"runs": self.runs,
|
||||
"runs": np.array(runs),
|
||||
"mean": mean,
|
||||
"std": std,
|
||||
"median": median,
|
||||
@@ -591,7 +626,7 @@ class StudyProcessor(Aggregator, HDF5Container):
|
||||
|
||||
self.rules[name] = Rule(
|
||||
partial(
|
||||
self._time_series,
|
||||
self.time_series,
|
||||
partial(
|
||||
self.get_global, glob_group + "/" + glob_name, unload_cells=True
|
||||
),
|
||||
@@ -635,9 +670,9 @@ class StudyProcessor(Aggregator, HDF5Container):
|
||||
|
||||
def fn(arg=None, **kwargs):
|
||||
if arg is None:
|
||||
return self._time_avg(src_name, group=group_src, **kwargs)
|
||||
return self.time_avg(src_name, group=group_src, **kwargs)
|
||||
else:
|
||||
return self._time_avg(
|
||||
return self.time_avg(
|
||||
src_name + "_" + str(arg), group=group_src, **kwargs
|
||||
)
|
||||
|
||||
@@ -790,7 +825,7 @@ class StudyProcessor(Aggregator, HDF5Container):
|
||||
# Read from outputs
|
||||
"time": Rule(
|
||||
partial(
|
||||
self._time_series, partial(self.get_global, "/globals/time_num")
|
||||
self.time_series, partial(self.get_global, "/globals/time_num")
|
||||
),
|
||||
group="/series",
|
||||
unit="unit_time",
|
||||
@@ -798,28 +833,28 @@ class StudyProcessor(Aggregator, HDF5Container):
|
||||
),
|
||||
"time_rho_prof": Rule(
|
||||
partial(
|
||||
self._time_series, partial(self.get_snap_value, "/profile/rho_prof")
|
||||
self.time_series, partial(self.get_snap_value, "/profile/rho_prof")
|
||||
),
|
||||
group="/series",
|
||||
dependencies={"time": None, "rho_prof": "__parent__"},
|
||||
),
|
||||
"time_coldens_pdf": Rule(
|
||||
partial(
|
||||
self._time_series, partial(self.get_snap_value, "/hist/pdf_coldens")
|
||||
self.time_series, partial(self.get_snap_value, "/hist/pdf_coldens")
|
||||
),
|
||||
group="/series",
|
||||
dependencies={"time": None, "pdf_coldens": "__parent__"},
|
||||
),
|
||||
"time_rho_pdf": Rule(
|
||||
partial(
|
||||
self._time_series, partial(self.get_snap_value, "/hist/rho_pdf")
|
||||
self.time_series, partial(self.get_snap_value, "/hist/rho_pdf")
|
||||
),
|
||||
group="/series",
|
||||
dependencies={"time": None},
|
||||
),
|
||||
"time_pdf_slope_coldens": Rule(
|
||||
partial(
|
||||
self._time_series,
|
||||
self.time_series,
|
||||
partial(
|
||||
self.get_attr,
|
||||
"slope",
|
||||
@@ -839,7 +874,7 @@ class StudyProcessor(Aggregator, HDF5Container):
|
||||
),
|
||||
# namelist
|
||||
"nml": Rule(
|
||||
lambda nml_key: self._compare(
|
||||
lambda nml_key: self.compare(
|
||||
partial(self.get_nml, nml_key), use_num=False
|
||||
),
|
||||
group="/comp",
|
||||
|
||||
Reference in New Issue
Block a user