# Source code for plasmapy.dispersion.analytical.two_fluid_

"""
Functionality for calculating various analytical solutions to the two
fluid dispersion relation.
"""

__all__ = ["two_fluid"]

import warnings
from numbers import Real

import astropy.units as u
import numpy as np
from astropy.constants.si import c

from plasmapy.formulary.frequencies import gyrofrequency, plasma_frequency
from plasmapy.formulary.speeds import Alfven_speed, ion_sound_speed
from plasmapy.particles import ParticleLike, particle_input
from plasmapy.utils.decorators import validate_quantities
from plasmapy.utils.exceptions import PhysicsWarning

[docs]
@particle_input
@validate_quantities(
B={"can_be_negative": False},
n_i={"can_be_negative": False},
T_e={"can_be_negative": False, "equivalencies": u.temperature_energy()},
T_i={"can_be_negative": False, "equivalencies": u.temperature_energy()},
)
def two_fluid(
B: u.Quantity[u.T],
ion: ParticleLike,
k: u.Quantity[u.rad / u.m],
n_i: u.Quantity[u.m**-3],
*,
T_e: u.Quantity[u.K],
T_i: u.Quantity[u.K],
gamma_e: float = 1,
gamma_i: float = 3,
mass_numb: int | None = None,
Z: float | None = None,
):
r"""
Using the solution provided by :cite:t:bellan:2012, calculate the
analytical solution to the two fluid, low-frequency
(:math:ω/kc ≪ 1) dispersion relation presented by
:cite:t:stringer:1963.  This dispersion relation also assumes a
uniform magnetic field :math:\mathbf{B}_0, no D.C. electric field
:math:\mathbf{E}_0=0, and quasi-neutrality.  For more information
see the **Notes** section below.

Parameters
----------
B : ~astropy.units.Quantity
The magnetic field magnitude in units convertible to T.

ion : |particle-like|
Representation of the ion species (e.g., 'p' for protons,
'D+' for deuterium, 'He-4 1+' for singly ionized
helium-4, etc.).

k : ~astropy.units.Quantity
Wavenumber in units convertible to rad/m. May be either single
valued or a 1D array of length :math:N.

n_i : ~astropy.units.Quantity
Ion number density in units convertible to m\ :sup:-3.

theta : ~astropy.units.Quantity
The angle of propagation of the wave with respect to the
magnetic field, :math:\cos^{-1}(k_z / k), in units must be
convertible to radians. May be either single valued or a 1D
array of size :math:M.

T_e : ~astropy.units.Quantity, |keyword-only|
The electron temperature in units of K or eV.

T_i : ~astropy.units.Quantity, |keyword-only|
The ion temperature in units of K or eV.

gamma_e : float or int, |keyword-only|, default: 1
The adiabatic index for electrons.  The default value assumes
that the electrons are able to equalize their temperature
rapidly enough that the electrons are effectively isothermal.

gamma_i : float or int, |keyword-only|, default: 3
The adiabatic index for ions. The default value assumes that ion
motion has only one degree of freedom, namely along magnetic
field lines.

mass_numb : int, |keyword-only|, optional
The mass number of an isotope corresponding to ion.

Z : real number, |keyword-only|, optional
The |charge number| corresponding to ion.

Returns
-------
omega : Dict[str, ~astropy.units.Quantity]
A dictionary of computed wave frequencies in units rad/s.  The
dictionary contains three keys: 'fast_mode' for the fast
mode, 'alfven_mode' for the Alfvén mode, and
'acoustic_mode' for the ion-acoustic mode.  The value for
each key will be a :math:N × M array.

Raises
------
TypeError
If applicable arguments are not instances of
~astropy.units.Quantity or cannot be converted into one.

|ParticleError|
If ion is not |particle-like|.

TypeError
If gamma_e, gamma_i, or Z are not a real number.

~astropy.units.UnitTypeError
If applicable arguments do not have units convertible to the
expected units.

ValueError
If any of B, k, n_i, T_e, or T_i is
negative.

ValueError
If k is negative or zero.

ValueError
If ion is not of category ion or element.

ValueError
If B, n_i, T_e, or T_I are not single valued
astropy.units.Quantity (i.e. an array).

ValueError
If k or theta are not single valued or a 1D array.

Warns
-----
: ~plasmapy.utils.exceptions.PhysicsWarning
When the computed wave frequencies violate the low-frequency
(:math:ω/kc ≪ 1) assumption of the dispersion relation.

Notes
-----
The complete dispersion equation presented by :cite:t:stringer:1963
(equation 1 of :cite:t:bellan:2012) is:

.. math::
\left( \cos^2 θ - Q \frac{ω^2}{k^2 {v_A}^2} \right) &
\left[
\left( \cos^2 θ - \frac{ω^2}{k^2 {c_s}^2} \right)
- Q \frac{ω^2}{k^2 {v_A}^2} \left(
1 - \frac{ω^2}{k^2 {c_s}^2}
\right)
\right] \\
&= \left(1 - \frac{ω^2}{k^2 {c_s}^2} \right)
\frac{ω^2}{{ω_{ci}}^2} \cos^2 θ

where

.. math::
Q &= 1 + k^2 c^2/{ω_{pe}}^2 \\
\cos θ &= \frac{k_z}{k} \\
\mathbf{B}_0 &= B_0 \mathbf{\hat{z}}

:math:ω is the wave frequency, :math:k is the wavenumber,
:math:v_A is the Alfvén velocity, :math:c_s is the sound speed,
:math:ω_{ci} is the ion gyrofrequency, and
:math:ω_{pe} is the electron plasma frequency. This relation
does additionally assume low-frequency waves
:math:ω/kc ≪ 1, no D.C. electric field
:math:\mathbf{E}_0=0 and quasi-neutrality.

Following section 5 of :cite:t:bellan:2012, the exact roots of the
above dispersion equation can be derived and expressed as one
analytical solution (equation 38 of :cite:t:bellan:2012):

.. math::
\frac{ω}{ω_{ci}} = \sqrt{
2 Λ \sqrt{-\frac{P}{3}} \cos\left(
\frac{1}{3} \cos^{-1}\left(
\frac{3q}{2p} \sqrt{-\frac{3}{p}}
\right)
- \frac{2 \pi}{3}j
\right)
+ \frac{Λ A}{3}
}

where :math:j = 0 represents the fast mode, :math:j = 1
represents the Alfvén mode, and :math:j = 2 represents the

.. math::
p &= \frac{3B-A^2}{3} \; , \; q = \frac{9AB-2A^3-27C}{27} \\
A &= \frac{Q + Q^2 β + Q α + α Λ}{Q^2} \;
, \; B = α \frac{1 + 2 Q β + Λ β}{Q^2} \;
, \; C = \frac{α^2 β}{Q^2} \\
α &= \cos^2 θ \;
, \; β = \left( \frac{c_s}{v_A}\right)^2 \;
, \; Λ = \left( \frac{k v_{A}}{ω_{ci}}\right)^2

Examples
--------
>>> import astropy.units as u
>>> from plasmapy.dispersion.analytical import two_fluid
>>> inputs = {
...     "k": 0.01 * u.rad / u.m,
...     "theta": 30 * u.deg,
...     "B": 8.3e-9 * u.T,
...     "n_i": 5e6 * u.m ** -3,
...     "T_e": 1.6e6 * u.K,
...     "T_i": 4.0e5 * u.K,
...     "ion": "p+",
... }
>>> omegas = two_fluid(**inputs)
>>> omegas
{'fast_mode': <Quantity 1520.57... rad / s>,
'alfven_mode': <Quantity 1261.75... rad / s>,
'acoustic_mode': <Quantity 0.688152... rad / s>}

>>> inputs = {
...     "k": [1e-7, 2e-7] * u.rad / u.m,
...     "theta": [10, 20] * u.deg,
...     "B": 8.3e-9 * u.T,
...     "n_i": 5e6 * u.m ** -3,
...     "T_e": 1.6e6 * u.K,
...     "T_i": 4.0e5 * u.K,
...     "ion": "He+",
... }
>>> omegas = two_fluid(**inputs)
>>> omegas['fast_mode']
<Quantity [[0.00767..., 0.00779... ],
[0.01534..., 0.01558...]] rad / s>
"""

# validate arguments
for arg_name in ("B", "n_i", "T_e", "T_i"):
val = locals()[arg_name].squeeze()
if val.shape != ():
raise ValueError(
f"Argument '{arg_name}' must a single value and not an array of "
f"shape {val.shape}."
)
locals()[arg_name] = val

# validate arguments
for arg_name in ("gamma_e", "gamma_i"):
if not isinstance(locals()[arg_name], Real):
raise TypeError(
f"Expected int or float for argument '{arg_name}', but got "
f"{type(locals()[arg_name])}."
)

# validate argument k
k = k.squeeze()
if k.ndim not in (0, 1):
raise ValueError(
f"Argument 'k' needs to be a single valued or 1D array astropy Quantity,"
f" got array of shape {k.shape}."
)
if np.any(k <= 0):
raise ValueError("Argument 'k' cannot be a or have negative values.")

# validate argument theta
theta = theta.squeeze()
if theta.ndim not in (0, 1):
raise ValueError(
f"Argument 'theta' needs to be a single valued or 1D array astropy "
f"Quantity, got array of shape {k.shape}."
)

# Calc needed plasma parameters
n_e = n_i * ion.charge_number
with warnings.catch_warnings():
warnings.simplefilter("ignore", category=PhysicsWarning)
c_s = ion_sound_speed(
T_e=T_e,
T_i=T_i,
ion=ion,
n_e=n_e,
gamma_e=gamma_e,
gamma_i=gamma_i,
Z=Z,
)
v_A = Alfven_speed(B, n_i, ion=ion)
omega_ci = gyrofrequency(B=B, particle=ion, signed=False)
omega_pe = plasma_frequency(n=n_e, particle="e-")

# Bellan2012JGR params equation 32
alpha = np.cos(theta.value) ** 2
beta = (c_s / v_A).to(u.dimensionless_unscaled).value ** 2
alphav, kv = np.meshgrid(alpha, k.value)  # create grid
Lambda = (kv * v_A.value / omega_ci.value) ** 2

# Bellan2012JGR params equation 2
Q = 1 + (kv * c.value / omega_pe.value) ** 2

# Bellan2012JGR params equation 35
A = ((1 + alphav) / Q) + beta + (alphav * Lambda / Q**2)
B = alphav * (1 + 2 * Q * beta + Lambda * beta) / Q**2
C = beta * (alphav / Q) ** 2

# Bellan2012JGR params equation 36
p = (3 * B - A**2) / 3
q = (9 * A * B - 2 * A**3 - 27 * C) / 27

# Bellan2012JGR params equation 38
R = 2 * Lambda * np.emath.sqrt(-p / 3)
S = 3 * q / (2 * p) * np.emath.sqrt(-3 / p)
T = Lambda * A / 3
omega = {}
for ind, wave_mode in enumerate(("fast_mode", "alfven_mode", "acoustic_mode")):
# The solution corresponding to equation 38
ω = omega_ci * np.emath.sqrt(
R * np.cos(1 / 3 * np.emath.arccos(S) - 2 * np.pi / 3 * ind) + T
)
omega[wave_mode] = ω.squeeze()

# check for violation of dispersion relation assumptions
# (i.e. low-frequency, ω/kc << 0.1)
wkc_max = np.max(ω.value / (kv * c.value))
if wkc_max > 0.1:
warnings.warn(
f"The {wave_mode} calculation produced a high-frequency wave (ω/kc == "
f"{wkc_max:.3f}), which violates the low-frequency (ω/kc << 1) "
f"assumption of the dispersion relation.",
PhysicsWarning,
)

return omega