From fa396178c66c8076542b5c627630ce0b43201d36 Mon Sep 17 00:00:00 2001 From: Noe Brucy Date: Mon, 28 Nov 2022 18:01:18 +0100 Subject: [PATCH] black without mortimer --- aggregator.py | 4 +- baseprocessor.py | 38 ++++---- galturb/galaxy.py | 11 +-- galturb/sectors_extraction.py | 11 +-- plotter.py | 112 ++++++++++++----------- scripts/mngseg.py | 9 +- scripts/pspec.py | 16 +--- snapshotprocessor.py | 166 ++++++++++++++++++---------------- studyprocessor.py | 120 +++++++++++++----------- turbox/turbox.py | 107 +++++++++++++--------- utils/extract_velcube.py | 27 +++--- utils/runselector.py | 19 ++-- utils/snapshotselector.py | 33 ++++--- utils/units.py | 2 +- 14 files changed, 371 insertions(+), 304 deletions(-) diff --git a/aggregator.py b/aggregator.py index e3f349d..8ef131c 100644 --- a/aggregator.py +++ b/aggregator.py @@ -8,6 +8,7 @@ from utils.mypool import MyPool try: from mpi4py.futures import MPIPoolExecutor + mpi_loaded = True except ModuleNotFoundError: mpi_loaded = False @@ -27,6 +28,7 @@ def _map_aux(fun, path, path_out, params, run_num, **kwargs): def _map_rule(snap, rule, **kwargs): return snap.process(rule, **kwargs) + class Aggregator: def get_snap_list(self, select=None): @@ -43,7 +45,7 @@ class Aggregator: return self.map(_map_rule, select, num_process, rule=func, **kwargs) snaps = self.get_snap_list(select) - + if num_process is None: num_process = self.params.process.num_process diff --git a/baseprocessor.py b/baseprocessor.py index 1f21f50..b0e263d 100644 --- a/baseprocessor.py +++ b/baseprocessor.py @@ -29,9 +29,8 @@ class Rule: kind="snapshot", unit=U.none, name="", - ): - self.name=name + self.name = name self.process_fn = process self.dependencies = dependencies self.group = group @@ -45,6 +44,7 @@ class Rule: else: return self.process_fn(**kwargs) + class BaseProcessor: """ Base class for processors, should not be instanciated @@ -56,7 +56,6 @@ class BaseProcessor: rules = {} solve_self_dep = True - def __init__(self, path, path_out=".", params=None, tag=None): if params is None: self.params = default_params() @@ -77,8 +76,8 @@ class BaseProcessor: # Initialize logger self.logger = logging.getLogger(self.log_id) self.logger.propagate = False - logging_format = '%(levelname)s | %(asctime)s | %(name)s.%(funcName)s:%(lineno)d | %(message)s' - formatter = logging.Formatter(logging_format, datefmt = '%H:%M:%S') + logging_format = "%(levelname)s | %(asctime)s | %(name)s.%(funcName)s:%(lineno)d | %(message)s" + formatter = logging.Formatter(logging_format, datefmt="%H:%M:%S") if not self.logger.hasHandlers(): stream = logging.StreamHandler(sys.stdout) @@ -95,7 +94,7 @@ class BaseProcessor: self.logger.setLevel(logging.WARNING) for handler in self.logger.handlers: - handler.setFormatter(formatter) + handler.setFormatter(formatter) def process( self, @@ -106,7 +105,7 @@ class BaseProcessor: skip_dep=False, select=None, **kwargs, - ): + ): self.overwrite_dep = overwrite_dep self.just_done = [] """ Process the rule 'to_process' @@ -126,7 +125,7 @@ class BaseProcessor: 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( @@ -166,7 +165,7 @@ class BaseProcessor: ------- The outbut of self._process_rule """ - updated = False + updated = False if not skip_dep: updated = self._solve_dependencies(name, rule, arg, overwrite, select) overwrite_rule = overwrite or updated @@ -221,9 +220,7 @@ class BaseProcessor: self.just_done.append(name_full) return data else: - self.logger.info( - "Data for {} is already computed.".format(name_full) - ) + self.logger.info("Data for {} is already computed.".format(name_full)) def def_rules(self): for rule in self.rules: @@ -302,7 +299,9 @@ class HDF5Container(BaseProcessor): if not (unit is None or unit_old is None or unit_old == U.none): value = value * unit_old.express(unit) except NoSuchNodeError: - self.logger.error(f"The value {node_name} is node available", stack_info=True) + self.logger.error( + f"The value {node_name} is node available", stack_info=True + ) raise finally: if not open_before: @@ -441,11 +440,13 @@ class HDF5Container(BaseProcessor): except TypeError: data = np.array([data]) - group_name = os.path.dirname(name_full) + group_name = os.path.dirname(name_full) if group_name in self.save: group = self.save.get_node(group_name) - if not isinstance(group, class_name_dict['Group']): - self.logger.warning(f"{group_name} already there and no a group, deleting") + if not isinstance(group, class_name_dict["Group"]): + self.logger.warning( + f"{group_name} already there and no a group, deleting" + ) self.save.remove_node(group) self.save.create_array( group_name, @@ -552,7 +553,8 @@ def oct_vect_getter(name, i, dset): def norm_getter(name, dset): - return np.sqrt(np.sum(dset[name] ** 2, axis=1)) + 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)) \ No newline at end of file + return np.sqrt(np.sum(dset[name] ** 2, axis=2)) diff --git a/galturb/galaxy.py b/galturb/galaxy.py index 8d2b7ae..28a24a1 100644 --- a/galturb/galaxy.py +++ b/galturb/galaxy.py @@ -9,7 +9,7 @@ def get_gas_dm_stars(pp): # Load arrays try: pp.load_parts(keys=["pos", "vel", "mass", "epoch"]) - except: + except KeyError: pp.load_parts(keys=["pos", "vel", "mass"]) pp.load_cells(keys=["pos", "vel", "dx", "rho"]) @@ -249,10 +249,9 @@ def load_wrapper(pp, fun): def allinone(pp, redo=False): def fun(pp): return analyse_disk(pp), analyse_rings(pp, [4, 5, 6, 7, 8]) - - + try: - assert(not redo) + assert not redo sectors = pd.read_csv("{pp.run}/disk_{pp.run}_{pp.num}.csv") disk = pd.read_csv(f"{pp.run}/disk_{pp.run}_{pp.num}.csv") @@ -260,7 +259,7 @@ def allinone(pp, redo=False): res = load_wrapper(pp, fun) disk = pd.DataFrame({key: [res[0][key]] for key in res[0]}) sectors = pd.DataFrame({key: res[1][key] for key in res[1]}) - sectors.to_csv(f"{pp.run}/sectors_{pp.run}_{pp.num}.csv") - disk.to_csv(f"{pp.run}/disk_{pp.run}_{pp.num}.csv") + sectors.to_csv(f"{pp.run}/sectors_{pp.run}_{pp.num}.csv") + disk.to_csv(f"{pp.run}/disk_{pp.run}_{pp.num}.csv") return disk, sectors diff --git a/galturb/sectors_extraction.py b/galturb/sectors_extraction.py index 8d2b7ae..28a24a1 100644 --- a/galturb/sectors_extraction.py +++ b/galturb/sectors_extraction.py @@ -9,7 +9,7 @@ def get_gas_dm_stars(pp): # Load arrays try: pp.load_parts(keys=["pos", "vel", "mass", "epoch"]) - except: + except KeyError: pp.load_parts(keys=["pos", "vel", "mass"]) pp.load_cells(keys=["pos", "vel", "dx", "rho"]) @@ -249,10 +249,9 @@ def load_wrapper(pp, fun): def allinone(pp, redo=False): def fun(pp): return analyse_disk(pp), analyse_rings(pp, [4, 5, 6, 7, 8]) - - + try: - assert(not redo) + assert not redo sectors = pd.read_csv("{pp.run}/disk_{pp.run}_{pp.num}.csv") disk = pd.read_csv(f"{pp.run}/disk_{pp.run}_{pp.num}.csv") @@ -260,7 +259,7 @@ def allinone(pp, redo=False): res = load_wrapper(pp, fun) disk = pd.DataFrame({key: [res[0][key]] for key in res[0]}) sectors = pd.DataFrame({key: res[1][key] for key in res[1]}) - sectors.to_csv(f"{pp.run}/sectors_{pp.run}_{pp.num}.csv") - disk.to_csv(f"{pp.run}/disk_{pp.run}_{pp.num}.csv") + sectors.to_csv(f"{pp.run}/sectors_{pp.run}_{pp.num}.csv") + disk.to_csv(f"{pp.run}/disk_{pp.run}_{pp.num}.csv") return disk, sectors diff --git a/plotter.py b/plotter.py index fb36b11..1e26bb0 100644 --- a/plotter.py +++ b/plotter.py @@ -104,16 +104,15 @@ def quiver(ax, map_h, map_v, extent, key_v=None, lognorm=False, label="", **kwar min_v = np.min(norm_v) if key_v is None: - key_v = (max_v + min_v) / 2.0 + key_v = (max_v + min_v) / 2.0 key = f"${key_v:g}$ {label}" if lognorm: - lognorm_v = np.log10(norm_v) - map_h *= lognorm_v/norm_v - map_v *= lognorm_v/norm_v - key_v = np.log10(key_v) - + lognorm_v = np.log10(norm_v) + map_h *= lognorm_v / norm_v + map_v *= lognorm_v / norm_v + key_v = np.log10(key_v) # plot vector field vec_field = ax.quiver(hh, vv, map_h, map_v, units="width", **kwargs) @@ -203,7 +202,7 @@ class Plotter(Aggregator, BaseProcessor): # Conversion table from namelist values (from ramses config file) into suitanle units value_convert = { - "sfr_avg_window": lambda x: 80 * x, # Myr + "sfr_avg_window": lambda x: 80 * x, # Myr "Bx": lambda x: x * 7.6189439, } @@ -241,7 +240,6 @@ class Plotter(Aggregator, BaseProcessor): # log info self.log_id = "plotter({})".format(tag) - super(Plotter, self).__init__(path, path_out, params, tag) # Select runs @@ -421,7 +419,7 @@ class Plotter(Aggregator, BaseProcessor): if ax is not None: onefigure = True if not movie: - plot_filename = self._find_filename(name_full) + plot_filename = self._find_filename(name_full) for i, (run, num) in enumerate(run_num): @@ -465,7 +463,7 @@ class Plotter(Aggregator, BaseProcessor): if self.params.astrophysix.generate: df = rule.datafile(name, arg) df[filetype] = plot_filename - + if plot_info is not None: df.plot_info = plot_info if num is not None: @@ -499,7 +497,7 @@ class Plotter(Aggregator, BaseProcessor): if self.params.plot.tight_layout and close: plt.tight_layout(pad=1) - if self.params.out.save: + if self.params.out.save: os.makedirs(os.path.dirname(plot_filename), exist_ok=True) plt.savefig(plot_filename) self.logger.info(f"{plot_filename} plotted") @@ -565,10 +563,10 @@ class Plotter(Aggregator, BaseProcessor): prop_label = self.label_convert[prop_name] else: prop_label = prop_name - try: - prop_value = self.study.get_nml(nml_key, run) + try: + prop_value = self.study.get_nml(nml_key, run) except KeyError: - return "" + 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: @@ -594,8 +592,8 @@ class Plotter(Aggregator, BaseProcessor): elif nml_key is not None: if not type(nml_key) == list: nml_key = [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 + 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 and len(label) > 0: @@ -831,7 +829,7 @@ class Plotter(Aggregator, BaseProcessor): s=title, color=overtext_color, transform=ax.transAxes, - **text_kwargs + **text_kwargs, ) else: plt.title(title) @@ -971,9 +969,7 @@ class Plotter(Aggregator, BaseProcessor): if sinks: try: self.current_processor.sinks() - data = pd.DataFrame( - self.current_processor.get_value("/datasets/sinks") - ) + data = pd.DataFrame(self.current_processor.get_value("/datasets/sinks")) part_pos = data[["x", "y", "z"]].values unit_length /= self.current_processor.lbox except KeyError: @@ -1038,6 +1034,8 @@ class Plotter(Aggregator, BaseProcessor): # Scatter plot scatter = plt.scatter(part_h, part_v, s=s, c=c, **kwargs) + return scatter + def _overlay_vector( self, name, @@ -1160,7 +1158,9 @@ class Plotter(Aggregator, BaseProcessor): else: 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) + nml_value = self.value_convert[os.path.basename(nml_color)]( + nml_value + ) try: color = colors[nml_value] except TypeError: @@ -1219,32 +1219,31 @@ class Plotter(Aggregator, BaseProcessor): def plot( self, - x:np.array, - y:np.array, - xlabel:str="", - ylabel:str="", - label:str="", - xscale:str="linear", - yscale:str="linear", - fit:str=None, - fitlabel:str=None, - smooth:float=0, + x: np.array, + y: np.array, + xlabel: str = "", + ylabel: str = "", + label: str = "", + xscale: str = "linear", + yscale: str = "linear", + fit: str = None, + fitlabel: str = None, + smooth: float = 0, nml_key=None, - run:str=None, - yerr:np.array=None, - grid:bool=False, - put_time:bool=False, + run: str = None, + yerr: np.array = None, + grid: bool = False, + put_time: bool = False, unit_time=U.Myr, colors=None, nml_color=None, - legend:bool=False, + legend: bool = False, **kwargs, ): """ Generic plot routine, with x, y two numpy arrauys """ - # Option to smooth data for readability (beware) if smooth > 0: y = gaussian_filter1d(y, sigma=smooth) @@ -1283,7 +1282,9 @@ class Plotter(Aggregator, BaseProcessor): else: 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) + nml_value = self.value_convert[os.path.basename(nml_color)]( + nml_value + ) try: color = colors[nml_value] except TypeError: @@ -1321,8 +1322,8 @@ class Plotter(Aggregator, BaseProcessor): def _plot( self, - name_x:str, - name_y:str, + name_x: str, + name_y: str, node_arg=None, xlabel=None, ylabel=None, @@ -1363,10 +1364,18 @@ class Plotter(Aggregator, BaseProcessor): # Find proper labels xlabel, xunit_old, xunit = self._ax_label_unit( - name_x, xlabel, xunit, xunit_coeff, put_units=put_units, + name_x, + xlabel, + xunit, + xunit_coeff, + put_units=put_units, ) ylabel, yunit_old, yunit = self._ax_label_unit( - name_y, ylabel, yunit, yunit_coeff, put_units=put_units, + name_y, + ylabel, + yunit, + yunit_coeff, + put_units=put_units, ) # Manage the different forms in which the data may be stored : @@ -1426,8 +1435,7 @@ class Plotter(Aggregator, BaseProcessor): "Errorbar may be meaningless when ytransform is used" ) - self.plot(x, y, yerr=yerr, xlabel=xlabel, - ylabel=ylabel, run=run, **kwargs) + self.plot(x, y, yerr=yerr, xlabel=xlabel, ylabel=ylabel, run=run, **kwargs) if subname_x: hdf5_x.close() @@ -1518,13 +1526,15 @@ class Plotter(Aggregator, BaseProcessor): This is where rules are defined """ self.rules = { - "plot_comp": PlotRule(lambda arg, **kwargs: self._plot(*arg, **kwargs), kind="comp" + "plot_comp": PlotRule( + lambda arg, **kwargs: self._plot(*arg, **kwargs), kind="comp" ), - "plot_run": PlotRule(lambda arg, **kwargs: self._plot(*arg, **kwargs), kind="run" + "plot_run": PlotRule( + lambda arg, **kwargs: self._plot(*arg, **kwargs), kind="run" ), - "plot_snapshot": PlotRule(lambda arg, **kwargs: self._plot(*arg, **kwargs) - ), - "plot_map": PlotRule(lambda mapname, **kwargs: self._plot_map(mapname, **kwargs) + "plot_snapshot": PlotRule(lambda arg, **kwargs: self._plot(*arg, **kwargs)), + "plot_map": PlotRule( + lambda mapname, **kwargs: self._plot_map(mapname, **kwargs) ), "coldens": PlotRule( partial( @@ -1955,7 +1965,7 @@ class Plotter(Aggregator, BaseProcessor): self._gen_from_log("fine_step_from_log", name) for name in ["time", "dt", "a", "mem_cells", "mem_parts"]: self._gen_from_log("fine_step_from_log", name_y=name, name_x="fine_step") - + self._gen_from_log("SN_momentum_from_log", name_x="time", name_y="SN_momentum") # Dict of overlays @@ -1963,7 +1973,7 @@ class Plotter(Aggregator, BaseProcessor): "g": partial(self._overlay_vector, "g"), "B": self._overlay_B, "vel": self._overlay_speed, - "speed": self._overlay_speed, + "speed": self._overlay_speed, "levels": self._overlay_levels, "contour": self._overlay_contour, "particles": self._overlay_particles, diff --git a/scripts/mngseg.py b/scripts/mngseg.py index b870017..f3ed002 100644 --- a/scripts/mngseg.py +++ b/scripts/mngseg.py @@ -6,7 +6,7 @@ import aplpy def make_images(im, wt, M, meanim, label): - """ show the Gaussian and coherent part of the image """ + """show the Gaussian and coherent part of the image""" total = np.sum(wt[:M, :, :], axis=0).real + meanim coherent = np.sum(wt[M : 2 * M, :, :], axis=0).real + meanim @@ -52,7 +52,7 @@ def make_images(im, wt, M, meanim, label): def scale_images(thingy, M, label, scale=14, mode="wt"): - """ visualize wt or S11a for a specific scale. Remark S11a = wt^2""" + """visualize wt or S11a for a specific scale. Remark S11a = wt^2""" total = thingy[scale, :, :].real coherent = thingy[M + scale, :, :].real Gaussian = thingy[2 * M + scale, :, :].real @@ -93,7 +93,7 @@ def scale_images(thingy, M, label, scale=14, mode="wt"): def plot_each_scale(S11a, wav_k, q, label, coherent=False, reso=1): - """ plot histogram at a certain scale """ + """plot histogram at a certain scale""" nsize = len(S11a[0, 0, :]) M = len(q) for scl in range(0, M): @@ -237,7 +237,7 @@ def load_results(label): def analyse_sim(im, load=False, scale_image=False): - """ Do the MnGseg analysis """ + """Do the MnGseg analysis""" meanim = np.mean(im) imzm = im - meanim M = nb_scale(im.shape) @@ -264,7 +264,6 @@ def analyse_sim(im, load=False, scale_image=False): # make images of the Gaussian and coherent part make_images(im, wt, M, meanim, label) - if scale_image: # (optional) make the image for each scale for s in range(fit_min, fit_max + 1): diff --git a/scripts/pspec.py b/scripts/pspec.py index 6883265..41508e8 100644 --- a/scripts/pspec.py +++ b/scripts/pspec.py @@ -13,7 +13,6 @@ import pymses.utils.misc import tables as T from numpy.fft import fftn, ifft from pymses.analysis import Camera, ScalarOperator, cube3d -import os __generator__ = "pspec.py" __version__ = "0.2" @@ -62,8 +61,6 @@ def degrade_cube(cube, lvl, integrate=False): return cube_new - - def calc_k(n, nbinsk, nbig, dkbig, dim=3, saxis=2): """Make cubes containing the wave vectors, a list of bins and the associated normalization factors @@ -433,9 +430,7 @@ parser = argparse.ArgumentParser( parser.add_argument( "repo", help="RAMSES output repository", type=str, default=".", nargs="?" ) -parser.add_argument( - "iouts", help="output numbers", type=int, default=[1], nargs="+" -) +parser.add_argument("iouts", help="output numbers", type=int, default=[1], nargs="+") parser.add_argument( "outfile", help="output file format (see below for fields)", @@ -525,12 +520,7 @@ def main(arg): read_lvl = None if True: # Load output ------------------------------------------------------------------ - # If ratarmount was used - if os.path.exists(f"{self.path}/output_{self.num:05}/output_{self.num:05}"): - path = f"{arg.repo}/output_{self.num:05}" - else: - path = arg.repo - ro = pymses.RamsesOutput(path, iout, order=arg.order) + ro = pymses.RamsesOutput(arg.repo, iout, order=arg.order) if arg.magnetic: amr = ro.amr_source(["rho", "vel", "Bl", "Br"]) else: @@ -576,7 +566,7 @@ def main(arg): ] cam = Camera( - center=[0.5, 0.5, 0.5], #arg.center, + center=[0.5, 0.5, 0.5], # arg.center, line_of_sight_axis="z", region_size=[arg.size, arg.size], distance=arg.size / 2.0, diff --git a/snapshotprocessor.py b/snapshotprocessor.py index 7fe37bb..6c93d6b 100644 --- a/snapshotprocessor.py +++ b/snapshotprocessor.py @@ -30,7 +30,7 @@ from pymses.analysis import ( cube3d, ) from pymses.filters import CellsToPoints, RegionFilter -from pymses.sources.hop.hop import HOP +from pymses.sources.hop.hop import HOP try: from fil_finder import FilFinder2D @@ -220,9 +220,9 @@ def pspec2d(map2D): pmap = pspec.pcube(fmap) # Use the power map and the fft to compute the powerspectrum # This is typically an histogram of k weighted by the fourier transform value - pspec, kbins, pspec2, fbins = pspec.pspectrum(pmap, kmap, kbins, 1, 0) + power, kbins, power2, fbins = pspec.pspectrum(pmap, kmap, kbins, 1, 0) # Return bin center and power spectrum - return 0.5 * (kbins[1:] + kbins[:-1]), pspec + return 0.5 * (kbins[1:] + kbins[:-1]), power def degrade_map(dmap, nnew, integrate=False): @@ -387,11 +387,12 @@ class SnapshotProcessor(HDF5Container): # 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) + 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) if self.params.astrophysix.generate: @@ -405,17 +406,16 @@ class SnapshotProcessor(HDF5Container): try: self.init_pymses() - except: + except IOError: self.logger.error("Pymses not initialized", exc_info=1) - self.def_rules() def init_pymses(self): # If ratarmount was used if os.path.exists(f"{self.path}/output_{self.num:05}/output_{self.num:05}"): path = f"{self.path}/output_{self.num:05}" - else: + else: path = self.path self._ro = pymses.RamsesOutput( path, @@ -494,21 +494,21 @@ class SnapshotProcessor(HDF5Container): far_cut_depth=distance, up_vector=ax_v, map_max_size=self.params.pymses.map_size, - ) + ) # Initialize HDF5 group - try: - self.open() - if "/maps" not in self.save: - self.save.create_group("/", "maps", "2D maps") - self.save.root.maps._v_attrs.center = center - self.save.root.maps._v_attrs.radius = self._radius - self.save.root.maps._v_attrs.im_extent = im_extent - except: - self.logger.error("Error in HDF5", exc_info=1) - raise - finally: - self.close() + try: + self.open() + if "/maps" not in self.save: + self.save.create_group("/", "maps", "2D maps") + self.save.root.maps._v_attrs.center = center + self.save.root.maps._v_attrs.radius = self._radius + self.save.root.maps._v_attrs.im_extent = im_extent + except Exception() as e: + self.logger.error("Error in HDF5", exc_info=1) + raise e + finally: + self.close() def load_data(self, points_src, filename, save, keys=None): """ @@ -613,13 +613,13 @@ class SnapshotProcessor(HDF5Container): return pos def getter_vect_r(self, dset, name_vect): - """ Radial component of a vector """ + """Radial component of a vector""" r = self.getter_pos_disk(dset)[:, :2] ur = np.transpose((np.transpose(r) / np.sqrt(np.sum(r ** 2, axis=1)))) return np.einsum("ij, ij -> i", dset[name_vect][:, :2], ur) def getter_vect_phi(self, dset, name_vect): - """ Azimuthal component of a vector """ + """Azimuthal component of a vector""" r = self.getter_pos_disk(dset)[:, :2] r_norm = np.sqrt(np.sum(r ** 2, axis=1)) @@ -638,7 +638,7 @@ class SnapshotProcessor(HDF5Container): return pos def oct_getter_vect_r(self, dset, name_vect): - """ Radial component of a vector """ + """Radial component of a vector""" r = self.oct_getter_pos_disk(dset)[:, :, :2] ur = np.transpose( (np.transpose(r, (2, 0, 1)) / np.sqrt(np.sum(r ** 2, axis=2))), (1, 2, 0) @@ -646,7 +646,7 @@ class SnapshotProcessor(HDF5Container): return np.einsum("ikj, ikj -> ik", dset[name_vect][:, :, :2], ur) def oct_getter_vect_phi(self, dset, name_vect): - """ Azimuthal component of a vector """ + """Azimuthal component of a vector""" r = self.oct_getter_pos_disk(dset)[:, :, :2] r_norm = np.sqrt(np.sum(r ** 2, axis=2)) @@ -660,7 +660,7 @@ class SnapshotProcessor(HDF5Container): return self.oct_getter_vect_r(dset, "vel") def oct_getter_vphi(self, dset): - """ Azimuthal velocity """ + """Azimuthal velocity""" return self.oct_getter_vect_phi(dset, "vel") def datacube(self, getter, level=None, unit=U.none): @@ -685,22 +685,19 @@ class SnapshotProcessor(HDF5Container): 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, - ) + 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 + 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): @@ -731,8 +728,9 @@ class SnapshotProcessor(HDF5Container): datamap = slicing.SliceMap(self._amr, self._cam[ax_los], op, z=z) return datamap.map.T - - def ax_avg(self, oct_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 @@ -847,7 +845,15 @@ class SnapshotProcessor(HDF5Container): self.unload_cells() return data - def vol_pdf(self, getter, bins=100, old_unit=None, unit=None, 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: @@ -1057,7 +1063,9 @@ 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=False, G1_units=False, coarsen_factor=1): + def _toomreQ_disk( + self, ax_los, omega_approx=False, G1_units=False, coarsen_factor=1 + ): """ Compute the Toomre Q parameter """ @@ -1404,8 +1412,8 @@ class SnapshotProcessor(HDF5Container): def _sinks(self): csv_name = f"{self.path}/output_{self.num:05}/sink_{self.num:05}.csv" - if not os.path.exists(csv_name): # If ratarmount was used - csv_name = f"{self.path}/output_{self.num:05}/output_{self.num:05}/sink_{self.num:05}.csv" + if not os.path.exists(csv_name): # If ratarmount was used + csv_name = f"{self.path}/output_{self.num:05}/output_{self.num:05}/sink_{self.num:05}.csv" f = open(csv_name) first_line = f.readline() @@ -1449,11 +1457,9 @@ class SnapshotProcessor(HDF5Container): def pspec(self, overwrite_file=False, **kwargs): outfile = self.pspec_filename if overwrite_file or not os.path.exists(self.pspec_filename): - pspec.pspec(repo=self.path, iouts=[self.num], outfile=outfile, **kwargs) + pspec.pspec(repo=self.path, iouts=[self.num], outfile=outfile, **kwargs) return np.array([self.pspec_filename]) - - def _write_particles(self): """Ensure particles are written in the hdf5 file""" if not os.path.exists(self.parts_filename) and not self.parts_loaded: @@ -1604,50 +1610,57 @@ class SnapshotProcessor(HDF5Container): """ # Selection of cells - mask = self.cells["rho"] * self.info["unit_density"].express(U.H_cc) > threshold_density + mask = ( + self.cells["rho"] * self.info["unit_density"].express(U.H_cc) + > threshold_density + ) ncells = np.sum(mask) # fill the matrice with ID, x,y,z and masses of particles cells_group = np.zeros((ncells, 6)) - cells_group[:,0] = np.arange(ncells) # index - position = self.cells["pos"][mask] * self.info["unit_length"].express(U.pc) - cells_group[:,1:4] = position # position - density = self.cells["rho"][mask] * self.info["unit_density"].express(U.Msun / U.pc**3) - size = self.cells["dx"][mask]*self.info["unit_length"].express(U.pc) - cells_group[:,4] = density * size**3 # mass + cells_group[:, 0] = np.arange(ncells) # index + position = self.cells["pos"][mask] * self.info["unit_length"].express(U.pc) + cells_group[:, 1:4] = position # position + density = self.cells["rho"][mask] * self.info["unit_density"].express( + U.Msun / U.pc ** 3 + ) + size = self.cells["dx"][mask] * self.info["unit_length"].express(U.pc) + cells_group[:, 4] = density * size ** 3 # mass # save file.txt head = str(ncells) np.savetxt( - self.filename[:-3] + '_hop.txt', + self.filename[:-3] + "_hop.txt", cells_group[:, :-1], - fmt='%10d %.10e %.10e %.10e %.10e', + fmt="%10d %.10e %.10e %.10e %.10e", header=head, - delimiter=' ', - comments=' ' - ) - + delimiter=" ", + comments=" ", + ) + # save file.den - f = open(self.filename[:-3] + '_hop.den','wb') - f.write(pack('I', ncells)) + f = open(self.filename[:-3] + "_hop.den", "wb") + f.write(pack("I", ncells)) self.cells["rho"][mask].astype("f").tofile(f) f.close() # exec HOP algo - h = HOP(self.filename[:-3] + '_hop.txt', os.path.dirname(self.filename)) - h.process_hop() + h = HOP(self.filename[:-3] + "_hop.txt", os.path.dirname(self.filename)) + h.process_hop() # get the igroup array group_ids = h.get_group_ids() # 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... + # this means that the particules of group 1 are written first then of group 2 etc... ind_sort = np.argsort(group_ids) cells_group = cells_group[ind_sort] cells_group[6] = group_ids[ind_sort] # Make it a pandas' DataFrame - cells_group = pd.DataFrame(cells_group, header=["id", "x", "y", "z", "mass", "group"]) + cells_group = pd.DataFrame( + cells_group, header=["id", "x", "y", "z", "mass", "group"] + ) self.clumps = cells_group @@ -1658,7 +1671,6 @@ class SnapshotProcessor(HDF5Container): cells_group = self.make_clump_hop() cells_group.groupby("group") - def def_rules(self): self.rules = { @@ -1714,8 +1726,7 @@ class SnapshotProcessor(HDF5Container): dependencies=["slice_rho"], unit=self.info["unit_pressure"] / self.info["unit_density"], ), - "levels": Rule(self._levels, "Max level within line of sight", "/maps" - ), + "levels": Rule(self._levels, "Max level within line of sight", "/maps"), "jeans": Rule( self._jeans, "Jeans length slice", @@ -1769,8 +1780,7 @@ class SnapshotProcessor(HDF5Container): }, ), "pspec": Rule(self.pspec, "Power spectrum", "/hdf5"), - "write_particles": Rule(self._write_particles, "Particles file", "/hdf5" - ), + "write_particles": Rule(self._write_particles, "Particles file", "/hdf5"), "write_cells": Rule(self._write_cells, "Cells file", "/hdf5"), "filaments": Rule( self._filaments, @@ -1986,8 +1996,10 @@ class SnapshotProcessor(HDF5Container): # Norm generic_rule( - field + "_norm", partial(norm_getter, field), self.unit_key[field], - oct_getter=partial(oct_norm_getter, field) + field + "_norm", + partial(norm_getter, field), + self.unit_key[field], + oct_getter=partial(oct_norm_getter, field), ) else: @@ -1998,8 +2010,6 @@ class SnapshotProcessor(HDF5Container): unit = self.rules[name].unit - - self.rules["rad_avg_" + name] = Rule( partial(self._rad_avg, name), "Azimuthal average of {}".format(name), diff --git a/studyprocessor.py b/studyprocessor.py index 32c4c0f..30f8786 100644 --- a/studyprocessor.py +++ b/studyprocessor.py @@ -38,12 +38,9 @@ class StudyProcessor(Aggregator, HDF5Container): Creates the basic structures needed for the outputs """ - # log id self.log_id = "study({})".format(tag) - - super(StudyProcessor, self).__init__(path, path_out, params, tag) # Open outfile @@ -60,7 +57,12 @@ class StudyProcessor(Aggregator, HDF5Container): # Select runs if selector is None: selector = RunSelector( - path, runs, nums, self.params.input.nml_filename, unit_time=unit_time, **kwargs + path, + runs, + nums, + self.params.input.nml_filename, + unit_time=unit_time, + **kwargs, ) # Save infos @@ -68,7 +70,6 @@ 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 @@ -90,8 +91,6 @@ class StudyProcessor(Aggregator, HDF5Container): unit_time=unit_time, ) - - # Save namelist and logs if self.params.out.copy_info: for run in self.runs: @@ -158,7 +157,16 @@ class StudyProcessor(Aggregator, HDF5Container): prop[run] = getter(run) return np.array(list(prop.values())) - def time_avg(self, name, start=None, end=None, span=None, unit_time=U.Myr, group="/series", select=None): + 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 @@ -273,18 +281,16 @@ class StudyProcessor(Aggregator, HDF5Container): return value 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 - - + 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] @@ -295,7 +301,7 @@ class StudyProcessor(Aggregator, HDF5Container): def _extract_sinks_from_log(self, series, log_filename, run): cmd_grep = f"grep 'Number of sink' {log_filename} -A 2" content = os.popen(cmd_grep).readlines() - block_err = [] # Block that will ill parsed + block_err = [] # Block that will ill parsed for i in range(0, len(content), 4): try: nb_sink = np.int(content[i].split("=")[1]) @@ -311,11 +317,11 @@ class StudyProcessor(Aggregator, HDF5Container): except (ValueError, IndexError): block_err.append(i) - + if len(block_err) > 0: self.logger.warning( - f"Errors encountered in parsing {log_filename} (grepped blocks {block_err})" - ) + f"Errors encountered in parsing {log_filename} (grepped blocks {block_err})" + ) return series def _extract_stellar_from_log(self, stellar_objects, log_filename, run): @@ -324,7 +330,7 @@ class StudyProcessor(Aggregator, HDF5Container): nb_stellar = list(map(lambda s: int(s.split()[1]), content)) line_numbers = list(map(lambda s: int(s.split(":")[0]), content)) current_line = 0 - block_err = [] # Block that will ill parsed + block_err = [] # Block that will ill parsed logfile = open(log_filename) for i in range(0, len(line_numbers)): try: @@ -334,7 +340,9 @@ class StudyProcessor(Aggregator, HDF5Container): for j in range(nb_stellar[i]): line_stellar = logfile.readline().split() current_line += 1 - while line_stellar[0] == "random": # random number outputs are ... random + while ( + line_stellar[0] == "random" + ): # random number outputs are ... random line_stellar = logfile.readline().split() current_line += 1 mass = float(line_stellar[3]) @@ -349,11 +357,11 @@ class StudyProcessor(Aggregator, HDF5Container): except (ValueError, IndexError): block_err.append(i) - + if len(block_err) > 0: self.logger.warning( - f"Errors encountered in parsing {log_filename} (grepped blocks {block_err})" - ) + f"Errors encountered in parsing {log_filename} (grepped blocks {block_err})" + ) logfile.close() return stellar_objects @@ -370,7 +378,7 @@ class StudyProcessor(Aggregator, HDF5Container): def _extract_fine_step_from_log(self, series, log_filename, run): cmd_grep = "grep 'Fine step' {} ".format(log_filename) content = os.popen(cmd_grep).readlines() - block_err = [] # Block that will ill parsed + block_err = [] # Block that will ill parsed for i in range(0, len(content)): try: data = content[i].replace("=", " ").split() @@ -388,11 +396,11 @@ class StudyProcessor(Aggregator, HDF5Container): series["mem_parts"][run].append(mempc2) except (ValueError, IndexError): block_err.append(i) - + if len(block_err) > 0: self.logger.warning( - f"Error encountered in parsing {log_filename} (grepped blocks {block_err})" - ) + f"Error encountered in parsing {log_filename} (grepped blocks {block_err})" + ) return series def _extract_coarse_step_from_log(self, series, log_filename, run): @@ -462,20 +470,19 @@ class StudyProcessor(Aggregator, HDF5Container): if content[i][1:5] == "Fine": data = content[i].replace("=", " ").split() time = np.float(data[4]) - elif content[i][1:3] == "SN" : + elif content[i][1:3] == "SN": series["time"][run].append(time) series["SN_momentum"][run].append(np.float(content[i].split()[-1])) else: raise ValueError("Wrong start of line") except (AssertionError, ValueError, IndexError): block_err.append(i) - + if len(block_err) > 0: self.logger.warning( - f"Error encountered in parsing {log_filename} (grepped blocks {block_err})" - ) + f"Error encountered in parsing {log_filename} (grepped blocks {block_err})" + ) return series - def get_logs(self, run): glob_str = f"{self.path}/{run}/{self.params.input.log_prefix}*" @@ -508,7 +515,7 @@ class StudyProcessor(Aggregator, HDF5Container): # Always prefer data from last log, assuming they come in the right order time = series["time"][run] time_new = time[size] - ind_overlap = np.searchsorted(time[:size], time_new, side='right') + ind_overlap = np.searchsorted(time[:size], time_new, side="right") for key in series: del series[key][run][ind_overlap:size] @@ -523,13 +530,13 @@ class StudyProcessor(Aggregator, HDF5Container): return series def get_coldens0(self, run): - mp = 1.4 * 1.66 * 10**(-24) * U.g + mp = 1.4 * 1.66 * 10 ** (-24) * U.g try: z0 = self.get_nml("galbox_params/height0", run) * U.pc - n0 = self.get_nml("galbox_params/dens0", run) * U.cm**(-3) + n0 = self.get_nml("galbox_params/dens0", run) * U.cm ** (-3) except KeyError: z0 = self.get_nml("cloud_params/height0", run) * U.pc - n0 = self.get_nml("cloud_params/dens0", run) * U.cm**(-3) + n0 = self.get_nml("cloud_params/dens0", run) * U.cm ** (-3) return (np.sqrt(2 * np.pi) * mp * z0 * n0).express(U.coldens) @@ -537,13 +544,13 @@ class StudyProcessor(Aggregator, HDF5Container): """ Sum of gas plus sink mass """ - - time_gas = self.get_value("/series/coarse_step_from_log/time") + + 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") + time_sink = self.get_value("/series/sinks_from_log/time") - total_mass = dict() + total_mass = {} for run in self.runs: if time_sink[run][-1] > time_gas[run][-1]: @@ -552,11 +559,13 @@ class StudyProcessor(Aggregator, HDF5Container): # 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 + m0 = self.get_coldens0(run) * surface # Initial mass in Msun offset = time_gas[run].size - time_sink[run].size - mass_gas[run] = m0 + m0*mass_gas[run] # convert in Msun + 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 + total_mass[run][offset:] = ( + mass_gas[run][offset:] + mass_sink[run] + ) # re add sink_mass return time_gas, total_mass, mass_gas def _ssfr_from_mass_sink(self, avg_window=None): @@ -841,9 +850,16 @@ class StudyProcessor(Aggregator, HDF5Container): }, ), "SN_momentum_from_log": Rule( - partial(self._from_log, ["time", "SN_momentum"], self._extract_SN_Mom_from_log), + partial( + self._from_log, + ["time", "SN_momentum"], + self._extract_SN_Mom_from_log, + ), group="/series", - unit={"time": "unit_time", "SN_momentum" : {"unit_mass" : 1, "unit_velocity" : 1}}, + unit={ + "time": "unit_time", + "SN_momentum": {"unit_mass": 1, "unit_velocity": 1}, + }, description={ "time": "Time", "SN_momentum": "Injected momentum", @@ -934,7 +950,7 @@ class StudyProcessor(Aggregator, HDF5Container): "turb_power", "time_rho_prof", "time_coldens_pdf", - "time_rho_pdf" + "time_rho_pdf", ]: self._gen_rule_avg(name) diff --git a/turbox/turbox.py b/turbox/turbox.py index 67f9526..ddf062e 100644 --- a/turbox/turbox.py +++ b/turbox/turbox.py @@ -7,7 +7,7 @@ """ import numpy as np -from plotter import U +from plotter import U from baseprocessor import Rule from matplotlib import pyplot as plt import pandas as pd @@ -15,8 +15,8 @@ import tables from scipy.stats import linregress -def get_pspec(pp, field:str, dim:int=3): - """Read power spectruù +def get_pspec(pp, field: str, dim: int = 3): + """Read power spectruù Parameters ---------- @@ -30,7 +30,7 @@ def get_pspec(pp, field:str, dim:int=3): ------- tupple (np.array, np.array) wave number and corresponding powers - """ + """ h5file = tables.File(pp.pspec_filename) path = f"/out_{pp.num:05}/d{dim}/{field}" node = h5file.get_node(path) @@ -38,17 +38,13 @@ def get_pspec(pp, field:str, dim:int=3): pspec = node.pspec.read() h5file.close() k = (kbins[:-1] + kbins[1:]) / 2 - return k, pspec + return k, pspec -span_resolution = { - 256: (0.8, 1.1), - 512: (0.5, 1.7), - 1024: (0.4, 1.7) -} +span_resolution = {256: (0.8, 1.1), 512: (0.5, 1.7), 1024: (0.4, 1.7)} -def get_pspec_slope(pp, field:str, resol:int, plotdebug:bool=False): +def get_pspec_slope(pp, field: str, resol: int, plotdebug: bool = False): """Get the slope of the Power specturm using linear regression in the selected range Parameters @@ -65,23 +61,35 @@ def get_pspec_slope(pp, field:str, resol:int, plotdebug:bool=False): Slope, square value of the correlation coefficient """ - # Trustworthy span od the power spectrum in log10(k) as a function of the resolution + # Trustworthy span od the power spectrum in log10(k) as a function of the resolution logkmin, logkmax = span_resolution[resol] k, power = get_pspec(pp, field) - logk, logpower = np.log10(k), np.log10(power) + logk, logpower = np.log10(k), np.log10(power) mask = (logk >= logkmin) & (logk < logkmax) results = linregress(logk[mask], logpower[mask]) if plotdebug: plt.figure() plt.plot(logk, logpower) - plt.plot(logk[mask], results.slope*logk[mask]+ results.intercept, lw=3, ls=":", color="k") - pp.logger.info(f"Fit results in get_slope({field}, {resol}): slope:{results.slope:.2f}, b:{results.intercept:.2f}, R2:{results.rvalue**2:.2f}") - if results.rvalue**2 < 0.8: - pp.logger.warning(f"Bad fit in get_slope({field}, {resol}) with {logkmin} <= logk < {logkmax}") + plt.plot( + logk[mask], + results.slope * logk[mask] + results.intercept, + lw=3, + ls=":", + color="k", + ) + pp.logger.info( + f"Fit results in get_slope({field}, {resol}): slope:{results.slope:.2f}" + + f", b:{results.intercept:.2f}, R2:{results.rvalue**2:.2f}" + ) + if results.rvalue ** 2 < 0.8: + pp.logger.warning( + f"Bad fit in get_slope({field}, {resol}) with {logkmin} <= logk < {logkmax}" + ) pp.logger.warning(f"log(k) is \n {logk[mask]}") pp.logger.warning(f"log(power) is \n {logpower[mask]}") - return results.slope, results.intercept, results.rvalue**2 + return results.slope, results.intercept, results.rvalue ** 2 + def build_suite(pl, redo=False, cs0=0.28834810480560674): """Compute an array of parameter for each run in the Plotter pl @@ -99,9 +107,9 @@ def build_suite(pl, redo=False, cs0=0.28834810480560674): pd.Dataframe dataframe with the properties of the simulation """ - - df = dict() - df["snapshots"] = pl.nums.values() + + df = {} + df["snapshots"] = pl.nums.values() df["n0"] = pl.study.get_nml("galbox_params/dens0").values() df["turbinit"] = pl.study.get_nml("galbox_params/turb").values() df["solver"] = pl.study.get_nml("hydro_params/riemann").values() @@ -113,33 +121,46 @@ def build_suite(pl, redo=False, cs0=0.28834810480560674): df["comp"] = pl.study.get_nml("turb_params/comp_frac").values() df["L"] = pl.study.get_nml("amr_params/boxlen").values() - df["T_turb"] = (np.array(list(pl.study.get_nml("turb_params/turb_T").values())) - * pl.study.info["unit_time"].express(U.Myr)) + df["T_turb"] = np.array( + list(pl.study.get_nml("turb_params/turb_T").values()) + ) * pl.study.info["unit_time"].express(U.Myr) df = pd.DataFrame(df, index=pl.runs) - + if redo: pl.study.avg_time_sigma("x", overwrite_dep=False) pl.study.avg_time_sigma("y", overwrite_dep=False) pl.study.avg_time_sigma("z", overwrite_dep=False) pl.study.time(overwrite=True) - - for ax in ["x", "y", "z"]: - df[f"sigma_{ax}"] = np.array(list(map( - lambda x : x.T[0], - [pl.study.get_value(f"/series/time_sigma_{ax}", - unit=U.km_s)[run] for run in pl.runs]))) - df["sigma_all"] = df[f"sigma_x"]**2 + df[f"sigma_y"]**2 + df[f"sigma_z"]**2 + for ax in ["x", "y", "z"]: + df[f"sigma_{ax}"] = np.array( + list( + map( + lambda x: x.T[0], + [ + pl.study.get_value(f"/series/time_sigma_{ax}", unit=U.km_s)[run] + for run in pl.runs + ], + ) + ) + ) + + df["sigma_all"] = df["sigma_x"] ** 2 + df["sigma_y"] ** 2 + df["sigma_z"] ** 2 df["sigma_all"] = list(map(np.sqrt, df["sigma_all"].values)) - df["Mach_all"] = list(map(lambda v: v/cs0, df["sigma_all"].values)) - df["time"] = list(map(lambda x : x.T[0], - [pl.study.get_value(f"/series/time", unit=U.Myr)[run] - for run in pl.runs])) - + df["Mach_all"] = list(map(lambda v: v / cs0, df["sigma_all"].values)) + df["time"] = list( + map( + lambda x: x.T[0], + [pl.study.get_value("/series/time", unit=U.Myr)[run] for run in pl.runs], + ) + ) + df["sigma"] = list(map(lambda l: np.mean(l), df["sigma_all"].values)) df["Mach"] = df["sigma"] / cs0 - df["turnover"] = (df["L"] * U.pc.express(U.km) / (2 * df["sigma"]))* U.s.express(U.Myr) - + df["turnover"] = (df["L"] * U.pc.express(U.km) / (2 * df["sigma"])) * U.s.express( + U.Myr + ) + return df @@ -148,13 +169,15 @@ def rho_pdf(pp): rho = pp.cells["rho"] * pp.info["unit_density"].express(U.H_cc) rho_0 = np.mean(rho) print(rho_0) - s = np.log(rho/rho_0) - values, edges = np.histogram(s, bins=300, range=(-15, 11), - density=True) + s = np.log(rho / rho_0) + values, edges = np.histogram(s, bins=300, range=(-15, 11), density=True) pp.unload_cells() centers = 0.5 * (edges[1:] + edges[:-1]) return (np.stack([values, centers]), {"logbins": True}) -rule_pdf=Rule(rho_pdf, "Density PDF", name="rho_pdf", group="/hist") + +rule_pdf = Rule(rho_pdf, "Density PDF", name="rho_pdf", group="/hist") + + def apply_rule_pdf(pp): return pp.process(rule_pdf, pp, overwrite=True) diff --git a/utils/extract_velcube.py b/utils/extract_velcube.py index a3f53b3..d936ebb 100644 --- a/utils/extract_velcube.py +++ b/utils/extract_velcube.py @@ -1,6 +1,5 @@ from snapshotprocessor import SnapshotProcessor, U -import pandas as pd -import os + def get_velocity_cubes(pp, unit=None): velcubes = [None, None, None] @@ -10,8 +9,9 @@ def get_velocity_cubes(pp, unit=None): velcubes[i] *= pp.info["unit_velocity"].express(unit) return velcubes + def get_density_cube(pp, unit=None): - dens_cube = pp.datacube(getter=lambda dset: dset["rho"]) + dens_cube = pp.datacube(getter=lambda dset: dset["rho"]) if unit is not None: dens_cube *= pp.info["unit_density"].express(unit) return dens_cube @@ -19,27 +19,32 @@ def get_density_cube(pp, unit=None): def write_data(filename, vel, dens): # write fields to ramses frig readable ascii file - f = open(filename, 'w') + f = open(filename, "w") dummy = 1 size = vel[0].shape[0] - f.write('{:8}{:13.5f}{:13.5f}{:13.5f}{:13.5f}\n'.format(size, dummy, dummy, dummy, dummy)) + f.write( + "{:8}{:13.5f}{:13.5f}{:13.5f}{:13.5f}\n".format( + size, dummy, dummy, dummy, dummy + ) + ) vx, vy, vz = vel # This strange order matches the one in the galbox condinit for z in range(size): for y in range(size): for x in range(size): - f.write('{:13.5f}{:13.5f}{:13.5f}{:13.5f}\n'.format(vx[x,y,z], vy[x,y,z], vz[x,y,z], dens[x, y, z])) + f.write( + "{:13.5f}{:13.5f}{:13.5f}{:13.5f}\n".format( + vx[x, y, z], vy[x, y, z], vz[x, y, z], dens[x, y, z] + ) + ) + - - def extract_from_pp(pp): vel = get_velocity_cubes(pp, unit=U.km_s) dens = get_density_cube(pp, unit=U.H_cc) write_data(f"{pp.path_out}/{pp.run}_{pp.num}_velrho.data", vel, dens) - + def extract(path, snap_number): pp = SnapshotProcessor(path, snap_number, params="../turbox_params.yml") extract_from_pp(pp) - - diff --git a/utils/runselector.py b/utils/runselector.py index 504c9f5..5e15331 100644 --- a/utils/runselector.py +++ b/utils/runselector.py @@ -283,7 +283,9 @@ class RunSelector: def load_info(self, run, num): info_filename_output = f"{self.path_in}/{run}/output_{num:05}/info_{num:05}.txt" # Path of the filename if ratarmount was used - info_filename_tarmount_output = f"{self.path_in}/{run}/output_{num:05}/output_{num:05}/info_{num:05}.txt" + info_filename_tarmount_output = ( + f"{self.path_in}/{run}/output_{num:05}/output_{num:05}/info_{num:05}.txt" + ) info_filename_folder = f"{self.path_in}/{run}/info/info_{num:05}.txt" if os.path.exists(info_filename_output): @@ -348,15 +350,17 @@ class RunSelector: return self.info[run][num]["time"] elif isinstance(unit_time, str): - + factor = self.get_nml_value(unit_time, run) + def get_time(num): - time_code = self.info[run][num]["time"] + time_code = self.info[run][num]["time"] return time_code / factor else: + def get_time(num): - time_code = self.info[run][num]["time"] + time_code = self.info[run][num]["time"] return time_code * self.info[run][num]["unit_time"].express(unit_time) # -- A function to search a given time using dichotomy @@ -477,7 +481,7 @@ class RunSelector: return nums def write_paths(self, prefix=None, filename="~/list_file"): - """ + """ Write the paths of the selected runs on a file Args: @@ -485,7 +489,7 @@ class RunSelector: filename (str, optional): F. Defaults to "~/list_file". """ if prefix is None: - prefix = self.path_in + prefix = self.path_in paths = [] for run in self.nums: for num in self.nums[run]: @@ -496,6 +500,3 @@ class RunSelector: f = open(os.path.expanduser(filename), "w") f.writelines(paths) f.close() - - - diff --git a/utils/snapshotselector.py b/utils/snapshotselector.py index fa60cf8..036de18 100644 --- a/utils/snapshotselector.py +++ b/utils/snapshotselector.py @@ -9,6 +9,7 @@ from utils.runselector import RunSelector from plotter import Plotter, U import os + def prep_mcons(study): study.coarse_step_from_log() @@ -22,36 +23,46 @@ def time_mcons(study, run, target=0.2): def find_nums(study, prep_function, time_function, time_min=0): - """ + """ Once other filter are applied, select one output based on the time given by time function Args: prep_function (study:studyProcessor -> None): prepare a study object - time_function (study:studyProcessor, run:str -> time:float): compute selected time from the object + time_function (study:studyProcessor, run:str -> time:float): compute selected time from the object """ nums = {} prep_function(study) for run in study.runs: time_target = max(time_min, time_function(study, run)) - rs = RunSelector(path_in=study.path, in_runs=run, time_min=time_min, time=time_target, unit_time=U.Myr) + rs = RunSelector( + path_in=study.path, + in_runs=run, + time_min=time_min, + time=time_target, + unit_time=U.Myr, + ) nums.update(rs.nums) return nums def write_paths(nums, path_from_home, filename="~/list_file"): paths = [] - for key in nums: - for num in self.nums[run]: - if os.path.exists("{prefix}/{run}/output_{num:05}/output_{num:05}\n"): - paths.append(f"{prefix}/{run}/output_{num:05}/output_{num:05}\n") - else: - paths.append(f"{prefix}/{run}/output_{num:05}\n") + for run in nums: + for num in nums[run]: + if os.path.exists( + "{path_from_home}/{run}/output_{num:05}/output_{num:05}\n" + ): + paths.append( + f"{path_from_home}/{run}/output_{num:05}/output_{num:05}\n" + ) + else: + paths.append(f"{path_from_home}/{run}/output_{num:05}\n") f = open(os.path.expanduser(filename), "w") f.writelines(paths) f.close() -if __name__ == '__main__': +if __name__ == "__main__": path_from_home = "simus/ismfeed/allmode" names = "n6_st_2e5_seed3_T5Myr_nsink1e3_comp*" @@ -63,4 +74,4 @@ if __name__ == '__main__': tag="select", ).study nums = find_nums(study, prep_mcons, time_mcons) - write_paths(nums, ".") \ No newline at end of file + write_paths(nums, ".") diff --git a/utils/units.py b/utils/units.py index bc92fb3..7898020 100644 --- a/utils/units.py +++ b/utils/units.py @@ -43,7 +43,7 @@ def unit_str(unit, base=None, prefix="", format=" [{unit}]"): return "" elif base is not None: coeff = unit.express(base) - return unit_str(base, prefix=convert_exp(coeff)+" ") + return unit_str(base, prefix=convert_exp(coeff) + " ") elif len(unit.latex) > 0: if "." in unit.latex or "^" in unit.latex: base_str = ".".join(map(parse_exp_unit, unit.name.split(".")))