From 76bd1b48f9edfc811d19c7b8426af909a555e749 Mon Sep 17 00:00:00 2001 From: Noe Brucy Date: Tue, 10 Dec 2019 17:05:31 +0100 Subject: [PATCH] Add Olivier and Cedric files for FFT --- pspec_new.py | 838 ++++++++++++++++++++++++++++++++++++++++++++++++++ pspec_read.py | 333 ++++++++++++++++++++ 2 files changed, 1171 insertions(+) create mode 100644 pspec_new.py create mode 100644 pspec_read.py diff --git a/pspec_new.py b/pspec_new.py new file mode 100644 index 0000000..ed4cd11 --- /dev/null +++ b/pspec_new.py @@ -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() diff --git a/pspec_read.py b/pspec_read.py new file mode 100644 index 0000000..7c00234 --- /dev/null +++ b/pspec_read.py @@ -0,0 +1,333 @@ +import tables as T +import numpy as np +import matplotlib.pyplot as P + + +################################################################ +# 3D +################################################################ + + +def pspec_rho(file, path_out, num, color="r", save=False): + + h5f = T.open_file(file, mode="r") + group = h5f.get_node("/out_" + format(num, "05") + "/d3/rho") + kbins = h5f.get_node(group, "kbins").read() + pspec = h5f.get_node(group, "pspec").read() + h5f.close() + + k = np.sqrt(kbins[1:] * kbins[:-1]) + pspec = (k ** 2) * pspec + if save: + P.figure(figsize=(8, 8)) + P.xlabel(r"$k$", fontsize=12) + P.ylabel(r"$k^2 \mathcal{P}(\rho)$", fontsize=12) + P.loglog(k, pspec, "-", color=color) + if save: + P.savefig(path_out + "pspec_rho_" + format(num, "05") + ".ps") + P.savefig(path_out + "pspec_rho_" + format(num, "05") + ".jpeg") + + +def pspec_B(file, path_out, num, color="r", save=False): + + h5f = T.open_file(file, mode="r") + group = h5f.get_node("/out_" + format(num, "05") + "/d3/B") + kbins = h5f.get_node(group, "kbins").read() + pspec = h5f.get_node(group, "pspec").read() + h5f.close() + + k = np.sqrt(kbins[1:] * kbins[:-1]) + pspec = k ** 2 * pspec + if save: + P.figure(figsize=(8, 8)) + P.xlabel(r"$k$", fontsize=12) + P.ylabel(r"$k^2 \mathcal{P}(B)$", fontsize=12) + P.loglog(k, pspec, "-", color=color) + if save: + P.savefig(path_out + "pspec_B_" + format(num, "05") + ".ps") + P.savefig(path_out + "pspec_B_" + format(num, "05") + ".jpeg") + + +def pspec_v(file, path_out, num, color="r", save=False): + + h5f = T.open_file(file, mode="r") + group = h5f.get_node("/out_" + format(num, "05") + "/d3/v") + kbins = h5f.get_node(group, "kbins").read() + pspec = h5f.get_node(group, "pspec").read() + h5f.close() + + k = np.sqrt(kbins[1:] * kbins[:-1]) + pspec = k ** 4 * pspec + if save: + P.figure(figsize=(8, 8)) + P.xlabel(r"$k$", fontsize=12) + P.ylabel(r"$k^4 \mathcal{P}(v)$", fontsize=12) + P.loglog(k, pspec, "-", color=color) + if save: + P.savefig(path_out + "pspec_v_" + format(num, "05") + ".ps") + P.savefig(path_out + "pspec_v_" + format(num, "05") + ".jpeg") + + +def pspec_cos(file, path_out, num, color="r", save=False): + + h5f = T.open_file(file, mode="r") + group = h5f.get_node("/out_" + format(num, "05") + "/d3/cos_vB") + kbins = h5f.get_node(group, "kbins").read() + pspec = h5f.get_node(group, "pspec").read() + h5f.close() + + k = np.sqrt(kbins[1:] * kbins[:-1]) + pspec = k ** 2 * pspec + if save: + P.figure(figsize=(8, 8)) + P.xlabel(r"$k$", fontsize=12) + P.ylabel(r"$k^2 \mathcal{P}(\cos(\vec{v}.\vec{B}))$", fontsize=12) + P.loglog(k, pspec, "-", color=color) + if save: + P.savefig(path_out + "pspec_cos_" + format(num, "05") + ".ps") + P.savefig(path_out + "pspec_cos_" + format(num, "05") + ".jpeg") + + +def pspec_logrho(file, path_out, num, color="r", save=False): + + h5f = T.open_file(file, mode="r") + group = h5f.get_node("/out_" + format(num, "05") + "/d3/logrho") + kbins = h5f.get_node(group, "kbins").read() + pspec = h5f.get_node(group, "pspec").read() + h5f.close() + + k = np.sqrt(kbins[1:] * kbins[:-1]) + pspec = k ** 2 * pspec + if save: + P.figure(figsize=(8, 8)) + P.xlabel(r"$k$", fontsize=12) + P.ylabel(r"$k^2 \mathcal{P}(\log(\rho))$", fontsize=12) + P.loglog(k, pspec, "-", color=color) + if save: + P.savefig(path_out + "pspec_logrho_" + format(num, "05") + ".ps") + P.savefig(path_out + "pspec_logrho_" + format(num, "05") + ".jpeg") + + +def pspec_kr(file, path_out, num, color="r", save=False): + + h5f = T.open_file(file, mode="r") + group = h5f.get_node("/out_" + format(num, "05") + "/d3/kr") + kbins = h5f.get_node(group, "kbins").read() + pspec = h5f.get_node(group, "pspec").read() + h5f.close() + + k = np.sqrt(kbins[1:] * kbins[:-1]) + pspec = (k ** 4) * pspec + if save: + P.figure(figsize=(8, 8)) + P.xlabel(r"$k$", fontsize=12) + P.ylabel(r"$k^4 \mathcal{P}(\rho^{1/3}v)$", fontsize=12) + P.loglog(k, pspec, "-", color=color) + if save: + P.savefig(path_out + "pspec_kr_" + format(num, "05") + ".ps") + P.savefig(path_out + "pspec_kr_" + format(num, "05") + ".jpeg") + + +def pspec_vc(file, path_out, num, color="r", save=False): + + h5f = T.open_file(file, mode="r") + group = h5f.get_node("/out_" + format(num, "05") + "/d3/vc") + kbins = h5f.get_node(group, "kbins").read() + pspec = h5f.get_node(group, "pspec").read() + h5f.close() + + k = np.sqrt(kbins[1:] * kbins[:-1]) + pspec = (k ** 4) * pspec + if save: + P.figure(figsize=(8, 8)) + P.xlabel(r"$k$", fontsize=12) + P.ylabel(r"$k^4 \mathcal{P}(v_c)$", fontsize=12) + P.loglog(k, pspec, "-", color=color) + if save: + P.savefig(path_out + "pspec_vc_" + format(num, "05") + ".ps") + P.savefig(path_out + "pspec_vc_" + format(num, "05") + ".jpeg") + + +def pspec_vs(file, path_out, num, color="r", save=False): + + h5f = T.open_file(file, mode="r") + group = h5f.get_node("/out_" + format(num, "05") + "/d3/vs") + kbins = h5f.get_node(group, "kbins").read() + pspec = h5f.get_node(group, "pspec").read() + h5f.close() + + k = np.sqrt(kbins[1:] * kbins[:-1]) + pspec = (k ** 4) * pspec + if save: + P.figure(figsize=(8, 8)) + P.xlabel(r"$k$", fontsize=12) + P.ylabel(r"$k^4 \mathcal{P}(v_s)$", fontsize=12) + P.loglog(k, pspec, "-", color=color) + if save: + P.savefig(path_out + "pspec_vs_" + format(num, "05") + ".ps") + P.savefig(path_out + "pspec_vs_" + format(num, "05") + ".jpeg") + + +################################################################### +# 2D +################################################################### + + +def pspec_rho_2D(file, path_out, num, plan, color="r", save=False): + + h5f = T.open_file(file, mode="r") + group = h5f.get_node("/out_" + format(num, "05") + "/d2/rho") + kbins = h5f.get_node(group, "kbins_" + str(plan)).read() + pspec = h5f.get_node(group, "pspec_" + str(plan)).read() + h5f.close() + + k = np.sqrt(kbins[1:] * kbins[:-1]) + pspec = (k ** 1) * pspec + if save: + P.figure(figsize=(8, 8)) + P.xlabel(r"$k$", fontsize=12) + P.ylabel(r"$k \mathcal{P}(\rho)$", fontsize=12) + P.loglog(k, pspec, "-", color=color) + if save: + P.savefig(path_out + "pspec_rho_2D_" + format(num, "05") + ".ps") + P.savefig(path_out + "pspec_rho_2D_" + format(num, "05") + ".jpeg") + + +def pspec_B_2D(file, path_out, num, plan, color="r", save=False): + + h5f = T.open_file(file, mode="r") + group = h5f.get_node("/out_" + format(num, "05") + "/d2/B") + kbins = h5f.get_node(group, "kbins_" + str(plan)).read() + pspec = h5f.get_node(group, "pspec_" + str(plan)).read() + h5f.close() + + k = np.sqrt(kbins[1:] * kbins[:-1]) + pspec = k ** 1 * pspec + if save: + P.figure(figsize=(8, 8)) + P.xlabel(r"$k$", fontsize=12) + P.ylabel(r"$k \mathcal{P}(B)$", fontsize=12) + P.loglog(k, pspec, "-", color=color) + if save: + P.savefig(path_out + "pspec_B_2D_" + format(num, "05") + ".ps") + P.savefig(path_out + "pspec_B_2D_" + format(num, "05") + ".jpeg") + + +def pspec_v_2D(file, path_out, num, plan, color="r", save=False): + + h5f = T.open_file(file, mode="r") + group = h5f.get_node("/out_" + format(num, "05") + "/d2/v") + kbins = h5f.get_node(group, "kbins_" + str(plan)).read() + pspec = h5f.get_node(group, "pspec_" + str(plan)).read() + h5f.close() + + k = np.sqrt(kbins[1:] * kbins[:-1]) + pspec = k ** 3 * pspec + if save: + P.figure(figsize=(8, 8)) + P.xlabel(r"$k$", fontsize=12) + P.ylabel(r"$k^3 \mathcal{P}(v)$", fontsize=12) + P.loglog(k, pspec, "-", color=color) + if save: + P.savefig(path_out + "pspec_v_2D_" + format(num, "05") + ".ps") + P.savefig(path_out + "pspec_v_2D_" + format(num, "05") + ".jpeg") + + +def pspec_cos_2D(file, path_out, num, plan, color="r", save=False): + + h5f = T.open_file(file, mode="r") + group = h5f.get_node("/out_" + format(num, "05") + "/d2/cos_vB") + kbins = h5f.get_node(group, "kbins_" + str(plan)).read() + pspec = h5f.get_node(group, "pspec_" + str(plan)).read() + h5f.close() + + k = np.sqrt(kbins[1:] * kbins[:-1]) + pspec = k * pspec + if save: + P.figure(figsize=(8, 8)) + P.xlabel(r"$k$", fontsize=12) + P.ylabel(r"$k \mathcal{P}(\cos(\vec{v}.\vec{B}))$", fontsize=12) + P.loglog(k, pspec, "-", color=color) + if save: + P.savefig(path_out + "pspec_cos_2D_" + format(num, "05") + ".ps") + P.savefig(path_out + "pspec_cos_2D_" + format(num, "05") + ".jpeg") + + +def pspec_logrho_2D(file, path_out, num, plan, color="r", save=False): + + h5f = T.open_file(file, mode="r") + group = h5f.get_node("/out_" + format(num, "05") + "/d2/logrho") + kbins = h5f.get_node(group, "kbins_" + str(plan)).read() + pspec = h5f.get_node(group, "pspec_" + str(plan)).read() + h5f.close() + + k = np.sqrt(kbins[1:] * kbins[:-1]) + pspec = k * pspec + if save: + P.figure(figsize=(8, 8)) + P.xlabel(r"$k$", fontsize=12) + P.ylabel(r"$k \mathcal{P}(\log(\rho))$", fontsize=12) + P.loglog(k, pspec, "-", color=color) + if save: + P.savefig(path_out + "pspec_logrho_2D_" + format(num, "05") + ".ps") + P.savefig(path_out + "pspec_logrho_2D_" + format(num, "05") + ".jpeg") + + +def pspec_kr_2D(file, path_out, num, plan, color="r", save=False): + + h5f = T.open_file(file, mode="r") + group = h5f.get_node("/out_" + format(num, "05") + "/d2/kr") + kbins = h5f.get_node(group, "kbins_" + str(plan)).read() + pspec = h5f.get_node(group, "pspec_" + str(plan)).read() + h5f.close() + + k = np.sqrt(kbins[1:] * kbins[:-1]) + pspec = (k ** 3) * pspec + if save: + P.figure(figsize=(8, 8)) + P.xlabel(r"$k$", fontsize=12) + P.ylabel(r"$k^3 \mathcal{P}(\rho^{1/3}v)$", fontsize=12) + P.loglog(k, pspec, "-", color=color) + if save: + P.savefig(path_out + "pspec_kr_2D_" + format(num, "05") + ".ps") + P.savefig(path_out + "pspec_kr_2D_" + format(num, "05") + ".jpeg") + + +def pspec_vc_2D(file, path_out, num, plan, color="r", save=False): + + h5f = T.open_file(file, mode="r") + group = h5f.get_node("/out_" + format(num, "05") + "/d2/vc") + kbins = h5f.get_node(group, "kbins_" + str(plan)).read() + pspec = h5f.get_node(group, "pspec_" + str(plan)).read() + h5f.close() + + k = np.sqrt(kbins[1:] * kbins[:-1]) + pspec = (k ** 3) * pspec + if save: + P.figure(figsize=(8, 8)) + P.xlabel(r"$k$", fontsize=12) + P.ylabel(r"$k^3 \mathcal{P}(v_c)$", fontsize=12) + P.loglog(k, pspec, "-", color=color) + if save: + P.savefig(path_out + "pspec_vc_2D" + format(num, "05") + ".ps") + P.savefig(path_out + "pspec_vc_2D" + format(num, "05") + ".jpeg") + + +def pspec_vs_2D(file, path_out, num, plan, color="r", save=False): + + h5f = T.open_file(file, mode="r") + group = h5f.get_node("/out_" + format(num, "05") + "/d2/vs") + kbins = h5f.get_node(group, "kbins_" + str(plan)).read() + pspec = h5f.get_node(group, "pspec_" + str(plan)).read() + h5f.close() + + k = np.sqrt(kbins[1:] * kbins[:-1]) + pspec = (k ** 3) * pspec + if save: + P.figure(figsize=(8, 8)) + P.xlabel(r"$k$", fontsize=12) + P.ylabel(r"$k^3 \mathcal{P}(v_s)$", fontsize=12) + P.loglog(k, pspec, "-", color=color) + if save: + P.savefig(path_out + "pspec_vs_2D" + format(num, "05") + ".ps") + P.savefig(path_out + "pspec_vs_2D" + format(num, "05") + ".jpeg")