Source code for torx.specializations.reax.rate_coefficient_m

"""
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 ... import reax_resources_dir
from torx import Quantity
from torx.autodoc_decorators_m import autodoc_class
from torx.fileio import filepath_resolver



@autodoc_class
class RateCoefficient:
    """Manages rate coefficient calculations based on AMJUEL."""



    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 = xr.DataArray(
            np.loadtxt(self.file_path, max_rows=9),
            dims=["T", "E"],
            coords={"T": np.arange(9), "E": np.arange(9)}
        )





    def compute(self,
                n: xr.DataArray,
                t: xr.DataArray,
                norm=None
    ) -> xr.DataArray:
        """Evaluate the normalized rate coefficient."""
        if isinstance(n, xr.DataArray) and isinstance(t, xr.DataArray):
            return self._compute_xarray(n, t, norm)

        # Raise an error for unsupported types
        raise TypeError("Unsupported types for compute. \
                         Expected xarray.DataArray.")


    def _compute_xarray(self,
                        ne: xr.DataArray,
                        te: xr.DataArray,
                        norm=None
    ) -> xr.DataArray:
        """Return rate coefficient as xarray."""
        # Set the rate normalization for the rate coefficient instance
        if norm:
            try:
                k_norm = norm.c_s0 / (norm.n0 * norm.R0)
            except AttributeError as e:
                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

        # 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_values = np.exp(rate_values) / rate_norm.magnitude
        rate_values.attrs["norm"] = rate_norm
        rate_values.attrs["description"] = self.file_name

        return rate_values