[pspec] first attempt to isolate magnetic field stuff

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):

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,13 +524,17 @@ def main(arg):
         if True:
             # Load output ------------------------------------------------------------------
             ro = pymses.RamsesOutput(arg.repo, iout, order=arg.order)
+            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 ---------------------------------------------------------------
+            if arg.magnetic:
                 cube_vars = [
                     lambda dset: dset["rho"],
                     lambda dset: dset["vel"][..., 0],
@@ -532,7 +544,6 @@ def main(arg):
                     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"],
@@ -543,6 +554,20 @@ def main(arg):
                     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,
                 line_of_sight_axis="z",
@@ -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,6 +592,7 @@ def main(arg):
         else:
             h5f = T.open_file("cube.hdf5", "r")
             read_lvl = np.asscalar(h5f.root.level.read())
+            if arg.magnetic:
                 cubes = {
                     "rho": h5f.root.rho.read(),
                     "vx": h5f.root.vx.read(),
@@ -572,8 +602,14 @@ def main(arg):
                     "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,6 +632,7 @@ def main(arg):
         cubes_k, kbins, knorm = calc_k(
             1 << clvl, arg.kbins, arg.kbinsbig, arg.dkbig, 3, saxis
         )
+        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(
@@ -605,35 +642,30 @@ def main(arg):
         print("Calculate derived quantities")
         # Additional quantities --------------------------------------------------------
         cubes["logrho"] = np.log(cubes["rho"])
+
+        if arg.magnetic:
             cubes["cos_vB"] = np.zeros_like(cubes["rho"])
-        vv = 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]
+            var_names.append("kr" + a)
+            if arg.magnetic:
                 cubes["cos_vB"] += cubes["v" + a] * cubes["B" + a]
-            vv += cubes["v" + a] ** 2
                 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",
+        for v in var_names:
-            "logrho",
-            "vx",
-            "vy",
-            "vz",
-            "krx",
-            "kry",
-            "krz",
-            "Bx",
-            "By",
-            "Bz",
-            "cos_vB",
-        ]:
             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["vz"] = pcube(fcubes["vz"])
+        if arg.magnetic:
             pcubes["B"] = pcube(fcubes["Bx"], fcubes["By"], fcubes["Bz"])
             pcubes["cos_vB"] = pcube(fcubes["cos_vB"])
-        pcubes["vz"] = pcube(fcubes["vz"])
 
         print("Compute 3D power spectra")
         # 3D power spectra -------------------------------------------------------------
-        for v in [
+        vspec = [
             "rho",
             "logrho",
             "v",
@@ -671,13 +704,16 @@ 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
             )
+            if arg.magnetic:
                 pspec_Bpar, _, _, _ = pspectrum(
                     pcubes[v], cubes_k["kBpar"], kbins, knorm_Bpar, 0
                 )
@@ -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)
+
+            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)
+                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)
+            # h5fd.create_group(outpath, "meta", createparents=True)
-            metagrp = h5fd.get_node(outpath, "meta")
+            # metagrp = h5fd.get_node(outpath, "meta")
-            h5fd.create_array(metagrp, "generator", __generator__)
+            # h5fd.create_array(metagrp, "generator", __generator__)
-            h5fd.create_array(metagrp, "version", __version__)
+            # h5fd.create_array(metagrp, "version", __version__)
 
             h5fd.close()
 
@@ -739,6 +781,7 @@ def main(arg):
         cubes_k, kbins, knorm = calc_k(
             1 << clvl, arg.kbins, arg.kbinsbig, arg.dkbig, 2, saxis
         )
+        if arg.magnetic:
             _, knorm_Bpar, knorm_Bperp = proj_B(
                 cubes_k, kbins, Bavg2, "B", dim=2, saxis=saxis, update=True
             )
@@ -746,7 +789,8 @@ def main(arg):
         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,33 +828,29 @@ 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["vz"] = pcube(fcubes2["vz"])
+
+        if arg.magnetic:
             pcubes2["B"] = pcube(fcubes2["Bx"], fcubes2["By"], fcubes2["Bz"])
             pcubes2["cos_vB"] = pcube(fcubes2["cos_vB"])
-        pcubes2["vz"] = pcube(fcubes2["vz"])
 
         print("Compute 2D power spectra")
         # 2D power spectra -------------------------------------------------------------
         ns = 2 ** clvl
         f = "_%%(i)0%dd" % (np.floor(np.log10(ns)) + 1)
-        for v in [
+        for v in var_names:
-            "rho",
-            "logrho",
-            "v",
-            "kr",
-            "vc",
-            "vs",
-            "krc",
-            "krs",
-            "B",
-            "cos_vB",
-            "vz",
-        ]:
             for i in range(ns):
                 pspec, kbins, pspec2, fbins = pspectrum(
                     pcubes2[v][:, :, i], cubes_k["k"][:, :, i], kbins, knorm, arg.fbins
                 )
+
+                if arg.magnetic:
                     pspec_Bpar, _, _, _ = pspectrum(
-                    pcubes2[v][:, :, i], cubes_k["kBpar"][:, :, i], kbins, knorm_Bpar, 0
+                        pcubes2[v][:, :, i],
+                        cubes_k["kBpar"][:, :, i],
+                        kbins,
+                        knorm_Bpar,
+                        0,
                     )
                     pspec_Bperp, kbins, _, _ = pspectrum(
                         pcubes2[v][:, :, i],
@@ -846,12 +885,15 @@ def main(arg):
                         pass
 
                 h5fd.create_array(outpath, "pspec" + suff, pspec, 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,6 +903,8 @@ def main(arg):
                         outpath, "fbins" + suff, fbins, createparents=True
                     )
                 h5fd.create_array(outpath, "norm" + suff, knorm, createparents=True)
+
+                if arg.magnetic:
                     h5fd.create_array(
                         outpath, "norm_Bpar" + suff, knorm_Bpar, createparents=True
                     )
@@ -873,10 +917,10 @@ def main(arg):
                 except T.NoSuchNodeError:
                     pass
 
-                h5fd.create_group(outpath, "meta", createparents=True)
+                # h5fd.create_group(outpath, "meta", createparents=True)
-                metagrp = h5fd.get_node(outpath, "meta")
+                # metagrp = h5fd.get_node(outpath, "meta")
-                h5fd.create_array(metagrp, "generator", __generator__)
+                # h5fd.create_array(metagrp, "generator", __generator__)
-                h5fd.create_array(metagrp, "version", __version__)
+                # h5fd.create_array(metagrp, "version", __version__)
 
                 h5fd.close()