Source code for struphy.initial.perturbations

#!/usr/bin/env python3
"Analytical perturbations."

import copy
import logging
from dataclasses import dataclass

import cunumpy as xp

from struphy.initial.base import Perturbation
from struphy.io.options import LiteralOptions
from struphy.utils.utils import check_option

logger = logging.getLogger("struphy")




@dataclass
class Noise(Perturbation):
    """White noise for FEEC coefficients.

    Parameters
    ----------
    direction: str
        The direction(s) of variation of the noise: 'e1', 'e2', 'e3', 'e1e2', etc.

    amp: float
        Noise amplitude.

    seed: int
        Seed for the random number generator.
    """

    direction: LiteralOptions.NoiseDirections = "e3"
    amp: float = 0.0001
    seed: int = None
    comp: int = 0
    given_in_basis: LiteralOptions.GivenInBasis = None

    def __post_init__(
        self,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        check_option(self.direction, LiteralOptions.NoiseDirections)

    def __call__(self):
        pass





class ModesSin(Perturbation):
    r"""Sinusoidal function in 3D.

    Parameters
    ----------
    ls : tuple[int]
        Mode numbers in x-direction (kx = l*2*pi/Lx).

    ms : tuple[int]
        Mode numbers in y-direction (ky = m*2*pi/Ly).

    ns : tuple[int]
        Mode numbers in z-direction (kz = n*2*pi/Lz).

    amps : tuple[float]
        Amplitude of each mode.

    theta : tuple | list
        Phase of each mode.

    pfuns : tuple | list[str]
        "Id" or "localize" define the profile functions.
        localize multiply the sinus by :math: `tanh((\eta_3 - 0.5)/\delta)/cosh((\eta_3 - 0.5)/\delta)`
        to localize it around 0.5. :math: `\delta` is given by the input parameter pfuns_params

    pfuns_params : tuple | list
        The parameter needed by the profile function

    Lx, Ly, Lz : float
        Domain lengths.

    given_in_basis : str
        In which basis the perturbation is represented (see base class).

    comp : int
        Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions)
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            u(x, y, z) =  \sum_{s} \chi_s(z) A_s \sin \left(l_s \frac{2\pi}{L_x} x + m_s \frac{2\pi}{L_y} y + n_s \frac{2\pi}{L_z} z + \theta_s \right) \,.

        where :math:`\chi_s(z)` is one of

        .. math::

            \chi_s(z) = \left\{
            \begin{aligned}
            1\,,
            \\[2mm]
             \tanh((z - 0.5)/\delta)/\cosh((z - 0.5)/\delta)\,,
            \end{aligned}
            \right.

        Can be used in logical space (use 'given_in_basis'), where :math:`x \to \eta_1,\, y\to \eta_2,\, z \to \eta_3`
        and :math:`L_x=L_y=L_z=1.0` (default).
        """


    def __init__(
        self,
        ls: tuple[int] = None,
        ms: tuple[int] = None,
        ns: tuple[int] = None,
        amps: tuple[float] = (1e-4,),
        theta: tuple[float] = None,
        pfuns=("Id",),
        pfuns_params=(0.0,),
        Lx=1.0,
        Ly=1.0,
        Lz=1.0,
        given_in_basis: LiteralOptions.GivenInBasis = None,
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        if ls is not None:
            n_modes = len(ls)
        elif ms is not None:
            n_modes = len(ms)
            ls = [0] * n_modes
        elif ns is not None:
            n_modes = len(ns)
            ls = [0] * n_modes
            ms = [0] * n_modes
        else:
            n_modes = 1
            ls = [0]
            ms = [0]
            ns = [0]

        if ms is None:
            ms = [0] * n_modes
        else:
            assert len(ms) == n_modes

        if ns is None:
            ns = [0] * n_modes
        else:
            assert len(ns) == n_modes

        if len(amps) == 1:
            amps = [amps[0]] * n_modes
        else:
            assert len(amps) == n_modes

        if theta is None:
            theta = [0] * n_modes

        if len(theta) == 1:
            theta = [theta[0]] * n_modes
        else:
            assert len(theta) == n_modes

        if len(pfuns) == 1:
            pfuns = [pfuns[0]] * n_modes
        else:
            assert len(pfuns) == n_modes

        if len(pfuns_params) == 1:
            pfuns_params = [pfuns_params[0]] * n_modes
        else:
            assert len(pfuns_params) == n_modes

        self.pfuns = []
        for pfun, params in zip(pfuns, pfuns_params):
            if pfun == "Id":
                self.pfuns += [lambda eta3: 1.0]
            elif pfun == "localize":
                self.pfuns += [
                    lambda eta3, p=params: xp.tanh((eta3 - 0.5) / p) / xp.cosh((eta3 - 0.5) / p),
                ]
            else:
                raise ValueError(f"Profile function {pfun} is not defined..")

        self.ls = tuple(ls)
        self.ms = tuple(ms)
        self.ns = tuple(ns)
        self.amps = tuple(amps)
        self.Lx = Lx
        self.Ly = Ly
        self.Lz = Lz
        self.theta = tuple(theta)
        self.pfuns = tuple(self.pfuns)

        # use the setters
        self.given_in_basis = given_in_basis
        self.comp = comp

    def __call__(self, x, y, z):
        val = 0.0

        for amp, l, m, n, t, pfun in zip(self.amps, self.ls, self.ms, self.ns, self.theta, self.pfuns):
            val += (
                amp
                * pfun(z)
                * xp.sin(
                    l * 2.0 * xp.pi / self.Lx * x + m * 2.0 * xp.pi / self.Ly * y + n * 2.0 * xp.pi / self.Lz * z + t,
                )
            )

        return val





class ModesCos(Perturbation):
    r"""Cosinusoidal function in 3D.

    Parameters
    ----------
    ls : tuple[int]
        Mode numbers in x-direction (kx = l*2*pi/Lx).

    ms : tuple[int]
        Mode numbers in y-direction (ky = m*2*pi/Ly).

    ns : tuple[int]
        Mode numbers in z-direction (kz = n*2*pi/Lz).

    amps : tuple[float]
        Amplitude of each mode.

    Lx, Ly, Lz : float
        Domain lengths.

    given_in_basis : str
        In which basis the perturbation is represented (see base class).

    comp : int
        Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions)

    perb_domain : tuple[tuple[float]]
        Subdomain in which the pertrubation is applied to: ((x_min, x_max), (y_min, y_max), (z_min, z_max)).
        None means apply perturbation to all domain in that direction
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            u(x, y, z) = \sum_{s} A_s \cos \left(l_s \frac{2\pi}{L_x} x + m_s \frac{2\pi}{L_y} y + n_s \frac{2\pi}{L_z} z \right) \,.

        Can be used in logical space (use 'given_in_basis'), where :math:`x \to \eta_1,\, y\to \eta_2,\, z \to \eta_3`
        and :math:`L_x=L_y=L_z=1.0` (default).
        """


    def __init__(
        self,
        ls: tuple[int] = None,
        ms: tuple[int] = None,
        ns: tuple[int] = None,
        amps: tuple[float] = (1e-4,),
        Lx=1.0,
        Ly=1.0,
        Lz=1.0,
        given_in_basis: LiteralOptions.GivenInBasis = None,
        comp: int = 0,
        perb_domain: tuple[tuple[float]] = (None, None, None),
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        if ls is not None:
            n_modes = len(ls)
        elif ms is not None:
            n_modes = len(ms)
            ls = [0] * n_modes
        elif ns is not None:
            n_modes = len(ns)
            ls = [0] * n_modes
            ms = [0] * n_modes
        else:
            n_modes = 1
            ls = [0]
            ms = [0]
            ns = [0]

        if ms is None:
            ms = [0] * n_modes
        else:
            assert len(ms) == n_modes

        if ns is None:
            ns = [0] * n_modes
        else:
            assert len(ns) == n_modes

        if len(amps) == 1:
            amps = [amps[0]] * n_modes
        else:
            assert len(amps) == n_modes

        self.ls = tuple(ls)
        self.ms = tuple(ms)
        self.ns = tuple(ns)
        self.amps = tuple(amps)
        self.Lx = Lx
        self.Ly = Ly
        self.Lz = Lz

        # use the setters
        self.given_in_basis = given_in_basis
        self.comp = comp
        self.perb_domain = perb_domain

    def __call__(self, x, y, z):

        if self.perb_domain != (None, None, None):
            val = xp.zeros_like(x)
            # find mask of particles within the sub domain
            mask = super()._mask_subdomain(x, y, z, perb_domain=self.perb_domain)

            # apply perturbation iff perb_domain not specified or (x,y,z) is within perb_domain
            for amp, l, m, n in zip(self.amps, self.ls, self.ms, self.ns):
                val[mask] += amp * xp.cos(
                    l * 2.0 * xp.pi / self.Lx * x[mask]
                    + m * 2.0 * xp.pi / self.Ly * y[mask]
                    + n * 2.0 * xp.pi / self.Lz * z[mask],
                )
        else:
            val = 0.0
            for amp, l, m, n in zip(self.amps, self.ls, self.ms, self.ns):
                val += amp * xp.cos(
                    l * 2.0 * xp.pi / self.Lx * x + m * 2.0 * xp.pi / self.Ly * y + n * 2.0 * xp.pi / self.Lz * z,
                )

        # logger.info( "Cos max value", val.max())
        return val





class CoaxialWaveguideElectric_r(Perturbation):
    r"""Initializes function for Coaxial Waveguide electric field in radial direction.

    Solutions taken from TUM master thesis of Alicia Robles Pérez:
    "Development of a Geometric Particle-in-Cell Method for Cylindrical Coordinate Systems", 2024

    Parameters
    ----------
    m : int
        Number of Modes
    a1, a2 : float
        inner and outer radius of Hollow Cylinder
    a, b : float
        Parameters of Electric field
    """

    def __init__(self, m=1, a1=1.0, a2=2.0, a=1, b=-0.28):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        self._m = m
        self._r1 = a1
        self._r2 = a2
        self._a = a
        self._b = b

        # use the setters
        self.given_in_basis = "norm"
        self.comp = 0

    def __call__(self, eta1, eta2, eta3):
        import scipy.special

        val = 0.0
        r = eta1 * (self._r2 - self._r1) + self._r1
        theta = eta2 * 2.0 * xp.pi

        val += (
            -self._m
            / r
            * xp.cos(self._m * theta)
            * (self._a * scipy.special.jv(self._m, r) + self._b * scipy.special.yn(self._m, r))
        )
        return val





class CoaxialWaveguideElectric_theta(Perturbation):
    r"""
    Initializes funtion for Coaxial Waveguide electric field in the azimuthal direction.

    Solutions taken from TUM master thesis of Alicia Robles Pérez:
    "Development of a Geometric Particle-in-Cell Method for Cylindrical Coordinate Systems", 2024

    Parameters
    ----------
    m : int
        Number of Modes
    a1, a2 : float
        inner and outer radius of Hollow Cylinder
    a, b : float
        Parameters of Electric field
    """

    def __init__(self, m=1, a1=1.0, a2=2.0, a=1, b=-0.28):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        self._m = m
        self._r1 = a1
        self._r2 = a2
        self._a = a
        self._b = b

        # use the setters
        self.given_in_basis = "norm"
        self.comp = 1

    def __call__(self, eta1, eta2, eta3):
        import scipy.special

        val = 0.0
        r = eta1 * (self._r2 - self._r1) + self._r1
        theta = eta2 * 2.0 * xp.pi

        val += (
            self._a * ((self._m / r) * scipy.special.jv(self._m, r) - scipy.special.jv(self._m + 1, r))
            + (self._b * ((self._m / r) * scipy.special.yn(self._m, r) - scipy.special.yn(self._m + 1, r)))
        ) * xp.sin(self._m * theta)
        return val





class CoaxialWaveguideMagnetic(Perturbation):
    r"""Initializes funtion for Coaxial Waveguide magnetic field in $z$-direction.

    Solutions taken from TUM master thesis of Alicia Robles Pérez:
    "Development of a Geometric Particle-in-Cell Method for Cylindrical Coordinate Systems", 2024

    Parameters
    ----------
    m : int
        Number of Modes
    a1, a2 : float
        inner and outer radius of Hollow Cylinder
    a, b : float
        Parameters of Electric field
    """

    def __init__(self, m=1, a1=1.0, a2=2.0, a=1, b=-0.28):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        self._m = m
        self._r1 = a1
        self._r2 = a2
        self._a = a
        self._b = b

        # use the setters
        self.given_in_basis = "norm"
        self.comp = 2

    def __call__(self, eta1, eta2, eta3):
        import scipy.special

        val = 0.0
        r = eta1 * (self._r2 - self._r1) + self._r1
        theta = eta2 * 2.0 * xp.pi
        z = eta3

        val += (self._a * scipy.special.jv(self._m, r) + self._b * scipy.special.yn(self._m, r)) * xp.cos(
            self._m * theta,
        )
        return val





class ModesCosCos(Perturbation):
    r"""Product of cosine modes in x and y, modulated by an optional localizing profile in z."""



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            u(x, y, z) = \sum_s A_s \, \chi_s(z)
            \cos \!\left(l_s \tfrac{2\pi}{L_x} x + \theta_{x,s}\right)
            \cos \!\left(m_s \tfrac{2\pi}{L_y} y + \theta_{y,s}\right)

        where :math:`\chi_s(z)` can be either 1 or localized in z.
        """


    def __init__(
        self,
        ls=None,
        ms=None,
        amps=(1e-4,),
        theta_x=None,
        theta_y=None,
        pfuns=("Id",),
        pfuns_params=(0.0,),
        Lx=1.0,
        Ly=1.0,
        Lz=1.0,
        given_in_basis: LiteralOptions.GivenInBasis = None,
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        if ls is not None:
            n_modes = len(ls)
        elif ms is not None:
            n_modes = len(ms)
            ls = [0] * n_modes
        else:
            n_modes = 1
            ls = [0]
            ms = [0]

        if ms is None:
            ms = [0] * n_modes
        else:
            assert len(ms) == n_modes

        if len(amps) == 1:
            amps = [amps[0]] * n_modes
        else:
            assert len(amps) == n_modes

        if theta_x is None:
            theta_x = [0] * n_modes
        elif len(theta_x) == 1:
            theta_x = [theta_x[0]] * n_modes
        else:
            assert len(theta_x) == n_modes

        if theta_y is None:
            theta_y = [0] * n_modes
        elif len(theta_y) == 1:
            theta_y = [theta_y[0]] * n_modes
        else:
            assert len(theta_y) == n_modes

        if len(pfuns) == 1:
            pfuns = [pfuns[0]] * n_modes
        else:
            assert len(pfuns) == n_modes

        if len(pfuns_params) == 1:
            pfuns_params = [pfuns_params[0]] * n_modes
        else:
            assert len(pfuns_params) == n_modes

        self._ls = ls
        self._ms = ms
        self._amps = amps
        self._theta_x = theta_x
        self._theta_y = theta_y
        self._Lx = Lx
        self._Ly = Ly
        self._Lz = Lz

        self._pfuns = []
        for pfun, params in zip(pfuns, pfuns_params):
            if pfun == "Id":
                self._pfuns += [lambda z: 1.0]
            elif pfun == "localize":
                self._pfuns += [lambda z, p=params: xp.tanh((z - 0.5) / p) / xp.cosh((z - 0.5) / p)]
            else:
                raise ValueError(f"Profile function {pfun} is not defined..")

        # use the setters
        self.given_in_basis = given_in_basis
        self.comp = comp

    def __call__(self, x, y, z):
        val = 0.0
        for amp, l, m, thx, thy, pfun in zip(self._amps, self._ls, self._ms, self._theta_x, self._theta_y, self._pfuns):
            val += (
                amp
                * pfun(z)
                * xp.cos(l * 2.0 * xp.pi / self._Lx * x + thx)
                * xp.cos(m * 2.0 * xp.pi / self._Ly * y + thy)
            )
        return val





class ModesSinSin(Perturbation):
    r"""Product of sine modes in x and y, modulated by an optional localizing profile in z."""



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            u(x, y, z) = \sum_s A_s \, \chi_s(z)
            \sin \!\left(l_s \tfrac{2\pi}{L_x} x + \theta_{x,s}\right)
            \sin \!\left(m_s \tfrac{2\pi}{L_y} y + \theta_{y,s}\right)

        where :math:`\chi_s(z)` can be either 1 or localized in z.
        """


    def __init__(
        self,
        ls=None,
        ms=None,
        amps=(1e-4,),
        theta_x=None,
        theta_y=None,
        pfuns=("Id",),
        pfuns_params=(0.0,),
        Lx=1.0,
        Ly=1.0,
        Lz=1.0,
        given_in_basis: LiteralOptions.GivenInBasis = None,
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        if ls is not None:
            n_modes = len(ls)
        elif ms is not None:
            n_modes = len(ms)
            ls = [0] * n_modes
        else:
            n_modes = 1
            ls = [0]
            ms = [0]

        if ms is None:
            ms = [0] * n_modes
        else:
            assert len(ms) == n_modes

        if len(amps) == 1:
            amps = [amps[0]] * n_modes
        else:
            assert len(amps) == n_modes

        if theta_x is None:
            theta_x = [0] * n_modes
        elif len(theta_x) == 1:
            theta_x = [theta_x[0]] * n_modes
        else:
            assert len(theta_x) == n_modes

        if theta_y is None:
            theta_y = [0] * n_modes
        elif len(theta_y) == 1:
            theta_y = [theta_y[0]] * n_modes
        else:
            assert len(theta_y) == n_modes

        if len(pfuns) == 1:
            pfuns = [pfuns[0]] * n_modes
        else:
            assert len(pfuns) == n_modes

        if len(pfuns_params) == 1:
            pfuns_params = [pfuns_params[0]] * n_modes
        else:
            assert len(pfuns_params) == n_modes

        self._ls = ls
        self._ms = ms
        self._amps = amps
        self._theta_x = theta_x
        self._theta_y = theta_y
        self._Lx = Lx
        self._Ly = Ly
        self._Lz = Lz

        self._pfuns = []
        for pfun, params in zip(pfuns, pfuns_params):
            if pfun == "Id":
                self._pfuns += [lambda z: 1.0]
            elif pfun == "localize":
                self._pfuns += [lambda z, p=params: xp.tanh((z - 0.5) / p) / xp.cosh((z - 0.5) / p)]
            else:
                raise ValueError(f"Profile function {pfun} is not defined..")

        # use the setters
        self.given_in_basis = given_in_basis
        self.comp = comp

    def __call__(self, x, y, z):
        val = 0.0
        for amp, l, m, thx, thy, pfun in zip(self._amps, self._ls, self._ms, self._theta_x, self._theta_y, self._pfuns):
            val += (
                amp
                * pfun(z)
                * xp.sin(l * 2.0 * xp.pi / self._Lx * x + thx)
                * xp.sin(m * 2.0 * xp.pi / self._Ly * y + thy)
            )
        return val





class ModesSinCos(Perturbation):
    r"""Product of a sine mode in x and a cosine mode in y, modulated by an optional localizing profile in z."""



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            u(x, y, z) = \sum_s A_s \, \chi_s(z)
            \sin \!\left(l_s \tfrac{2\pi}{L_x} x + \theta_{x,s}\right)
            \cos \!\left(m_s \tfrac{2\pi}{L_y} y + \theta_{y,s}\right)

        where :math:`\chi_s(z)` can be either 1 or localized in z.
        """


    def __init__(
        self,
        ls=None,
        ms=None,
        amps=(1e-4,),
        theta_x=None,
        theta_y=None,
        pfuns=("Id",),
        pfuns_params=(0.0,),
        Lx=1.0,
        Ly=1.0,
        Lz=1.0,
        given_in_basis: LiteralOptions.GivenInBasis = None,
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        # number of modes
        if ls is not None:
            n_modes = len(ls)
        elif ms is not None:
            n_modes = len(ms)
            ls = [0] * n_modes
        else:
            n_modes = 1
            ls = [0]
            ms = [0]

        if ms is None:
            ms = [0] * n_modes
        else:
            assert len(ms) == n_modes

        if len(amps) == 1:
            amps = [amps[0]] * n_modes
        else:
            assert len(amps) == n_modes

        if theta_x is None:
            theta_x = [0] * n_modes
        elif len(theta_x) == 1:
            theta_x = [theta_x[0]] * n_modes
        else:
            assert len(theta_x) == n_modes

        if theta_y is None:
            theta_y = [0] * n_modes
        elif len(theta_y) == 1:
            theta_y = [theta_y[0]] * n_modes
        else:
            assert len(theta_y) == n_modes

        if len(pfuns) == 1:
            pfuns = [pfuns[0]] * n_modes
        else:
            assert len(pfuns) == n_modes

        if len(pfuns_params) == 1:
            pfuns_params = [pfuns_params[0]] * n_modes
        else:
            assert len(pfuns_params) == n_modes

        # store
        self._ls = ls
        self._ms = ms
        self._amps = amps
        self._theta_x = theta_x
        self._theta_y = theta_y
        self._Lx = Lx
        self._Ly = Ly
        self._Lz = Lz

        self._pfuns = []
        for pfun, params in zip(pfuns, pfuns_params):
            if pfun == "Id":
                self._pfuns += [lambda z: 1.0]
            elif pfun == "localize":
                self._pfuns += [lambda z, p=params: xp.tanh((z - 0.5) / p) / xp.cosh((z - 0.5) / p)]
            else:
                raise ValueError(f"Profile function {pfun} is not defined..")

        # use the setters
        self.given_in_basis = given_in_basis
        self.comp = comp

    def __call__(self, x, y, z):
        val = 0.0
        for amp, l, m, thx, thy, pfun in zip(self._amps, self._ls, self._ms, self._theta_x, self._theta_y, self._pfuns):
            val += (
                amp
                * pfun(z)
                * xp.sin(l * 2.0 * xp.pi / self._Lx * x + thx)
                * xp.cos(m * 2.0 * xp.pi / self._Ly * y + thy)
            )
        return val





class ModesCosSin(Perturbation):
    r"""Product of a cosine mode in x and a sine mode in y, modulated by an optional localizing profile in z."""



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            u(x, y, z) = \sum_s A_s \, \chi_s(z)
            \cos \!\left(l_s \tfrac{2\pi}{L_x} x + \theta_{x,s}\right)
            \sin \!\left(m_s \tfrac{2\pi}{L_y} y + \theta_{y,s}\right)

        where :math:`\chi_s(z)` can be either 1 or localized in z.
        """


    def __init__(
        self,
        ls=None,
        ms=None,
        amps=(1e-4,),
        theta_x=None,
        theta_y=None,
        pfuns=("Id",),
        pfuns_params=(0.0,),
        Lx=1.0,
        Ly=1.0,
        Lz=1.0,
        given_in_basis: LiteralOptions.GivenInBasis = None,
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        # number of modes
        if ls is not None:
            n_modes = len(ls)
        elif ms is not None:
            n_modes = len(ms)
            ls = [0] * n_modes
        else:
            n_modes = 1
            ls = [0]
            ms = [0]

        if ms is None:
            ms = [0] * n_modes
        else:
            assert len(ms) == n_modes

        if len(amps) == 1:
            amps = [amps[0]] * n_modes
        else:
            assert len(amps) == n_modes

        if theta_x is None:
            theta_x = [0] * n_modes
        elif len(theta_x) == 1:
            theta_x = [theta_x[0]] * n_modes
        else:
            assert len(theta_x) == n_modes

        if theta_y is None:
            theta_y = [0] * n_modes
        elif len(theta_y) == 1:
            theta_y = [theta_y[0]] * n_modes
        else:
            assert len(theta_y) == n_modes

        if len(pfuns) == 1:
            pfuns = [pfuns[0]] * n_modes
        else:
            assert len(pfuns) == n_modes

        if len(pfuns_params) == 1:
            pfuns_params = [pfuns_params[0]] * n_modes
        else:
            assert len(pfuns_params) == n_modes

        # store
        self._ls = ls
        self._ms = ms
        self._amps = amps
        self._theta_x = theta_x
        self._theta_y = theta_y
        self._Lx = Lx
        self._Ly = Ly
        self._Lz = Lz

        self._pfuns = []
        for pfun, params in zip(pfuns, pfuns_params):
            if pfun == "Id":
                self._pfuns += [lambda z: 1.0]
            elif pfun == "localize":
                self._pfuns += [lambda z, p=params: xp.tanh((z - 0.5) / p) / xp.cosh((z - 0.5) / p)]
            else:
                raise ValueError(f"Profile function {pfun} is not defined..")

        # use the setters
        self.given_in_basis = given_in_basis
        self.comp = comp

    def __call__(self, x, y, z):
        val = 0.0
        for amp, l, m, thx, thy, pfun in zip(self._amps, self._ls, self._ms, self._theta_x, self._theta_y, self._pfuns):
            val += (
                amp
                * pfun(z)
                * xp.cos(l * 2.0 * xp.pi / self._Lx * x + thx)
                * xp.sin(m * 2.0 * xp.pi / self._Ly * y + thy)
            )
        return val





class TorusModesSin(Perturbation):
    r"""Sinusoidal function in the periodic coordinates of a Torus.

    Parameters
    ----------
    ms : tuple | list[int]
        Poloidal mode numbers.

    ns : tuple | list[int]
        Toroidal mode numbers.

    pfuns : tuple | list[str]
        "sin", "exp" or "d_exp" to define the profile functions.

    amps : tuple | list[float]
        Amplitudes of each mode (m_i, n_i).

    pfun_params : tuple | list
        Provides :math:`[r_0, \sigma]` parameters for each "exp" and "d_exp" profile function, and l_s for "sin".

    given_in_basis : str
        In which basis the perturbation is represented (see base class).

    comp : int
        Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions)
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            u(\eta_1, \eta_2, \eta_3) = \sum_{s} \chi_s(\eta_1) A_s \sin(m_s\,2\pi \eta_2 + n_s\,2\pi \eta_3) \,,

        where :math:`\chi_s(\eta_1)` is one of

        .. math::

            \chi_s(\eta_1) = \left\{
            \begin{aligned}
            &\sin(l_s\pi\eta_1)\,,
            \\[2mm]
            &\frac{1}{\sqrt{2 \pi \sigma^2}} \exp \left(- \frac{(\eta_1 - r_0)^2}{2 \sigma^2} \right) \,,
            \\[2mm]
            & - \frac{\eta_1 - r_0}{\sigma^2 \sqrt{2 \pi \sigma^2}} \exp \left(- \frac{(\eta_1 - r_0)^2}{2 \sigma^2} \right) \,.
            \end{aligned}
            \right.

        for ``"sin"``, ``"exp"`` and ``"d_exp"``, respectively; ``"d_exp"`` is the derivative of ``"exp"``.

        Can only be used in logical space (use 'given_in_basis').
        """


    def __init__(
        self,
        ms=None,
        ns=None,
        amps=(1e-4,),
        pfuns=("sin",),
        pfun_params=None,
        given_in_basis: LiteralOptions.GivenInBasis = None,
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        if given_in_basis is not None:
            assert "physical" not in given_in_basis

        if ms is not None:
            n_modes = len(ms)
        elif ns is not None:
            n_modes = len(ns)
            ms = [0] * n_modes
        else:
            n_modes = 1
            ms = [1]
            ns = [0]

        if ns is None:
            ns = [0] * n_modes
        else:
            assert len(ns) == n_modes

        if len(amps) == 1:
            amps = [amps[0]] * n_modes
        else:
            assert len(amps) == n_modes

        if len(pfuns) == 1:
            pfuns = [pfuns[0]] * n_modes
        else:
            assert len(pfuns) == n_modes

        if pfun_params is None:
            pfun_params = [None] * n_modes

        self._ms = ms
        self._ns = ns
        self._amps = amps

        self._pfuns = []
        for pfun, params in zip(pfuns, pfun_params):
            if pfun == "sin":
                if params is None:
                    ls = 1
                else:
                    ls = params
                self._pfuns += [lambda eta1, ls=ls: xp.sin(ls * xp.pi * eta1)]
            elif pfun == "exp":
                self._pfuns += [
                    lambda eta1, p=params: (
                        xp.exp(-((eta1 - p[0]) ** 2) / (2 * p[1] ** 2)) / xp.sqrt(2 * xp.pi * p[1] ** 2)
                    ),
                ]
            elif pfun == "d_exp":
                self._pfuns += [
                    lambda eta1, p=params: (
                        -(eta1 - p[0])
                        / p[1] ** 2
                        * xp.exp(-((eta1 - p[0]) ** 2) / (2 * p[1] ** 2))
                        / xp.sqrt(2 * xp.pi * p[1] ** 2)
                    ),
                ]
            else:
                raise ValueError(f"Profile function {pfun} is not defined..")

        # use the setters
        self.given_in_basis = given_in_basis
        self.comp = comp

    def __call__(self, eta1, eta2, eta3):
        val = 0.0
        for mi, ni, pfun, amp in zip(self._ms, self._ns, self._pfuns, self._amps):
            val += (
                amp
                * pfun(eta1)
                * xp.sin(
                    mi * 2.0 * xp.pi * eta2 + ni * 2.0 * xp.pi * eta3,
                )
            )

        return val





class TorusModesCos(Perturbation):
    r"""Cosinusoidal function in the periodic coordinates of a Torus.

    Parameters
    ----------
    ms : tuple[int]
        Poloidal mode numbers.

    ns : tuple[int]
        Toroidal mode numbers.

    pfuns : tuple[str]
        "sin", "cos", "exp" or "d_exp" to define the profile functions.

    amps : tuple[float]
        Amplitudes of each mode (m_i, n_i).

    pfun_params : tuple | list
        Provides :math:`[r_0, \sigma]` parameters for each "exp" and "d_exp" profile function, and l_s for "sin".

    given_in_basis : str
        In which basis the perturbation is represented (see base class).

    comp : int
        Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions)
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            u(\eta_1, \eta_2, \eta_3) = \sum_{s} \chi_s(\eta_1) A_s \cos(m_s\,2\pi \eta_2 + n_s\,2\pi \eta_3) \,,

        where :math:`\chi_s(\eta_1)` is one of

        .. math::

            \chi_s(\eta_1) = \left\{
            \begin{aligned}
            &\sin(l_s\pi\eta_1)\,,
            \\[2mm]
            &\cos(\pi\eta_1)\,,
            \\[2mm]
            &\frac{1}{\sqrt{2 \pi \sigma^2}} \exp \left(- \frac{(\eta_1 - r_0)^2}{2 \sigma^2} \right) \,,
            \\[2mm]
            & - \frac{\eta_1 - r_0}{\sigma^2 \sqrt{2 \pi \sigma^2}} \exp \left(- \frac{(\eta_1 - r_0)^2}{2 \sigma^2} \right) \,.
            \end{aligned}
            \right.

        for ``"sin"``, ``"cos"``, ``"exp"`` and ``"d_exp"``, respectively; ``"d_exp"`` is the derivative of ``"exp"``.

        Can only be used in logical space (use 'given_in_basis').
        """


    def __init__(
        self,
        ms: tuple = (2,),
        ns: tuple = (1,),
        amps: tuple = (0.1,),
        pfuns: tuple = ("sin",),
        pfun_params=None,
        given_in_basis: LiteralOptions.GivenInBasis = None,
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        if given_in_basis is not None:
            assert "physical" not in given_in_basis

        if ms is not None:
            n_modes = len(ms)
        elif ns is not None:
            n_modes = len(ns)
            ms = [0] * n_modes
        else:
            n_modes = 1
            ms = [1]
            ns = [0]

        if ns is None:
            ns = [0] * n_modes
        else:
            assert len(ns) == n_modes

        if len(amps) == 1:
            amps = [amps[0]] * n_modes
        else:
            assert len(amps) == n_modes

        if len(pfuns) == 1:
            pfuns = [pfuns[0]] * n_modes
        else:
            assert len(pfuns) == n_modes

        if pfun_params is None:
            pfun_params = [None] * n_modes

        self._ms = ms
        self._ns = ns
        self._amps = amps

        self._pfuns = []
        for pfun, params in zip(pfuns, pfun_params):
            if pfun == "sin":
                if params is None:
                    ls = 1
                else:
                    ls = params
                self._pfuns += [lambda eta1, ls=ls: xp.sin(ls * xp.pi * eta1)]
            elif pfun == "cos":
                self._pfuns += [lambda eta1: xp.cos(xp.pi * eta1)]
            elif pfun == "exp":
                self._pfuns += [
                    lambda eta1, p=params: (
                        xp.exp(-((eta1 - p[0]) ** 2) / (2 * p[1] ** 2)) / xp.sqrt(2 * xp.pi * p[1] ** 2)
                    ),
                ]
            elif pfun == "d_exp":
                self._pfuns += [
                    lambda eta1, p=params: (
                        -(eta1 - p[0])
                        / p[1] ** 2
                        * xp.exp(-((eta1 - p[0]) ** 2) / (2 * p[1] ** 2))
                        / xp.sqrt(2 * xp.pi * p[1] ** 2)
                    ),
                ]
            else:
                raise ValueError(
                    'Profile function must be "sin" or "cos" or "exp".',
                )

        # use the setters
        self.given_in_basis = given_in_basis
        self.comp = comp

    def __call__(self, eta1, eta2, eta3):
        val = 0.0
        for mi, ni, pfun, amp in zip(self._ms, self._ns, self._pfuns, self._amps):
            val += (
                amp
                * pfun(eta1)
                * xp.cos(
                    mi * 2.0 * xp.pi * eta2 + ni * 2.0 * xp.pi * eta3,
                )
            )

        return val





class Shear_x(Perturbation):
    r"""Double shear layer in eta1 (-1 in outer regions, 1 in inner regions).

    Parameters
    ----------
    amps : float
        Amplitude of the velocity on each side.

    delta : float
        Characteristic size of the shear layer

    given_in_basis : str
        In which basis the perturbation is represented (see base class).

    comp : int
        Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions)
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            u(\eta_1, \eta_2, \eta_3) = A(-\tanh((\eta_1 - 0.25)/\delta)+\tanh((\eta_1 - 0.75)/\delta) - 1) \,.

        Can only be used in logical space.
        """


    def __init__(
        self,
        amp=1e-4,
        delta=1 / 15,
        given_in_basis: LiteralOptions.GivenInBasis = None,
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        if given_in_basis is not None:
            assert "physical" not in given_in_basis, f"Perturbation {self.__name__} can only be used in logical space."

        self._amp = amp
        self._delta = delta

        # use the setters
        self.given_in_basis = given_in_basis
        self.comp = comp

    def __call__(self, e1, e2, e3):
        val = self._amp * (-xp.tanh((e1 - 0.75) / self._delta) + xp.tanh((e1 - 0.25) / self._delta) - 1)

        return val





class Shear_y(Perturbation):
    r"""Double shear layer in eta2 (-1 in outer regions, 1 in inner regions).

    Parameters
    ----------
    amps : float
        Amplitude of the velocity on each side.

    delta : float
        Characteristic size of the shear layer

    given_in_basis : str
        In which basis the perturbation is represented (see base class).

    comp : int
        Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions)
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            u(\eta_1, \eta_2, \eta_3) = A(-\tanh((\eta_2 - 0.25)/\delta) + \tanh((\eta_2 - 0.75)/\delta) - 1) \,.

        Can only be used in logical space.
        """


    def __init__(
        self,
        amp=1e-4,
        delta=1 / 15,
        given_in_basis: LiteralOptions.GivenInBasis = None,
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        if given_in_basis is not None:
            assert "physical" not in given_in_basis, f"Perturbation {self.__name__} can only be used in logical space."

        self._amp = amp
        self._delta = delta

        # use the setters
        self.given_in_basis = given_in_basis
        self.comp = comp

    def __call__(self, e1, e2, e3):
        val = self._amp * (-xp.tanh((e2 - 0.75) / self._delta) + xp.tanh((e2 - 0.25) / self._delta) - 1)

        return val





class Shear_z(Perturbation):
    r"""Double shear layer in eta3 (-1 in outer regions, 1 in inner regions).

    Parameters
    ----------
    amps : float
        Amplitude of the velocity on each side.

    delta : float
        Characteristic size of the shear layer

    given_in_basis : str
        In which basis the perturbation is represented (see base class).

    comp : int
        Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions)
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            u(\eta_1, \eta_2, \eta_3) = A(-\tanh((\eta_3 - 0.25)/\delta) + \tanh((\eta_3 - 0.75)/\delta) - 1) \,.

        Can only be used in logical space.
        """


    def __init__(
        self,
        amp=1e-4,
        delta=1 / 15,
        given_in_basis: LiteralOptions.GivenInBasis = None,
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        if given_in_basis is not None:
            assert "physical" not in given_in_basis, f"Perturbation {self.__name__} can only be used in logical space."

        self._amp = amp
        self._delta = delta

        # use the setters
        self.given_in_basis = given_in_basis
        self.comp = comp

    def __call__(self, e1, e2, e3):
        val = self._amp * (-xp.tanh((e3 - 0.75) / self._delta) + xp.tanh((e3 - 0.25) / self._delta) - 1)

        return val





class Erf_z(Perturbation):
    r"""Shear layer in eta3 (-1 in lower regions, 1 in upper regions).

    Parameters
    ----------
    amp : float
        Amplitude of the velocity on each side.

    delta : float
        Characteristic size of the shear layer

    given_in_basis : str
        In which basis the perturbation is represented (see base class).

    comp : int
        Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions)
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            u(\eta_1, \eta_2, \eta_3) = A \, erf((\eta_3 - 0.5)/\delta) \,.

        Can only be used in logical space.
        """


    def __init__(
        self,
        amp=1e-4,
        delta=1 / 15,
        given_in_basis: LiteralOptions.GivenInBasis = None,
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        if given_in_basis is not None:
            assert "physical" not in given_in_basis, f"Perturbation {self.__name__} can only be used in logical space."

        self._amp = amp
        self._delta = delta

        # use the setters
        self.given_in_basis = given_in_basis
        self.comp = comp

    def __call__(self, e1, e2, e3):
        from scipy.special import erf

        val = self._amp * erf((e3 - 0.5) / self._delta)

        return val





class Gaussian(Perturbation):
    r"""Gaussian bump (e.g. a localized pulse as initial condition for wave propagation).

    Parameters
    ----------
    amp : float
        Amplitude.

    center : tuple[float]
        Center :math:`(x_0, y_0, z_0)` of the bump.

    sigma : tuple[float]
        Widths :math:`(\sigma_x, \sigma_y, \sigma_z)` per direction; ``None`` for a direction without variation.

    given_in_basis : str
        In which basis the perturbation is represented (see base class). The coordinates are physical with
        ``"physical"`` and logical otherwise.

    comp : int
        Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions)
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            u(x, y, z) = A \, \exp\left(-\frac{(x - x_0)^2}{2 \sigma_x^2} - \frac{(y - y_0)^2}{2 \sigma_y^2}
            - \frac{(z - z_0)^2}{2 \sigma_z^2}\right)\,,

        where a direction with :math:`\sigma = ` ``None`` is left out.
        """


    def __init__(
        self,
        amp: float = 1e-4,
        center: tuple = (0.5, 0.5, 0.5),
        sigma: tuple = (0.1, 0.1, None),
        given_in_basis: LiteralOptions.GivenInBasis = None,
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        assert len(center) == 3 and len(sigma) == 3, "center and sigma need three entries (one per direction)."
        assert any(s is not None for s in sigma), "At least one width must be given."

        self._amp = amp
        self._center = center
        self._sigma = sigma

        # use the setters
        self.given_in_basis = given_in_basis
        self.comp = comp

    def __call__(self, x, y, z):
        arg = 0.0 * x
        for coord, c, s in zip((x, y, z), self._center, self._sigma):
            if s is not None:
                arg = arg + (coord - c) ** 2 / (2 * s**2)

        return self._amp * xp.exp(-arg)





class RestelliAnalyticSolutionVelocity(Perturbation):
    r"""Analytic solution :math:`u=u_e` of the system:

    Parameters
    ----------
    a : float
        Minor radius of torus (default: 1.).
    R0 : float
        Major radius of torus (default: 2.).
    B0 : float
        On-axis (r=0) toroidal magnetic field (default: 10.).
    Bp : float
        Poloidal magnetic field (default: 12.5).
    alpha : float
        (default: 0.1)
    beta : float
        (default: 1.0)
    comp : int
        Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions)

    References
    ----------
    [1] Juan Vicente Gutiérrez-Santacreu, Omar Maj, Marco Restelli: Finite element discretization of a Stokes-like model arising
    in plasma physics, Journal of Computational Physics 2018.
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            \partial_t u = - \nabla \phi + u \times B + \nu \Delta u + f \,,\\
            0 = \nabla \phi- u_e \times B + \nu_e \Delta u_e + f_e \,, \\
            \nabla \cdot (u-u_e) = 0 \,.

        where :math:`f` is defined as follows: 

        .. math::

            f = \nu \omega \,, 
            \\[2mm]
            \omega = \left[0, \alpha \frac{R_0 - 4R}{a R_0 R} - \beta \frac{B_p}{B_0}\frac{R_0^2}{a R^3}, 0 \right] \,, 
            \\[2mm]
            R = \sqrt{x^2 + y^2} \,.

        Can only be defined in Cartesian coordinates. 
        The solution is given by:

        .. math::
            \alpha \frac{R}{a R_0} \left[\begin{array}{c} -z \\ R-R_0 \\ 0 \end{array} \right] + \beta \frac{B_p}{B_0} \frac{R_0}{aR} \left[\begin{array}{c} z \\ -(R-R_0) \\ \frac{B_0}{B_p} a \end{array} \right] \,,
            \\[2mm]
            R = \sqrt{x^2 + y^2} \,.
        """


    def __init__(
        self,
        a=1.0,
        R0=2.0,
        B0=10.0,
        Bp=12.5,
        alpha=0.1,
        beta=1.0,
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        self._a = a
        self._R0 = R0
        self._B0 = B0
        self._Bp = Bp
        self._alpha = alpha
        self._beta = beta

        # use the setters
        self.given_in_basis = "physical"
        self.comp = comp

    # equilibrium ion velocity
    def __call__(self, x, y, z):
        """Velocity of ions and electrons."""
        R = xp.sqrt(x**2 + y**2)
        R = xp.where(R == 0.0, 1e-9, R)
        phi = xp.arctan2(-y, x)
        ustarR = (
            self._alpha * R / (self._a * self._R0) * (-z)
            + self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * z
        )
        ustarZ = self._alpha * R / (self._a * self._R0) * (R - self._R0) + self._beta * self._Bp * self._R0 / (
            self._B0 * self._a * R
        ) * (-(R - self._R0))
        ustarphi = self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * self._B0 * self._a / self._Bp

        # form normalized to cylindrical coordinates:
        uR = ustarR
        uphi = ustarphi / R
        uZ = ustarZ

        # from cylindrical to cartesian:

        if self.comp == 0:
            ux = xp.cos(phi) * uR - R * xp.sin(phi) * uphi
            return ux
        elif self.comp == 1:
            uy = -xp.sin(phi) * uR - R * xp.cos(phi) * uphi
            return uy
        elif self.comp == 2:
            uz = uZ
            return uz
        else:
            raise ValueError(f"Invalid component '{self._comp}'. Must be 0, 1, or 2.")





class RestelliAnalyticSolutionVelocity_2(Perturbation):
    r"""Analytic solution :math:`u=u_e` of the system:

    Parameters
    ----------
    a : float
        Minor radius of torus (default: 1.).
    R0 : float
        Major radius of torus (default: 2.).
    B0 : float
        On-axis (r=0) toroidal magnetic field (default: 10.).
    Bp : float
        Poloidal magnetic field (default: 12.5).
    alpha : float
        (default: 0.1)
    beta : float
        (default: 1.0)
    comp : int
        Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions)

    References
    ----------
    [1] Juan Vicente Gutiérrez-Santacreu, Omar Maj, Marco Restelli: Finite element discretization of a Stokes-like model arising
    in plasma physics, Journal of Computational Physics 2018.
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            \partial_t u = - \nabla \phi + u \times B + \nu \Delta u + f \,,\\
            0 = \nabla \phi- u_e \times B + \nu_e \Delta u_e + f_e \,, \\
            \nabla \cdot (u-u_e) = 0 \,.

        where :math:`f` is defined as follows: 

        .. math::

            f = \nu \omega \,, 
            \\[2mm]
            \omega = \left[0, \alpha \frac{R_0 - 4R}{a R_0 R} - \beta \frac{B_p}{B_0}\frac{R_0^2}{a R^3}, 0 \right] \,, 
            \\[2mm]
            R = \sqrt{x^2 + y^2} \,.

        Can only be defined in Cartesian coordinates. 
        The solution is given by:

        .. math::
            \alpha \frac{R}{a R_0} \left[\begin{array}{c} -z \\ R-R_0 \\ 0 \end{array} \right] + \beta \frac{B_p}{B_0} \frac{R_0}{aR} \left[\begin{array}{c} z \\ -(R-R_0) \\ \frac{B_0}{B_p} a \end{array} \right] \,,
            \\[2mm]
            R = \sqrt{x^2 + y^2} \,.
        """


    def __init__(
        self,
        a=1.0,
        R0=2.0,
        B0=10.0,
        Bp=12.5,
        alpha=0.1,
        beta=1.0,
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        self._a = a
        self._R0 = R0
        self._B0 = B0
        self._Bp = Bp
        self._alpha = alpha
        self._beta = beta

        # use the setter
        self.given_in_basis = "physical"
        self.comp = comp

    # equilibrium ion velocity
    def __call__(self, x, y, z):
        """Velocity of ions and electrons."""
        R = xp.sqrt(x**2 + y**2)
        R = xp.where(R == 0.0, 1e-9, R)
        phi = xp.arctan2(-y, x)
        ustarR = (
            self._alpha * R / (self._a * self._R0) * (-z)
            + self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * z
        )
        ustarZ = self._alpha * R / (self._a * self._R0) * (R - self._R0) + self._beta * self._Bp * self._R0 / (
            self._B0 * self._a * R
        ) * (-(R - self._R0))
        ustarphi = self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * self._B0 * self._a / self._Bp

        # form normalized to cylindrical coordinates:
        uR = ustarR
        uphi = ustarphi / R
        uZ = ustarZ

        # from cylindrical to cartesian:

        if self.comp == 0:
            ux = xp.cos(phi) * uR - R * xp.sin(phi) * uphi
            return ux
        elif self.comp == 1:
            uy = -xp.sin(phi) * uR - R * xp.cos(phi) * uphi
            return uy
        elif self.comp == 2:
            uz = uZ
            return uz
        else:
            raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.")





class RestelliAnalyticSolutionVelocity_3(Perturbation):
    r"""Analytic solution :math:`u=u_e` of the system:

    Parameters
    ----------
    a : float
        Minor radius of torus (default: 1.).
    R0 : float
        Major radius of torus (default: 2.).
    B0 : float
        On-axis (r=0) toroidal magnetic field (default: 10.).
    Bp : float
        Poloidal magnetic field (default: 12.5).
    alpha : float
        (default: 0.1)
    beta : float
        (default: 1.0)
    comp : int
        Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions)

    References
    ----------
    [1] Juan Vicente Gutiérrez-Santacreu, Omar Maj, Marco Restelli: Finite element discretization of a Stokes-like model arising
    in plasma physics, Journal of Computational Physics 2018.
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            \partial_t u = - \nabla \phi + u \times B + \nu \Delta u + f \,,\\
            0 = \nabla \phi- u_e \times B + \nu_e \Delta u_e + f_e \,, \\
            \nabla \cdot (u-u_e) = 0 \,.

        where :math:`f` is defined as follows: 

        .. math::

            f = \nu \omega \,, 
            \\[2mm]
            \omega = \left[0, \alpha \frac{R_0 - 4R}{a R_0 R} - \beta \frac{B_p}{B_0}\frac{R_0^2}{a R^3}, 0 \right] \,, 
            \\[2mm]
            R = \sqrt{x^2 + y^2} \,.

        Can only be defined in Cartesian coordinates. 
        The solution is given by:

        .. math::
            \alpha \frac{R}{a R_0} \left[\begin{array}{c} -z \\ R-R_0 \\ 0 \end{array} \right] + \beta \frac{B_p}{B_0} \frac{R_0}{aR} \left[\begin{array}{c} z \\ -(R-R_0) \\ \frac{B_0}{B_p} a \end{array} \right] \,,
            \\[2mm]
            R = \sqrt{x^2 + y^2} \,.
        """


    def __init__(
        self,
        a=1.0,
        R0=2.0,
        B0=10.0,
        Bp=12.5,
        alpha=0.1,
        beta=1.0,
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        self._a = a
        self._R0 = R0
        self._B0 = B0
        self._Bp = Bp
        self._alpha = alpha
        self._beta = beta

        # use the setters
        self.given_in_basis = "physical"
        self.comp = comp

    # equilibrium ion velocity
    def __call__(self, x, y, z):
        """Velocity of ions and electrons."""
        R = xp.sqrt(x**2 + y**2)
        R = xp.where(R == 0.0, 1e-9, R)
        phi = xp.arctan2(-y, x)
        ustarR = (
            self._alpha * R / (self._a * self._R0) * (-z)
            + self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * z
        )
        ustarZ = self._alpha * R / (self._a * self._R0) * (R - self._R0) + self._beta * self._Bp * self._R0 / (
            self._B0 * self._a * R
        ) * (-(R - self._R0))
        ustarphi = self._beta * self._Bp * self._R0 / (self._B0 * self._a * R) * self._B0 * self._a / self._Bp

        # form normalized to cylindrical coordinates:
        uR = ustarR
        uphi = ustarphi / R
        uZ = ustarZ

        # from cylindrical to cartesian:

        if self.comp == 0:
            ux = xp.cos(phi) * uR - R * xp.sin(phi) * uphi
            return ux
        elif self.comp == 1:
            uy = -xp.sin(phi) * uR - R * xp.cos(phi) * uphi
            return uy
        elif self.comp == 2:
            uz = uZ
            return uz
        else:
            raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.")





class RestelliAnalyticSolutionPotential(Perturbation):
    r"""Analytic solution :math:`\phi` of the system:

    Parameters
    ----------
    a : float
        Minor radius of torus (default: 1.).
    R0 : float
        Major radius of torus (default: 2.).
    B0 : float
        On-axis (r=0) toroidal magnetic field (default: 10.).
    Bp : float
        Poloidal magnetic field (default: 12.5).
    alpha : float
        (default: 0.1)
    beta : float
        (default: 1.0)

    References
    ----------
    [1] Juan Vicente Gutiérrez-Santacreu, Omar Maj, Marco Restelli: Finite element discretization of a Stokes-like model arising
    in plasma physics, Journal of Computational Physics 2018.
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            \partial_t u = - \nabla \phi + u \times B + \nu \Delta u + f \,,\\
            0 = \nabla \phi- u_e \times B + \nu_e \Delta u_e + f_e \,, \\
            \nabla \cdot (u-u_e) = 0 \,.

        where :math:`f` is defined as follows: 

        .. math::

            f = \nu \omega \,, 
            \\[2mm]
            \omega = \left[0, \alpha \frac{R_0 - 4R}{a R_0 R} - \beta \frac{B_p}{B_0}\frac{R_0^2}{a R^3}, 0 \right] \,, 
            \\[2mm]
            R = \sqrt{x^2 + y^2} \,.

        Can only be defined in Cartesian coordinates. 
        The solution is given by:

        .. math::
            \phi = \frac{1}{2} a B_0 \alpha \left( \frac{(R-R_0)^2+z^2}{a^2} - \frac{2}{3} \right)
            \\[2mm]
            R = \sqrt{x^2 + y^2} \,.
        """


    def __init__(self, a=1.0, R0=2.0, B0=10.0, Bp=12.5, alpha=0.1, beta=1.0):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        self._a = a
        self._R0 = R0
        self._B0 = B0
        self._Bp = Bp
        self._alpha = alpha
        self._beta = beta

        # use the setter
        self.given_in_basis = "physical"

    # equilibrium potential
    def __call__(self, x, y, z):
        """Equilibrium potential."""
        R = xp.sqrt(x**2 + y**2)
        pp = 0.5 * self._a * self._B0 * self._alpha * (((R - self._R0) ** 2 + z**2) / self._a**2 - 2.0 / 3.0)

        return pp





class ManufacturedSolutionVelocity(Perturbation):
    r"""Analytic solutions :math:`u` and :math:`u_e` of the system:

    Parameters
    ----------
    species : string
        'Ions' or 'Electrons'.
    comp : string
        Which component of the solution ('0', '1' or '2').
    dimension: string
        Defines the manufactured solution to be selected ('1D' or '2D').
    b0 : float
        Magnetic field (default: 1.0).
    comp : int
        Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions)
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            \partial_t u = - \nabla \phi + u \times B + \nu \Delta u + f \,,\\
            0 = \nabla \phi- u_e \times B + \nu_e \Delta u_e + f_e \,, \\
            \nabla \cdot (u-u_e) = 0 \,.

        Can only be defined in Cartesian coordinates. 
        The solution in 1D is given by:

        .. math::
            u =  \left[\begin{array}{c} sin(2 \pi x) + 1.0 \\ 0 \\ 0 \end{array} \right] \,,
            u_e =  \left[\begin{array}{c} sin(2 \pi x) \\ 0 \\ 0 \end{array} \right] \,.

        The solution in 2D is given by:

        .. math::
            u =  \left[\begin{array}{c} -sin(2 \pi x) sin(2 \pi y) \\ -cos(2 \pi y) cos(2 \pi y) \\ 0 \end{array} \right] \,,
            u_e =  \left[\begin{array}{c} -sin(4 \pi x) sin(4 \pi y) \\ -cos(4 \pi y) cos(4 \pi y) \\ 0 \end{array} \right] \,.
        """


    def __init__(
        self,
        species="Ions",
        dimension="1D",
        b0=1.0,
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        self._b = b0
        self._species = species
        self._dimension = dimension

        # use the setters
        self.given_in_basis = "physical"
        self.comp = comp

    # equilibrium ion velocity
    def __call__(self, x, y, z):
        if self._species == "Ions":
            """Velocity of ions."""
            """x component"""
            if self._dimension == "2D":
                ux = -xp.sin(2 * xp.pi * x) * xp.sin(2 * xp.pi * y)
            elif self._dimension == "1D":
                ux = xp.sin(2 * xp.pi * x) + 1.0

            """y component"""
            if self._dimension == "2D":
                uy = -xp.cos(2 * xp.pi * x) * xp.cos(2 * xp.pi * y)
            elif self._dimension == "1D":
                uy = xp.cos(2 * xp.pi * x)

            """z component"""
            uz = 0.0 * x

            if self.comp == 0:
                return ux
            elif self.comp == 1:
                return uy
            elif self.comp == 2:
                return uz
            else:
                raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.")

        elif self._species == "Electrons":
            """Velocity of electrons."""
            """x component"""
            if self._dimension == "2D":
                ux = -xp.sin(4 * xp.pi * x) * xp.sin(4 * xp.pi * y)
            elif self._dimension == "1D":
                ux = xp.sin(2.0 * xp.pi * x)

            """y component"""
            if self._dimension == "2D":
                uy = -xp.cos(4 * xp.pi * x) * xp.cos(4 * xp.pi * y)
            elif self._dimension == "1D":
                uy = xp.cos(2 * xp.pi * x)

            """z component"""
            uz = 0.0 * x

            if self.comp == 0:
                return ux
            if self.comp == 1:
                return uy
            if self.comp == 2:
                return uz
            else:
                raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.")

        else:
            raise ValueError(f"Invalid species '{self._species}'. Must be 'Ions' or 'Electrons'.")





class ManufacturedSolutionPotential(Perturbation):
    r"""Analytic solution :math:`\phi` of the system:

    Parameters
    ----------
    dimension: string
        Defines the manufactured solution to be selected ('1D' or '2D').
    b0 : float
        Magnetic field (default: 1.0).
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            \partial_t u = - \nabla \phi + u \times B + \nu \Delta u + f \,,\\
            0 = \nabla \phi- u_e \times B + \nu_e \Delta u_e + f_e \,, \\
            \nabla \cdot (u-u_e) = 0 \,.

        where :math:`f` is defined as follows: 

        .. math::

            f = \left[1 - b_0 cos(x) - \nu sin(y), 1 - b_0 sin(y) + \nu cos(x) , 0 \right] \,, 
            \\[2mm]
            f_e = \left[-1 + 0.5 b_0 cos(x) - \nu_e 0.5 sin(y), -1 + 0.5 b_0 sin(y) + \nu_e cos(x) , 0 \right] \,.

        Can only be defined in Cartesian coordinates. 
        The solution in 1D is given by:

        .. math::
            \phi =  sin(2\pi x) \,.

        The solution in 2D is given by:

        .. math::
            \phi =  cos(2\pi x) + sin(2\pi y) \,.
        """


    def __init__(self, dimension="1D", b0=1.0):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        self._ab = b0
        self._dimension = dimension

        # use the setter
        self.given_in_basis = "physical"

    # equilibrium ion velocity
    def __call__(self, x, y, z):
        """Potential."""
        if self._dimension == "2D":
            phi = xp.cos(2 * xp.pi * x) + xp.sin(2 * xp.pi * y)
        elif self._dimension == "1D":
            phi = xp.sin(2.0 * xp.pi * x)

        return phi





class ManufacturedSolutionVelocity_2(Perturbation):
    r"""Analytic solutions :math:`u` and :math:`u_e` of the system:

    Parameters
    ----------
    species : string
        'Ions' or 'Electrons'.
    dimension: string
        Defines the manufactured solution to be selected ('1D' or '2D').
    b0 : float
        Magnetic field (default: 1.0).
    comp : int
        Which component (0, 1 or 2) of vector is perturbed (=0 for scalar-valued functions)
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            \partial_t u = - \nabla \phi + u \times B + \nu \Delta u + f \,,\\
            0 = \nabla \phi- u_e \times B + \nu_e \Delta u_e + f_e \,, \\
            \nabla \cdot (u-u_e) = 0 \,.

        Can only be defined in Cartesian coordinates. 
        The solution in 1D is given by:

        .. math::
            u =  \left[\begin{array}{c} sin(2 \pi x) + 1.0 \\ 0 \\ 0 \end{array} \right] \,,
            u_e =  \left[\begin{array}{c} sin(2 \pi x) \\ 0 \\ 0 \end{array} \right] \,.

        The solution in 2D is given by:

        .. math::
            u =  \left[\begin{array}{c} -sin(2 \pi x) sin(2 \pi y) \\ -cos(2 \pi y) cos(2 \pi y) \\ 0 \end{array} \right] \,,
            u_e =  \left[\begin{array}{c} -sin(4 \pi x) sin(4 \pi y) \\ -cos(4 \pi y) cos(4 \pi y) \\ 0 \end{array} \right] \,.
        """


    def __init__(
        self,
        species="Ions",
        dimension="1D",
        b0=1.0,
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        self._b = b0
        self._species = species
        self._dimension = dimension

        # use the setters
        self.given_in_basis = "physical"
        self.comp = comp

    # equilibrium ion velocity
    def __call__(self, x, y, z):
        if self._species == "Ions":
            """Velocity of ions."""
            """x component"""
            if self._dimension == "2D":
                ux = -xp.sin(2 * xp.pi * x) * xp.sin(2 * xp.pi * y)
            elif self._dimension == "1D":
                ux = xp.sin(2 * xp.pi * x) + 1.0

            """y component"""
            if self._dimension == "2D":
                uy = -xp.cos(2 * xp.pi * x) * xp.cos(2 * xp.pi * y)
            elif self._dimension == "1D":
                uy = xp.cos(2 * xp.pi * x)

            """z component"""
            uz = 0.0 * x

            if self.comp == 0:
                return ux
            elif self.comp == 1:
                return uy
            elif self.comp == 2:
                return uz
            else:
                raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.")

        elif self._species == "Electrons":
            """Velocity of electrons."""
            """x component"""
            if self._dimension == "2D":
                ux = -xp.sin(4 * xp.pi * x) * xp.sin(4 * xp.pi * y)
            elif self._dimension == "1D":
                ux = xp.sin(2.0 * xp.pi * x)

            """y component"""
            if self._dimension == "2D":
                uy = -xp.cos(4 * xp.pi * x) * xp.cos(4 * xp.pi * y)
            elif self._dimension == "1D":
                uy = xp.cos(2 * xp.pi * x)

            """z component"""
            uz = 0.0 * x

            if self.comp == 0:
                return ux
            if self.comp == 1:
                return uy
            if self.comp == 2:
                return uz
            else:
                raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.")

        else:
            raise ValueError(f"Invalid species '{self._species}'. Must be 'Ions' or 'Electrons'.")





class ITPA_density(Perturbation):
    r"""ITPA radial density profile in `A. Könies et al. 2018  <https://iopscience.iop.org/article/10.1088/1741-4326/aae4e6>`_

    Parameters
    ----------
    n0 : float
        ITPA profile density

    c : tuple | list
        4 ITPA profile coefficients
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            n(\eta_1) = n_0*c_3\exp\left[-\frac{c_2}{c_1}\tanh\left(\frac{\eta_1 - c_0}{c_2}\right)\right]\,.
        """


    def __init__(
        self,
        n0: float = 0.00720655,
        c: tuple = (0.491230, 0.298228, 0.198739, 0.521298),
        comp: int = 0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        assert len(c) == 4

        self._n0 = n0
        self._c = c

        # use the setters
        self.given_in_basis = "physical"
        self.comp = comp

    def __call__(self, eta1, eta2=None, eta3=None):
        val = 0.0

        if self._c[2] == 0.0:
            val = self._n0 * self._c[3] - 0 * eta1
        else:
            val = (
                self._n0
                * self._c[3]
                * xp.exp(
                    -self._c[2] / self._c[1] * xp.tanh((eta1 - self._c[0]) / self._c[2]),
                )
            )

        return val





class TokamakManufacturedSolutionVelocity(Perturbation):
    r"""Analytic solution :math:`u=u_e` of the system:

    Parameters
    ----------
    comp : string
        Which component of the solution ('0', '1' or '2').
    a : float
        Minor radius of torus (default: 1.).
    R0 : float
        Major radius of torus (default: 2.).
    B0 : float
        On-axis (r=0) toroidal magnetic field (default: 10.).
    Bp : float
        Poloidal magnetic field (default: 12.5).
    alpha : float
        (default: 0.1)
    beta : float
        (default: 1.0)

    References
    ----------
    [1] Juan Vicente Gutiérrez-Santacreu, Omar Maj, Marco Restelli: Finite element discretization of a Stokes-like model arising
    in plasma physics, Journal of Computational Physics 2018.
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            \partial_t u = - \nabla \phi + u \times B + \nu \Delta u + f \,,\\
            0 = \nabla \phi- u_e \times B + \nu_e \Delta u_e + f_e \,, \\
            \nabla \cdot (u-u_e) = 0 \,.

        where :math:`f` is defined as follows: 

        .. math::

            f = \left[\begin{array}{c} \alpha \frac{B_0}{a}(R-R_0) - \alpha \frac{1}{a R_0} \frac{R_0 B_0 Z}{R} + \nu \alpha \frac{1}{a R_0} \frac{R_0}{R^2}   \\
                     \alpha \frac{1}{a R_0} (R-R_0) \frac{R_0 B_0}{R} + \alpha \frac{B_0Z}{a} \\
                    \alpha \frac{1}{a R_0} \frac{R_0 B_p}{a R^2} \left( (R-R_0)^2 + Z^2\right)  \end{array} \right] \,, 
            \\[2mm]
            f = \left[\begin{array}{c} -\alpha \frac{B_0}{a}(R-R_0) + \alpha \frac{1}{a R_0} \frac{R_0 B_0 Z}{R} + \nu_e \alpha \frac{1}{a R_0} \frac{R_0}{R^2}   \\
                     -\alpha \frac{1}{a R_0} (R-R_0) \frac{R_0 B_0}{R} - \alpha \frac{B_0 Z}{a} \\
                    -\alpha \frac{1}{a R_0} \frac{ R_0 B_p}{a R^2} \left( (R-R_0)^2 + Z^2\right)  \end{array} \right] \,, 
            \\[2mm]
            R = \sqrt{x^2 + y^2} \,.

        Can only be defined in Cartesian coordinates. 
        The solution is given by:

        .. math::
            \mathbf{u} = \alpha \frac{1}{a R_0} \left[\begin{array}{c} R-R_0 \\ z \\ 0 \end{array} \right]  \,,
            \\[2mm]
            R = \sqrt{x^2 + y^2} \,.
        """


    def __init__(
        self,
        comp=0,
        a=1.0,
        R0=2.0,
        B0=10.0,
        Bp=12.5,
        alpha=0.1,
        beta=1.0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        self._comp = comp
        self._a = a
        self._R0 = R0
        self._B0 = B0
        self._Bp = Bp
        self._alpha = alpha
        self._beta = beta

        # use the setters
        self.given_in_basis = "physical"
        self.comp = comp

    # equilibrium ion velocity
    def __call__(self, x, y, z):
        """Velocity of ions and electrons."""
        R = xp.sqrt(x**2 + y**2)
        R = xp.where(R == 0.0, 1e-9, R)
        phi = xp.arctan2(-y, x)
        A = self._alpha / (self._a * self._R0)
        C = self._beta * self._Bp * self._R0 / (self._B0 * self._a)

        uR = A * (R - self._R0)
        uZ = A * z
        uphi = 0

        # from cylindrical to cartesian:

        if self.comp == 0:
            ux = xp.cos(phi) * uR - R * xp.sin(phi) * uphi
            return ux
        elif self.comp == 1:
            uy = -xp.sin(phi) * uR - R * xp.cos(phi) * uphi
            return uy
        elif self.comp == 2:
            uz = uZ
            return uz
        else:
            raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.")





class TokamakManufacturedSolutionVelocity_1(Perturbation):
    r"""Analytic solution :math:`u=u_e` of the system:

    Parameters
    ----------
    comp : string
        Which component of the solution ('0', '1' or '2').
    a : float
        Minor radius of torus (default: 1.).
    R0 : float
        Major radius of torus (default: 2.).
    B0 : float
        On-axis (r=0) toroidal magnetic field (default: 10.).
    Bp : float
        Poloidal magnetic field (default: 12.5).
    alpha : float
        (default: 0.1)
    beta : float
        (default: 1.0)

    References
    ----------
    [1] Juan Vicente Gutiérrez-Santacreu, Omar Maj, Marco Restelli: Finite element discretization of a Stokes-like model arising
    in plasma physics, Journal of Computational Physics 2018.
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            \partial_t u = - \nabla \phi + u \times B + \nu \Delta u + f \,,\\
            0 = \nabla \phi- u_e \times B + \nu_e \Delta u_e + f_e \,, \\
            \nabla \cdot (u-u_e) = 0 \,.

        where :math:`f` is defined as follows: 

        .. math::

            f = \left[\begin{array}{c} \alpha \frac{B_0}{a}(R-R_0) - \alpha \frac{1}{a R_0} \frac{R_0 B_0 Z}{R} + \nu \alpha \frac{1}{a R_0} \frac{R_0}{R^2}   \\
                     \alpha \frac{1}{a R_0} (R-R_0) \frac{R_0 B_0}{R} + \alpha \frac{B_0Z}{a} \\
                    \alpha \frac{1}{a R_0} \frac{R_0 B_p}{a R^2} \left( (R-R_0)^2 + Z^2\right)  \end{array} \right] \,, 
            \\[2mm]
            f = \left[\begin{array}{c} -\alpha \frac{B_0}{a}(R-R_0) + \alpha \frac{1}{a R_0} \frac{R_0 B_0 Z}{R} + \nu_e \alpha \frac{1}{a R_0} \frac{R_0}{R^2}   \\
                     -\alpha \frac{1}{a R_0} (R-R_0) \frac{R_0 B_0}{R} - \alpha \frac{B_0 Z}{a} \\
                    -\alpha \frac{1}{a R_0} \frac{ R_0 B_p}{a R^2} \left( (R-R_0)^2 + Z^2\right)  \end{array} \right] \,, 
            \\[2mm]
            R = \sqrt{x^2 + y^2} \,.

        Can only be defined in Cartesian coordinates. 
        The solution is given by:

        .. math::
            \mathbf{u} = \alpha \frac{1}{a R_0} \left[\begin{array}{c} R-R_0 \\ z \\ 0 \end{array} \right]  \,,
            \\[2mm]
            R = \sqrt{x^2 + y^2} \,.
        """


    def __init__(
        self,
        comp=0,
        a=1.0,
        R0=2.0,
        B0=10.0,
        Bp=12.5,
        alpha=0.1,
        beta=1.0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        self._comp = comp
        self._a = a
        self._R0 = R0
        self._B0 = B0
        self._Bp = Bp
        self._alpha = alpha
        self._beta = beta

        # use the setters
        self.given_in_basis = "physical"
        self.comp = comp

    # equilibrium ion velocity
    def __call__(self, x, y, z):
        """Velocity of ions and electrons."""
        R = xp.sqrt(x**2 + y**2)
        R = xp.where(R == 0.0, 1e-9, R)
        phi = xp.arctan2(-y, x)
        A = self._alpha / (self._a * self._R0)
        C = self._beta * self._Bp * self._R0 / (self._B0 * self._a)

        uR = A * (R - self._R0)
        uZ = A * z
        uphi = 0

        # from cylindrical to cartesian:

        if self.comp == 0:
            ux = xp.cos(phi) * uR - R * xp.sin(phi) * uphi
            return ux
        elif self.comp == 1:
            uy = -xp.sin(phi) * uR - R * xp.cos(phi) * uphi
            return uy
        elif self.comp == 2:
            uz = uZ
            return uz
        else:
            raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.")





class TokamakManufacturedSolutionVelocity_2(Perturbation):
    r"""Analytic solution :math:`u=u_e` of the system:

    Parameters
    ----------
    comp : string
        Which component of the solution ('0', '1' or '2').
    a : float
        Minor radius of torus (default: 1.).
    R0 : float
        Major radius of torus (default: 2.).
    B0 : float
        On-axis (r=0) toroidal magnetic field (default: 10.).
    Bp : float
        Poloidal magnetic field (default: 12.5).
    alpha : float
        (default: 0.1)
    beta : float
        (default: 1.0)

    References
    ----------
    [1] Juan Vicente Gutiérrez-Santacreu, Omar Maj, Marco Restelli: Finite element discretization of a Stokes-like model arising
    in plasma physics, Journal of Computational Physics 2018.
    """



    @classmethod
    def doc_formula(cls):
        r"""
        .. math::

            \partial_t u = - \nabla \phi + u \times B + \nu \Delta u + f \,,\\
            0 = \nabla \phi- u_e \times B + \nu_e \Delta u_e + f_e \,, \\
            \nabla \cdot (u-u_e) = 0 \,.

        where :math:`f` is defined as follows: 

        .. math::

            f = \left[\begin{array}{c} \alpha \frac{B_0}{a}(R-R_0) - \alpha \frac{1}{a R_0} \frac{R_0 B_0 Z}{R} + \nu \alpha \frac{1}{a R_0} \frac{R_0}{R^2}   \\
                     \alpha \frac{1}{a R_0} (R-R_0) \frac{R_0 B_0}{R} + \alpha \frac{B_0Z}{a} \\
                    \alpha \frac{1}{a R_0} \frac{R_0 B_p}{a R^2} \left( (R-R_0)^2 + Z^2\right)  \end{array} \right] \,, 
            \\[2mm]
            f = \left[\begin{array}{c} -\alpha \frac{B_0}{a}(R-R_0) + \alpha \frac{1}{a R_0} \frac{R_0 B_0 Z}{R} + \nu_e \alpha \frac{1}{a R_0} \frac{R_0}{R^2}   \\
                     -\alpha \frac{1}{a R_0} (R-R_0) \frac{R_0 B_0}{R} - \alpha \frac{B_0 Z}{a} \\
                    -\alpha \frac{1}{a R_0} \frac{ R_0 B_p}{a R^2} \left( (R-R_0)^2 + Z^2\right)  \end{array} \right] \,, 
            \\[2mm]
            R = \sqrt{x^2 + y^2} \,.

        Can only be defined in Cartesian coordinates. 
        The solution is given by:

        .. math::
            \mathbf{u} = \alpha \frac{1}{a R_0} \left[\begin{array}{c} R-R_0 \\ z \\ 0 \end{array} \right]  \,,
            \\[2mm]
            R = \sqrt{x^2 + y^2} \,.
        """


    def __init__(
        self,
        comp=0,
        a=1.0,
        R0=2.0,
        B0=10.0,
        Bp=12.5,
        alpha=0.1,
        beta=1.0,
    ):
        # use setter to store input parameters
        self.params = copy.deepcopy(locals())

        self._comp = comp
        self._a = a
        self._R0 = R0
        self._B0 = B0
        self._Bp = Bp
        self._alpha = alpha
        self._beta = beta

        # use the setters
        self.given_in_basis = "physical"
        self.comp = comp

    # equilibrium ion velocity
    def __call__(self, x, y, z):
        """Velocity of ions and electrons."""
        R = xp.sqrt(x**2 + y**2)
        R = xp.where(R == 0.0, 1e-9, R)
        phi = xp.arctan2(-y, x)
        A = self._alpha / (self._a * self._R0)
        C = self._beta * self._Bp * self._R0 / (self._B0 * self._a)

        uR = A * (R - self._R0)
        uZ = A * z
        uphi = 0

        # from cylindrical to cartesian:

        if self.comp == 0:
            ux = xp.cos(phi) * uR - R * xp.sin(phi) * uphi
            return ux
        elif self.comp == 1:
            uy = -xp.sin(phi) * uR - R * xp.cos(phi) * uphi
            return uy
        elif self.comp == 2:
            uz = uZ
            return uz
        else:
            raise ValueError(f"Invalid component '{self._comp}'. Must be '0', '1', or '2'.")