From bc43269763b7440e7104a81a9d94011935dbf9dd Mon Sep 17 00:00:00 2001 From: Noe Brucy Date: Thu, 25 Mar 2021 14:16:07 +0100 Subject: [PATCH] [pspec] first attempt to isolate magnetic field stuff --- postprocessor.py | 4 +- pspec_new.py | 306 +++++++++++++++++++++++++++-------------------- 2 files changed, 177 insertions(+), 133 deletions(-) diff --git a/postprocessor.py b/postprocessor.py index ad3cddc..6926b53 100644 --- a/postprocessor.py +++ b/postprocessor.py @@ -1210,9 +1210,9 @@ class PostProcessor(HDF5Container): return sinks_dict - def _pspec(self): + def _pspec(self, **kwargs): outfile = self.pspec_filename - pspec_new.pspec(repo=self.path, iouts=[self.num], outfile=outfile) + pspec_new.pspec(repo=self.path, iouts=[self.num], outfile=outfile, **kwargs) return True def _particles(self): diff --git a/pspec_new.py b/pspec_new.py index e5589c6..0d3f773 100644 --- a/pspec_new.py +++ b/pspec_new.py @@ -3,7 +3,7 @@ from __future__ import division, print_function import argparse -import sys + import textwrap from builtins import range @@ -11,7 +11,7 @@ import numpy as np import pymses import pymses.utils.misc import tables as T -from i_utils import args, tools +from i_utils import args from numpy.fft import fftn, ifft from pymses.analysis import Camera, ScalarOperator, cube3d @@ -492,6 +492,14 @@ parser.add_argument( default="z", ) +parser.add_argument( + "-m", + "--magnetic", + help="MHD simulation", + type=bool, + default=True, +) + def main(arg): add_pspec2 = False @@ -516,32 +524,49 @@ def main(arg): if True: # Load output ------------------------------------------------------------------ ro = pymses.RamsesOutput(arg.repo, iout, order=arg.order) - amr = ro.amr_source(["rho", "vel", "Bl", "Br"]) + if arg.magnetic: + amr = ro.amr_source(["rho", "vel", "Bl", "Br"]) + else: + amr = ro.amr_source(["rho", "vel"]) if clvl == 0: clvl = ro.info["levelmin"] read_lvl = ro.info["levelmin"] # Extract cubes --------------------------------------------------------------- - cube_vars = [ - lambda dset: dset["rho"], - lambda dset: dset["vel"][..., 0], - lambda dset: dset["vel"][..., 1], - lambda dset: dset["vel"][..., 2], - lambda dset: 0.5 * (dset["Bl"][..., 0] + dset["Br"][..., 0]), - lambda dset: 0.5 * (dset["Bl"][..., 1] + dset["Br"][..., 1]), - lambda dset: 0.5 * (dset["Bl"][..., 2] + dset["Br"][..., 2]), - ] + if arg.magnetic: + cube_vars = [ + lambda dset: dset["rho"], + lambda dset: dset["vel"][..., 0], + lambda dset: dset["vel"][..., 1], + lambda dset: dset["vel"][..., 2], + lambda dset: 0.5 * (dset["Bl"][..., 0] + dset["Br"][..., 0]), + lambda dset: 0.5 * (dset["Bl"][..., 1] + dset["Br"][..., 1]), + lambda dset: 0.5 * (dset["Bl"][..., 2] + dset["Br"][..., 2]), + ] + cube_units = [ + ro.info["unit_density"], + ro.info["unit_velocity"], + ro.info["unit_velocity"], + ro.info["unit_velocity"], + ro.info["unit_mag"], + ro.info["unit_mag"], + ro.info["unit_mag"], + ] - cube_units = [ - ro.info["unit_density"], - ro.info["unit_velocity"], - ro.info["unit_velocity"], - ro.info["unit_velocity"], - ro.info["unit_mag"], - ro.info["unit_mag"], - ro.info["unit_mag"], - ] + else: + cube_vars = [ + lambda dset: dset["rho"], + lambda dset: dset["vel"][..., 0], + lambda dset: dset["vel"][..., 1], + lambda dset: dset["vel"][..., 2], + ] + cube_units = [ + ro.info["unit_density"], + ro.info["unit_velocity"], + ro.info["unit_velocity"], + ro.info["unit_velocity"], + ] cam = Camera( center=arg.center, @@ -554,7 +579,11 @@ def main(arg): ) cubes = {} - for i, v in enumerate(["rho", "vx", "vy", "vz", "Bx", "By", "Bz"]): + if arg.magnetic: + var_names = ["rho", "vx", "vy", "vz", "Bx", "By", "Bz"] + else: + var_names = ["rho", "vx", "vy", "vz"] + for i, v in enumerate(var_names): operator = ScalarOperator(cube_vars[i], cube_units[i]) extractor = cube3d.CubeExtractor(amr, operator) cubes[v] = extractor.process( @@ -563,17 +592,24 @@ def main(arg): else: h5f = T.open_file("cube.hdf5", "r") read_lvl = np.asscalar(h5f.root.level.read()) - cubes = { - "rho": h5f.root.rho.read(), - "vx": h5f.root.vx.read(), - "vy": h5f.root.vy.read(), - "vz": h5f.root.vz.read(), - "Bx": h5f.root.Bx.read(), - "By": h5f.root.By.read(), - "Bz": h5f.root.Bz.read(), - } + if arg.magnetic: + cubes = { + "rho": h5f.root.rho.read(), + "vx": h5f.root.vx.read(), + "vy": h5f.root.vy.read(), + "vz": h5f.root.vz.read(), + "Bx": h5f.root.Bx.read(), + "By": h5f.root.By.read(), + "Bz": h5f.root.Bz.read(), + } + else: + cubes = { + "rho": h5f.root.rho.read(), + "vx": h5f.root.vx.read(), + "vy": h5f.root.vy.read(), + "vz": h5f.root.vz.read(), + } h5f.close() - if clvl > read_lvl: print( "WARNING: adjusting cube level (%d) to match data level (%d)" @@ -585,7 +621,7 @@ def main(arg): if clvl < read_lvl: print("Degrade cubes") - for v in ["rho", "vx", "vy", "vz", "Bx", "By", "Bz"]: + for v in var_names: # Don't average, simply integrate the magnetic field cube_tmp = degrade_cube(cubes[v], clvl, v.startswith("B")) del cubes[v] @@ -596,44 +632,40 @@ def main(arg): cubes_k, kbins, knorm = calc_k( 1 << clvl, arg.kbins, arg.kbinsbig, arg.dkbig, 3, saxis ) - Bavg = avg_vect(cubes, "B", dim=3) - Bavg2 = avg_vect(cubes, "B", dim=2, saxis=saxis) - _, knorm_Bpar, knorm_Bperp = proj_B( - cubes_k, kbins, Bavg, "B", dim=3, update=True - ) + if arg.magnetic: + Bavg = avg_vect(cubes, "B", dim=3) + Bavg2 = avg_vect(cubes, "B", dim=2, saxis=saxis) + _, knorm_Bpar, knorm_Bperp = proj_B( + cubes_k, kbins, Bavg, "B", dim=3, update=True + ) print("Calculate derived quantities") # Additional quantities -------------------------------------------------------- cubes["logrho"] = np.log(cubes["rho"]) - cubes["cos_vB"] = np.zeros_like(cubes["rho"]) - vv = np.zeros_like(cubes["rho"]) - BB = np.zeros_like(cubes["rho"]) + + if arg.magnetic: + cubes["cos_vB"] = np.zeros_like(cubes["rho"]) + BB = np.zeros_like(cubes["rho"]) + vv = np.zeros_like(cubes["rho"]) for a in ["x", "y", "z"]: cubes["kr" + a] = cubes["rho"] ** (1.0 / 3.0) * cubes["v" + a] - cubes["cos_vB"] += cubes["v" + a] * cubes["B" + a] - vv += cubes["v" + a] ** 2 - BB += cubes["B" + a] ** 2 - cubes["cos_vB"] /= vv * BB - del vv - del BB + var_names.append("kr" + a) + if arg.magnetic: + cubes["cos_vB"] += cubes["v" + a] * cubes["B" + a] + BB += cubes["B" + a] ** 2 + vv += cubes["v" + a] ** 2 + var_names.append("logrho") + if arg.magnetic: + cubes["cos_vB"] /= vv * BB + var_names.append("cos_vB") + del vv + del BB print("3D FFT") # 3D Fourier transform --------------------------------------------------------- fcubes = {} - for v in [ - "rho", - "logrho", - "vx", - "vy", - "vz", - "krx", - "kry", - "krz", - "Bx", - "By", - "Bz", - "cos_vB", - ]: + + for v in var_names: fcubes[v] = fftn(cubes[v]) fcubes.update(cubes_k) @@ -656,13 +688,14 @@ def main(arg): pcubes["vs"] = pcube(fcubes["vsx"], fcubes["vsy"], fcubes["vsz"]) pcubes["krc"] = pcube(fcubes["krc"]) pcubes["krs"] = pcube(fcubes["krsx"], fcubes["krsy"], fcubes["krsz"]) - pcubes["B"] = pcube(fcubes["Bx"], fcubes["By"], fcubes["Bz"]) - pcubes["cos_vB"] = pcube(fcubes["cos_vB"]) pcubes["vz"] = pcube(fcubes["vz"]) + if arg.magnetic: + pcubes["B"] = pcube(fcubes["Bx"], fcubes["By"], fcubes["Bz"]) + pcubes["cos_vB"] = pcube(fcubes["cos_vB"]) print("Compute 3D power spectra") # 3D power spectra ------------------------------------------------------------- - for v in [ + vspec = [ "rho", "logrho", "v", @@ -671,19 +704,22 @@ def main(arg): "vs", "krc", "krs", - "B", - "cos_vB", "vz", - ]: + ] + if arg.magnetic: + vspec += ["B", "cos_vB"] + + for v in vspec: pspec, kbins, pspec2, fbins = pspectrum( pcubes[v], cubes_k["k"], kbins, knorm, arg.fbins ) - pspec_Bpar, _, _, _ = pspectrum( - pcubes[v], cubes_k["kBpar"], kbins, knorm_Bpar, 0 - ) - pspec_Bperp, kbins, _, _ = pspectrum( - pcubes[v], cubes_k["kBperp"], kbins, knorm_Bperp, 0 - ) + if arg.magnetic: + pspec_Bpar, _, _, _ = pspectrum( + pcubes[v], cubes_k["kBpar"], kbins, knorm_Bpar, 0 + ) + pspec_Bperp, kbins, _, _ = pspectrum( + pcubes[v], cubes_k["kBperp"], kbins, knorm_Bperp, 0 + ) # Save 3D power spectra params = {"iout": iout, "varname": v, "dim": 3} @@ -709,25 +745,31 @@ def main(arg): pass h5fd.create_array(outpath, "pspec", pspec, createparents=True) - h5fd.create_array(outpath, "pspec_Bpar", pspec_Bpar, createparents=True) - h5fd.create_array(outpath, "pspec_Bperp", pspec_Bperp, createparents=True) h5fd.create_array(outpath, "kbins", kbins, createparents=True) if add_pspec2: h5fd.create_array(outpath, "pspec2", pspec2, createparents=True) h5fd.create_array(outpath, "fbins", fbins, createparents=True) h5fd.create_array(outpath, "norm", knorm, createparents=True) - h5fd.create_array(outpath, "norm_Bpar", knorm_Bpar, createparents=True) - h5fd.create_array(outpath, "norm_Bperp", knorm_Bperp, createparents=True) + + if arg.magnetic: + h5fd.create_array(outpath, "pspec_Bpar", pspec_Bpar, createparents=True) + h5fd.create_array( + outpath, "pspec_Bperp", pspec_Bperp, createparents=True + ) + h5fd.create_array(outpath, "norm_Bpar", knorm_Bpar, createparents=True) + h5fd.create_array( + outpath, "norm_Bperp", knorm_Bperp, createparents=True + ) try: h5fd.remove_node(outpath, "meta", recursive=True) except T.NoSuchNodeError: pass - h5fd.create_group(outpath, "meta", createparents=True) - metagrp = h5fd.get_node(outpath, "meta") - h5fd.create_array(metagrp, "generator", __generator__) - h5fd.create_array(metagrp, "version", __version__) + # h5fd.create_group(outpath, "meta", createparents=True) + # metagrp = h5fd.get_node(outpath, "meta") + # h5fd.create_array(metagrp, "generator", __generator__) + # h5fd.create_array(metagrp, "version", __version__) h5fd.close() @@ -739,14 +781,16 @@ def main(arg): cubes_k, kbins, knorm = calc_k( 1 << clvl, arg.kbins, arg.kbinsbig, arg.dkbig, 2, saxis ) - _, knorm_Bpar, knorm_Bperp = proj_B( - cubes_k, kbins, Bavg2, "B", dim=2, saxis=saxis, update=True - ) + if arg.magnetic: + _, knorm_Bpar, knorm_Bperp = proj_B( + cubes_k, kbins, Bavg2, "B", dim=2, saxis=saxis, update=True + ) print("Project 3D -> 2D") # 3D -> 2D --------------------------------------------------------------------- fcubes2 = {} - for v in [ + + vars2D = [ "rho", "logrho", "vx", @@ -763,12 +807,11 @@ def main(arg): "krsx", "krsy", "krsz", - "Bx", - "By", - "Bz", - "cos_vB", "vz", - ]: + ] + if arg.magnetic: + vars2D += ["Bx", "By", "Bz", "cos_vB"] + for v in vars2D: fcubes2[v] = ifft(fcubes[v], axis=saxis) # Memory cleanup --------------------------------------------------------------- @@ -785,41 +828,37 @@ def main(arg): pcubes2["vs"] = pcube(fcubes2["vsx"], fcubes2["vsy"], fcubes2["vsz"]) pcubes2["krc"] = pcube(fcubes2["krc"]) pcubes2["krs"] = pcube(fcubes2["krsx"], fcubes2["krsy"], fcubes2["krsz"]) - pcubes2["B"] = pcube(fcubes2["Bx"], fcubes2["By"], fcubes2["Bz"]) - pcubes2["cos_vB"] = pcube(fcubes2["cos_vB"]) pcubes2["vz"] = pcube(fcubes2["vz"]) + if arg.magnetic: + pcubes2["B"] = pcube(fcubes2["Bx"], fcubes2["By"], fcubes2["Bz"]) + pcubes2["cos_vB"] = pcube(fcubes2["cos_vB"]) + print("Compute 2D power spectra") # 2D power spectra ------------------------------------------------------------- ns = 2 ** clvl f = "_%%(i)0%dd" % (np.floor(np.log10(ns)) + 1) - for v in [ - "rho", - "logrho", - "v", - "kr", - "vc", - "vs", - "krc", - "krs", - "B", - "cos_vB", - "vz", - ]: + for v in var_names: for i in range(ns): pspec, kbins, pspec2, fbins = pspectrum( pcubes2[v][:, :, i], cubes_k["k"][:, :, i], kbins, knorm, arg.fbins ) - pspec_Bpar, _, _, _ = pspectrum( - pcubes2[v][:, :, i], cubes_k["kBpar"][:, :, i], kbins, knorm_Bpar, 0 - ) - pspec_Bperp, kbins, _, _ = pspectrum( - pcubes2[v][:, :, i], - cubes_k["kBperp"][:, :, i], - kbins, - knorm_Bperp, - 0, - ) + + if arg.magnetic: + pspec_Bpar, _, _, _ = pspectrum( + pcubes2[v][:, :, i], + cubes_k["kBpar"][:, :, i], + kbins, + knorm_Bpar, + 0, + ) + pspec_Bperp, kbins, _, _ = pspectrum( + pcubes2[v][:, :, i], + cubes_k["kBperp"][:, :, i], + kbins, + knorm_Bperp, + 0, + ) # Save 2D power spectra suff = f % {"i": i} @@ -846,12 +885,15 @@ def main(arg): pass h5fd.create_array(outpath, "pspec" + suff, pspec, createparents=True) - h5fd.create_array( - outpath, "pspec_Bpar" + suff, pspec_Bpar, createparents=True - ) - h5fd.create_array( - outpath, "pspec_Bperp" + suff, pspec_Bperp, createparents=True - ) + + if arg.magnetic: + h5fd.create_array( + outpath, "pspec_Bpar" + suff, pspec_Bpar, createparents=True + ) + h5fd.create_array( + outpath, "pspec_Bperp" + suff, pspec_Bperp, createparents=True + ) + h5fd.create_array(outpath, "kbins" + suff, kbins, createparents=True) if add_pspec2: h5fd.create_array( @@ -861,22 +903,24 @@ def main(arg): outpath, "fbins" + suff, fbins, createparents=True ) h5fd.create_array(outpath, "norm" + suff, knorm, createparents=True) - h5fd.create_array( - outpath, "norm_Bpar" + suff, knorm_Bpar, createparents=True - ) - h5fd.create_array( - outpath, "norm_Bperp" + suff, knorm_Bperp, createparents=True - ) + + if arg.magnetic: + h5fd.create_array( + outpath, "norm_Bpar" + suff, knorm_Bpar, createparents=True + ) + h5fd.create_array( + outpath, "norm_Bperp" + suff, knorm_Bperp, createparents=True + ) try: h5fd.remove_node(outpath, "meta", recursive=True) except T.NoSuchNodeError: pass - h5fd.create_group(outpath, "meta", createparents=True) - metagrp = h5fd.get_node(outpath, "meta") - h5fd.create_array(metagrp, "generator", __generator__) - h5fd.create_array(metagrp, "version", __version__) + # h5fd.create_group(outpath, "meta", createparents=True) + # metagrp = h5fd.get_node(outpath, "meta") + # h5fd.create_array(metagrp, "generator", __generator__) + # h5fd.create_array(metagrp, "version", __version__) h5fd.close()