Possible to call rule directly

This commit is contained in:
Noe Brucy
2020-01-16 17:35:13 +01:00
parent 76bd1b48f9
commit beb6203d06
6 changed files with 626 additions and 340 deletions
+106 -84
View File
@@ -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))