Add preliminary power spectrum integration

pspec_new.py

@@ -15,7 +15,6 @@ from pymses.analysis import cube3d
 import pymses.utils.misc
 
 import tables as T
-import bunch
 
 from i_utils import args
 from i_utils import tools
@@ -464,6 +463,45 @@ parser.add_argument("-s", "--size", help="cube size", type=float, default=1.0)
 parser.add_argument(
     "-l", "--level", help="cube level (default: levelMIN)", type=int, default=0
 )
+parser.add_argument(
+    "repo", help="RAMSES output repository", type=str, default=".", nargs="?"
+)
+parser.add_argument(
+    "iouts", help="output numbers", type=args.selection, default=":", nargs="?"
+)
+parser.add_argument(
+    "outfile",
+    help="output file format (see below for fields)",
+    type=str,
+    default=".",
+    nargs="?",
+)
+parser.add_argument(
+    "-n",
+    "--nodename",
+    help="node name format (see below for fields)",
+    type=str,
+    default="/out_%(iout)05d/d%(dim)d/%(varname)s",
+)
+parser.add_argument(
+    "-O",
+    "--order",
+    help="byte order (= for native)",
+    type=str,
+    default="=",
+    choices=["<", ">", "="],
+)
+parser.add_argument(
+    "-c",
+    "--center",
+    help="coordinates of the center",
+    type=args.center,
+    default=[0.5, 0.5, 0.5],
+)
+parser.add_argument("-s", "--size", help="cube size", type=float, default=1.0)
+parser.add_argument(
+    "-l", "--level", help="cube level (default: levelMIN)", type=int, default=0
+)
 parser.add_argument(
     "-k", "--kbins", help="number of wave number bins", type=int, default=100
 )
@@ -504,7 +542,7 @@ def main(arg):
         parser.error("Invalid slicing axis: %r" % arg.sliceaxis)
 
     # Output numbers
-    iouts = tools.select_outputs(arg.repo, arg.iouts)
+    iouts = arg.iouts  # tools.select_outputs(arg.repo, arg.iouts)
 
     for iout in iouts:
         print("Output %d" % iout)
@@ -547,7 +585,7 @@ def main(arg):
                 distance=arg.size / 2.0,
                 far_cut_depth=arg.size / 2.0,
                 up_vector="y",
-                map_max_size=256,
+                map_max_size=2 ** clvl,
             )
 
             cubes = {}
@@ -555,7 +593,7 @@ def main(arg):
                 operator = ScalarOperator(cube_vars[i], cube_units[i])
                 extractor = cube3d.CubeExtractor(amr, operator)
                 cubes[v] = extractor.process(
-                    cam, cube_size=arg.size, resolution=256
+                    cam, cube_size=arg.size, resolution=2 ** clvl
                 ).data
         else:
             h5f = T.open_file("cube.hdf5", "r")
@@ -577,7 +615,6 @@ def main(arg):
                 % (clvl, read_lvl)
             )
             clvl = read_lvl
-
         # Degrade cubes ----------------------------------------------------------------
         if clvl < read_lvl:
             print("Degrade cubes")
@@ -655,10 +692,23 @@ def main(arg):
         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"])
 
         print("Compute 3D power spectra")
         # 3D power spectra -------------------------------------------------------------
-        for v in ["rho", "logrho", "v", "kr", "vc", "vs", "krc", "krs", "B", "cos_vB"]:
+        for v in [
+            "rho",
+            "logrho",
+            "v",
+            "kr",
+            "vc",
+            "vs",
+            "krc",
+            "krs",
+            "B",
+            "cos_vB",
+            "vz",
+        ]:
             pspec, kbins, pspec2, fbins = pspectrum(
                 pcubes[v], cubes_k["k"], kbins, knorm, arg.fbins
             )
@@ -751,6 +801,7 @@ def main(arg):
             "By",
             "Bz",
             "cos_vB",
+            "vz",
         ]:
             fcubes2[v] = ifft(fcubes[v], axis=saxis)
 
@@ -770,12 +821,25 @@ def main(arg):
         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"])
 
         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"]:
+        for v in [
+            "rho",
+            "logrho",
+            "v",
+            "kr",
+            "vc",
+            "vs",
+            "krc",
+            "krs",
+            "B",
+            "cos_vB",
+            "vz",
+        ]:
             for i in xrange(ns):
                 pspec, kbins, pspec2, fbins = pspectrum(
                     pcubes2[v][:, :, i], cubes_k["k"][:, :, i], kbins, knorm, arg.fbins
@@ -857,4 +921,7 @@ if __name__ == "__main__":
 
 
 def pspec(**kwargs):
-    main(bunch.bunchify(kwargs))
+    arg = parser.parse_args("")
+    for kwarg in kwargs:
+        setattr(arg, kwarg, kwargs[kwarg])
+    main(arg)