Source code for torx.specializations.grillix.atomic_data_m

"""Allows interfacing with GRILLIX atomic data."""

from numba import njit
import numpy as np
from pathlib import Path
import torx
import xarray as xr
from torx.autodoc_decorators_m import autodoc_class

@njit
def poly8x8_eval(pc, x, y):
    """Return double polynomial of degree 8x8 evaluated at x and y."""
    pcx = np.zeros(9)
    for m in range(9):
        # collapse to single polynomial
        pcx[m] = poly8_eval(pc[m, :], y)

    result = poly8_eval(pcx, x)

    return result


@njit
def poly8_eval(pc, x):
    """Return polynomial of degree 8 evaluated at x and y."""
    val = 0.0
    for n in range(8, 0, -1):
        val = (val + pc[n]) * x

    val = val + pc[0]
    return val



@autodoc_class
class AtomicData:
    """Reads in and evaluates rate coefficients for neutrals."""



    def __init__(self, run_filepath, neutrals_params_path="params_neutrals.in"):
        """Initialize the atomic data."""
        self.filepath = Path(run_filepath)

        assert (self.filepath / neutrals_params_path).exists()

        self.params = torx.fileio.read_fortran_namelist(
            self.filepath / neutrals_params_path
        )
        self.ne_ref = self.params["neut_model"]["ne_ref"]
        self.te_ref = self.params["neut_model"]["te_ref"]
        self.k_ref = self.params["neut_model"]["k_ref"]
        self.w_ref = self.params["neut_model"]["w_ref"]
        self.s_cx = self.params["neut_model"]["s_cx"]
        self.mach_limit = self.params["neut_model"]["mach_limit"]
        self.pot_iz = self.params["neut_model"]["pot_iz"]

        self._k_iz = np.loadtxt(
            (
                self.filepath / self.params["neut_data_paths"]["path_k_iz"]
            ).absolute()
        )
        self._k_rec = np.loadtxt(
            (
                self.filepath / self.params["neut_data_paths"]["path_k_rec"]
            ).absolute()
        )
        self._w_iz = np.loadtxt(
            (
                self.filepath / self.params["neut_data_paths"]["path_w_iz"]
            ).absolute()
        )
        self._p_rec = np.loadtxt(
            (
                self.filepath / self.params["neut_data_paths"]["path_p_rec"]
            ).absolute()
        )


    def _evaluate_rate_coeff(
        self, ne: np.ndarray, te: np.ndarray, rate: np.ndarray
    ) -> np.ndarray:
        """Evaluate the normalized rate coefficient."""
        # Convert inputs to log values, as required for fit (see AMJUEL manual)

        te_fit = np.log(te * self.te_ref)
        ne_fit = np.log(
            ne * self.ne_ref * 1.0e-8
        )  # apply AMJUEL electron density scaling

        return np.exp(poly8x8_eval(rate, te_fit, ne_fit)) / self.k_ref



    def evaluate_rate_coeff(
        self, density: xr.DataArray, electron_temp: xr.DataArray, rate: np.ndarray
    ) -> xr.DataArray:
        """Evaluate the normalized rate coefficient (parallel)."""
        return xr.apply_ufunc(
            self._evaluate_rate_coeff,
            density,
            electron_temp,
            kwargs=dict(rate=rate),
            vectorize=True,
            dask="parallelized",
        )




    def k_iz(self, density: xr.DataArray, electron_temp: xr.DataArray) -> xr.DataArray:
        """Evaluate the k_iz rate coefficient."""
        return self.evaluate_rate_coeff(density, electron_temp, rate=self._k_iz)




    def k_rec(self, density: xr.DataArray, electron_temp: xr.DataArray) -> xr.DataArray:
        """Evaluate the k_rec rate coefficient."""
        return self.evaluate_rate_coeff(density, electron_temp, rate=self._k_rec)




    def w_iz(self, density: xr.DataArray, electron_temp: xr.DataArray) -> xr.DataArray:
        """Evaluate the w_iz rate coefficient."""
        return self.evaluate_rate_coeff(density, electron_temp, rate=self._w_iz)




    def p_rec(self, density: xr.DataArray, electron_temp: xr.DataArray) -> xr.DataArray:
        """Evaluate the p_rec rate coefficient."""
        return self.evaluate_rate_coeff(density, electron_temp, rate=self._p_rec)