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

Add Olivier and Cedric files for FFT

Browse Source
This commit is contained in:
Noe Brucy
2019-12-10 17:05:31 +01:00
parent 8e5813f6e2
commit 76bd1b48f9
2 changed files with 1171 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pspec_new.py
+838
View File
@@ -0,0 +1,838 @@
"""Compute power spectra"""
from __future__ import division
import argparse
import sys
import textwrap
import numpy as np
from numpy.fft import fftn, ifft
import pymses
from pymses.analysis import amr2cube
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
from utils import args
from utils import tools
__generator__ = "pspec_new.py"
__version__ = "0.2"
def degrade_cube(cube, lvl, integrate=False):
"""Degrade a data cube to a coarser resolution
Parameters:
-----------
cube
(numpy.ndarray, ndim=3) input data cube, lengths in each direction must
be equal, and a power of 2
lvl
(int) level of the output cube (2**lvl values in each direction)
integrate
(bool, default False) if True, values are added instead of averaged
Return value:
-------------
degraded_cube
(numpy.ndarray, ndim=3) output data cube, 2**lvl values in each
direction
"""
assert cube.ndim == 3
nold = cube.shape[0]
nnew = 1 << lvl
assert nnew <= nold
nsum, rem = divmod(nold, nnew)
assert rem == 0
# For each direction, we split the corresponding axis (length nold) into
# 2 axes, the first one being the axis for the output cube (length nnew),
# and the second one the axis to average or integrate over (fine cells).
# reshape creates a view, no copy is involved
v = cube.reshape((nnew, nsum, nnew, nsum, nnew, nsum), order="C")
# Even indexes are kept, odd ones are to be summed
cube_new = np.einsum("iajbkc->ijk", v)
if not integrate:
cube_new /= nsum
return cube_new
def calc_k(n, nbinsk, nbig, dkbig, dim=3, saxis=2):
"""Make cubes containing the wave vectors, a list of bins and the
associated normalization factors
The wave vector bins are first linear, then logarithmic.
Parameters:
-----------
n
(int) number of points in each direction
nbinsk
(int) number of wave vector bins
nbig
(int) number of linear bins
dkbig
(float) linear bin width
dim
(int, default 3) dimension of the Fourier transform to be performed,
should be either 2 or 3
saxis
(int, default 2) for 2D Fourier transforms, the slicing axis (i.e. the
axis where the transform is NOT done), should be 0, 1 or 2
Return value:
-------------
cubes_k
(dict) Dictionary containing the wave vector cubes:
'kx', 'ky', 'kz'
Components of the wave vectors (in 2D, the one corresponding to
saxis is always 0)
'k'
Norm of the wave vector
kbins
(numpy.array) Wave vector bins (length = nbinsk + 1): bin `i` is
between `kbins[i]` and `kbins[i+1]`.
norm
(numpy.array) Number of points of the Fourier space in each wave vector
bin (length = nbinsk)
"""
k_alias = np.arange(n, dtype=np.float64)
k = np.where(k_alias >= n // 2, k_alias - n, k_alias)
if dim == 3:
kx, ky, kz = np.meshgrid(k, k, k, indexing="ij")
else:
a = [k, k, k]
a[saxis] = np.zeros_like(k)
kx, ky, kz = np.meshgrid(a[0], a[1], a[2], indexing="ij")
cube_k = np.sqrt(kx ** 2 + ky ** 2 + kz ** 2)
cubes_k = {"kx": kx, "ky": ky, "kz": kz, "k": cube_k}
kmin = 1.0
kmax = n // 2
nbins2 = nbinsk - nbig
kmid = kmin + nbig * dkbig
kbins = np.concatenate(
(
kmin + (kmid - kmin) * (np.arange(nbig, dtype=np.float) / nbig),
kmid * (kmax / kmid) ** (np.arange(nbins2 + 1, dtype=np.float) / nbins2),
)
)
norm, _ = np.histogram(cube_k, bins=kbins)
return cubes_k, kbins, norm
def helmholtz(cubes, var, update=False):
r"""Perform a Helmholtz decomposition of a vector field
The Helmholtz decomposition of a vector field $\vec{v}$ is the couple of
vector fields $\left(\vec{v}_c, \vec{v}_s\right)$, where $\vec{\nabla}
\cdot \vec{v}_s = 0$ and $\vec{\nabla} \cdot \vec{v}_c = \vec{0}$. 'c'
stands for compressive and 's' for solenoidal.
This routine uses the Fourier-domain projection method, where
$\hat{\vec{v}}_c$ is proportional to the wave vector $\vec{k}$, and
$\hat{\vec{v}}_s$ is orthogonal to $\vec{k}$. For the compressive
component, only the projection of the field on the wave vector is computed
Parameters:
-----------
cubes
(dict) 3D Fourier space data cubes containing the vector field (see
`var`) and the wave numbers (`'kx'`, `'ky'`, `'kz'`)
var
(str) name of the vector field, the components are assumed to
correspond to the `var + dim` key, where `dim` is `'x'`, `'y'` or `'z'`
update
(bool) if True, `cubes` is updated with the data in `hcubes`
Return value:
-------------
hcubes
(dict) the components of the compressive and solenoidal fields, in
Fourier space:
var + 'c'
(numpy.ndarray, ndim=3) the projection of the compressive component
on the unit wave vector
var + 's' + dim
(numpy.ndarray, ndim=3) the solenoidal component, for each `dim` in
`['x', 'y', 'z']`
"""
knorm = np.zeros_like(cubes["k"])
vpar = np.zeros_like(cubes[var + "x"])
for d in ["x", "y", "z"]:
vpar += cubes[var + d] * cubes["k" + d]
knorm += cubes["k" + d] ** 2
knorm = np.sqrt(knorm)
# Prevent NaNs
knorm[0, 0, 0] = 1.0
vpar[0, 0, 0] = 0.0
vpar /= knorm
hcubes = {}
hcubes[var + "c"] = vpar
for d in ["x", "y", "z"]:
hcubes[var + "s" + d] = cubes[var + d] - vpar * cubes["k" + d] / knorm
if update:
cubes.update(hcubes)
return hcubes
def avg_vect(cubes, var, dim=3, saxis=2):
"""Average a vector in the cube
In 3D, the output is a (3,)-array (global average).
In 2D, it is a (n, 3)-array (average by slices).
Parameters:
-----------
cubes
(dict) 3D data cubes containing the vector field (see `var`)
var
(str) name of the vector field, the components are assumed to
correspond to the `var + dim` key, where `dim` is `'x'`, `'y'` or `'z'`
dim
(int, default 3) dimension of the Fourier transform to be performed,
should be either 2 or 3
saxis
(int, default 2) for 2D Fourier transforms, the slicing axis (i.e. the
axis where the transform is NOT done), should be 0, 1 or 2
Return value:
-------------
avg
(numpy.ndarray) averaged vector, shape is (3,) in 3D, (n, 3) in 2D (n =
number of points along the slicing axis)
"""
if dim == 3:
vavg = np.zeros(3, dtype=np.float64)
for i, d in enumerate(["x", "y", "z"]):
vavg[i] = cubes[var + d].mean()
else:
vavg = np.zeros((cubes[var + "x"].shape[saxis], 3), dtype=np.float64)
axes = [0, 1, 2]
axes.remove(saxis)
for i, d in enumerate(["x", "y", "z"]):
vavg[:, i] = cubes[var + d].mean(axis=tuple(axes))
return vavg
def proj_B(cubes_k, kbins, vec, var="", dim=3, saxis=2, update=False):
r"""Project wave vectors parallel and perpendicular to a given vector
If the mean field value is zero, the parallel component is set to 0 and the
perpendicular component is the wave vector itself.
In 2D, `vec` can be either a (n, 3) or a (3,) shaped array, the additional
dimension is taken along the slice axis.
Parameters:
-----------
cubes_k
(dict) 3D data cube containing the wave numbers (`'kx'`, `'ky'`, `'kz'`)
kbins
(numpy.array) wave vector bin edges (length nbins + 1), see `calc_k`
vec
(numpy.array) 3D: vector to project on, 2D: vector or array of vectors
var
(str) name of the vector in the output
dim
(int, default 3) dimension of the Fourier transform to be performed,
should be either 2 or 3
saxis
(int, default 2) for 2D Fourier transforms, the slicing axis (i.e. the
axis where the transform is NOT done), should be 0, 1 or 2
update
(bool, default False) if True, `cubes_k` is updated with the data in
`proj_cubes`
Return value:
-------------
proj_cubes
(dict) the magnitudes of the projected wave vectors:
'k' + var + 'par'
(numpy.ndarray, ndim=3) the projection of the wave vector parallel
to the mean field
'k' + var + 'perp'
(numpy.ndarray, ndim=3) the projection of the wave vector
perpendicular to the mean field
knorm_par
(numpy.array) number of wave vectors in each bin for the parallel
component (length nbins), see `calc_k`
knorm_perp
(numpy.array) number of wave vectors in each bin for the perpendicular
component (length nbins), see `calc_k`
"""
# we want vec_z[..., saxis] to be set to 0 without changing the argument
vec_z = vec.copy()
if dim == 2:
vec_z[..., saxis] = 0.0
# we need vec_z to be indexed correctly: we create dummy axes in the FT
# directions
vind = [np.newaxis, np.newaxis, np.newaxis]
if dim == 2:
vind[saxis] = slice(None)
vind = tuple(vind)
vec_z = vec_z[vind + (slice(None),)]
vnorm = np.sqrt(np.sum(vec_z ** 2, axis=-1))
kpar = np.zeros_like(cubes_k["k"])
for i, d in enumerate(["x", "y", "z"]):
kpar += cubes_k["k" + d] * vec_z[..., i]
mask = np.logical_not(np.isclose(vnorm, 0))
# Broadcast to the right shape for bool indexing in kpar
mask_b = np.broadcast_to(mask, kpar.shape)
vnorm_b = np.broadcast_to(vnorm, kpar.shape)
# Normalize kpar and vnorm, correctly taking care of zeros
kpar[mask_b] /= vnorm_b[mask_b]
kpar[~mask_b] = 0
vec_z[mask, :] /= vnorm[mask, np.newaxis]
vec_z[~mask] = 0
proj_cubes = {}
proj_cubes["k" + var + "par"] = kpar
kvp = "k" + var + "perp"
proj_cubes[kvp] = np.zeros_like(kpar)
for i, d in enumerate(["x", "y", "z"]):
# note that vec_z[..., saxis] is 0 by construction
proj_cubes[kvp] += (cubes_k["k" + d] - kpar * vec_z[..., i]) ** 2
proj_cubes[kvp] = np.sqrt(proj_cubes[kvp])
if update:
cubes_k.update(proj_cubes)
norm_par, _ = np.histogram(proj_cubes["k" + var + "par"], bins=kbins)
norm_perp, _ = np.histogram(proj_cubes["k" + var + "perp"], bins=kbins)
return proj_cubes, norm_par, norm_perp
def pcube(cube, *others):
"""Compute the power associated to a Fourier space data cube
The power is the sum of square magnitude of each component.
Parameters:
-----------
cube1, cube2, ...
(numpy.ndarray, ndim=3) data cubes for each component
Return value:
-------------
pcube
(numpy.ndarray, ndim=3) power at each point of the Fourier space
"""
pcube = np.real(cube * np.conj(cube))
for c in others:
pcube += np.real(c * np.conj(c))
return pcube
def pspectrum(pcube, kcube, kbins, norm, nbinsf):
"""Compute the power spectrum associated to a power cube
Parameters:
-----------
pcube
(numpy.ndarray, ndim=3) power cube, see `pcube`
kcube
(numpy.ndarray, ndim=3) wave vector magnitude, see `calc_k`
kbins
(numpy.array) wave vector bin edges (length nbins + 1), see `calc_k`
norm
(numpy.array) number of wave vectors in each bin (length nbins), see
`calc_k`
nbinsf
(int) number of power bins for a 2d-bin power spectrum, computed only
if nbinsf > 1
Return value:
-------------
pspec
(numpy.array) power spectrum (length nbins)
kbins
(numpy.array) vector bin edges (the input parameter)
pspec2
(numpy.ndarray, ndim=2 | None) 2d-bin power spectrum (shape nbins,
nbinsf), or None if `nbinsf <= 1`
fbins
(numpy.array | None) power bin edges (length nbinsf + 1), or None if
`nbinsf <= 1`
"""
# Flatten
k_pts = kcube.flatten()
f_pts = pcube.flatten()
# Drop zeros, NaNs and infinities
mask = np.logical_and(f_pts > 0.0, np.isfinite(f_pts))
k_pts = k_pts[mask]
f_pts = f_pts[mask]
fbins = None
pspec2 = None
if nbinsf > 1:
# Fourier coefficients binning
fmin = f_pts.min()
fmax = f_pts.max()
if fmin == fmax:
fmin *= 0.99
fmax *= 1.01
fbins = fmin * (fmax / fmin) ** (np.arange(nbinsf + 1, dtype=np.float) / nbinsf)
# Binned power spectrum
pspec2, _, _ = np.histogram2d(k_pts, f_pts, bins=(kbins, fbins))
pspec2 /= norm[:, np.newaxis]
# Averaged power spectrum
pow_unnorm, _ = np.histogram(k_pts, bins=kbins, weights=f_pts)
pspec = pow_unnorm / norm
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(
"""
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",
)
arg = parser.parse_args()
add_pspec2 = False
if arg.fbins > 0:
add_pspec2 = True
# Cube level (0 means levelmin)
clvl = arg.level
# Slicing axis for 2D data (x = 0, y = 1, z = 2)
try:
saxis = {"x": 0, "y": 1, "z": 2}[arg.sliceaxis.strip()]
except KeyError:
parser.error("Invalid slicing axis: %r" % arg.sliceaxis)
# Output numbers
iouts = tools.select_outputs(arg.repo, arg.iouts)
for iout in iouts:
print("Output %d" % iout)
print("Load data")
read_lvl = None
if True:
# Load output ------------------------------------------------------------------
ro = pymses.RamsesOutput(arg.repo, iout, order=arg.order)
amr = ro.amr_source(["rho", "vel", "Bl", "Br"])
if clvl == 0:
clvl = ro.info["levelmin"]
read_lvl = max(clvl, 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["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]),
]
cubes_arr = amr2cube(amr, cube_vars, xmin, xmax, read_lvl)
cubes = {}
for i, v in enumerate(["rho", "vx", "vy", "vz", "Bx", "By", "Bz"]):
cubes[v] = cubes_arr[i, ...].copy()
del cubes_arr
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(),
}
h5f.close()
if clvl > read_lvl:
print(
"WARNING: adjusting cube level (%d) to match data level (%d)"
% (clvl, read_lvl)
)
clvl = read_lvl
# Degrade cubes ----------------------------------------------------------------
if clvl < read_lvl:
print("Degrade cubes")
for v in ["rho", "vx", "vy", "vz", "Bx", "By", "Bz"]:
# Don't average, simply integrate the magnetic field
cube_tmp = degrade_cube(cubes[v], clvl, v.startswith("B"))
del cubes[v]
cubes[v] = cube_tmp
print("Calculate wave numbers")
# Wave numbers -----------------------------------------------------------------
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
)
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"])
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
print("3D FFT")
# 3D Fourier transform ---------------------------------------------------------
fcubes = {}
for v in [
"rho",
"logrho",
"vx",
"vy",
"vz",
"krx",
"kry",
"krz",
"Bx",
"By",
"Bz",
"cos_vB",
]:
fcubes[v] = fftn(cubes[v])
fcubes.update(cubes_k)
# Memory cleanup ---------------------------------------------------------------
del cubes
print("Calculate Fourier-space derived quantities")
# Additional quantities (Fourier domain) ---------------------------------------
helmholtz(fcubes, "v", update=True)
helmholtz(fcubes, "kr", update=True)
print("Compute power cubes")
# Power cubes ------------------------------------------------------------------
pcubes = {}
pcubes["rho"] = pcube(fcubes["rho"])
pcubes["logrho"] = pcube(fcubes["logrho"])
pcubes["v"] = pcube(fcubes["vx"], fcubes["vy"], fcubes["vz"])
pcubes["kr"] = pcube(fcubes["krx"], fcubes["kry"], fcubes["krz"])
pcubes["vc"] = pcube(fcubes["vc"])
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"])
print("Compute 3D power spectra")
# 3D power spectra -------------------------------------------------------------
for v in ["rho", "logrho", "v", "kr", "vc", "vs", "krc", "krs", "B", "cos_vB"]:
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
)
# Save 3D power spectra
params = {"iout": iout, "varname": v, "dim": 3}
outfile = arg.outfile % params
outpath = arg.nodename % params
h5fd = T.open_file(outfile, mode="a")
for n in [
"pspec",
"kbins",
"pspec2",
"fbins",
"norm",
"pspec_Bpar",
"pspec_Bperp",
"norm_Bpar",
"norm_Bperp",
]:
try:
h5fd.remove_node(outpath, n, recursive=True)
except T.NoSuchNodeError:
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)
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.close()
# Memory cleanup ---------------------------------------------------------------
del pcubes
print("Calculate 2D wave numbers")
# 2D wave numbers --------------------------------------------------------------
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
)
print("Project 3D -> 2D")
# 3D -> 2D ---------------------------------------------------------------------
fcubes2 = {}
for v in [
"rho",
"logrho",
"vx",
"vy",
"vz",
"krx",
"kry",
"krz",
"vc",
"vsx",
"vsy",
"vsz",
"krc",
"krsx",
"krsy",
"krsz",
"Bx",
"By",
"Bz",
"cos_vB",
]:
fcubes2[v] = ifft(fcubes[v], axis=saxis)
# Memory cleanup ---------------------------------------------------------------
del fcubes
print("Compute 2D power cubes")
# 2D power cubes ---------------------------------------------------------------
pcubes2 = {}
pcubes2["rho"] = pcube(fcubes2["rho"])
pcubes2["logrho"] = pcube(fcubes2["logrho"])
pcubes2["v"] = pcube(fcubes2["vx"], fcubes2["vy"], fcubes2["vz"])
pcubes2["kr"] = pcube(fcubes2["krx"], fcubes2["kry"], fcubes2["krz"])
pcubes2["vc"] = pcube(fcubes2["vc"])
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"])
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 i in xrange(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,
)
# Save 2D power spectra
suff = f % {"i": i}
params = {"iout": iout, "varname": v, "dim": 2}
outfile = arg.outfile % params
outpath = arg.nodename % params
h5fd = T.open_file(outfile, mode="a")
for n in [
"pspec",
"kbins",
"pspec2",
"fbins",
"norm",
"pspec_Bpar",
"pspec_Bperp",
"norm_Bpar",
"norm_Bperp",
]:
try:
h5fd.remove_node(outpath, n + suff, recursive=True)
except T.NoSuchNodeError:
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
)
h5fd.create_array(outpath, "kbins" + suff, kbins, createparents=True)
if add_pspec2:
h5fd.create_array(
outpath, "pspec2" + suff, pspec2, createparents=True
)
h5fd.create_array(
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
)
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.close()