Source code for viscid.calculator.plasma

#!/usr/bin/env python
""" some equations that are useful for plasmas """

from __future__ import print_function, division

import numpy as np
try:
    import numexpr as ne
    _HAS_NUMEXPR = True
except ImportError:
    _HAS_NUMEXPR = False

from viscid import field
from viscid import logger
# from viscid.calculator import calc

__all__ = ["calc_psi", "calc_beta"]


[docs]def calc_psi(B, rev=False): """Calc Flux function (only valid in 2d) Parameters: B (VectorField): magnetic field, should only have two spatial dimensions so we can infer the symmetry dimension rev (bool): since this integration doesn't like going through undefined regions (like within 1 earth radius of the origin for openggcm), you can use this to start integrating from the opposite corner. Returns: ScalarField: 2-D scalar flux function Raises: ValueError: If B has <> 2 spatial dimensions """ # TODO: if this is painfully slow, i bet just putting this exact # code in a cython module would make it a bunch faster, the problem # being that the loops are in python instead of some broadcasting # numpy type thing B = B.slice_reduce(":") # try to guess if a dim of a 3D field is invariant reduced_axes = [] if B.nr_sdims > 2: slcs = [slice(None)] * B.nr_sdims for i, nxi in enumerate(B.sshape): if nxi <= 2: slcs[i] = 0 reduced_axes.append(B.crds.axes[i]) slcs.insert(B.nr_comp, slice(None)) B = B[slcs] # ok, so the above didn't work... just nip out the smallest dim? if B.nr_sdims == 3: slcs = [slice(None)] * B.nr_sdims i = np.argmin(B.sshape) slcs[i] = 0 reduced_axes.append(B.crds.axes[i]) logger.warning("Tried to get the flux function of a 3D field. " "I can't do that, so I'm\njust ignoring the {0} " "dimension".format(reduced_axes[-1])) slcs.insert(B.nr_comp, slice(None)) B = B[slcs] if B.nr_sdims != 2: raise ValueError("flux function only implemented for 2D fields") comps = "" for comp in "xyz": if comp in B.crds.axes: comps += comp # ex: comps = "yz", comp_inds = [1, 2] comp_inds = [dict(x=0, y=1, z=2)[comp] for comp in comps] # Note: what follows says y, z, but it has been generalized # to any two directions, so hy isn't necessarily hy, but it's # easier to see at a glance if it's correct using a specific # example ycc, zcc = B.get_crds(comps) comp_views = B.component_views() hy, hz = comp_views[comp_inds[0]], comp_views[comp_inds[1]] dy = ycc[1:] - ycc[:-1] dz = zcc[1:] - zcc[:-1] ny, nz = len(ycc), len(zcc) A = np.empty((ny, nz), dtype=B.dtype) if rev: A[-1, -1] = 0.0 for i in range(ny - 2, -1, -1): A[i, -1] = A[i + 1, -1] - dy[i] * 0.5 * (hz[i, -1] + hz[i + 1, -1]) for j in range(nz - 2, -1, -1): A[:, j] = A[:, j + 1] + dz[j] * 0.5 * (hy[:, j + 1] + hy[:, j]) else: A[0, 0] = 0.0 for i in range(1, ny): A[i, 0] = A[i - 1, 0] + dy[i - 1] * 0.5 * (hz[i, 0] + hz[i - 1, 0]) for j in range(1, nz): A[:, j] = A[:, j - 1] - dz[j - 1] * 0.5 * (hy[:, j - 1] + hy[:, j]) psi = field.wrap_field(A, B.crds, name="psi", center=B.center, pretty_name=r"$\psi$", parents=[B]) if reduced_axes: slc = "..., " + ", ".join("{0}=None".format(ax) for ax in reduced_axes) psi = psi[slc] return psi
[docs]def calc_beta(pp, B, scale=1.0): """Calc plasma beta (2.0 * p / B^2) Parameters: pp (ScalarField or ndarray): pressure B (VectorField): magnetic field scale (float, optional): overall scale factor Returns: ScalarField: Plasma beta Note: For OpenGGCM, where pp is in pPa and B is in nT, scale should be 40.0. """ two = np.array([2.0], dtype=pp.dtype) bx, by, bz = B.component_views() if _HAS_NUMEXPR: ldict = dict(scale=scale, bx=bx, by=by, bz=bz, pp=pp, two=two) result = ne.evaluate("scale * two * pp / sqrt(bx**2 + by**2 + bz**2)", local_dict=ldict) else: result = scale * two * pp / np.sqrt(bx**2 + by**2 + bz**2) context = dict(name="beta", pretty_name=r"$\beta_{pl}$") return pp.wrap(result, context=context)
## ## EOF ##