"""
Base class for calculating rates based on the AMJUEL database.
Contains ionization / recombination reaction and cooling rates.
The RateCoefficient class mimics the REAX module rate_coefficients_m
to perform the same computation a posteriori, i.e. based on the xr.DataArray
that stores the GENE-X and GRILLIX simulation data.
"""
import numpy as np
import xarray as xr
from pathlib import Path
from numba import njit, prange
from ... import reax_resources_dir
from torx.units import Quantity
from torx.decorators import autodoc_class
from torx.fileio import filepath_resolver
@njit(parallel=True, fastmath=True)
def _rate_values_numba(ne, te, coeff, ne_norm_value, te_norm_value):
out = np.empty_like(ne)
for p in prange(ne.size):
ne_fit = np.log(ne[p] * ne_norm_value * 1.0e-8)
te_fit = np.log(te[p] * te_norm_value)
val = 0.0
for i in range(8, -1, -1):
row = 0.0
for j in range(8, -1, -1):
row = row * ne_fit + coeff[i, j]
val = val * te_fit + row
out[p] = np.exp(val)
return out
[docs]
@autodoc_class
class RateCoefficient:
"""Manages rate coefficient calculations based on AMJUEL."""
[docs]
def __init__(self,
coeff_type: int,
file_name: str,
folder_path: Path = None
):
"""
Initialize the RateCoefficient object.
Given the coeff_type (0 for reaction or 1 for cooling).
By default, the file stored in TorX is used, but one can use another
file by additionally specifying a folder_path.
"""
if coeff_type not in (0, 1):
raise ValueError("Attribute 'coeff_type' must be 0 (reaction rate) \
or 1 (cooling rate).")
self.coeff_type = coeff_type
# Load fit coefficient
if folder_path is None:
folder_path = Path(reax_resources_dir)
self.file_path = Path(filepath_resolver(folder_path, file_name))
self.file_name = file_name
# AMJUEL fit have T-index as row and E-index as column
self.fit_coeff_np = np.loadtxt(self.file_path, max_rows=9)
self.fit_coeff = xr.DataArray(
self.fit_coeff_np,
dims=["T", "E"],
coords={"T": np.arange(9), "E": np.arange(9)}
)
[docs]
def compute(self,
n: xr.DataArray,
t: xr.DataArray,
norm=None,
use_numba=True,
) -> xr.DataArray:
"""Evaluate the normalized rate coefficient."""
if isinstance(n, xr.DataArray) and isinstance(t, xr.DataArray):
if use_numba:
return self._compute_numba(n, t, norm)
else:
return self._compute(n, t, norm)
# Raise an error for unsupported types
raise TypeError("Unsupported types for compute. \
Expected xarray.DataArray.")
def _get_raw_rate_norm(self):
"""Return rate normalization of the raw AMJUEL data."""
if self.coeff_type == 0:
return Quantity(1.0, "cm^3/s")
elif self.coeff_type == 1:
return Quantity(1.0, "eV * cm^3/s")
def _get_rate_norm(self,
ne: xr.DataArray,
te: xr.DataArray,
norm=None
):
"""Return rate normalization of the input data."""
if norm is not None:
try:
k_norm = norm.c_s0 / (norm.n0 * norm.R0)
except AttributeError:
print("Incompatible norm provided, \
using default k_norm (cm^3/s).")
k_norm = Quantity(1.0, "cm^3/s")
else:
print("No norm provided, using default k_norm (cm^3/s).")
k_norm = Quantity(1.0, "cm^3/s")
# Extract normalization constants from attribute
if hasattr(ne, "norm"):
ne_norm = ne.norm
elif norm is not None:
ne_norm = norm.n0
else:
ne_norm = Quantity(1.0, "1/cm^3")
if hasattr(te, "norm"):
te_norm = te.norm
elif norm is not None:
te_norm = norm.Te0
else:
te_norm = Quantity(1.0, "eV")
if self.coeff_type == 0:
rate_norm = k_norm
elif self.coeff_type == 1:
rate_norm = te_norm * k_norm
return ne_norm, te_norm, rate_norm
def _compute_numba(self,
ne: xr.DataArray,
te: xr.DataArray,
norm=None
) -> xr.DataArray:
"""
Return evaluated rate coefficient as xarray.
Employs raw numpy with numba rather than dask-parallelized
xarrays to improve performance on repeated computations.
"""
ne_norm, te_norm, rate_norm = self._get_rate_norm(ne, te, norm)
# Convert norms to magnitudes for computation
ne_norm_value = ne_norm.to("1/cm^3").magnitude
te_norm_value = te_norm.to("eV").magnitude
# Only broadcast if needed
if ne.dims != te.dims or ne.shape != te.shape:
ne, te = xr.broadcast(ne, te)
sh = ne.shape
ne_np = np.asarray(ne.to_numpy()).ravel()
te_np = np.asarray(te.to_numpy()).ravel()
if ne_np.shape != te_np.shape:
raise ValueError("ne and te must have the same shape after \
broadcasting.")
te_fit = np.log(te_np * te_norm_value)
ne_fit = np.log(ne_np * ne_norm_value * 1.0e-8)
# Numba accelerated version
rate_values = _rate_values_numba(
ne_np,
te_np,
self.fit_coeff_np,
ne_norm_value,
te_norm_value
)
# Convert the raw output to the same unit as rate_norm
rate_norm_raw = self._get_raw_rate_norm().to(rate_norm.units)
rate_values *= (rate_norm_raw / rate_norm).magnitude
rate_values = xr.DataArray(
rate_values.reshape(sh),
coords=ne.coords,
dims=ne.dims
)
rate_values.attrs["norm"] = rate_norm
rate_values.attrs["description"] = self.file_name
return rate_values
def _compute(self,
ne: xr.DataArray,
te: xr.DataArray,
norm=None
) -> xr.DataArray:
"""
Return evaluated rate coefficient as xarray, dask-parallelized variant.
The output of this method is identical to _compute_numba.
"""
ne_norm, te_norm, rate_norm = self._get_rate_norm(ne, te, norm)
# Convert norms to magnitudes for computation
ne_norm_value = ne_norm.to("1/cm^3").magnitude
te_norm_value = te_norm.to("eV").magnitude
te_fit = np.log(te * te_norm_value)
ne_fit = np.log(ne * ne_norm_value * 1.0e-8)
# dask-compatible equivalent of
# np.polynomial.polynomial.polyval2d(te_fit, ne_fit, self.fit)
i_vals = np.arange(9)
j_vals = np.arange(9)
n_powers = xr.concat([ne_fit ** i for i in i_vals], dim="E")
t_powers = xr.concat([te_fit ** j for j in j_vals], dim="T")
rate_values = xr.dot(self.fit_coeff, n_powers, dims="E")
rate_values = xr.dot(rate_values, t_powers, dims="T")
rate_raw = np.exp(rate_values)
# Convert the raw output to the same unit as rate_norm
rate_norm_raw = self._get_raw_rate_norm()
conversion_factor = (rate_norm_raw.to(rate_norm.units)
/ rate_norm).magnitude
rate_values = rate_raw * conversion_factor
rate_values.attrs["norm"] = rate_norm
rate_values.attrs["description"] = self.file_name
return rate_values
