"Base classes for MHD equilibria."
import logging
from abc import ABCMeta, abstractmethod
import cunumpy as xp
from matplotlib import pyplot as plt
from pyevtk.hl import gridToVTK
from struphy.geometry.base import Domain
from struphy.utils.utils import (
__class_with_params_repr_no_defaults__,
__dataclass_repr_no_defaults__,
all_class_params_are_default,
)
logger = logging.getLogger("struphy")
[docs]
class FluidEquilibrium(metaclass=ABCMeta):
"""
Abstract base class for callable fluid equilibria on arbitrary domains.
This class provides a unified interface for representing fluid equilibrium states,
including velocity, pressure, and number density fields. It supports coordinate
transformations between logical (reference) and physical domains through the Domain
object, enabling computations on mapped domains.
Attributes
----------
params : dict
Dictionary of parameters passed to the class initialization. Automatically
strips 'self' and '__class__' entries.
domain : Domain
Domain object that characterizes the mapping from the logical cube [0, 1]^3
to the physical domain. Enables coordinate transformations and differential
form conversions (0-forms, 1-forms, 2-forms, 3-forms).
Implementation Requirements
---------------------------
Child classes must provide at least one method from each of these pairs:
* Velocity: either ``u_xyz`` (Cartesian) or override ``uv`` (logical coordinates)
* Pressure: either ``p_xyz`` (Cartesian) or override ``p0`` (0-form on logical domain)
* Number Density: either ``n_xyz`` (Cartesian) or override ``n0`` (0-form on logical domain)
Derived Quantities
------------------
The class automatically provides derived fields computed from the basic fields:
* Temperature: ``t0``, ``t3`` (from p/n)
* Thermal velocity: ``vth0``, ``vth3`` (from temperature)
* Entropy density: ``s0_monoatomic``, ``s3_monoatomic``, ``s0_diatomic``, ``s3_diatomic``
Differential Forms
-------------------
Vector fields (velocity) are available as:
* ``uv``: contravariant components on logical domain
* ``u1``: 1-form components
* ``u2``: 2-form components
* ``u_cart``: Cartesian components with physical coordinates
Scalar fields are available as 0-forms (point values) and 3-forms (densities):
* ``p0``, ``p3``: pressure
* ``n0``, ``n3``: number density
* ``t0``, ``t3``: temperature
* ``q0``, ``q3``: square root of pressure
Notes
-----
The class uses abstract methods to enforce implementation in child classes.
Subclasses should override coordinate-appropriate base methods (CartesianFluidEquilibrium
or LogicalFluidEquilibrium) to simplify implementation.
"""
@property
def params(self) -> dict:
"""Parameters passed to __init__() of the class in equils.py, as dictionary."""
if not hasattr(self, "_params"):
self._params = {}
return self._params
@params.setter
def params(self, new):
assert isinstance(new, dict)
if "self" in new:
new.pop("self")
if "__class__" in new:
new.pop("__class__")
self._params = new
@property
def domain(self):
"""Domain object that characterizes the mapping from the logical to the physical domain."""
assert hasattr(self, "_domain"), (
"Domain for FluidEquilibrium not set; do obj.domain = ... to have access to all transformations."
)
return self._domain
@domain.setter
def domain(self, new_domain):
assert isinstance(new_domain, Domain) or new_domain is None
self._domain = new_domain
def __repr__(self) -> str:
out = f"{self.__class__.__name__}(\n"
for k, v in self.params.items():
out += " " * 4
out += f"{k}={v},\n"
out += ")"
return out
def __repr_no_defaults__(self):
return __class_with_params_repr_no_defaults__(self)
@property
def is_default(self):
return all_class_params_are_default(self)
###########################
# Vector-valued callables #
###########################
[docs]
def u1(self, *etas, squeeze_out=False):
"""1-form components of velocity on logical cube [0, 1]^3."""
return self.domain.transform(
self.uv(*etas, squeeze_out=False),
*etas,
kind="v_to_1",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
[docs]
def u2(self, *etas, squeeze_out=False):
"""2-form components of velocity on logical cube [0, 1]^3."""
return self.domain.transform(
self.uv(*etas, squeeze_out=False),
*etas,
kind="v_to_2",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
[docs]
def uv(self, *etas, squeeze_out=False):
"""Contra-variant components of velocity on logical cube [0, 1]^3."""
xyz = self.domain(*etas, squeeze_out=False)
return self.domain.pull(self.u_xyz(xyz[0], xyz[1], xyz[2]), *etas, kind="v", squeeze_out=squeeze_out)
[docs]
def u_cart(self, *etas, squeeze_out=False):
"""Cartesian components of velocity evaluated on logical cube [0, 1]^3. Returns also (x,y,z)."""
out = self.domain.push(
self.uv(*etas, squeeze_out=False),
*etas,
kind="v",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
return out, self.domain(*etas, squeeze_out=squeeze_out)
###########################
# Scalar-valued callables #
###########################
[docs]
def p0(self, *etas, squeeze_out=False):
"""0-form pressure on logical cube [0, 1]^3."""
xyz = self.domain(*etas, squeeze_out=False)
return self.domain.pull(self.p_xyz(xyz[0], xyz[1], xyz[2]), *etas, kind="0", squeeze_out=squeeze_out)
[docs]
def p3(self, *etas, squeeze_out=False):
"""3-form pressure on logical cube [0, 1]^3."""
return self.domain.transform(
self.p0(*etas, squeeze_out=False),
*etas,
kind="0_to_3",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
[docs]
def n0(self, *etas, squeeze_out=False):
"""0-form number density on logical cube [0, 1]^3."""
xyz = self.domain(*etas, squeeze_out=False)
return self.domain.pull(self.n_xyz(xyz[0], xyz[1], xyz[2]), *etas, kind="0", squeeze_out=squeeze_out)
[docs]
def n3(self, *etas, squeeze_out=False):
"""3-form number density on logical cube [0, 1]^3."""
return self.domain.transform(
self.n0(*etas, squeeze_out=False),
*etas,
kind="0_to_3",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
[docs]
def t0(self, *etas, squeeze_out=False):
"""0-form temperature on logical cube [0, 1]^3."""
return self.p0(*etas, squeeze_out=squeeze_out) / self.n0(*etas, squeeze_out=squeeze_out)
[docs]
def t3(self, *etas, squeeze_out=False):
"""3-form temperature on logical cube [0, 1]^3."""
return self.domain.transform(
self.t0(*etas, squeeze_out=False),
*etas,
kind="0_to_3",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
[docs]
def vth0(self, *etas, squeeze_out=False):
"""0-form thermal velocity on logical cube [0, 1]^3."""
return xp.sqrt(self.t0(*etas, squeeze_out=squeeze_out))
[docs]
def vth3(self, *etas, squeeze_out=False):
"""3-form thermal velocity on logical cube [0, 1]^3."""
return self.domain.transform(
self.vth0(*etas, squeeze_out=False),
*etas,
kind="0_to_3",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
[docs]
def q0(self, *etas, squeeze_out=False):
"""0-form square root of the pressure on logical cube [0, 1]^3."""
# xyz = self.domain(*etas, squeeze_out=False)
p = self.p0(*etas)
q = xp.sqrt(p)
return self.domain.pull(q, *etas, kind="0", squeeze_out=squeeze_out)
[docs]
def q3(self, *etas, squeeze_out=False):
"""3-form square root of the pressure on logical cube [0, 1]^3."""
return self.domain.transform(
self.q0(*etas, squeeze_out=False),
*etas,
kind="0_to_3",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
[docs]
def s0_monoatomic(self, *etas, squeeze_out=False):
"""0-form entropy density on logical cube [0, 1]^3.
Hard coded assumption : gamma = 5/3 (monoatomic perfect gaz)
"""
# xyz = self.domain(*etas, squeeze_out=False)
p = self.p0(*etas)
n = self.n0(*etas)
s = n * xp.log(p / (2 / 3 * xp.power(n, 5 / 3)))
return self.domain.pull(s, *etas, kind="0", squeeze_out=squeeze_out)
[docs]
def s3_monoatomic(self, *etas, squeeze_out=False):
"""3-form entropy density on logical cube [0, 1]^3.
Hard coded assumption : gamma = 5/3 (monoatomic perfect gaz)
"""
return self.domain.transform(
self.s0_monoatomic(*etas, squeeze_out=False),
*etas,
kind="0_to_3",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
[docs]
def s0_diatomic(self, *etas, squeeze_out=False):
"""0-form entropy density on logical cube [0, 1]^3.
Hard coded assumption : gamma = 7/5 (diatomic perfect gaz)
"""
# xyz = self.domain(*etas, squeeze_out=False)
p = self.p0(*etas)
n = self.n0(*etas)
s = n * xp.log(p / (2 / 5 * xp.power(n, 7 / 5)))
return self.domain.pull(s, *etas, kind="0", squeeze_out=squeeze_out)
[docs]
def s3_diatomic(self, *etas, squeeze_out=False):
"""3-form entropy density on logical cube [0, 1]^3.
Hard coded assumption : gamma = 5/3 (monoatomic perfect gaz)
"""
return self.domain.transform(
self.s0_diatomic(*etas, squeeze_out=False),
*etas,
kind="0_to_3",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
#########################################################
# Single components (for input of commuting projectors) #
#########################################################
[docs]
def u1_1(self, *etas, squeeze_out=False):
return self.u1(*etas, squeeze_out=squeeze_out)[0]
[docs]
def u1_2(self, *etas, squeeze_out=False):
return self.u1(*etas, squeeze_out=squeeze_out)[1]
[docs]
def u1_3(self, *etas, squeeze_out=False):
return self.u1(*etas, squeeze_out=squeeze_out)[2]
[docs]
def u2_1(self, *etas, squeeze_out=False):
return self.u2(*etas, squeeze_out=squeeze_out)[0]
[docs]
def u2_2(self, *etas, squeeze_out=False):
return self.u2(*etas, squeeze_out=squeeze_out)[1]
[docs]
def u2_3(self, *etas, squeeze_out=False):
return self.u2(*etas, squeeze_out=squeeze_out)[2]
[docs]
def uv_1(self, *etas, squeeze_out=False):
return self.uv(*etas, squeeze_out=squeeze_out)[0]
[docs]
def uv_2(self, *etas, squeeze_out=False):
return self.uv(*etas, squeeze_out=squeeze_out)[1]
[docs]
def uv_3(self, *etas, squeeze_out=False):
return self.uv(*etas, squeeze_out=squeeze_out)[2]
[docs]
def u_cart_1(self, *etas, squeeze_out=False):
return self.u_cart(*etas, squeeze_out=squeeze_out)[0][0]
[docs]
def u_cart_2(self, *etas, squeeze_out=False):
return self.u_cart(*etas, squeeze_out=squeeze_out)[0][1]
[docs]
def u_cart_3(self, *etas, squeeze_out=False):
return self.u_cart(*etas, squeeze_out=squeeze_out)[0][2]
[docs]
def to_dict(self) -> dict:
return {
"type": self.__class__.__name__,
"params": self.params,
}
[docs]
@classmethod
def from_dict(cls, dct: dict | None) -> "FluidEquilibrium":
if dct is None:
return None
equil = cls.get_equil_by_name(dct["type"])
return equil(**dct["params"])
[docs]
@classmethod
def get_equil_by_name(cls, equil_name: str) -> type["FluidEquilibrium"]:
from struphy.fields_background import equils
try:
equil_class: FluidEquilibrium = getattr(equils, equil_name)
if not issubclass(equil_class, FluidEquilibrium):
raise TypeError(f"{equil_name} is not a FluidEquilibrium subclass.")
else:
return equil_class
except AttributeError:
raise ModuleNotFoundError(f"{equil_name} not found in equils.")
def __eq__(self, other: "FluidEquilibrium") -> bool:
assert isinstance(other, FluidEquilibrium), f"Cannot compare FluidEquilibrium with {type(other)}."
return self.to_dict() == other.to_dict()
[docs]
class CartesianFluidEquilibrium(FluidEquilibrium):
r"""
Specialization for equilibria defined in Cartesian coordinates.
Child classes must implement the abstract methods ``u_xyz``, ``p_xyz``, and ``n_xyz``,
which return velocity, pressure, and number density in Cartesian physical space (x, y, z).
The base class automatically handles coordinate transformations and differential form conversions.
"""
[docs]
@abstractmethod
def u_xyz(self, x, y, z):
"""Cartesian velocity in physical space.
Must return the components as a tuple."""
pass
[docs]
@abstractmethod
def p_xyz(self, x, y, z):
"""Equilibrium pressure in physical space."""
pass
[docs]
@abstractmethod
def n_xyz(self, x, y, z):
"""Equilibrium number density in physical space."""
pass
@FluidEquilibrium.domain.setter
def domain(self, new_domain):
super(CartesianFluidEquilibrium, type(self)).domain.fset(self, new_domain)
[docs]
class LogicalFluidEquilibrium(FluidEquilibrium):
r"""
Specialization for equilibria defined on the logical cube [0, 1]^3.
Child classes must implement the abstract methods ``uv``, ``p0``, and ``n0``,
which return contravariant velocity, 0-form pressure, and 0-form number density
on the logical reference domain. Useful for direct implementation when physical
coordinates are obtained via coordinate transformation through the domain mapping.
"""
[docs]
@abstractmethod
def uv(self, *etas, squeeze_out=False):
"""Contra-variant (vector field) velocity on logical cube [0, 1]^3.
Must return the components as a tuple.
"""
pass
[docs]
@abstractmethod
def p0(self, *etas, squeeze_out=False):
"""0-form pressure on logical cube [0, 1]^3."""
pass
[docs]
@abstractmethod
def n0(self, *etas, squeeze_out=False):
"""0-form density on logical cube [0, 1]^3."""
pass
@FluidEquilibrium.domain.setter
def domain(self, new_domain):
super(LogicalFluidEquilibrium, type(self)).domain.fset(self, new_domain)
[docs]
class NumericalFluidEquilibrium(LogicalFluidEquilibrium):
r"""
Specialization for equilibria with numerically computed domain mappings.
Child classes must provide a ``numerical_domain`` property that returns a Domain object
representing the mapping from the logical cube [0, 1]^3 to the physical domain.
This class overrides the domain property to use the numerically computed mapping.
"""
@property
@abstractmethod
def numerical_domain(self):
"""Numerically computed mapping from the logical cube [0, 1]^3 to the physical domain
in the form of a :class:`~struphy.geometry.base.Domain` object."""
pass
@property
def domain(self):
return self.numerical_domain
[docs]
class FluidEquilibriumWithB(FluidEquilibrium):
"""
Extension of FluidEquilibrium with magnetic field and its gradient.
Child classes must implement either Cartesian (``b_xyz``, ``gradB_xyz``) or
logical (``bv``, ``gradB1``) methods for magnetic field and its gradient.
Provides methods for 1-form and 2-form transformations of the magnetic field.
"""
@FluidEquilibrium.domain.setter
def domain(self, new_domain):
super(FluidEquilibriumWithB, type(self)).domain.fset(self, new_domain)
###########################
# Vector-valued callables #
###########################
[docs]
def b1(self, *etas, squeeze_out=False):
"""1-form components of magnetic field on logical cube [0, 1]^3."""
return self.domain.transform(
self.bv(*etas, squeeze_out=False),
*etas,
kind="v_to_1",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
[docs]
def b2(self, *etas, squeeze_out=False):
"""2-form components of magnetic field on logical cube [0, 1]^3."""
return self.domain.transform(
self.bv(*etas, squeeze_out=False),
*etas,
kind="v_to_2",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
[docs]
def bv(self, *etas, squeeze_out=False):
"""Contra-variant components of magnetic field on logical cube [0, 1]^3."""
xyz = self.domain(*etas, squeeze_out=False)
return self.domain.pull(self.b_xyz(xyz[0], xyz[1], xyz[2]), *etas, kind="v", squeeze_out=squeeze_out)
[docs]
def b_cart(self, *etas, squeeze_out=False):
"""Cartesian components of magnetic field evaluated on logical cube [0, 1]^3. Returns also (x,y,z)."""
b_out = self.domain.push(
self.bv(*etas, squeeze_out=False),
*etas,
kind="v",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
return b_out, self.domain(*etas, squeeze_out=squeeze_out)
[docs]
def unit_b1(self, *etas, squeeze_out=False):
"""Unit vector components of magnetic field (1-form) on logical cube [0, 1]^3."""
return self.domain.pull(self.unit_b_cart(*etas, squeeze_out=False)[0], *etas, kind="1", squeeze_out=squeeze_out)
[docs]
def unit_b2(self, *etas, squeeze_out=False):
"""Unit vector components of magnetic field (2-form) on logical cube [0, 1]^3."""
return self.domain.pull(self.unit_b_cart(*etas, squeeze_out=False)[0], *etas, kind="2", squeeze_out=squeeze_out)
[docs]
def unit_bv(self, *etas, squeeze_out=False):
"""Unit vector components of magnetic field (contra-variant) on logical cube [0, 1]^3."""
return self.domain.pull(self.unit_b_cart(*etas, squeeze_out=False)[0], *etas, kind="v", squeeze_out=squeeze_out)
[docs]
def unit_b_cart(self, *etas, squeeze_out=False):
"""Unit vector Cartesian components of magnetic field evaluated on logical cube [0, 1]^3. Returns also (x,y,z)."""
b, xyz = self.b_cart(*etas, squeeze_out=squeeze_out)
absB = self.absB0(*etas, squeeze_out=squeeze_out)
out = xp.array([b[0] / absB, b[1] / absB, b[2] / absB], dtype=float)
return out, xyz
[docs]
def gradB1(self, *etas, squeeze_out=False):
"""1-form components of gradient of magnetic field strength evaluated on logical cube [0, 1]^3. Returns also (x,y,z)."""
xyz = self.domain(*etas, squeeze_out=False)
return self.domain.pull(self.gradB_xyz(xyz[0], xyz[1], xyz[2]), *etas, kind="1", squeeze_out=squeeze_out)
[docs]
def gradB2(self, *etas, squeeze_out=False):
"""2-form components of gradient of magnetic field strength evaluated on logical cube [0, 1]^3. Returns also (x,y,z)."""
return self.domain.transform(
self.gradB1(*etas, squeeze_out=False),
*etas,
kind="1_to_2",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
[docs]
def gradBv(self, *etas, squeeze_out=False):
"""Contra-variant components of gradient of magnetic field strength evaluated on logical cube [0, 1]^3. Returns also (x,y,z)."""
return self.domain.transform(
self.gradB1(*etas, squeeze_out=False),
*etas,
kind="1_to_v",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
[docs]
def gradB_cart(self, *etas, squeeze_out=False):
"""Cartesian components of gradient of magnetic field strength evaluated on logical cube [0, 1]^3. Returns also (x,y,z)."""
gradB_out = self.domain.push(
self.gradB1(*etas, squeeze_out=False),
*etas,
kind="1",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
return gradB_out, self.domain(*etas)
[docs]
def a1(self, *etas, squeeze_out=False):
"""1-form components of vector potential on logical cube [0, 1]^3."""
avail_list = ["HomogenSlab"]
assert self.__class__.__name__ in avail_list, (
f'Vector potential currently available only for {avail_list}, but mhd_equil is "{self.__class__.__name__}".'
)
return self.domain.transform(
self.a2(*etas, squeeze_out=False),
*etas,
kind="2_to_1",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
[docs]
def a2(self, *etas, squeeze_out=False):
"""2-form components of vector potential on logical cube [0, 1]^3."""
avail_list = ["HomogenSlab"]
assert self.__class__.__name__ in avail_list, (
f'Vector potential currently available only for {avail_list}, but mhd_equil is "{self.__class__.__name__}".'
)
xyz = self.domain(*etas, squeeze_out=False)
return self.domain.pull(self.a_xyz(xyz[0], xyz[1], xyz[2]), *etas, kind="2", squeeze_out=squeeze_out)
[docs]
def av(self, *etas, squeeze_out=False):
"""Contra-variant components of vector potneital on logical cube [0, 1]^3."""
avail_list = ["HomogenSlab"]
assert self.__class__.__name__ in avail_list, (
f'Vector potential currently available only for {avail_list}, but mhd_equil is "{self.__class__.__name__}".'
)
return self.domain.transform(
self.a2(*etas, squeeze_out=False),
*etas,
kind="2_to_v",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
###########################
# Scalar-valued callables #
###########################
[docs]
def absB0(self, *etas, squeeze_out=False):
"""0-form absolute value of magnetic field on logical cube [0, 1]^3."""
b, xyz = self.b_cart(*etas, squeeze_out=squeeze_out)
return xp.sqrt(b[0] ** 2 + b[1] ** 2 + b[2] ** 2)
[docs]
def absB3(self, *etas, squeeze_out=False):
"""3-form absolute value of magnetic field on logical cube [0, 1]^3."""
return self.domain.transform(
self.absB0(*etas, squeeze_out=False),
*etas,
kind="0_to_3",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
[docs]
def u_para0(self, *etas, squeeze_out=False):
"""0-form parallel velocity on logical cube [0, 1]^3."""
tmp_uv = self.uv(*etas, squeeze_out=squeeze_out)
tmp_unit_b1 = self.unit_b1(*etas, squeeze_out=squeeze_out)
return sum([ji * bi for ji, bi in zip(tmp_uv, tmp_unit_b1)])
[docs]
def u_para3(self, *etas, squeeze_out=False):
"""3-form parallel velocity on logical cube [0, 1]^3."""
return self.domain.transform(
self.u_para0(*etas, squeeze_out=False),
*etas,
kind="0_to_3",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
#########################################################
# Single components (for input of commuting projectors) #
#########################################################
[docs]
def b1_1(self, *etas, squeeze_out=False):
return self.b1(*etas, squeeze_out=squeeze_out)[0]
[docs]
def b1_2(self, *etas, squeeze_out=False):
return self.b1(*etas, squeeze_out=squeeze_out)[1]
[docs]
def b1_3(self, *etas, squeeze_out=False):
return self.b1(*etas, squeeze_out=squeeze_out)[2]
[docs]
def b2_1(self, *etas, squeeze_out=False):
return self.b2(*etas, squeeze_out=squeeze_out)[0]
[docs]
def b2_2(self, *etas, squeeze_out=False):
return self.b2(*etas, squeeze_out=squeeze_out)[1]
[docs]
def b2_3(self, *etas, squeeze_out=False):
return self.b2(*etas, squeeze_out=squeeze_out)[2]
[docs]
def bv_1(self, *etas, squeeze_out=False):
return self.bv(*etas, squeeze_out=squeeze_out)[0]
[docs]
def bv_2(self, *etas, squeeze_out=False):
return self.bv(*etas, squeeze_out=squeeze_out)[1]
[docs]
def bv_3(self, *etas, squeeze_out=False):
return self.bv(*etas, squeeze_out=squeeze_out)[2]
[docs]
def b_cart_1(self, *etas, squeeze_out=False):
return self.b_cart(*etas, squeeze_out=squeeze_out)[0][0]
[docs]
def b_cart_2(self, *etas, squeeze_out=False):
return self.b_cart(*etas, squeeze_out=squeeze_out)[0][1]
[docs]
def b_cart_3(self, *etas, squeeze_out=False):
return self.b_cart(*etas, squeeze_out=squeeze_out)[0][2]
[docs]
def unit_b1_1(self, *etas, squeeze_out=False):
return self.unit_b1(*etas, squeeze_out=squeeze_out)[0]
[docs]
def unit_b1_2(self, *etas, squeeze_out=False):
return self.unit_b1(*etas, squeeze_out=squeeze_out)[1]
[docs]
def unit_b1_3(self, *etas, squeeze_out=False):
return self.unit_b1(*etas, squeeze_out=squeeze_out)[2]
[docs]
def unit_b2_1(self, *etas, squeeze_out=False):
return self.unit_b2(*etas, squeeze_out=squeeze_out)[0]
[docs]
def unit_b2_2(self, *etas, squeeze_out=False):
return self.unit_b2(*etas, squeeze_out=squeeze_out)[1]
[docs]
def unit_b2_3(self, *etas, squeeze_out=False):
return self.unit_b2(*etas, squeeze_out=squeeze_out)[2]
[docs]
def unit_bv_1(self, *etas, squeeze_out=False):
return self.unit_bv(*etas, squeeze_out=squeeze_out)[0]
[docs]
def unit_bv_2(self, *etas, squeeze_out=False):
return self.unit_bv(*etas, squeeze_out=squeeze_out)[1]
[docs]
def unit_bv_3(self, *etas, squeeze_out=False):
return self.unit_bv(*etas, squeeze_out=squeeze_out)[2]
[docs]
def gradB1_1(self, *etas, squeeze_out=False):
return self.gradB1(*etas, squeeze_out=squeeze_out)[0]
[docs]
def gradB1_2(self, *etas, squeeze_out=False):
return self.gradB1(*etas, squeeze_out=squeeze_out)[1]
[docs]
def gradB1_3(self, *etas, squeeze_out=False):
return self.gradB1(*etas, squeeze_out=squeeze_out)[2]
[docs]
def gradB2_1(self, *etas, squeeze_out=False):
return self.gradB2(*etas, squeeze_out=squeeze_out)[0]
[docs]
def gradB2_2(self, *etas, squeeze_out=False):
return self.gradB2(*etas, squeeze_out=squeeze_out)[1]
[docs]
def gradB2_3(self, *etas, squeeze_out=False):
return self.gradB2(*etas, squeeze_out=squeeze_out)[2]
[docs]
def gradBv_1(self, *etas, squeeze_out=False):
return self.gradBv(*etas, squeeze_out=squeeze_out)[0]
[docs]
def gradBv_2(self, *etas, squeeze_out=False):
return self.gradBv(*etas, squeeze_out=squeeze_out)[1]
[docs]
def gradBv_3(self, *etas, squeeze_out=False):
return self.gradBv(*etas, squeeze_out=squeeze_out)[2]
[docs]
def a1_1(self, *etas, squeeze_out=False):
return self.a1(*etas, squeeze_out=squeeze_out)[0]
[docs]
def a1_2(self, *etas, squeeze_out=False):
return self.a1(*etas, squeeze_out=squeeze_out)[1]
[docs]
def a1_3(self, *etas, squeeze_out=False):
return self.a1(*etas, squeeze_out=squeeze_out)[2]
[docs]
def a2_1(self, *etas, squeeze_out=False):
return self.a2(*etas, squeeze_out=squeeze_out)[0]
[docs]
def a2_2(self, *etas, squeeze_out=False):
return self.a2(*etas, squeeze_out=squeeze_out)[1]
[docs]
def a2_3(self, *etas, squeeze_out=False):
return self.a2(*etas, squeeze_out=squeeze_out)[2]
[docs]
def av_1(self, *etas, squeeze_out=False):
return self.av(*etas, squeeze_out=squeeze_out)[0]
[docs]
def av_2(self, *etas, squeeze_out=False):
return self.av(*etas, squeeze_out=squeeze_out)[1]
[docs]
def av_3(self, *etas, squeeze_out=False):
return self.av(*etas, squeeze_out=squeeze_out)[2]
###########
# Methods #
###########
[docs]
def parallel_component(self, *etas, squeeze_out=False):
raise NotImplementedError()
[docs]
def perpendicular_component(self, *etas, squeeze_out=False):
raise NotImplementedError()
[docs]
class CartesianFluidEquilibriumWithB(CartesianFluidEquilibrium, FluidEquilibriumWithB):
r"""
Specialization for fluid equilibria with magnetic field in Cartesian coordinates.
Child classes must implement the abstract methods ``b_xyz`` and ``gradB_xyz``,
which return magnetic field and its gradient strength in Cartesian physical space.
"""
[docs]
@abstractmethod
def b_xyz(self, x, y, z):
"""Cartesian magnetic field in physical space.
Must return the components as a tuple."""
pass
[docs]
@abstractmethod
def gradB_xyz(self, x, y, z):
"""Cartesian gradient of magnetic field strength in physical space. Must return the components as a tuple."""
pass
@CartesianFluidEquilibrium.domain.setter
def domain(self, new_domain):
super(CartesianFluidEquilibriumWithB, type(self)).domain.fset(self, new_domain)
[docs]
class LogicalFluidEquilibriumWithB(LogicalFluidEquilibrium, FluidEquilibriumWithB):
r"""
Specialization for fluid equilibria with magnetic field on the logical cube [0, 1]^3.
Child classes must implement the abstract methods ``bv`` (contravariant magnetic field)
and ``gradB1`` (1-form gradient of magnetic field strength) on the logical domain.
"""
[docs]
@abstractmethod
def bv(self, *etas, squeeze_out=False):
"""Contra-variant (vector field) magnetic field on logical cube [0, 1]^3.
Must return the components as a tuple.
"""
pass
[docs]
@abstractmethod
def gradB1(self, *etas, squeeze_out=False):
"""Co-variant (1-from) gradient of magnetic field strength on logical cube [0, 1]^3.
Must return the components as a tuple.
"""
pass
@LogicalFluidEquilibrium.domain.setter
def domain(self, new_domain):
super(LogicalFluidEquilibriumWithB, type(self)).domain.fset(self, new_domain)
[docs]
class NumericalFluidEquilibriumWithB(LogicalFluidEquilibriumWithB):
r"""
Specialization for fluid equilibria with magnetic field and numerically computed domain mappings.
Child classes must provide a ``numerical_domain`` property that returns a Domain object.
This class overrides the domain property to use the numerically computed mapping.
"""
@property
@abstractmethod
def numerical_domain(self):
"""Numerically computed mapping from the logical cube [0, 1]^3 to the physical domain
in the form of a :class:`~struphy.geometry.base.Domain` object."""
pass
@property
def domain(self):
return self.numerical_domain
[docs]
class MHDequilibrium(FluidEquilibriumWithB):
"""
Extension of FluidEquilibriumWithB with current density field.
Child classes must implement either Cartesian (``j_xyz``) or logical (``jv``) methods
for current density. The velocity field is derived from current density as j/n,
overriding the base FluidEquilibrium. Provides methods for 1-form and 2-form
transformations of the current density.
"""
@FluidEquilibriumWithB.domain.setter
def domain(self, new_domain):
super(MHDequilibrium, type(self)).domain.fset(self, new_domain)
###########################
# Vector-valued callables #
###########################
[docs]
def j1(self, *etas, squeeze_out=False):
"""1-form components of current on logical cube [0, 1]^3."""
return self.domain.transform(
self.jv(*etas, squeeze_out=False),
*etas,
kind="v_to_1",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
[docs]
def j2(self, *etas, squeeze_out=False):
"""2-form components of current on logical cube [0, 1]^3."""
return self.domain.transform(
self.jv(*etas, squeeze_out=False),
*etas,
kind="v_to_2",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
[docs]
def jv(self, *etas, squeeze_out=False):
"""Contra-variant components of current on logical cube [0, 1]^3."""
xyz = self.domain(*etas, squeeze_out=False)
return self.domain.pull(self.j_xyz(xyz[0], xyz[1], xyz[2]), *etas, kind="v", squeeze_out=squeeze_out)
[docs]
def j_cart(self, *etas, squeeze_out=False):
"""Cartesian components of current evaluated on logical cube [0, 1]^3. Returns also (x,y,z)."""
j_out = self.domain.push(
self.jv(*etas, squeeze_out=False),
*etas,
kind="v",
a_kwargs={"squeeze_out": False},
squeeze_out=squeeze_out,
)
return j_out, self.domain(*etas, squeeze_out=squeeze_out)
[docs]
def u1(self, *etas, squeeze_out=False):
"""1-form components of mean velocity evaluated on logical cube [0, 1]^3. Returns also (x,y,z)."""
return self.j1(
*etas,
squeeze_out=squeeze_out,
) / self.n0(
*etas,
squeeze_out=squeeze_out,
)
[docs]
def u2(self, *etas, squeeze_out=False):
"""2-form components of mean velocity evaluated on logical cube [0, 1]^3. Returns also (x,y,z)."""
return self.j2(
*etas,
squeeze_out=squeeze_out,
) / self.n0(
*etas,
squeeze_out=squeeze_out,
)
[docs]
def uv(self, *etas, squeeze_out=False):
"""Contra-variant components of mean velocity evaluated on logical cube [0, 1]^3. Returns also (x,y,z)."""
return self.jv(
*etas,
squeeze_out=squeeze_out,
) / self.n0(
*etas,
squeeze_out=squeeze_out,
)
[docs]
def u_cart(self, *etas, squeeze_out=False):
"""Cartesian components of mean velocity evaluated on logical cube [0, 1]^3. Returns also (x,y,z)."""
u_out = self.j_cart(
*etas,
squeeze_out=squeeze_out,
)[0] / self.n0(
*etas,
squeeze_out=squeeze_out,
)
return u_out, self.domain(*etas, squeeze_out=squeeze_out)
[docs]
def curl_unit_b1(self, *etas, squeeze_out=False):
"""1-form components of curl of unit magnetic field evaluated on logical cube [0, 1]^3. Returns also (x,y,z)."""
return self.domain.pull(
self.curl_unit_b_cart(*etas, squeeze_out=False)[0],
*etas,
kind="1",
squeeze_out=squeeze_out,
)
[docs]
def curl_unit_b2(self, *etas, squeeze_out=False):
"""2-form components of curl of unit magnetic field evaluated on logical cube [0, 1]^3. Returns also (x,y,z)."""
return self.domain.pull(
self.curl_unit_b_cart(*etas, squeeze_out=False)[0],
*etas,
kind="2",
squeeze_out=squeeze_out,
)
[docs]
def curl_unit_bv(self, *etas, squeeze_out=False):
"""Contra-variant components of curl of unit magnetic field evaluated on logical cube [0, 1]^3. Returns also (x,y,z)."""
return self.domain.pull(
self.curl_unit_b_cart(*etas, squeeze_out=False)[0],
*etas,
kind="v",
squeeze_out=squeeze_out,
)
[docs]
def curl_unit_b_cart(self, *etas, squeeze_out=False):
"""Cartesian components of curl of unit magnetic field evaluated on logical cube [0, 1]^3. Returns also (x,y,z)."""
b, xyz = self.b_cart(*etas, squeeze_out=squeeze_out)
j, xyz = self.j_cart(*etas, squeeze_out=squeeze_out)
gradB, xyz = self.gradB_cart(*etas, squeeze_out=squeeze_out)
absB = self.absB0(*etas, squeeze_out=squeeze_out)
out = xp.array(
[
j[0] / absB + (b[1] * gradB[2] - b[2] * gradB[1]) / absB**2,
j[1] / absB + (b[2] * gradB[0] - b[0] * gradB[2]) / absB**2,
j[2] / absB + (b[0] * gradB[1] - b[1] * gradB[0]) / absB**2,
],
dtype=float,
)
return out, xyz
###########################
# Scalar-valued callables #
###########################
[docs]
def absJ0(self, *etas, squeeze_out=False):
"""0-form absolute value of current on logical cube [0, 1]^3."""
j, xyz = self.j_cart(*etas, squeeze_out=squeeze_out)
return xp.sqrt(j[0] ** 2 + j[1] ** 2 + j[2] ** 2)
[docs]
def curl_unit_b_dot_b0(self, *etas, squeeze_out=False):
r"""0-form of :math:`(\nabla \times \mathbf b_0) \times \mathbf b_0` evaluated on logical cube [0, 1]^3."""
curl_b, xyz = self.curl_unit_b_cart(*etas, squeeze_out=squeeze_out)
b, xyz = self.unit_b_cart(*etas, squeeze_out=squeeze_out)
out = curl_b[0] * b[0] + curl_b[1] * b[1] + curl_b[2] * b[2]
return out
###################
# Single components
###################
[docs]
def j1_1(self, *etas, squeeze_out=False):
return self.j1(*etas, squeeze_out=squeeze_out)[0]
[docs]
def j1_2(self, *etas, squeeze_out=False):
return self.j1(*etas, squeeze_out=squeeze_out)[1]
[docs]
def j1_3(self, *etas, squeeze_out=False):
return self.j1(*etas, squeeze_out=squeeze_out)[2]
[docs]
def j2_1(self, *etas, squeeze_out=False):
return self.j2(*etas, squeeze_out=squeeze_out)[0]
[docs]
def j2_2(self, *etas, squeeze_out=False):
return self.j2(*etas, squeeze_out=squeeze_out)[1]
[docs]
def j2_3(self, *etas, squeeze_out=False):
return self.j2(*etas, squeeze_out=squeeze_out)[2]
[docs]
def jv_1(self, *etas, squeeze_out=False):
return self.jv(*etas, squeeze_out=squeeze_out)[0]
[docs]
def jv_2(self, *etas, squeeze_out=False):
return self.jv(*etas, squeeze_out=squeeze_out)[1]
[docs]
def jv_3(self, *etas, squeeze_out=False):
return self.jv(*etas, squeeze_out=squeeze_out)[2]
[docs]
def j_cart_1(self, *etas, squeeze_out=False):
return self.j_cart(*etas, squeeze_out=squeeze_out)[0][0]
[docs]
def j_cart_2(self, *etas, squeeze_out=False):
return self.j_cart(*etas, squeeze_out=squeeze_out)[0][1]
[docs]
def j_cart_3(self, *etas, squeeze_out=False):
return self.j_cart(*etas, squeeze_out=squeeze_out)[0][2]
[docs]
def curl_unit_b1_1(self, *etas, squeeze_out=False):
return self.curl_unit_b1(*etas, squeeze_out=squeeze_out)[0]
[docs]
def curl_unit_b1_2(self, *etas, squeeze_out=False):
return self.curl_unit_b1(*etas, squeeze_out=squeeze_out)[1]
[docs]
def curl_unit_b1_3(self, *etas, squeeze_out=False):
return self.curl_unit_b1(*etas, squeeze_out=squeeze_out)[2]
[docs]
def curl_unit_b2_1(self, *etas, squeeze_out=False):
return self.curl_unit_b2(*etas, squeeze_out=squeeze_out)[0]
[docs]
def curl_unit_b2_2(self, *etas, squeeze_out=False):
return self.curl_unit_b2(*etas, squeeze_out=squeeze_out)[1]
[docs]
def curl_unit_b2_3(self, *etas, squeeze_out=False):
return self.curl_unit_b2(*etas, squeeze_out=squeeze_out)[2]
[docs]
def curl_unit_bv_1(self, *etas, squeeze_out=False):
return self.curl_unit_bv(*etas, squeeze_out=squeeze_out)[0]
[docs]
def curl_unit_bv_2(self, *etas, squeeze_out=False):
return self.curl_unit_bv(*etas, squeeze_out=squeeze_out)[1]
[docs]
def curl_unit_bv_3(self, *etas, squeeze_out=False):
return self.curl_unit_bv(*etas, squeeze_out=squeeze_out)[2]
##########
# Plotting
##########
[docs]
def show(self, n1=16, n2=33, n3=21, n_planes=5):
"""Generate vtk files of equilibirum and do some 2d plots with matplotlib.
Parameters
----------
n1, n2, n3 : int
Evaluation points of mapping in each direcion.
n_planes : int
Number of planes to show perpendicular to eta3."""
import struphy
torus_mappings = (
"Tokamak",
"GVECunit",
"DESCunit",
"IGAPolarTorus",
"HollowTorus",
)
e1 = xp.linspace(0.0001, 1, n1)
e2 = xp.linspace(0, 1, n2)
e3 = xp.linspace(0, 1, n3)
if self.domain.__class__.__name__ in ("GVECunit", "DESCunit"):
if n_planes > 1:
jump = (n3 - 1) / (n_planes - 1)
else:
jump = 0
else:
n_planes = 1
jump = 0
x, y, z = self.domain(e1, e2, e3)
logger.info("Evaluation of mapping done.")
det_df = self.domain.jacobian_det(e1, e2, e3)
p = self.p0(e1, e2, e3)
logger.info("Computation of pressure done.")
# ori 240624
n_dens = self.n0(e1, e2, e3)
logger.info("Computation of density done.")
absB = self.absB0(e1, e2, e3)
logger.info("Computation of abs(B) done.")
j_cart, xyz = self.j_cart(e1, e2, e3)
logger.info("Computation of current density done.")
absJ = xp.sqrt(j_cart[0] ** 2 + j_cart[1] ** 2 + j_cart[2] ** 2)
_path = struphy.__path__[0] + "/fields_background/mhd_equil/gvec/output/"
gridToVTK(
_path + "vtk/gvec_equil",
x,
y,
z,
pointData={"det_df": det_df, "pressure": p, "absB": absB},
)
logger.info("Generation of vtk files done.")
# show params
logger.info("\nEquilibrium parameters:")
for key, val in self.params.items():
logger.info(f"{key}: {val}")
logger.info("\nMapping parameters:")
for key, val in self.domain.params.items():
if key not in {"cx", "cy", "cz"}:
logger.info(f"{key}: {val}")
# poloidal plane grid
fig = plt.figure(figsize=(13, xp.ceil(n_planes / 2) * 6.5))
for n in range(n_planes):
xpp = x[:, :, int(n * jump)].squeeze()
yp = y[:, :, int(n * jump)].squeeze()
zp = z[:, :, int(n * jump)].squeeze()
if self.domain.__class__.__name__ in torus_mappings:
pc1 = xp.sqrt(xpp**2 + yp**2)
pc2 = zp
l1 = "R"
l2 = "Z"
else:
pc1 = xpp
pc2 = yp
l1 = "x"
l2 = "y"
ax = fig.add_subplot(int(xp.ceil(n_planes / 2)), 2, n + 1)
for i in range(pc1.shape[0]):
for j in range(pc1.shape[1] - 1):
if i < pc1.shape[0] - 1:
ax.plot(
[pc1[i, j], pc1[i + 1, j]],
[pc2[i, j], pc2[i + 1, j]],
"b",
linewidth=0.6,
)
if j < pc1.shape[1] - 1:
ax.plot(
[pc1[i, j], pc1[i, j + 1]],
[pc2[i, j], pc2[i, j + 1]],
"b",
linewidth=0.6,
)
ax.scatter(pc1[0, 0], pc2[0, 0], 20, "red", zorder=10)
# ax.scatter(pc1[0, 32], pc2[0, 32], 20, 'red', zorder=10)
ax.set_xlabel(l1)
ax.set_ylabel(l2)
ax.axis("equal")
ax.set_title(
r"Poloidal plane at $\eta_3$={0:4.3f}".format(e3[int(n * jump)]),
)
# top view
e1 = xp.linspace(0, 1, n1) # radial coordinate in [0, 1]
e2 = xp.linspace(0, 1, 3) # poloidal angle in [0, 1]
e3 = xp.linspace(0, 1, n3) # toroidal angle in [0, 1]
xt, yt, zt = self.domain(e1, e2, e3)
fig = plt.figure(figsize=(13, 2 * 6.5))
ax = fig.add_subplot()
for m in range(2):
xpp = xt[:, m, :].squeeze()
yp = yt[:, m, :].squeeze()
zp = zt[:, m, :].squeeze()
if self.domain.__class__.__name__ in torus_mappings:
tc1 = xpp
tc2 = yp
l1 = "x"
l2 = "y"
else:
tc1 = xpp
tc2 = zp
l1 = "x"
l2 = "z"
for i in range(tc1.shape[0]):
for j in range(tc1.shape[1] - 1):
if i < tc1.shape[0] - 1:
ax.plot(
[tc1[i, j], tc1[i + 1, j]],
[tc2[i, j], tc2[i + 1, j]],
"b",
linewidth=0.6,
)
if j < tc1.shape[1] - 1:
if i == 0:
ax.plot(
[tc1[i, j], tc1[i, j + 1]],
[tc2[i, j], tc2[i, j + 1]],
"r",
linewidth=1,
)
else:
ax.plot(
[tc1[i, j], tc1[i, j + 1]],
[tc2[i, j], tc2[i, j + 1]],
"b",
linewidth=0.6,
)
ax.set_xlabel(l1)
ax.set_ylabel(l2)
ax.axis("equal")
ax.set_title("Device top view")
# Jacobian determinant
fig = plt.figure(figsize=(13, xp.ceil(n_planes / 2) * 6.5))
for n in range(n_planes):
xpp = x[:, :, int(n * jump)].squeeze()
yp = y[:, :, int(n * jump)].squeeze()
zp = z[:, :, int(n * jump)].squeeze()
if self.domain.__class__.__name__ in torus_mappings:
pc1 = xp.sqrt(xpp**2 + yp**2)
pc2 = zp
l1 = "R"
l2 = "Z"
else:
pc1 = xpp
pc2 = yp
l1 = "x"
l2 = "y"
detp = det_df[:, :, int(n * jump)].squeeze()
ax = fig.add_subplot(int(xp.ceil(n_planes / 2)), 2, n + 1)
map = ax.contourf(pc1, pc2, detp, 30)
ax.set_xlabel(l1)
ax.set_ylabel(l2)
ax.axis("equal")
ax.set_title(
r"Jacobian determinant at $\eta_3$={0:4.3f}".format(e3[int(n * jump)]),
)
fig.colorbar(map, ax=ax, location="right")
# pressure
fig = plt.figure(figsize=(15, xp.ceil(n_planes / 2) * 6.5))
for n in range(n_planes):
xpp = x[:, :, int(n * jump)].squeeze()
yp = y[:, :, int(n * jump)].squeeze()
zp = z[:, :, int(n * jump)].squeeze()
if self.domain.__class__.__name__ in torus_mappings:
pc1 = xp.sqrt(xpp**2 + yp**2)
pc2 = zp
l1 = "R"
l2 = "Z"
else:
pc1 = xpp
pc2 = yp
l1 = "x"
l2 = "y"
pp = p[:, :, int(n * jump)].squeeze()
ax = fig.add_subplot(int(xp.ceil(n_planes / 2)), 2, n + 1)
map = ax.contourf(pc1, pc2, pp, 30)
ax.set_xlabel(l1)
ax.set_ylabel(l2)
ax.axis("equal")
ax.set_title(
r"Pressure at $\eta_3$={0:4.3f}".format(e3[int(n * jump)]),
)
fig.colorbar(map, ax=ax, location="right")
# density
fig = plt.figure(figsize=(15, xp.ceil(n_planes / 2) * 6.5))
for n in range(n_planes):
xpp = x[:, :, int(n * jump)].squeeze()
yp = y[:, :, int(n * jump)].squeeze()
zp = z[:, :, int(n * jump)].squeeze()
if self.domain.__class__.__name__ in torus_mappings:
pc1 = xp.sqrt(xpp**2 + yp**2)
pc2 = zp
l1 = "R"
l2 = "Z"
else:
pc1 = xpp
pc2 = yp
l1 = "x"
l2 = "y"
nn = n_dens[:, :, int(n * jump)].squeeze()
ax = fig.add_subplot(int(xp.ceil(n_planes / 2)), 2, n + 1)
map = ax.contourf(pc1, pc2, nn, 30)
ax.set_xlabel(l1)
ax.set_ylabel(l2)
ax.axis("equal")
ax.set_title(
r"Equilibrium density at $\eta_3$={0:4.3f}".format(e3[int(n * jump)]),
)
fig.colorbar(map, ax=ax, location="right")
# magnetic field strength
fig = plt.figure(figsize=(15, xp.ceil(n_planes / 2) * 6.5))
for n in range(n_planes):
xpp = x[:, :, int(n * jump)].squeeze()
yp = y[:, :, int(n * jump)].squeeze()
zp = z[:, :, int(n * jump)].squeeze()
if self.domain.__class__.__name__ in torus_mappings:
pc1 = xp.sqrt(xpp**2 + yp**2)
pc2 = zp
l1 = "R"
l2 = "Z"
else:
pc1 = xpp
pc2 = yp
l1 = "x"
l2 = "y"
ab = absB[:, :, int(n * jump)].squeeze()
ax = fig.add_subplot(int(xp.ceil(n_planes / 2)), 2, n + 1)
map = ax.contourf(pc1, pc2, ab, 30)
ax.set_xlabel(l1)
ax.set_ylabel(l2)
ax.axis("equal")
ax.set_title(
r"Magnetic field strength at $\eta_3$={0:4.3f}".format(e3[int(n * jump)]),
)
fig.colorbar(map, ax=ax, location="right")
# current density
fig = plt.figure(figsize=(15, xp.ceil(n_planes / 2) * 6.5))
for n in range(n_planes):
xpp = x[:, :, int(n * jump)].squeeze()
yp = y[:, :, int(n * jump)].squeeze()
zp = z[:, :, int(n * jump)].squeeze()
if self.domain.__class__.__name__ in torus_mappings:
pc1 = xp.sqrt(xpp**2 + yp**2)
pc2 = zp
l1 = "R"
l2 = "Z"
else:
pc1 = xpp
pc2 = yp
l1 = "x"
l2 = "y"
ab = absJ[:, :, int(n * jump)].squeeze()
ax = fig.add_subplot(int(xp.ceil(n_planes / 2)), 2, n + 1)
map = ax.contourf(pc1, pc2, ab, 30)
ax.set_xlabel(l1)
ax.set_ylabel(l2)
ax.axis("equal")
ax.set_title(
r"Current density (abs) at $\eta_3$={0:4.3f}".format(e3[int(n * jump)]),
)
fig.colorbar(map, ax=ax, location="right")
plt.show()
[docs]
class CartesianMHDequilibrium(MHDequilibrium):
r"""
Specialization for MHD equilibria in Cartesian coordinates.
Child classes must implement the abstract methods ``b_xyz``, ``j_xyz``, ``p_xyz``,
``n_xyz``, and ``gradB_xyz`` in Cartesian physical space.
"""
[docs]
@abstractmethod
def b_xyz(self, x, y, z):
"""Cartesian magnetic field in physical space.
Must return the components as a tuple."""
pass
[docs]
@abstractmethod
def j_xyz(self, x, y, z):
"""Cartesian current (curl of magnetic field) in physical space.
Must return the components as a tuple."""
pass
[docs]
@abstractmethod
def p_xyz(self, x, y, z):
"""Equilibrium pressure in physical space."""
pass
[docs]
@abstractmethod
def n_xyz(self, x, y, z):
"""Equilibrium number density in physical space."""
pass
[docs]
@abstractmethod
def gradB_xyz(self, x, y, z):
"""Cartesian gradient of magnetic field strength in physical space.
Must return the components as a tuple."""
pass
@MHDequilibrium.domain.setter
def domain(self, new_domain):
super(CartesianMHDequilibrium, type(self)).domain.fset(self, new_domain)
[docs]
class AxisymmMHDequilibrium(CartesianMHDequilibrium):
r"""
Base class for ideal axisymmetric MHD equilibria based on a poloidal flux function
:math:`\psi(R, Z)` and a toroidal field function :math:`g_{tor}(R, Z)`
in a cylindrical coordinate system :math:`(R, \phi, Z)`.
The magnetic field and current density are then given by
.. math::
\mathbf B = \nabla \psi \times \nabla \phi + g_{tor} \nabla \phi\,,\qquad \mathbf j = \nabla \times \mathbf B\,.
The pressure and density profiles need to be implemented by child classes.
"""
[docs]
@abstractmethod
def psi(self, R, Z, dR=0, dZ=0):
"""Poloidal flux function per radian. First AND second derivatives dR=0,1,2 and dZ=0,1,2 must be implemented."""
pass
[docs]
@abstractmethod
def g_tor(self, R, Z, dR=0, dZ=0):
"""Toroidal field function. First derivatives dR=0,1 and dZ=0,1 must be implemented."""
pass
@property
@abstractmethod
def psi_range(self):
"""Psi on-axis and at plasma boundary returned as list [psi_axis, psi_boundary]."""
pass
@property
@abstractmethod
def psi_axis_RZ(self):
"""Location of magnetic axis in R-Z-coordinates returned as list [psi_axis_R, psi_axis_Z]."""
pass
[docs]
@abstractmethod
def p_xyz(self, x, y, z):
"""Equilibrium pressure in physical space."""
pass
[docs]
@abstractmethod
def n_xyz(self, x, y, z):
"""Equilibrium number density in physical space."""
pass
[docs]
def b_xyz(self, x, y, z):
"""Cartesian B-field components calculated as BR = -(dpsi/dZ)/R, BPhi = g_tor/R, BZ = (dpsi/dR)/R."""
R, Phi, Z = self.inverse_map(x, y, z)
# at phi = 0°
BR = -self.psi(R, Z, dZ=1) / R
BP = self.g_tor(R, Z) / R
BZ = self.psi(R, Z, dR=1) / R
# push-forward to Cartesian components
Bx = BR * xp.cos(Phi) - BP * xp.sin(Phi)
By = BR * xp.sin(Phi) + BP * xp.cos(Phi)
Bz = 1 * BZ
return Bx, By, Bz
[docs]
def j_xyz(self, x, y, z):
"""Cartesian current density components calculated as curl(B)."""
R, Phi, Z = self.inverse_map(x, y, z)
# at phi = 0° (j = curl(B))
jR = -self.g_tor(R, Z, dZ=1) / R
jP = -self.psi(R, Z, dZ=2) / R + self.psi(R, Z, dR=1) / R**2 - self.psi(R, Z, dR=2) / R
jZ = self.g_tor(R, Z, dR=1) / R
# push-forward to Cartesian components
jx = jR * xp.cos(Phi) - jP * xp.sin(Phi)
jy = jR * xp.sin(Phi) + jP * xp.cos(Phi)
jz = 1 * jZ
return jx, jy, jz
[docs]
def gradB_xyz(self, x, y, z):
"""Cartesian gradient |B| components calculated as grad(sqrt(BR**2 + BPhi**2 + BZ**2))."""
R, Phi, Z = self.inverse_map(x, y, z)
RabsB = xp.sqrt(
self.psi(R, Z, dZ=1) ** 2 + self.g_tor(R, Z) ** 2 + self.psi(R, Z, dR=1) ** 2,
)
# at phi = 0° (gradB = grad(absB))
gradBR = (
-RabsB / R**2
+ (
self.psi(R, Z, dZ=1)
* self.psi(
R,
Z,
dR=1,
dZ=1,
)
+ self.psi(R, Z, dR=1) * self.psi(R, Z, dR=2)
)
/ RabsB
/ R
)
gradBP = 0.0
gradBZ = (
(self.psi(R, Z, dZ=1) * self.psi(R, Z, dZ=2) + self.psi(R, Z, dR=1) * self.psi(R, Z, dR=1, dZ=1))
/ RabsB
/ R
)
# push-forward to Cartesian components
gradBx = gradBR * xp.cos(Phi) - gradBP * xp.sin(Phi)
gradBy = gradBR * xp.sin(Phi) + gradBP * xp.cos(Phi)
gradBz = 1 * gradBZ
return gradBx, gradBy, gradBz
[docs]
@staticmethod
def inverse_map(x, y, z):
"""Inverse cylindrical mapping."""
R = xp.sqrt(x**2 + y**2)
P = xp.arctan2(y, x)
Z = 1 * z
return R, P, Z
@CartesianMHDequilibrium.domain.setter
def domain(self, new_domain):
super(AxisymmMHDequilibrium, type(self)).domain.fset(self, new_domain)
[docs]
class LogicalMHDequilibrium(MHDequilibrium):
r"""
Specialization for MHD equilibria on the logical cube [0, 1]^3.
Child classes must implement the abstract methods ``bv``, ``jv``, ``p0``, ``n0``,
and ``gradB1`` on the logical reference domain.
"""
[docs]
@abstractmethod
def bv(self, *etas, squeeze_out=False):
"""Contra-variant (vector field) magnetic field on logical cube [0, 1]^3.
Must return the components as a tuple.
"""
pass
[docs]
@abstractmethod
def jv(self, *etas, squeeze_out=False):
"""Contra-variant (vector field) current density (=curl B) on logical cube [0, 1]^3.
Must return the components as a tuple.
"""
pass
[docs]
@abstractmethod
def p0(self, *etas, squeeze_out=False):
"""0-form pressure on logical cube [0, 1]^3.
Must return the components as a tuple.
"""
pass
[docs]
@abstractmethod
def n0(self, *etas, squeeze_out=False):
"""0-form density on logical cube [0, 1]^3."""
pass
[docs]
@abstractmethod
def gradB1(self, *etas, squeeze_out=False):
"""1-form gradient of magnetic field strength strength on logical cube [0, 1]^3.
Must return the components as a tuple.
"""
pass
@MHDequilibrium.domain.setter
def domain(self, new_domain):
super(LogicalMHDequilibrium, type(self)).domain.fset(self, new_domain)
[docs]
class NumericalMHDequilibrium(LogicalMHDequilibrium):
r"""
Specialization for MHD equilibria with numerically computed domain mappings.
Child classes must provide a ``numerical_domain`` property that returns a Domain object.
This class overrides the domain property to use the numerically computed mapping.
"""
@property
@abstractmethod
def numerical_domain(self):
"""Numerically computed mapping from the logical cube [0, 1]^3 to the physical domain
in the form of a :class:`~struphy.geometry.base.Domain` object."""
pass
@property
def domain(self) -> Domain:
return self.numerical_domain
