Source code for struphy.models.linear_extended_mh_duniform

import copy

from feectools.ddm.mpi import mpi as MPI
from feectools.linalg.block import BlockVector

from struphy.io.options import BaseUnits, LiteralOptions
from struphy.models.base import StruphyModel
from struphy.models.scalars import BilinearEnergyFEEC, FunctionScalarFEEC, Scalars, VolumeFormEnergyFEEC
from struphy.models.species import (
    FieldSpecies,
    FluidSpecies,
)
from struphy.models.variables import FEECVariable
from struphy.polar.basic import PolarVector
from struphy.propagators.base import Propagator
from struphy.propagators.hall import Hall
from struphy.propagators.magnetosonic_uniform import MagnetosonicUniform
from struphy.propagators.shear_alfven_b1 import ShearAlfvenB1

rank = MPI.COMM_WORLD.Get_rank()




class LinearExtendedMHDuniform(StruphyModel):
    """Linear extended MHD with zero-flow equilibrium, for uniform background conditions only.

    Parameters
    ----------
    base_units: BaseUnits
        Base units for normalization (default: BaseUnits())
    charge_number: int
        Charge number (in units of the positive elementary charge) of the ion species (default: 1)
    mass_number: float
        Mass number (in units of Proton mass) of the ion species (default: 1.0)
    epsilon: float, optional
        Normalized cyclotron period: 1 / (cyclotron frequency × time unit). If None, computed from units and charge/mass numbers.
    """

    @classmethod
    def model_type(cls) -> LiteralOptions.ModelTypes:
        return "Fluid"

    ## species

    class EMFields(FieldSpecies):
        def __init__(self):
            self.b_field = FEECVariable(space="Hcurl")
            self.init_variables()

    class MHD(FluidSpecies):
        def __init__(
            self,
            charge_number: int = 1,
            mass_number: float = 1.0,
            epsilon: float = None,
        ):
            self.density = FEECVariable(space="L2")
            self.velocity = FEECVariable(space="Hdiv")
            self.pressure = FEECVariable(space="L2")
            self.init_variables(
                charge_number=charge_number,
                mass_number=mass_number,
                epsilon=epsilon,
            )

    ## propagators

    class Propagators:
        def __init__(self, epsilon_from=None):
            self.shear_alf = ShearAlfvenB1()
            self.hall = Hall(epsilon_from=epsilon_from)
            self.mag_sonic = MagnetosonicUniform()

    ## abstract methods

    def __init__(
        self,
        base_units: BaseUnits = BaseUnits(),
        charge_number: int = 1,
        mass_number: float = 1.0,
        epsilon: float = None,
    ):

        # 0. store input parameters
        self.params = copy.deepcopy(locals())

        # 1. instantiate all species
        self.em_fields = self.EMFields()
        self.mhd = self.MHD(
            charge_number=charge_number,
            mass_number=mass_number,
            epsilon=epsilon,
        )

        # 2. derive units (must be done after instantiating species to access charge and mass numbers)
        self.setup_equation_params(base_units=base_units)

        # 3. instantiate all propagators
        self.propagators = self.Propagators(epsilon_from=self.mhd)

        # 4. assign variables to propagators
        self.propagators.shear_alf.variables.u = self.mhd.velocity
        self.propagators.shear_alf.variables.b = self.em_fields.b_field

        self.propagators.hall.variables.b = self.em_fields.b_field

        self.propagators.mag_sonic.variables.n = self.mhd.density
        self.propagators.mag_sonic.variables.u = self.mhd.velocity
        self.propagators.mag_sonic.variables.p = self.mhd.pressure

        # 5. define scalars to be tracked during simulation
        kinetic_energy = BilinearEnergyFEEC(self.mhd.velocity, bilinear_form_name="M2n")
        pressure_energy = VolumeFormEnergyFEEC(self.mhd.pressure, normalization=1.0 / (5.0 / 3.0 - 1.0))
        magnetic_energy = BilinearEnergyFEEC(self.em_fields.b_field)
        background_pressure = FunctionScalarFEEC(self._compute_en_p_eq)
        background_magnetic = FunctionScalarFEEC(self._compute_en_B_eq)
        total_magnetic = FunctionScalarFEEC(self._compute_en_B_tot)
        helicity = FunctionScalarFEEC(self._compute_helicity)
        self.scalars = Scalars(
            en_U=kinetic_energy,
            en_p=pressure_energy,
            en_B=magnetic_energy,
            en_p_eq=background_pressure,
            en_B_eq=background_magnetic,
            en_B_tot=total_magnetic,
            en_tot=kinetic_energy + pressure_energy + magnetic_energy,
            helicity=helicity,
        )

    @property
    def bulk_species(self):
        return self.mhd

    @property
    def velocity_scale(self):
        return "alfvén"

    def allocate_helpers(self):
        self._b_eq = Propagator.projected_equil.b1
        self._a_eq = Propagator.projected_equil.a1
        self._p_eq = Propagator.projected_equil.p3

        self._ones = Propagator.projected_equil.p3.space.zeros()
        if isinstance(self._ones, PolarVector):
            self._ones.tp[:] = 1.0
        else:
            self._ones[:] = 1.0

        self._tmp_b1: BlockVector = Propagator.derham.V1.zeros()
        self._tmp_b2: BlockVector = Propagator.derham.V1.zeros()

        # adjust coupling parameters
        epsilon = self.mhd.equation_params.epsilon

        if abs(epsilon - 1) < 1e-6:
            self.mhd.equation_params.epsilon = 1.0

    def _compute_helicity(self):
        u = self.mhd.velocity.spline.vector
        p = self.mhd.pressure.spline.vector
        b = self.em_fields.b_field.spline.vector

        b1 = Propagator.mass_ops.M1.dot(b, out=self._tmp_b1)
        return 2.0 * self._a_eq.inner(b1)

    def _compute_en_B_eq(self):
        b1 = Propagator.mass_ops.M1.dot(self._b_eq, apply_bc=False, out=self._tmp_b1)
        return self._b_eq.inner(b1) / 2.0

    def _compute_en_p_eq(self):
        return self._p_eq.inner(self._ones) / (5.0 / 3.0 - 1.0)

    def _compute_en_B_tot(self):
        b = self.em_fields.b_field.spline.vector
        b1 = self._b_eq.copy(out=self._tmp_b1)
        self._tmp_b1 += b

        b2 = Propagator.mass_ops.M1.dot(b1, apply_bc=False, out=self._tmp_b2)
        return b1.inner(b2) / 2.0

    # default parameters
    def generate_default_parameter_file(self, path=None, prompt=True):
        params_path = super().generate_default_parameter_file(path=path, prompt=prompt)
        new_file = []
        with open(params_path, "r") as f:
            for line in f:
                new_file += [line]

        with open(params_path, "w") as f:
            for line in new_file:
                f.write(line)



    @classmethod
    def doc_pde(cls):
        r"""**PDEs solved by model:**

        Continuity:

        .. math::

            \frac{\partial \tilde{\rho}}{\partial t} + \nabla \cdot (\rho_0 \tilde{\mathbf{U}}) = 0

        Momentum:

        .. math::

            \rho_0 \frac{\partial \tilde{\mathbf{U}}}{\partial t} + \nabla \tilde{p} = (\nabla \times \tilde{\mathbf{B}}) \times \mathbf{B}_0

        Pressure:

        .. math::

            \frac{\partial \tilde{p}}{\partial t} + \frac{5}{3} p_0 \nabla \cdot \tilde{\mathbf{U}} = 0

        Induction:

        .. math::

            \frac{\partial \tilde{\mathbf{B}}}{\partial t} - \nabla \times \left( \tilde{\mathbf{U}} \times \mathbf{B}_0 - \frac{1}{\varepsilon} \frac{\nabla \times \tilde{\mathbf{B}}}{\rho_0} \times \mathbf{B}_0 \right) = 0

        """




    @classmethod
    def doc_normalization(cls):
        r"""All velocities are normalized by the Alfvén speed,

        .. math::

            \hat U = \hat v_A,

        and Hall effects enter through
        :math:`\varepsilon = 1/(\hat\Omega_c\hat t)`."""




    @classmethod
    def doc_scalar_quantities(cls):
        r"""**The following scalars are tracked during simulation:**

        - Flow kinetic energy: ``en_U``
        - Thermal pressure energy: ``en_p``
        - Magnetic perturbation energy: ``en_B``
        - Background pressure energy: ``en_p_eq``
        - Background magnetic energy: ``en_B_eq``
        - Total magnetic energy: ``en_B_tot``
        - Total perturbation energy: ``en_tot``
        - Magnetic helicity-like invariant: ``helicity``"""




    @classmethod
    def doc_discretization(cls):
        """Time integration is performed by the following propagators (in sequence):

        1. :class:`~struphy.propagators.shear_alfven_b1.ShearAlfvenB1`
        2. :class:`~struphy.propagators.hall.Hall`
        3. :class:`~struphy.propagators.magnetosonic_uniform.MagnetosonicUniform`
        """
        doc = rf"""**1. ShearAlfvenB1:**

{ShearAlfvenB1.__doc__}

**2. Hall:**

{Hall.__doc__}

**3. MagnetosonicUniform:**

{MagnetosonicUniform.__doc__}
"""
        return doc




    @classmethod
    def doc_long_description(cls):
        r"""LinearExtendedMHDuniform extends the linear MHD model with Hall
        dynamics for a uniform equilibrium. It is targeted at wave propagation
        and mode studies where the equilibrium is simple but finite Hall
        corrections are important."""




    @classmethod
    def doc_examples(cls):
        r"""Create and initialize a linear extended-MHD model:

        .. code-block:: python

            from struphy.models import LinearExtendedMHDuniform

            model = LinearExtendedMHDuniform()
            model.em_fields.b_field
            model.mhd.density
            model.mhd.velocity
            model.mhd.pressure
        """




    @classmethod
    def doc_use_cases(cls):
        r"""This model is appropriate for:

        - linear Hall-MHD wave studies
        - uniform-equilibrium benchmarks for extended MHD
        - comparison against ideal-MHD dispersion in the Hall-corrected regime"""




    @classmethod
    def doc_cannot_be_used_for(cls):
        r"""This model is not suitable for:

        - nonlinear extended-MHD turbulence
        - non-uniform equilibrium configurations
        - kinetic ion/electron effects beyond the Hall correction
        - dissipation-dominated problems with explicit viscosity or resistivity"""