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
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
+175 -131
View File
@@ -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,32 +524,49 @@ def main(arg):
if True:
# Load output ------------------------------------------------------------------
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:
clvl = ro.info["levelmin"]
read_lvl = ro.info["levelmin"]
# Extract cubes ---------------------------------------------------------------
cube_vars = [
lambda dset: dset["rho"],
lambda dset: dset["vel"][..., 0],
lambda dset: dset["vel"][..., 1],
lambda dset: dset["vel"][..., 2],
lambda dset: 0.5 * (dset["Bl"][..., 0] + dset["Br"][..., 0]),
lambda dset: 0.5 * (dset["Bl"][..., 1] + dset["Br"][..., 1]),
lambda dset: 0.5 * (dset["Bl"][..., 2] + dset["Br"][..., 2]),
]
if arg.magnetic:
cube_vars = [
lambda dset: dset["rho"],
lambda dset: dset["vel"][..., 0],
lambda dset: dset["vel"][..., 1],
lambda dset: dset["vel"][..., 2],
lambda dset: 0.5 * (dset["Bl"][..., 0] + dset["Br"][..., 0]),
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 = [
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"],
]
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,
@@ -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,17 +592,24 @@ def main(arg):
else:
h5f = T.open_file("cube.hdf5", "r")
read_lvl = np.asscalar(h5f.root.level.read())
cubes = {
"rho": h5f.root.rho.read(),
"vx": h5f.root.vx.read(),
"vy": h5f.root.vy.read(),
"vz": h5f.root.vz.read(),
"Bx": h5f.root.Bx.read(),
"By": h5f.root.By.read(),
"Bz": h5f.root.Bz.read(),
}
if arg.magnetic:
cubes = {
"rho": h5f.root.rho.read(),
"vx": h5f.root.vx.read(),
"vy": h5f.root.vy.read(),
"vz": h5f.root.vz.read(),
"Bx": h5f.root.Bx.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()
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,44 +632,40 @@ def main(arg):
cubes_k, kbins, knorm = calc_k(
1 << clvl, arg.kbins, arg.kbinsbig, arg.dkbig, 3, saxis
)
Bavg = avg_vect(cubes, "B", dim=3)
Bavg2 = avg_vect(cubes, "B", dim=2, saxis=saxis)
_, knorm_Bpar, knorm_Bperp = proj_B(
cubes_k, kbins, Bavg, "B", dim=3, update=True
)
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(
cubes_k, kbins, Bavg, "B", dim=3, update=True
)
print("Calculate derived quantities")
# Additional quantities --------------------------------------------------------
cubes["logrho"] = np.log(cubes["rho"])
cubes["cos_vB"] = np.zeros_like(cubes["rho"])
vv = np.zeros_like(cubes["rho"])
BB = np.zeros_like(cubes["rho"])
if arg.magnetic:
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"]:
cubes["kr" + a] = cubes["rho"] ** (1.0 / 3.0) * cubes["v" + a]
cubes["cos_vB"] += cubes["v" + a] * cubes["B" + a]
vv += cubes["v" + a] ** 2
BB += cubes["B" + a] ** 2
cubes["cos_vB"] /= vv * BB
del vv
del BB
var_names.append("kr" + a)
if arg.magnetic:
cubes["cos_vB"] += cubes["v" + a] * cubes["B" + a]
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",
"logrho",
"vx",
"vy",
"vz",
"krx",
"kry",
"krz",
"Bx",
"By",
"Bz",
"cos_vB",
]:
for v in var_names:
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["B"] = pcube(fcubes["Bx"], fcubes["By"], fcubes["Bz"])
pcubes["cos_vB"] = pcube(fcubes["cos_vB"])
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")
# 3D power spectra -------------------------------------------------------------
for v in [
vspec = [
"rho",
"logrho",
"v",
@@ -671,19 +704,22 @@ 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
)
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
)
if arg.magnetic:
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
)
# Save 3D power spectra
params = {"iout": iout, "varname": v, "dim": 3}
@@ -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)
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:
h5fd.remove_node(outpath, "meta", recursive=True)
except T.NoSuchNodeError:
pass
h5fd.create_group(outpath, "meta", createparents=True)
metagrp = h5fd.get_node(outpath, "meta")
h5fd.create_array(metagrp, "generator", __generator__)
h5fd.create_array(metagrp, "version", __version__)
# h5fd.create_group(outpath, "meta", createparents=True)
# metagrp = h5fd.get_node(outpath, "meta")
# h5fd.create_array(metagrp, "generator", __generator__)
# h5fd.create_array(metagrp, "version", __version__)
h5fd.close()
@@ -739,14 +781,16 @@ def main(arg):
cubes_k, kbins, knorm = calc_k(
1 << clvl, arg.kbins, arg.kbinsbig, arg.dkbig, 2, saxis
)
_, knorm_Bpar, knorm_Bperp = proj_B(
cubes_k, kbins, Bavg2, "B", dim=2, saxis=saxis, update=True
)
if arg.magnetic:
_, knorm_Bpar, knorm_Bperp = proj_B(
cubes_k, kbins, Bavg2, "B", dim=2, saxis=saxis, update=True
)
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,41 +828,37 @@ 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["B"] = pcube(fcubes2["Bx"], fcubes2["By"], fcubes2["Bz"])
pcubes2["cos_vB"] = pcube(fcubes2["cos_vB"])
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")
# 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",
"vz",
]:
for v in var_names:
for i in range(ns):
pspec, kbins, pspec2, fbins = pspectrum(
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
)
pspec_Bperp, kbins, _, _ = pspectrum(
pcubes2[v][:, :, i],
cubes_k["kBperp"][:, :, i],
kbins,
knorm_Bperp,
0,
)
if arg.magnetic:
pspec_Bpar, _, _, _ = pspectrum(
pcubes2[v][:, :, i],
cubes_k["kBpar"][:, :, i],
kbins,
knorm_Bpar,
0,
)
pspec_Bperp, kbins, _, _ = pspectrum(
pcubes2[v][:, :, i],
cubes_k["kBperp"][:, :, i],
kbins,
knorm_Bperp,
0,
)
# Save 2D power spectra
suff = f % {"i": i}
@@ -846,12 +885,15 @@ def main(arg):
pass
h5fd.create_array(outpath, "pspec" + suff, pspec, createparents=True)
h5fd.create_array(
outpath, "pspec_Bpar" + suff, pspec_Bpar, createparents=True
)
h5fd.create_array(
outpath, "pspec_Bperp" + suff, pspec_Bperp, 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,22 +903,24 @@ def main(arg):
outpath, "fbins" + suff, fbins, createparents=True
)
h5fd.create_array(outpath, "norm" + suff, knorm, createparents=True)
h5fd.create_array(
outpath, "norm_Bpar" + suff, knorm_Bpar, createparents=True
)
h5fd.create_array(
outpath, "norm_Bperp" + suff, knorm_Bperp, createparents=True
)
if arg.magnetic:
h5fd.create_array(
outpath, "norm_Bpar" + suff, knorm_Bpar, createparents=True
)
h5fd.create_array(
outpath, "norm_Bperp" + suff, knorm_Bperp, createparents=True
)
try:
h5fd.remove_node(outpath, "meta", recursive=True)
except T.NoSuchNodeError:
pass
h5fd.create_group(outpath, "meta", createparents=True)
metagrp = h5fd.get_node(outpath, "meta")
h5fd.create_array(metagrp, "generator", __generator__)
h5fd.create_array(metagrp, "version", __version__)
# h5fd.create_group(outpath, "meta", createparents=True)
# metagrp = h5fd.get_node(outpath, "meta")
# h5fd.create_array(metagrp, "generator", __generator__)
# h5fd.create_array(metagrp, "version", __version__)
h5fd.close()