noe/pipeline
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

[pspec] first attempt to isolate magnetic field stuff

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