Add possibility to launch an external rule

+ some plotter improvements
+ some namelist improvements
+ add datacube extraction
This commit is contained in:
Noe Brucy
2022-08-23 10:10:42 +02:00
parent fcfd7a8dab
commit cdafe5cb61
5 changed files with 207 additions and 167 deletions
+7 -3
View File
@@ -2,8 +2,10 @@
import numpy as np
from functools import partial
from baseprocessor import Rule
import snapshotprocessor
try:
from mpi4py.futures import MPIPoolExecutor
@@ -25,9 +27,8 @@ def _map_aux(fun, path, path_out, params, run_num, **kwargs):
return fun(snap, **kwargs)
def _map_rule(snap, rule, arg, overwrite, overwrite_dep):
return snap.process(rule, arg, overwrite, overwrite_dep)
def _map_rule(snap, rule, **kwargs):
return snap.process(rule, **kwargs)
class Aggregator:
def get_snap_list(self, select=None):
@@ -41,6 +42,9 @@ class Aggregator:
def map(self, func, select=None, num_process=None, **kwargs):
if isinstance(func, Rule):
return self.map(_map_rule, select, num_process, rule=func, **kwargs)
snaps = self.get_snap_list(select)
if num_process is None:
+55 -15
View File
@@ -14,19 +14,19 @@ from tables import HDF5ExtError
from params import default_params, load_params
from units import U
class Rule:
def __init__(
self,
postproc,
process,
description="",
group="",
dependencies=[],
kind="snapshot",
unit=U.none,
name="",
):
self.postproc = postproc
self.name=name
self.process_fn = process
self.dependencies = dependencies
self.group = group
@@ -39,8 +39,6 @@ class Rule:
return self.process_fn(arg, **kwargs)
else:
return self.process_fn(**kwargs)
class BaseProcessor:
"""
Base class for processors, should not be instanciated
@@ -82,20 +80,39 @@ class BaseProcessor:
arg=None,
overwrite=False,
overwrite_dep=False,
skip_dep=False,
select=None,
**kwargs,
):
"""
Process the rule `to_process`
"""
):
self.overwrite_dep = overwrite_dep
self.just_done = []
""" Process the rule 'to_process'
Parameters
----------
to_process : str of Rule
name of the rule to process or Rule object with nonempty rule.name
arg : optional
argument to give to the rule
overwrite : bool, optional
Force redo if already done
overwrite_dep : bool, optional
Force redoing of the dependencies even if already done
skip_dep : bool, optional
Skip the dependency checks (assume they are already done)
select : dict, optional
Select object (see RunSelector) to only select some run/snapshot
"""
if to_process in self.rules:
rule = self.rules[to_process]
return self._solve_and_process_rule(
to_process, rule, arg, overwrite, select, **kwargs
to_process, rule, arg, overwrite, skip_dep, select, **kwargs
)
elif isinstance(to_process, Rule):
rule = to_process
return self._solve_and_process_rule(
rule.name, rule, arg, overwrite, skip_dep, select, **kwargs
)
else:
self._log(
@@ -106,9 +123,30 @@ class BaseProcessor:
)
def _solve_and_process_rule(
self, name, rule, arg, overwrite=False, select=None, **kwargs
self, name, rule, arg, overwrite=False, skip_dep=False, select=None, **kwargs
):
updated = self._solve_dependencies(name, rule, arg, overwrite, select)
"""Resolve dependencies and proceed in the processing of a rule
Parameters
----------
name : str
name of the rule
rule : Rule
rule object
overwrite : bool, optional
Force redo if already done
skip_dep : bool, optional
Skip the dependency checks (assume they are already done)
select : dict, optional
Select object (see RunSelector) to only select some run/snapshot
Returns
-------
The outbut of self._process_rule
"""
updated = False
if not skip_dep:
updated = self._solve_dependencies(name, rule, arg, overwrite, select)
overwrite_rule = overwrite or updated
return self._process_rule(name, rule, arg, overwrite_rule, select, **kwargs)
@@ -147,7 +185,7 @@ class BaseProcessor:
return overwrite
def _process_rule(self, name, rule, arg, overwrite=False, select=None, **kwargs):
if arg is not None:
if arg is not None and not isinstance(arg, BaseProcessor):
name_full = rule.group + "/" + name + "_" + str(arg)
else:
name_full = rule.group + "/" + name
@@ -418,7 +456,6 @@ class HDF5Container(BaseProcessor):
name = transform_name + "_" + rule_src_name
self.rules[name] = Rule(
self,
fn,
group=group,
unit=unit,
@@ -469,3 +506,6 @@ def oct_vect_getter(name, i, dset):
def norm_getter(name, dset):
return np.sqrt(np.sum(dset[name] ** 2, axis=1))
def oct_norm_getter(name, dset):
return np.sqrt(np.sum(dset[name] ** 2, axis=2))
+42 -64
View File
@@ -155,8 +155,7 @@ def line_integral_convolution(ax, map_h, map_v, extent, **kwargs):
class PlotRule(Rule):
"""
The rule class, speficic to plot.
Add an extra method, plot, that take the reference to an open hdf5 file (from pytables)
The rule class, specific to plot.
"""
def datafile(self, name, arg):
@@ -197,10 +196,16 @@ class Plotter(Aggregator, BaseProcessor):
"comp_frac": "$\chi$",
}
# Conversion table from namelist values (from amses config file) into LaTex strings
value_convert = {
# Conversion table from namelist values (from ramses config file) into LaTex strings
value_str = {
"sfr_avg_window": lambda x: "${:g}$ Myr".format(80 * x),
"Bx": lambda x: "${:g}$ $\\mu G$".format(7.6189439 * x),
"Bx": lambda x: "${:.1f}$ $\\mu G$".format(7.6189439 * x),
}
# Conversion table from namelist values (from ramses config file) into suitanle units
value_convert = {
"sfr_avg_window": lambda x: 80 * x, # Myr
"Bx": lambda x: x * 7.6189439,
}
def __init__(
@@ -560,11 +565,16 @@ class Plotter(Aggregator, BaseProcessor):
prop_label = self.label_convert[prop_name]
else:
prop_label = prop_name
prop_value = self.study.get_nml(nml_key, run)
if prop_name in self.value_convert:
try:
prop_value = self.study.get_nml(nml_key, run)
except KeyError:
return ""
if prop_name in self.value_str:
prop_value_str = self.value_str[prop_name](prop_value)
elif prop_name in self.value_convert:
prop_value_str = self.value_convert[prop_name](prop_value)
elif type(prop_value) in [int, float]:
prop_value_str = convert_exp(prop_value, digits=5)
prop_value_str = convert_exp(prop_value, digits=4)
else:
prop_value_str = str(prop_value)
return r"{} = {}".format(prop_label, prop_value_str)
@@ -584,7 +594,9 @@ class Plotter(Aggregator, BaseProcessor):
elif nml_key is not None:
if not type(nml_key) == list:
nml_key = [nml_key]
label_run = ", ".join(map(get_label_nml, nml_key))
lbl_list = map(get_label_nml, nml_key) # get namelist value
lbl_list = filter(lambda x: len(x) > 0, lbl_list) # Remove void labels
label_run = ", ".join(lbl_list)
if label is not None:
label_run = label + " (" + label_run + ")"
@@ -667,6 +679,7 @@ class Plotter(Aggregator, BaseProcessor):
colorbar=True,
embeded=False,
text_embeded=None,
text_kwargs={},
colorbar_embeded=None,
axes_indicator=None,
overtext_color="w",
@@ -814,10 +827,11 @@ class Plotter(Aggregator, BaseProcessor):
if text_embeded:
ax.text(
x=0.05,
y=0.95,
y=0.91,
s=title,
color=overtext_color,
transform=ax.transAxes,
**text_kwargs
)
else:
plt.title(title)
@@ -1021,7 +1035,7 @@ class Plotter(Aggregator, BaseProcessor):
c = c(data)[mask]
# Scatter plot
plt.scatter(part_h, part_v, s=s, c=c, **kwargs)
scatter = plt.scatter(part_h, part_v, s=s, c=c, **kwargs)
def _overlay_vector(
self,
@@ -1143,11 +1157,13 @@ class Plotter(Aggregator, BaseProcessor):
).express(unit_time)
color = colors(time)
else:
nml = self.study.get_nml(nml_color, run)
nml_value = self.study.get_nml(nml_color, run)
if os.path.basename(nml_color) in self.value_convert:
nml_value = self.value_convert[ os.path.basename(nml_color)](nml_value)
try:
color = colors[nml]
color = colors[nml_value]
except TypeError:
color = colors(nml)
color = colors(nml_value)
# Actual plot
if kind == "bar":
@@ -1209,6 +1225,7 @@ class Plotter(Aggregator, BaseProcessor):
label=None,
xunit=None,
yunit=None,
put_units=True,
xunit_coeff=1.0,
yunit_coeff=1.0,
xtransform=None,
@@ -1255,10 +1272,10 @@ class Plotter(Aggregator, BaseProcessor):
# Find proper labels
xlabel, xunit_old, xunit = self._ax_label_unit(
name_x, xlabel, xunit, xunit_coeff
name_x, xlabel, xunit, xunit_coeff, put_units=put_units,
)
ylabel, yunit_old, yunit = self._ax_label_unit(
name_y, ylabel, yunit, yunit_coeff
name_y, ylabel, yunit, yunit_coeff, put_units=put_units,
)
# If relevent, get time
@@ -1350,11 +1367,13 @@ class Plotter(Aggregator, BaseProcessor):
).express(unit_time)
color = colors(time)
else:
nml = self.study.get_nml(nml_color, run)
nml_value = self.study.get_nml(nml_color, run)
if os.path.basename(nml_color) in self.value_convert:
nml_value = self.value_convert[os.path.basename(nml_color)](nml_value)
try:
color = colors[nml]
color = colors[nml_value]
except TypeError:
color = colors(nml)
color = colors(nml_value)
if yerr is None or yerr_kind is None:
(base_line,) = plt.plot(x, y, label=label, color=color, **kwargs)
else:
@@ -1469,7 +1488,6 @@ class Plotter(Aggregator, BaseProcessor):
else:
name_rule = name_y + "_" + name_x
self.rules[name_rule] = PlotRule(
self,
partial(
self._plot,
"/series/" + logrule + "/" + name_x,
@@ -1485,20 +1503,15 @@ class Plotter(Aggregator, BaseProcessor):
This is where rules are defined
"""
self.rules = {
"plot": PlotRule(
self, lambda arg, **kwargs: self._plot(*arg, **kwargs), kind="comp"
"plot": PlotRule(lambda arg, **kwargs: self._plot(*arg, **kwargs), kind="comp"
),
"plot_run": PlotRule(
self, lambda arg, **kwargs: self._plot(*arg, **kwargs), kind="run"
"plot_run": PlotRule(lambda arg, **kwargs: self._plot(*arg, **kwargs), kind="run"
),
"plot_snapshot": PlotRule(
self, lambda arg, **kwargs: self._plot(*arg, **kwargs)
"plot_snapshot": PlotRule(lambda arg, **kwargs: self._plot(*arg, **kwargs)
),
"plot_map": PlotRule(
self, lambda mapname, **kwargs: self._plot_map(mapname, **kwargs)
"plot_map": PlotRule(lambda mapname, **kwargs: self._plot_map(mapname, **kwargs)
),
"coldens": PlotRule(
self,
partial(
self._plot_map,
"coldens",
@@ -1509,7 +1522,6 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["coldens"],
),
"slice_T": PlotRule(
self,
partial(
self._plot_map,
"T",
@@ -1519,19 +1531,16 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["T"],
),
"alpha_disk": PlotRule(
self,
partial(self._plot_map, "alpha_disk", label=r"$\alpha$"),
"Map of the Shakura&Sunaev alpha parameter for disks",
dependencies=["alpha_disk"],
),
"alpha_grav": PlotRule(
self,
partial(self._plot_map, "alpha_grav", label=r"$\alpha_g$"),
"Map of the grav Shakura&Sunaev alpha parameter for disks",
dependencies=["alpha_grav"],
),
"coldens_l": PlotRule(
self,
partial(
self._plot_map,
"coldens",
@@ -1543,7 +1552,6 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["coldens", "levels"],
),
"slice_rho_v": PlotRule(
self,
partial(
self._plot_map,
"slice_rho",
@@ -1555,7 +1563,6 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["slice_rho"],
),
"jeans_ratio": PlotRule(
self,
partial(
self._plot_map,
"jeans_ratio",
@@ -1568,7 +1575,6 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["jeans_ratio", "levels"],
),
"Q": PlotRule(
self,
partial(
self._plot_map,
"Q",
@@ -1581,25 +1587,21 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["Q"],
),
"rho_pdf": PlotRule(
self,
partial(self._plot_hist, "rho_pdf"),
"$\rho$-PDF",
dependencies=["rho_pdf"],
),
"rho_pdf_mw": PlotRule(
self,
partial(self._plot_hist, "rho_pdf_mw"),
"Mass weighted $\rho$-PDF",
dependencies=["rho_pdf_mw"],
),
"cos_pdf": PlotRule(
self,
partial(self._plot_hist, "cos_pdf"),
"cos-PDF",
dependencies=["cos_pdf", "mwa_speed"],
),
"avg_coldens_pdf": PlotRule(
self,
partial(
self._plot_hist,
"avg_time_coldens_pdf_z",
@@ -1612,19 +1614,16 @@ class Plotter(Aggregator, BaseProcessor):
dependencies={"avg_time_coldens_pdf": "z"},
),
"T_pdf": PlotRule(
self,
partial(self._plot_hist, "T_pdf"),
"T-PDF on a 2D slice",
dependencies=["T_pdf"],
),
"P_pdf": PlotRule(
self,
partial(self._plot_hist, "P_pdf"),
"P-PDF on a 2D slice ",
dependencies=["P_pdf"],
),
"Brho": PlotRule(
self,
partial(
self._plot,
"/datasets/Brho/rho",
@@ -1636,7 +1635,6 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["Brho"],
),
"Ek_Eb_rho": PlotRule(
self,
partial(
self._plot,
"/datasets/Ek_Eb_rho/rho",
@@ -1648,14 +1646,12 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["Ek_Eb_rho", "mwa_speed"],
),
"rho_prof": PlotRule(
self,
partial(self._plot, "/profile/axis", "/profile/rho_prof"),
"Density profile",
dependencies=["axis", "rho_prof"],
),
"pspec": PlotRule(self, self._pspec, dependencies={"pspec": None}),
"sbeta": PlotRule(
self,
partial(
self._plot,
"/comp/nml_cloud_params/beta_cool",
@@ -1669,7 +1665,6 @@ class Plotter(Aggregator, BaseProcessor):
},
),
"sbeta_onavg": PlotRule(
self,
partial(
self._plot,
"/comp/sbeta_onavg/beta",
@@ -1681,7 +1676,6 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["sbeta_onavg"],
),
"sink_mass": PlotRule(
self,
partial(
self._plot,
"/series/sinks_from_log/time",
@@ -1694,7 +1688,6 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["sinks_from_log"],
),
"ssm": PlotRule(
self,
partial(
self._plot,
"/series/sinks_from_log/time",
@@ -1707,7 +1700,6 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["ssm"],
),
"assfr": PlotRule(
self,
partial(
self._plot,
"/series/sfr_from_log/time",
@@ -1720,7 +1712,6 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["sfr_from_log"],
),
"issfr": PlotRule(
self,
partial(
self._plot,
"/series/sinks_from_log/time",
@@ -1733,7 +1724,6 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["issfr"],
),
"turb_rms": PlotRule(
self,
partial(
self._plot,
"/series/rms_from_log/time",
@@ -1745,7 +1735,6 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["rms_from_log"],
),
"turb_energy": PlotRule(
self,
partial(
self._plot,
"/series/rms_from_log/time",
@@ -1757,7 +1746,6 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["rms_from_log"],
),
"turb_power": PlotRule(
self,
partial(
self._plot,
"/series/rms_from_log/time",
@@ -1769,7 +1757,6 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["turb_power"],
),
"sigma": PlotRule(
self,
partial(
self._plot,
"/series/time",
@@ -1783,7 +1770,6 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["time_sigma"],
),
"mwa_B_int": PlotRule(
self,
partial(
self._plot,
"/series/time",
@@ -1796,7 +1782,6 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["time_mwa_B_int"],
),
"mass": PlotRule(
self,
partial(
self._plot,
"/series/time",
@@ -1809,7 +1794,6 @@ class Plotter(Aggregator, BaseProcessor):
dependencies=["time_mass"],
),
"max_fluct_coldens": PlotRule(
self,
partial(
self._plot,
"/series/time",
@@ -1879,7 +1863,6 @@ class Plotter(Aggregator, BaseProcessor):
for name in averageables:
self.rules["rad_" + name] = PlotRule(
self,
partial(
self._plot,
"/radial/radial_centers",
@@ -1890,7 +1873,6 @@ class Plotter(Aggregator, BaseProcessor):
)
self.rules["dispersion_rad_" + name] = PlotRule(
self,
partial(
self._plot,
"/radial/radial_centers",
@@ -1901,28 +1883,24 @@ class Plotter(Aggregator, BaseProcessor):
)
self.rules["avg_map_" + name] = PlotRule(
self,
partial(self._plot_map, "avg_map_" + name),
"Map of the radial average of {}".format(name),
dependencies=["avg_map_" + name],
)
self.rules["mwavg_map_" + name] = PlotRule(
self,
partial(self._plot_map, "mwavg_map_" + name),
"Map of the mass weighted radial average of {}".format(name),
dependencies=["avg_map_" + name],
)
self.rules["fluct_" + name] = PlotRule(
self,
partial(self._plot_map, "fluct_" + name, cmap="RdBu_r"),
"Fluctuation of {}".format(name),
dependencies=["fluct_" + name],
)
self.rules["pdf_" + name] = PlotRule(
self,
partial(self._plot_hist, "pdf_" + name, ylog=True),
"Probability density function of {} fluctuations".format(name),
dependencies=["pdf_" + name],
+69 -58
View File
@@ -27,6 +27,7 @@ from pymses.analysis import (
raytracing,
slicing,
splatting,
cube3d,
)
from pymses.filters import CellsToPoints, RegionFilter
from pymses.sources.hop.hop import HOP
@@ -41,6 +42,7 @@ from baseprocessor import (
HDF5Container,
Rule,
norm_getter,
oct_norm_getter,
simple_getter,
vect_getter,
oct_vect_getter,
@@ -379,6 +381,14 @@ class SnapshotProcessor(HDF5Container):
else:
self.fil = None
# Convert time unit
if isinstance(unit_time, str):
factor = self.get_nml(unit_time)
unit_time = U.Unit(
name=os.path.basename(unit_time),
base_unit=self.info["unit_time"],
coeff=factor)
time_in_right_unit = self.time * self.info["unit_time"].express(unit_time)
self.snapshot = Snapshot(
name=str(self.num),
@@ -636,6 +646,46 @@ class SnapshotProcessor(HDF5Container):
""" Azimuthal velocity """
return self.oct_getter_vect_phi(dset, "vel")
def datacube(self, getter, level=None, unit=U.none):
"""
Return a datacube of the source box.
Parameters
----------
getter : callable
A callable that extract the wanted data from a pymses dataset
unit : U.Unit
Unit of the resulting dataset
Returns
-------
A numpy array containing the slice
"""
unit = self._get_units(unit)
operator = ScalarOperator(getter, unit)
extractor = cube3d.CubeExtractor(self._amr, operator)
if level is None:
level = self.get_nml("amr_params/levelmin")
size = 1.0
cam = Camera(
center=self.params.pymses.center,
line_of_sight_axis="z",
region_size=[size, size],
distance=size / 2.0,
far_cut_depth=size / 2.0,
up_vector="y",
map_max_size=2 ** level,
)
cube = extractor.process(
cam, cube_size=1.0, resolution=2 ** level
).data
return cube
def slice(self, getter, ax_los="z", z=0.0, unit=U.none):
"""
Slice process function.
@@ -664,7 +714,8 @@ class SnapshotProcessor(HDF5Container):
datamap = slicing.SliceMap(self._amr, self._cam[ax_los], op, z=z)
return datamap.map.T
def ax_avg(self, getter, ax_los, unit=U.none, mass_weighted=True, surf_qty=False):
def ax_avg(self, oct_getter, ax_los, unit=U.none, mass_weighted=True, surf_qty=False):
"""
Map of the average of a quantity (given by getter) along an axis (ax_los)
Returns 2D array if getter returns a scalar quantity
@@ -673,12 +724,12 @@ class SnapshotProcessor(HDF5Container):
"""
unit = self._get_units(unit)
if surf_qty:
op = ScalarOperator(getter, unit)
op = ScalarOperator(oct_getter, unit)
else:
if mass_weighted:
def num(dset):
value = getter(dset)
value = oct_getter(dset)
rho = getter_rho(dset)
return rho * value
@@ -688,7 +739,7 @@ class SnapshotProcessor(HDF5Container):
else: # Volume weighted
def num(dset):
value = getter(dset)
value = oct_getter(dset)
return value
def denum(dset):
@@ -779,9 +830,11 @@ class SnapshotProcessor(HDF5Container):
self.unload_cells()
return data
def vol_pdf(self, getter, bins=100, logbins=False, weight_func=vol_func):
def vol_pdf(self, getter, bins=100, old_unit=None, unit=None, logbins=False, weight_func=vol_func):
self.load_cells()
data = getter(self.cells)
if old_unit is not None and unit is not None:
data *= old_unit.express(unit)
if logbins:
data = np.log10(data)
weights = weight_func(self.cells)
@@ -987,7 +1040,7 @@ class SnapshotProcessor(HDF5Container):
dmap_omega = rt_omega.process(self._cam[ax_los]).map.T
return dmap_omega
def _toomreQ_disk(self, ax_los, omega_approx=True, G1_units=True, coarsen_factor=1):
def _toomreQ_disk(self, ax_los, omega_approx=False, G1_units=False, coarsen_factor=1):
"""
Compute the Toomre Q parameter
"""
@@ -1347,6 +1400,8 @@ class SnapshotProcessor(HDF5Container):
for key in header:
if units[header.index(key)] == "m":
df[key] = df[key] * self.info["unit_mass"].express(U.Msun)
if "M" in df:
df["msink"] = df["M"]
else:
header = [
"id",
@@ -1369,6 +1424,7 @@ class SnapshotProcessor(HDF5Container):
"Teff",
]
df = pd.read_csv(csv_name, header=None, names=header)
return {key: df[key].values for key in df}
def _pspec(self, overwrite_file=False, **kwargs):
@@ -1566,7 +1622,7 @@ class SnapshotProcessor(HDF5Container):
# sort it and apply the sorting to the coordinates
# this means that the particules of group 1 are written first then of group 2 etc...
ind_sort = np.argsort(group_ids)
cells_group = val_mat[ind_sort]
cells_group = cells_group[ind_sort]
cells_group[6] = group_ids[ind_sort]
# Make it a pandas' DataFrame
@@ -1587,14 +1643,12 @@ class SnapshotProcessor(HDF5Container):
self.rules = {
# Base rules
"coldens": Rule(
self,
self._coldens,
"Column density",
"/maps",
unit=self.info["unit_density"] * self.info["unit_length"],
),
"T_mwavg": Rule(
self,
partial(
self.ax_avg,
getter_T,
@@ -1606,7 +1660,6 @@ class SnapshotProcessor(HDF5Container):
unit=self.info["unit_pressure"] / self.info["unit_density"],
),
"omega_mwavg": Rule(
self,
partial(self._omega_average),
"Ax mass-weighted averaged azimuthal of angular frequency",
"/maps",
@@ -1614,7 +1667,6 @@ class SnapshotProcessor(HDF5Container):
/ (self.info["unit_length"] / self.lbox),
),
"alpha_disk": Rule(
self,
self._alpha_disk,
"Map of the Shakura&Sunaev alpha parameter for disks",
"/maps",
@@ -1627,7 +1679,6 @@ class SnapshotProcessor(HDF5Container):
],
),
"alpha_grav": Rule(
self,
self._alpha_grav,
"Map of the graviational contrib to\
Shakura&Sunaev alpha parameter for disks",
@@ -1636,39 +1687,33 @@ class SnapshotProcessor(HDF5Container):
dependencies=["avg_map_coldens", "avg_map_T_mwavg"],
),
"T": Rule(
self,
self._temperature,
"Temperature slice",
"/maps",
dependencies=["slice_rho"],
unit=self.info["unit_pressure"] / self.info["unit_density"],
),
"levels": Rule(
self, self._levels, "Max level within line of sight", "/maps"
"levels": Rule(self._levels, "Max level within line of sight", "/maps"
),
"jeans": Rule(
self,
self._jeans,
"Jeans length slice",
"/maps",
dependencies=["slice_rho", "T"],
),
"jeans_ratio": Rule(
self,
self._jeans_ratio,
"Jeans' length divided by the max resolution",
"/maps",
dependencies=["jeans", "levels"],
),
"Q": Rule(
self,
self._toomreQ_disk,
"Toomre Q parameter for a Keplerian disk",
"/maps",
dependencies=["coldens", "T_mwavg", "omega_mwavg"],
),
"sinks": Rule(
self,
self._sinks,
group="/datasets",
unit={
@@ -1702,20 +1747,17 @@ class SnapshotProcessor(HDF5Container):
"vz_gas": "unit_velocity",
},
),
"pspec": Rule(self, self._pspec, "Power spectrum", "/hdf5"),
"write_particles": Rule(
self, self._write_particles, "Particles file", "/hdf5"
"pspec": Rule(self._pspec, "Power spectrum", "/hdf5"),
"write_particles": Rule(self._write_particles, "Particles file", "/hdf5"
),
"write_cells": Rule(self, self._write_cells, "Cells file", "/hdf5"),
"write_cells": Rule(self._write_cells, "Cells file", "/hdf5"),
"filaments": Rule(
self,
self._filaments,
"Filaments",
"/datasets",
dependencies={self.params.filaments.datamap: "z"},
),
"filaments_forces": Rule(
self,
self._filaments_forces,
"Filaments",
"/datasets",
@@ -1730,14 +1772,12 @@ class SnapshotProcessor(HDF5Container):
),
# Helpers
"radial_bins": Rule(
self,
self._radial_bins,
"Radial bins",
"/radial",
unit=self.info["unit_length"] / self.lbox,
),
"radial_centers": Rule(
self,
self._radial_centers,
"Centers of radial bins",
"/radial",
@@ -1745,21 +1785,18 @@ class SnapshotProcessor(HDF5Container):
unit=self.info["unit_length"] / self.lbox,
),
"rr": Rule(
self,
self._rr,
"Coordinate map",
"/maps",
unit=self.info["unit_length"] / self.lbox,
),
"bins_on_map": Rule(
self,
self._bins_on_map,
"Convert map coordinates to bins",
"/maps",
dependencies=["radial_bins", "rr"],
),
"B_int": Rule(
self,
self._B_int,
"Magnetic intensity slice",
"/maps",
@@ -1768,14 +1805,12 @@ class SnapshotProcessor(HDF5Container):
),
# PDF
"rho_pdf": Rule(
self,
partial(self.vol_pdf, partial(simple_getter, "rho"), logbins=True),
"Global rho-PDF",
"/hist",
unit=self.info["unit_density"],
),
"rho_pdf_mw": Rule(
self,
partial(
self.vol_pdf,
partial(simple_getter, "rho"),
@@ -1787,14 +1822,12 @@ class SnapshotProcessor(HDF5Container):
unit=self.info["unit_density"],
),
"T_pdf": Rule(
self,
partial(self.vol_pdf, getter_T, logbins=True),
"Global T-PDF",
"/hist",
unit=self.info["unit_pressure"] / self.info["unit_density"],
),
"cos_pdf": Rule(
self,
partial(self.cos_vfluct_B),
"Global cos fluctuation-PDF",
"/hist",
@@ -1802,14 +1835,12 @@ class SnapshotProcessor(HDF5Container):
unit=U.none,
),
"Brho": Rule(
self,
self._Brho,
"Average of B as a function of rho",
"/datasets",
unit={"rho": self.info["unit_density"], "B": self.info["unit_mag"]},
),
"Ek_Eb_rho": Rule(
self,
self._Ek_Eb_rho,
"Average of Ek/Eb as a function of rho",
"/datasets",
@@ -1818,14 +1849,12 @@ class SnapshotProcessor(HDF5Container):
),
# Profiles
"axis": Rule(
self,
partial(self._get_axis),
"Axis",
"/profile",
unit=self.info["unit_length"],
),
"rho_prof": Rule(
self,
partial(self._plane_avg_uniform, partial(simple_getter, "rho")),
"Rho profile",
"/profile",
@@ -1834,28 +1863,24 @@ class SnapshotProcessor(HDF5Container):
),
# globals
"time_num": Rule(
self,
lambda: self.info["time"],
"Time",
"/globals",
unit=self.info["unit_time"],
),
"mass": Rule(
self,
partial(self.sum, mass_func, mass_weighted=False),
"Total mass",
"/globals",
unit=self.info["unit_density"] * self.info["unit_length"] ** 3,
),
"mwa_speed": Rule(
self,
partial(self.vol_avg, partial(simple_getter, "vel")),
"Mass weighted speed average",
"/globals",
unit=self.info["unit_velocity"],
),
"mwa_sigma": Rule(
self,
self._mwa_sigma,
"Mass weighted speed average",
"/globals",
@@ -1863,7 +1888,6 @@ class SnapshotProcessor(HDF5Container):
unit=self.info["unit_velocity"],
),
"mwa_B_int": Rule(
self,
partial(self.vol_avg, getter_B_int),
"Mass weighted Magnetic intensity average",
"/globals",
@@ -1888,7 +1912,6 @@ class SnapshotProcessor(HDF5Container):
oct_getter = getter
self.rules["slice_" + name] = Rule(
self,
partial(self.slice, getter, z=0.0, unit=unit),
"{} slice".format(name),
"/maps",
@@ -1896,7 +1919,6 @@ class SnapshotProcessor(HDF5Container):
)
self.rules[name + "_mwavg"] = Rule(
self,
partial(self.ax_avg, oct_getter, mass_weighted=True, unit=unit),
"Ax mass-weighted averaged {}".format(name),
"/maps",
@@ -1904,7 +1926,6 @@ class SnapshotProcessor(HDF5Container):
)
self.rules[name + "_avg"] = Rule(
self,
partial(self.ax_avg, oct_getter, mass_weighted=False, unit=unit),
"Ax averaged {}".format(name),
"/maps",
@@ -1944,7 +1965,8 @@ class SnapshotProcessor(HDF5Container):
# Norm
generic_rule(
field + "_norm", partial(norm_getter, field), self.unit_key[field]
field + "_norm", partial(norm_getter, field), self.unit_key[field],
oct_getter=partial(oct_norm_getter, field)
)
else:
@@ -1958,7 +1980,6 @@ class SnapshotProcessor(HDF5Container):
self.rules["rad_avg_" + name] = Rule(
self,
partial(self._rad_avg, name),
"Azimuthal average of {}".format(name),
"/radial",
@@ -1966,7 +1987,6 @@ class SnapshotProcessor(HDF5Container):
unit=unit,
)
self.rules["rad_mwavg_" + name] = Rule(
self,
partial(self._rad_avg, name, mass_weighted=True),
"Mass weighted azimuthal average of {}".format(name),
"/radial",
@@ -1974,7 +1994,6 @@ class SnapshotProcessor(HDF5Container):
unit=unit,
)
self.rules["rad_avg_map_" + name] = Rule(
self,
partial(self._rad_avg_map, name),
"Interpolated map of azimuthal average of {}".format(name),
"/maps",
@@ -1982,7 +2001,6 @@ class SnapshotProcessor(HDF5Container):
unit=unit,
)
self.rules["rad_mwavg_map_" + name] = Rule(
self,
partial(self._rad_avg_map, name, mass_weighted=True),
"Interpolated map of azimuthal average of {}".format(name),
"/maps",
@@ -1990,14 +2008,12 @@ class SnapshotProcessor(HDF5Container):
unit=unit,
)
self.rules["fluct_" + name] = Rule(
self,
partial(self._fluct_map, name),
"Fluctuation wrt to average of {}".format(name),
"/maps",
dependencies=[name, "rad_avg_map_" + name],
)
self.rules["dispersion_rad_" + name] = Rule(
self,
partial(self._dispersion_rad, name),
"radial RMS of {}".format(name),
"/radial",
@@ -2005,21 +2021,18 @@ class SnapshotProcessor(HDF5Container):
unit=unit,
)
self.rules["mwfluct_" + name] = Rule(
self,
partial(self._fluct_map, name, mass_weigthed=True),
"Fluctuation wrt to mass weighted average of {}".format(name),
"/maps",
dependencies=[name, "rad_mwavg_map_" + name],
)
self.rules["fluct_pdf_" + name] = Rule(
self,
partial(self._rad_fluct_pdf, name),
"Probability density function of {} fluctuations".format(name),
"/hist",
dependencies=["rr", "fluct_" + name],
)
self.rules["fluct_mwpdf_" + name] = Rule(
self,
partial(self._rad_fluct_pdf, name, mass_weigthed=True),
"Probability density function of {} mass weighted fluctuations".format(
name
@@ -2028,7 +2041,6 @@ class SnapshotProcessor(HDF5Container):
dependencies=["rr", "mwfluct_" + name],
)
self.rules["fit_fluct_pdf_" + name] = Rule(
self,
partial(self._fit_pdf, name),
"Fit the PDF of {} fluctuations".format(name),
"/hist",
@@ -2038,7 +2050,6 @@ class SnapshotProcessor(HDF5Container):
unit_bin = self.rules[name_bin].unit
if name_bin is not name:
self.rules["mbb_" + name + "_" + name_bin] = Rule(
self,
partial(self.mean_by_bins, name_bin, name),
"Mean of {} by bins of {}".format(name, name_bin),
"/hist",
+32 -25
View File
@@ -63,10 +63,13 @@ class StudyProcessor(Aggregator, HDF5Container):
self.runs = selector.runs
self.nums = selector.nums
run0 = self.runs[0]
self.info = selector.info[run0][self.nums[run0][0]]
self.namelist = selector.namelist
# Get postprocesor objets for each run and infos on them
self.snaps = {}
self.info = {}
for run in self.runs:
path_run = path + "/" + run
path_out_run = self.path_out + "/" + run
@@ -82,9 +85,7 @@ class StudyProcessor(Aggregator, HDF5Container):
unit_time=unit_time,
)
run0 = self.runs[0]
self.info = selector.info[run0][self.nums[run0][0]]
self.namelist = selector.namelist
# Save namelist and logs
if self.params.out.copy_info:
@@ -234,8 +235,19 @@ class StudyProcessor(Aggregator, HDF5Container):
value = snap.get_value(node_name)
return value
def get_nml(self, nml_key, run):
return self.namelist[run][nml_key]
def get_nml(self, nml_key=None, run=None):
if run is not None:
if nml_key is not None:
return self.namelist[run][nml_key]
else:
return self.namelist[run]
else:
if nml_key is not None:
return {run : self.namelist[run][nml_key] for run in self.runs}
else:
return self.namelist
def get_pdf_slope(self, name, run, num):
snap = self.snaps[run][num]
@@ -464,15 +476,19 @@ class StudyProcessor(Aggregator, HDF5Container):
time_gas = self.get_value("/series/coarse_step_from_log/time")
mass_gas = self.get_value("/series/coarse_step_from_log/mcons")
mass_sink = self.get_value("/series/sinks_from_log/mass_sink")
time_sink = self.get_value("/series/sinks_from_log/time")
total_mass = dict()
for run in self.runs:
if time_sink[run][-1] > time_gas[run][-1]:
time_sink[run] = time_sink[run][:-1]
mass_sink[run] = mass_sink[run][:-1]
# A bit specific ... needs to be generalized (TODO)
info = self.snaps[run][self.nums[run][0]].info
surface = (info["unit_length"].express(U.pc)) ** 2
m0 = self.get_coldens0(run) * surface # Initial mass in Msun
offset = time_gas[run].size - mass_sink[run].size
offset = time_gas[run].size - time_sink[run].size
mass_gas[run] = m0 + m0*mass_gas[run] # convert in Msun
total_mass[run] = mass_gas[run].copy()
total_mass[run][offset:] = mass_gas[run][offset:] + mass_sink[run] # re add sink_mass
@@ -574,7 +590,6 @@ class StudyProcessor(Aggregator, HDF5Container):
name = "time_" + glob_name
self.rules[name] = Rule(
self,
partial(
self._time_series,
partial(
@@ -627,7 +642,6 @@ class StudyProcessor(Aggregator, HDF5Container):
)
self.rules[name] = Rule(
self,
fn,
group="/comp",
unit=units,
@@ -640,7 +654,6 @@ class StudyProcessor(Aggregator, HDF5Container):
self.rules = {
# Read from log
"sinks_from_log": Rule(
self,
partial(
self._from_log,
["time", "mass_sink", "nb_sink"],
@@ -655,7 +668,6 @@ class StudyProcessor(Aggregator, HDF5Container):
},
),
"issfr": Rule(
self,
self._ssfr_from_mass_sink,
group="/series/sinks_from_log",
unit=U.ssfr,
@@ -663,7 +675,6 @@ class StudyProcessor(Aggregator, HDF5Container):
dependencies=["sinks_from_log"],
),
"ssm": Rule(
self,
self._surfacic_sink_mass,
group="/series/sinks_from_log",
unit=U.Msun / U.pc ** 2,
@@ -671,7 +682,6 @@ class StudyProcessor(Aggregator, HDF5Container):
dependencies=["sinks_from_log"],
),
"sfr_from_log": Rule(
self,
partial(self._from_log, ["time", "sfr"], self._extract_sfr_from_log),
group="/series",
unit={"time": U.year, "sfr": U.ssfr},
@@ -681,7 +691,6 @@ class StudyProcessor(Aggregator, HDF5Container):
},
),
"rms_from_log": Rule(
self,
partial(
self._from_log,
["time", "dt", "turb_rms", "turb_energy"],
@@ -706,7 +715,6 @@ class StudyProcessor(Aggregator, HDF5Container):
},
),
"coarse_step_from_log": Rule(
self,
partial(
self._from_log,
[
@@ -738,7 +746,6 @@ class StudyProcessor(Aggregator, HDF5Container):
},
),
"fine_step_from_log": Rule(
self,
partial(
self._from_log,
["time", "fine_step", "dt", "a", "mem_cells", "mem_parts"],
@@ -755,7 +762,6 @@ class StudyProcessor(Aggregator, HDF5Container):
},
),
"stellar_from_log": Rule(
self,
partial(
self._from_log,
["time", "mass", "lifetime", "id"],
@@ -770,7 +776,6 @@ class StudyProcessor(Aggregator, HDF5Container):
},
),
"turb_power": Rule(
self,
self._turb_power,
group="/series/rms_from_log",
unit={
@@ -784,7 +789,6 @@ class StudyProcessor(Aggregator, HDF5Container):
),
# Read from outputs
"time": Rule(
self,
partial(
self._time_series, partial(self.get_global, "/globals/time_num")
),
@@ -793,7 +797,6 @@ class StudyProcessor(Aggregator, HDF5Container):
dependencies=["time_num"],
),
"time_rho_prof": Rule(
self,
partial(
self._time_series, partial(self.get_snap_value, "/profile/rho_prof")
),
@@ -801,15 +804,20 @@ class StudyProcessor(Aggregator, HDF5Container):
dependencies={"time": None, "rho_prof": "__parent__"},
),
"time_coldens_pdf": Rule(
self,
partial(
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")
),
group="/series",
dependencies={"time": None},
),
"time_pdf_slope_coldens": Rule(
self,
partial(
self._time_series,
partial(
@@ -822,7 +830,6 @@ class StudyProcessor(Aggregator, HDF5Container):
dependencies={"time": None, "fit_pdf_coldens": "z"},
),
"sbeta_onavg": Rule(
self,
partial(self._sbeta_onavg),
group="/comp",
dependencies={
@@ -832,7 +839,6 @@ class StudyProcessor(Aggregator, HDF5Container):
),
# namelist
"nml": Rule(
self,
lambda nml_key: self._compare(
partial(self.get_nml, nml_key), use_num=False
),
@@ -854,6 +860,7 @@ class StudyProcessor(Aggregator, HDF5Container):
"turb_power",
"time_rho_prof",
"time_coldens_pdf",
"time_rho_pdf"
]:
self._gen_rule_avg(name)