from __future__ import annotations
from typing import TYPE_CHECKING, Tuple
if TYPE_CHECKING:
from pylbo.data_containers import LegolasDataSet
import numpy as np
from pylbo.utilities.logger import pylboLogger
from pylbo.visualisation.legend_handler import LegendHandler
[docs]CONTINUA_NAMES = {
SLOW_MIN: r"$\Omega_S^-",
SLOW_PLUS: r"$\Omega_S^+",
ALFVEN_MIN: r"$\Omega_A^-",
ALFVEN_PLUS: r"$\Omega_A^+",
THERMAL: r"$\Omega_T",
DOPPLER: r"$\Omega_0",
}
[docs]CONTINUA_COLORS = ["red", "red", "cyan", "cyan", "green", "grey"]
[docs]_DEFAULT_ZERO_TOL = 1e-12
[docs]def _is_zero(values: np.ndarray, zero_tol: float = _DEFAULT_ZERO_TOL) -> bool:
return np.all(np.isclose(values, 0, atol=zero_tol))
[docs]def _is_nonadiabatic(ds: LegolasDataSet) -> bool:
zeroes = np.zeros_like(ds.grid_gauss)
dLdT = ds.equilibria.get("dLdT", zeroes)
dLdrho = ds.equilibria.get("dLdrho", zeroes)
has_cooling = not _is_zero(dLdT) or not _is_zero(dLdrho)
has_conduction = not _is_zero(ds.equilibria.get("kappa_para", zeroes))
return has_cooling or has_conduction
[docs]def _get_parallel_wave_vector(ds: LegolasDataSet) -> float:
bg = ds.equilibria
k2 = ds.parameters.get("k2", 0)
k3 = ds.parameters.get("k3", 0)
B0 = bg["B0"]
eps = ds.scale_factor
return 0 if _is_zero(B0) else (k2 * bg["B02"] / eps + k3 * bg["B03"]) / B0
[docs]def _get_squared_sound_speed(ds: LegolasDataSet) -> np.ndarray:
return ds.gamma * ds.equilibria["T0"]
[docs]def _get_squared_isothermal_sound_speed(ds: LegolasDataSet) -> np.ndarray:
return ds.equilibria["T0"]
[docs]def _get_squared_Alfven_speed(ds: LegolasDataSet) -> np.ndarray:
return ds.equilibria["B0"] ** 2 / ds.equilibria["rho0"]
[docs]def calculate_continua(ds: LegolasDataSet) -> dict:
"""
Calculates the different continua for a given dataset.
The Alfvén and flow continua are always analytical. Depending on the background
and physical effects the slow and thermal continua are either all analytical, or
coupled through a third-order polynomial.
In case of the latter this polynomical is numerically solved through numpy.roots.
Parameters
----------
ds : ~pylbo.data_containers.LegolasDataSet
The Legolas dataset.
Returns
-------
dict, None
Dictionary containing the various continua. The keys are the names of the
continua, the values are the corresponding frequencies as numpy arrays.
Returns `None` if the dataset has no background.
"""
if not ds.has_background:
return None
doppler = get_doppler_shift(ds)
alfven2 = get_squared_alfven_continuum(ds)
slowneg, slowpos, thermal = _get_thermal_and_slow_continua(ds)
continua = {
DOPPLER: doppler,
SLOW_MIN: slowneg,
SLOW_PLUS: slowpos,
THERMAL: thermal,
ALFVEN_MIN: -np.sqrt(alfven2),
ALFVEN_PLUS: np.sqrt(alfven2),
}
# correct for doppler shift
for name in CONTINUA_NAMES.keys():
if name == DOPPLER:
continue
if not _is_zero(continua[name]):
continua[name] += doppler
return continua
[docs]def get_squared_alfven_continuum(ds: LegolasDataSet) -> np.ndarray:
"""
Calculates the squared Alfvén continuum.
Returns
-------
np.ndarray
The squared Alfvén continuum.
"""
bg = ds.equilibria
k2 = ds.parameters.get("k2", 0)
k3 = ds.parameters.get("k3", 0)
eps = ds.scale_factor
return (1 / bg["rho0"]) * (k2 * bg["B02"] / eps + k3 * bg["B03"]) ** 2
[docs]def get_doppler_shift(ds: LegolasDataSet) -> np.ndarray:
"""
Calculates the Doppler shift as the dot product between the wave vector and the
background velocity.
Returns
-------
np.ndarray
The Doppler shift.
"""
bg = ds.equilibria
k2 = ds.parameters.get("k2", 0)
k3 = ds.parameters.get("k3", 0)
eps = ds.scale_factor
return k2 * bg["v02"] / eps + k3 * bg["v03"]
[docs]def get_squared_slow_continuum(ds: LegolasDataSet) -> np.ndarray:
"""
Calculates the squared slow continuum.
Returns
-------
np.ndarray
The squared slow continuum.
"""
bg = ds.equilibria
p = bg["rho0"] * bg["T0"]
alfven_sq = get_squared_alfven_continuum(ds)
return (ds.gamma * p / (ds.gamma * p + bg["B0"] ** 2)) * alfven_sq
[docs]def _get_thermal_and_slow_continua(
ds: LegolasDataSet,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
Calculates the thermal and slow continua.
Returns
-------
Tuple[np.ndarray, np.ndarray, np.ndarray]
A tuple containing the slow-, slow+ and thermal continua, in this order.
"""
zeroes = np.zeros_like(ds.grid_gauss)
bg = ds.equilibria
T0 = bg["T0"]
if _is_zero(T0):
# for pressureless cases (no T) there is no slow/thermal continuum
return (zeroes, zeroes, zeroes)
slow_sq = get_squared_slow_continuum(ds)
if not _is_nonadiabatic(ds):
return (-np.sqrt(slow_sq), np.sqrt(slow_sq), zeroes)
if _is_zero(slow_sq):
# if slow continuum vanishes, thermal continuum is analytical
thermal = _get_thermal_continuum_analytical(ds)
return (zeroes, zeroes, thermal)
# for standard cases we have a third-order polynomial
return _get_slow_and_thermal_continuum_coupled(ds)
[docs]def _get_thermal_continuum_analytical(ds: LegolasDataSet) -> np.ndarray:
"""
Calculates the thermal continuum analytically, when the slow continuum is zero.
Returns
-------
np.ndarray
The thermal continuum.
"""
zeroes = np.zeros_like(ds.grid_gauss)
bg = ds.equilibria
L0 = bg.get("L0", zeroes)
dLdT = bg.get("dLdT", zeroes)
dLdrho = bg.get("dLdrho", zeroes)
rho0 = bg["rho0"]
cs2 = _get_squared_sound_speed(ds)
ca2 = _get_squared_Alfven_speed(ds)
ci2 = _get_squared_isothermal_sound_speed(ds)
gamma_1 = ds.gamma - 1
return 1j * gamma_1 * (L0 + rho0 * dLdrho - (ca2 + ci2) * dLdT) / (cs2 + ca2)
[docs]def _get_slow_and_thermal_continuum_coupled(
ds: LegolasDataSet,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
cs2 = _get_squared_sound_speed(ds)
ca2 = _get_squared_Alfven_speed(ds)
ci2 = _get_squared_isothermal_sound_speed(ds)
kpara = _get_parallel_wave_vector(ds)
gamma_1 = ds.gamma - 1
zeroes = np.zeros_like(ds.grid_gauss)
L0 = ds.equilibria.get("L0", zeroes)
dLdT = ds.equilibria.get("dLdT", zeroes)
dLdrho = ds.equilibria.get("dLdrho", zeroes)
kappa_para = ds.equilibria.get("kappa_para", zeroes)
rho0 = ds.equilibria["rho0"]
# coeffi means the coefficient corresponding to the term omega^i
coeff3 = rho0 * (cs2 + ca2) * 1j / gamma_1
coeff2 = rho0 * (L0 + rho0 * dLdrho) - (kappa_para * kpara**2 + rho0 * dLdT) * (
ci2 + ca2
)
coeff1 = -rho0 * cs2 * ca2 * kpara**2 * 1j / gamma_1
coeff0 = -(
(rho0 * (L0 + rho0 * dLdrho) - (kappa_para * kpara**2 + rho0 * dLdT) * ci2)
* ca2
* kpara**2
)
return _solve_coupled_continuum_polynomial(coeff3, coeff2, coeff1, coeff0)
[docs]def _solve_coupled_continuum_polynomial(
coeff3: np.ndarray, coeff2: np.ndarray, coeff1: np.ndarray, coeff0: np.ndarray
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
Solves the third-order polynomial that couples the slow and thermal continua.
The thermal continuum corresponds to the purely imaginary solution.
Parameters
----------
coeff3 : np.ndarray
The coefficient corresponding to the term :math:`\\omega^3`.
coeff2 : np.ndarray
The coefficient corresponding to the term :math:`\\omega^2`.
coeff1 : np.ndarray
The coefficient corresponding to the term :math:`\\omega^1`.
coeff0 : np.ndarray
The coefficient corresponding to the constant term.
Returns
-------
Tuple[np.ndarray, np.ndarray, np.ndarray]
A tuple containing the slow-, slow+ and thermal continua, in this order.
"""
thermal = np.zeros_like(coeff3, dtype=complex)
slowneg = np.zeros_like(thermal, dtype=complex)
slowpos = np.zeros_like(thermal, dtype=complex)
for i, (c3, c2, c1, c0) in enumerate(zip(coeff3, coeff2, coeff1, coeff0)):
roots = np.roots([c3, c2, c1, c0])
slowneg[i], slowpos[i], thermal[i] = _extract_solutions_from_roots(roots, i)
return slowneg, slowpos, thermal
[docs]def _log_slow_continuum_zero_warning(roots: np.ndarray, i: int):
pylboLogger.warning(
f"encountered index = {i} where the slow continuum has a "
f"real value close to zero. \nFound thermal-slow roots: {roots}"
)
[docs]def _log_assumed_thermal_continuum(root: complex):
pylboLogger.warning(
f"Assuming that the largest imaginary root {root} is the thermal continuum."
)
[docs]class ContinuaHandler(LegendHandler):
"""
Handler to draw continua regions on the plots and make them interactive.
Parameters
----------
interactive : bool
If `True`, makes the legend pickable and continuum plotting interactive.
Attributes
----------
continua_names : list
The list of continua names
"""
def __init__(self, interactive):
super().__init__(interactive)
self.continua_names = list(CONTINUA_NAMES.keys())
self.continua_latex = list(CONTINUA_NAMES.values())
self._continua_colors = CONTINUA_COLORS
self.marker = "."
self.markersize = 6
@property
[docs] def continua_colors(self):
"""
Returns the list of continua colors.
Returns
-------
The continua colors as a list.
"""
return self._continua_colors
@continua_colors.setter
def continua_colors(self, colors):
"""
Setter for the continua colors attribute.
Parameters
----------
colors : list, numpy.ndarray
The colors to use when plotting the continua as a list of strings.
Raises
------
ValueError
If a wrong argument is passed or if it is of improper length.
"""
if colors is None:
return
if not isinstance(colors, (list, np.ndarray)):
raise ValueError(
f"continua_colors should be an array/list but got {type(colors)}"
)
if not len(colors) == len(CONTINUA_COLORS):
raise ValueError(
f"continua_colors should be of length {len(CONTINUA_COLORS)}"
)
self._continua_colors = colors