Possible to call rule directly
This commit is contained in:
+106
-84
@@ -10,16 +10,15 @@ import numpy as np
|
||||
from numpy.fft import fftn, ifft
|
||||
|
||||
import pymses
|
||||
from pymses.analysis import amr2cube
|
||||
from pymses.analysis import ScalarOperator, Camera
|
||||
from pymses.analysis import cube3d
|
||||
import pymses.utils.misc
|
||||
|
||||
# Don't use multiprocessing, it will crash with level 10 cubes
|
||||
pymses.utils.misc.NUMBER_OF_PROCESSES_LIMIT = 1
|
||||
|
||||
import tables as T
|
||||
import bunch
|
||||
|
||||
from utils import args
|
||||
from utils import tools
|
||||
from i_utils import args
|
||||
from i_utils import tools
|
||||
|
||||
__generator__ = "pspec_new.py"
|
||||
__version__ = "0.2"
|
||||
@@ -420,85 +419,78 @@ def pspectrum(pcube, kcube, kbins, norm, nbinsf):
|
||||
return pspec, kbins, pspec2, fbins
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Command-line parser ----------------------------------------------------------
|
||||
parser = argparse.ArgumentParser(
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
description="Compute 2D and 3D power spectra",
|
||||
epilog=textwrap.dedent(
|
||||
"""
|
||||
# Command-line parser ----------------------------------------------------------
|
||||
parser = argparse.ArgumentParser(
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
description="Compute 2D and 3D power spectra",
|
||||
epilog=textwrap.dedent(
|
||||
"""
|
||||
In the output file and node name formats, you can use the formatting
|
||||
fields:
|
||||
* %(iout)d - output number
|
||||
* %(varname)s - variable name
|
||||
* %(dim)d - dimension (3 for cube, 2 for slices)
|
||||
"""
|
||||
),
|
||||
)
|
||||
parser.add_argument("repo", help="RAMSES output repository", type=str, default=".")
|
||||
parser.add_argument(
|
||||
"iouts", help="output numbers", type=args.selection, default=":"
|
||||
)
|
||||
parser.add_argument(
|
||||
"outfile", help="output file format (see below for fields)", type=str
|
||||
)
|
||||
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
|
||||
)
|
||||
parser.add_argument(
|
||||
"-f",
|
||||
"--fbins",
|
||||
help="number of Fourier bins for k-power 2D histogram (0 to disable)",
|
||||
type=int,
|
||||
default=0,
|
||||
)
|
||||
parser.add_argument(
|
||||
"-K", "--kbinsbig", help="number of big wave number bins", type=int, default=9
|
||||
)
|
||||
parser.add_argument(
|
||||
"-d",
|
||||
"--dkbig",
|
||||
help="width of the big wave number bins",
|
||||
type=float,
|
||||
default=1.0,
|
||||
)
|
||||
parser.add_argument(
|
||||
"-S",
|
||||
"--sliceaxis",
|
||||
help="slicing axis",
|
||||
type=str,
|
||||
choices=["x", "y", "z"],
|
||||
default="z",
|
||||
)
|
||||
),
|
||||
)
|
||||
parser.add_argument("repo", help="RAMSES output repository", type=str, default=".")
|
||||
parser.add_argument("iouts", help="output numbers", type=args.selection, default=":")
|
||||
parser.add_argument(
|
||||
"outfile", help="output file format (see below for fields)", type=str
|
||||
)
|
||||
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
|
||||
)
|
||||
parser.add_argument(
|
||||
"-f",
|
||||
"--fbins",
|
||||
help="number of Fourier bins for k-power 2D histogram (0 to disable)",
|
||||
type=int,
|
||||
default=0,
|
||||
)
|
||||
parser.add_argument(
|
||||
"-K", "--kbinsbig", help="number of big wave number bins", type=int, default=9
|
||||
)
|
||||
parser.add_argument(
|
||||
"-d", "--dkbig", help="width of the big wave number bins", type=float, default=1.0
|
||||
)
|
||||
parser.add_argument(
|
||||
"-S",
|
||||
"--sliceaxis",
|
||||
help="slicing axis",
|
||||
type=str,
|
||||
choices=["x", "y", "z"],
|
||||
default="z",
|
||||
)
|
||||
|
||||
arg = parser.parse_args()
|
||||
|
||||
def main(arg):
|
||||
add_pspec2 = False
|
||||
if arg.fbins > 0:
|
||||
add_pspec2 = True
|
||||
@@ -525,13 +517,11 @@ if __name__ == "__main__":
|
||||
|
||||
if clvl == 0:
|
||||
clvl = ro.info["levelmin"]
|
||||
read_lvl = max(clvl, ro.info["levelmin"])
|
||||
read_lvl = ro.info["levelmin"]
|
||||
|
||||
# Extract cubes ---------------------------------------------------------------
|
||||
xmin = [x - arg.size / 2.0 for x in arg.center]
|
||||
xmax = [x + arg.size / 2.0 for x in arg.center]
|
||||
cube_vars = [
|
||||
"rho",
|
||||
lambda dset: dset["rho"],
|
||||
lambda dset: dset["vel"][..., 0],
|
||||
lambda dset: dset["vel"][..., 1],
|
||||
lambda dset: dset["vel"][..., 2],
|
||||
@@ -539,11 +529,34 @@ if __name__ == "__main__":
|
||||
lambda dset: 0.5 * (dset["Bl"][..., 1] + dset["Br"][..., 1]),
|
||||
lambda dset: 0.5 * (dset["Bl"][..., 2] + dset["Br"][..., 2]),
|
||||
]
|
||||
cubes_arr = amr2cube(amr, cube_vars, xmin, xmax, read_lvl)
|
||||
|
||||
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"],
|
||||
]
|
||||
|
||||
cam = Camera(
|
||||
center=arg.center,
|
||||
line_of_sight_axis="z",
|
||||
region_size=[arg.size, arg.size],
|
||||
distance=arg.size / 2.0,
|
||||
far_cut_depth=arg.size / 2.0,
|
||||
up_vector="y",
|
||||
map_max_size=256,
|
||||
)
|
||||
|
||||
cubes = {}
|
||||
for i, v in enumerate(["rho", "vx", "vy", "vz", "Bx", "By", "Bz"]):
|
||||
cubes[v] = cubes_arr[i, ...].copy()
|
||||
del cubes_arr
|
||||
operator = ScalarOperator(cube_vars[i], cube_units[i])
|
||||
extractor = cube3d.CubeExtractor(amr, operator)
|
||||
cubes[v] = extractor.process(
|
||||
cam, cube_size=arg.size, resolution=256
|
||||
).data
|
||||
else:
|
||||
h5f = T.open_file("cube.hdf5", "r")
|
||||
read_lvl = np.asscalar(h5f.root.level.read())
|
||||
@@ -836,3 +849,12 @@ if __name__ == "__main__":
|
||||
h5fd.create_array(metagrp, "version", __version__)
|
||||
|
||||
h5fd.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
arg = parser.parse_args()
|
||||
main(arg)
|
||||
|
||||
|
||||
def pspec(**kwargs):
|
||||
main(bunch.bunchify(kwargs))
|
||||
|
||||
Reference in New Issue
Block a user