Improve interface to radas data

.github/workflows/workflow_actions.yml

@@ -25,7 +25,7 @@ jobs:
       - uses: actions/checkout@v4
 
       - name: Install Poetry
-        run: curl -sSL https://install.python-poetry.org | python - --version 1.8.2
+        run: curl -sSL https://install.python-poetry.org | python - --version 2.1.3
 
       - name: Set up Python 3.11
         uses: actions/setup-python@v5
@@ -45,7 +45,7 @@ jobs:
 
       - name: Make radas data
         if: steps.radas.outputs.cache-hit != 'true'
-        run: poetry run radas
+        run: poetry run radas -c radas_config.yaml
 
       - name: Upload radas artifacts
         uses: actions/upload-artifact@v4
@@ -71,7 +71,7 @@ jobs:
         run: sudo apt-get update && sudo apt-get install pandoc
 
       - name: Install Poetry
-        run: curl -sSL https://install.python-poetry.org | python - --version 1.8.2
+        run: curl -sSL https://install.python-poetry.org | python - --version 2.1.3
 
       - name: Set up Python ${{ matrix.python-version }}
         uses: actions/setup-python@v5
@@ -138,7 +138,7 @@ jobs:
       - uses: actions/checkout@v4
 
       - name: Install Poetry
-        run: curl -sSL https://install.python-poetry.org | python - --version 1.8.2
+        run: curl -sSL https://install.python-poetry.org | python - --version 2.1.3
 
       - name: Poetry build
         run: poetry build

Dockerfile

@@ -28,4 +28,4 @@ RUN poetry config virtualenvs.create false
 RUN poetry install --without dev
 
 # 6. Run radas to get the atomic data files
-RUN poetry run radas
+RUN poetry run radas -c radas_config.yaml

cfspopcon/file_io.py

@@ -131,7 +131,7 @@ class _RoundingFloat(float):
     From: https://stackoverflow.com/questions/54370322/how-to-limit-the-number-of-float-digits-jsonencoder-produces
     """
 
-    __repr__ = staticmethod(lambda x: f"{x:#.10g}")  # type:ignore[assignment,unused-ignore]
+    __repr__ = staticmethod(lambda x: f"{x:#.6g}")  # type:ignore[assignment,unused-ignore]
 
 
 class _ModifyJSONFloatRepr:

cfspopcon/formulas/atomic_data/__init__.py

@@ -1,5 +1,6 @@
 """Interface to atomic data files."""
 
 from .atomic_data import AtomicData, read_atomic_data
+from .coeff_interpolator import CoeffInterpolator
 
-__all__ = ["AtomicData", "read_atomic_data"]
+__all__ = ["AtomicData", "CoeffInterpolator", "read_atomic_data"]

cfspopcon/formulas/atomic_data/coeff_interpolator.py

@@ -0,0 +1,138 @@
+"""Class definition for CoeffInterpolator which is used to interpolate atomic data."""
+
+from functools import partial
+
+import numpy as np
+import xarray as xr
+from scipy.interpolate import RectBivariateSpline
+
+from cfspopcon.unit_handling import Quantity, get_units, magnitude_in_units, ureg, wraps_ufunc
+
+
+class CoeffInterpolator(RectBivariateSpline):
+    """An extension of a 2D spline interpolator which handles interpolations in log-space, for interpolating atomic data."""
+
+    tiny = np.finfo(np.float64).tiny
+
+    def log10_with_floor(self, x: xr.DataArray | np.ndarray | float) -> xr.DataArray | np.ndarray | float:
+        """Return the log of x if x > 0, and otherwise return the log of the smallest representable float."""
+        return np.log10(np.maximum(x, self.tiny))  # type: ignore[no-any-return]
+
+    def __init__(
+        self,
+        coeff: xr.DataArray,
+        reference_electron_density: Quantity = 1.0 * ureg.m**-3,
+        reference_electron_temp: Quantity = 1.0 * ureg.eV,
+    ) -> None:
+        """Builds a bivariate spline interpolator for the provided DataArray of coefficients.
+
+        The interpolator performs interpolations on the log of the provided coefficient, for the log of the
+        provided temperature and density, since the atomic data arrays are provided with log-spaced values.
+
+        Parameters:
+        - coeff (xr.DataArray): The xarray DataArray containing the data to interpolate, indexed by electron temperature and density.
+        - reference_electron_density (Quantity): the normalization value for the electron density coordinates.
+        - reference_electron_temp (Quantity): the normalization value for the electron temp coordinates.
+
+        Returns:
+        - CoeffInterpolator: The bivariate spline interpolator object with eval, vector_eval and grid_eval methods.
+        """
+        if not np.all(coeff.dim_electron_density > 0.0):
+            raise ValueError("Encountered null or negative values in electron density for CoeffInterpolator")
+
+        if not np.all(coeff.dim_electron_temp > 0.0):
+            raise ValueError("Encountered null or negative values in electron temp for CoeffInterpolator")
+
+        if not np.all(coeff >= 0.0):
+            raise ValueError("Encountered negative values in coeff for CoeffInterpolator")
+
+        self.units = get_units(coeff)
+        coeff_magnitude = magnitude_in_units(coeff.transpose("dim_electron_temp", "dim_electron_density"), self.units)
+
+        self.electron_density_units = get_units(reference_electron_density)
+        self.electron_temp_units = get_units(reference_electron_temp)
+        # Get the normalization of the electron density and electron temp. Usually the reference values are 1.0 * units.
+        # This prevents an error if there is a scalar prefactor i.e. 2.0 * units
+        electron_density_norm = magnitude_in_units(reference_electron_density, self.electron_density_units)
+        electron_temp_norm = magnitude_in_units(reference_electron_temp, self.electron_temp_units)
+
+        electron_temp_coord = coeff.dim_electron_temp / electron_temp_norm
+        electron_density_coord = coeff.dim_electron_density / electron_density_norm
+
+        self.coeff_is_zero = np.all(coeff == 0.0)
+
+        if self.coeff_is_zero:
+            super().__init__(
+                self.log10_with_floor(electron_temp_coord),  # type: ignore[arg-type]
+                self.log10_with_floor(electron_density_coord),  # type: ignore[arg-type]
+                np.zeros_like(coeff_magnitude),
+            )
+        else:
+            if not np.all(coeff > 0.0):
+                raise ValueError("Encountered mix of null and positive values in coeff for CoeffInterpolator")
+
+            super().__init__(
+                self.log10_with_floor(electron_temp_coord),  # type: ignore[arg-type]
+                self.log10_with_floor(electron_density_coord),  # type: ignore[arg-type]
+                self.log10_with_floor(coeff_magnitude),  # type: ignore[arg-type]
+            )
+
+        self.electron_temp_coord = electron_temp_coord
+        self.electron_density_coord = electron_density_coord
+
+        self.max_temp = np.max(electron_temp_coord).item()
+        self.min_temp = np.min(electron_temp_coord).item()
+        self.max_density = np.max(electron_density_coord).item()
+        self.min_density = np.min(electron_density_coord).item()
+
+        call_config = dict(
+            input_units=dict(
+                electron_density=self.electron_density_units,
+                electron_temp=self.electron_temp_units,
+                allow_extrap=None,
+                grid_from_inputs=None,
+            ),
+            return_units=dict(coeff=self.units),
+            pass_as_kwargs=("allow_extrap", "grid_from_inputs"),
+        )
+
+        self.eval = wraps_ufunc(**call_config)(self.__call__)  # type: ignore[arg-type]
+        self.vector_eval = wraps_ufunc(**call_config)(partial(self.__call__, grid_from_inputs=False))  # type: ignore[arg-type]
+        self.grid_eval = wraps_ufunc(
+            **call_config,  # type: ignore[arg-type]
+            input_core_dims=(("dim_electron_density",), ("dim_electron_temp",)),
+            output_core_dims=(("dim_electron_temp", "dim_electron_density"),),
+        )(partial(self.__call__, grid_from_inputs=True))
+
+    def __call__(  # type: ignore[override]
+        self,
+        electron_density: float | np.ndarray,
+        electron_temp: float | np.ndarray,
+        allow_extrap: bool = False,
+        grid_from_inputs: bool = False,
+    ) -> np.ndarray:
+        """Interpolate the atomic data to the requested point, handling the log-spaced interpolation but not handling units.
+
+        For unit-aware interpolation, use the eval, vector_eval or grid_eval methods which are dynamically constructed in the __init__ function.
+
+        NOTE: Unlike RectBivariateSpline, the "grid_from_inputs" option (mapping to 'grid' in RectBivariateSpline) is false by default.
+        """
+        if allow_extrap:
+            electron_temp = np.minimum(electron_temp, self.max_temp)
+            electron_temp = np.maximum(electron_temp, self.min_temp)
+            electron_density = np.minimum(electron_density, self.max_density)
+            electron_density = np.maximum(electron_density, self.min_density)
+        else:
+            assert (
+                (np.min(electron_density) >= self.min_density)
+                & (np.max(electron_density) <= self.max_density)
+                & (np.min(electron_temp) >= self.min_temp)
+                & (np.max(electron_temp) <= self.max_temp)
+            ), "CoeffInterpolator error: values off grid and allow_extrap=False. "
+
+        values = super().__call__(np.log10(electron_temp), np.log10(electron_density), grid=grid_from_inputs)
+
+        if self.coeff_is_zero:
+            return values
+        else:
+            return np.power(10, values)

cfspopcon/formulas/impurities/core_radiator_conc.py

@@ -5,8 +5,10 @@
 
 from ... import named_options
 from ...algorithm_class import Algorithm
+from ...helpers import get_item
 from ...unit_handling import Unitfull
 from .. import radiated_power
+from ..atomic_data import AtomicData
 from .impurity_array_helpers import extend_impurity_concentration_array, make_impurity_concentration_array
 
 
@@ -80,7 +82,7 @@ def calc_core_seeded_impurity_concentration(
     radiated_power_method: named_options.RadiationMethod,
     radiated_power_scalar: Unitfull,
     core_impurity_species: named_options.AtomicSpecies,
-    atomic_data: xr.DataArray,
+    atomic_data: xr.DataArray | AtomicData,
 ) -> Unitfull:
     """Calculate the concentration of a core radiator required to increase the radiated power by a desired amount.
 
@@ -110,7 +112,7 @@ def calc_core_seeded_impurity_concentration(
         electron_temp_profile=electron_temp_profile,
         electron_density_profile=electron_density_profile,
         plasma_volume=plasma_volume,
-        atomic_data=atomic_data.item(),
+        atomic_data=get_item(atomic_data),
     )
 
     P_radiated_per_unit_concentration = radiated_power_scalar * radiated_power.impurity_radiated_power.calc_impurity_radiated_power(

cfspopcon/formulas/impurities/edge_radiator_conc.py

@@ -4,11 +4,12 @@
 
 import numpy as np
 import xarray as xr
-from scipy.interpolate import InterpolatedUnivariateSpline  # type:ignore[import-untyped]
+from scipy.interpolate import InterpolatedUnivariateSpline
 
 from ...algorithm_class import Algorithm
+from ...helpers import get_item
 from ...named_options import AtomicSpecies
-from ...unit_handling import Unitfull, convert_units, magnitude, ureg, wraps_ufunc
+from ...unit_handling import Unitfull, magnitude, magnitude_in_units, ureg, wraps_ufunc
 from ..atomic_data import AtomicData
 from .impurity_array_helpers import extend_impurity_concentration_array
 
@@ -31,7 +32,7 @@ def calc_edge_impurity_concentration(
     lengyel_overestimation_factor: Unitfull,
     edge_impurity_enrichment: Unitfull,
     impurity_concentration: xr.DataArray,
-    atomic_data: xr.DataArray,
+    atomic_data: xr.DataArray | AtomicData,
     reference_electron_density: Unitfull = 1.0 * ureg.n20,
     reference_ne_tau: Unitfull = 1.0 * ureg.n20 * ureg.ms,
 ) -> tuple[Unitfull, ...]:
@@ -56,7 +57,7 @@ def calc_edge_impurity_concentration(
         :term:`edge_impurity_concentration`, :term:`edge_impurity_concentration_in_core`, :term:`impurity_concentration`
     """
     L_int_integrator = build_L_int_integrator(
-        atomic_data=atomic_data.item(),
+        atomic_data=get_item(atomic_data),
         impurity_species=edge_impurity_species,
         reference_electron_density=reference_electron_density,
         reference_ne_tau=reference_ne_tau,
@@ -101,22 +102,14 @@ def build_L_int_integrator(
     Returns:
         L_int_integrator
     """
-    if isinstance(impurity_species, xr.DataArray):
-        impurity_species = impurity_species.item()
+    Lz_curve = atomic_data.get_noncoronal_Lz_interpolator(get_item(impurity_species), ne_tau=reference_ne_tau)
 
-    electron_density_ref = magnitude(convert_units(reference_electron_density, ureg.m**-3))
-    ne_tau_ref = magnitude(convert_units(reference_ne_tau, ureg.m**-3 * ureg.s))
+    _electron_temp = Lz_curve.electron_temp_coord
+    _electron_density = magnitude_in_units(reference_electron_density, Lz_curve.electron_density_units)
+    assert np.size(_electron_density) == 1
 
-    Lz_curve = (
-        atomic_data.get_dataset(impurity_species)
-        .equilibrium_Lz.sel(dim_electron_density=electron_density_ref, method="nearest", tolerance=1e-6 * electron_density_ref)
-        .sel(dim_ne_tau=ne_tau_ref, method="nearest", tolerance=1e-6 * ne_tau_ref)
-    )
-
-    electron_temp = Lz_curve.dim_electron_temp
-    Lz_sqrt_Te = Lz_curve * np.sqrt(electron_temp)
-
-    interpolator = InterpolatedUnivariateSpline(electron_temp, magnitude(Lz_sqrt_Te))
+    Lz_sqrt_Te = Lz_curve(electron_density=_electron_density, electron_temp=_electron_temp) * np.sqrt(_electron_temp)
+    interpolator = InterpolatedUnivariateSpline(_electron_temp, magnitude(Lz_sqrt_Te))  # type: ignore[arg-type]
 
     def L_int(start_temp: float, stop_temp: float) -> float:
         integrated_Lz: float = interpolator.integral(start_temp, stop_temp)

cfspopcon/formulas/impurities/impurity_charge_state.py

@@ -4,8 +4,8 @@
 import xarray as xr
 
 from ...algorithm_class import Algorithm
-from ...named_options import AtomicSpecies
-from ...unit_handling import Unitfull, ureg, wraps_ufunc
+from ...helpers import get_item
+from ...unit_handling import Unitfull
 from ..atomic_data import AtomicData
 
 
@@ -14,7 +14,7 @@ def calc_impurity_charge_state(
     average_electron_density: Unitfull,
     average_electron_temp: Unitfull,
     impurity_concentration: xr.DataArray,
-    atomic_data: xr.DataArray,
+    atomic_data: AtomicData | xr.DataArray,
 ) -> Unitfull:
     """Calculate the impurity charge state for each species in impurity_concentration.
 
@@ -27,49 +27,15 @@ def calc_impurity_charge_state(
     Returns:
         :term:`impurity_charge_state`
     """
-    if isinstance(atomic_data, xr.DataArray):
-        atomic_data = atomic_data.item()
+    atomic_data = get_item(atomic_data)
 
-    return _calc_impurity_charge_state(average_electron_density, average_electron_temp, impurity_concentration.dim_species, atomic_data)
+    def calc_mean_z(impurity_concentration: xr.DataArray) -> xr.DataArray:
+        species = get_item(impurity_concentration.dim_species)
+        interpolator = atomic_data.get_coronal_Z_interpolator(species)
+        mean_z = interpolator.eval(electron_density=average_electron_density, electron_temp=average_electron_temp, allow_extrap=True)
 
+        mean_z = np.minimum(mean_z, atomic_data.datasets[species].atomic_number)
+        mean_z = np.maximum(mean_z, 0.0)
+        return mean_z  # type:ignore[no-any-return]
 
-@wraps_ufunc(
-    return_units=dict(mean_charge_state=ureg.dimensionless),
-    input_units=dict(
-        average_electron_density=ureg.m**-3,
-        average_electron_temp=ureg.eV,
-        impurity_species=None,
-        atomic_data=None,
-    ),
-    pass_as_kwargs=("atomic_data",),
-)
-def _calc_impurity_charge_state(
-    average_electron_density: float,
-    average_electron_temp: float,
-    impurity_species: AtomicSpecies,
-    atomic_data: AtomicData,
-) -> float:
-    """Calculate the impurity charge state of the specified impurity species.
-
-    Args:
-        average_electron_density: [m^-3] :term:`glossary link<average_electron_density>`
-        average_electron_temp: [eV] :term:`glossary link<average_electron_temp>`
-        impurity_species: [] :term:`glossary link<impurity_species>`
-        atomic_data: :term:`glossary link<atomic_data>`
-
-    Returns:
-        :term:`impurity_charge_state`
-    """
-    average_electron_temp, average_electron_density = atomic_data.nearest_neighbour_off_grid(  # type:ignore[assignment]
-        impurity_species,
-        average_electron_temp,
-        average_electron_density,
-    )
-    interpolator = atomic_data.coronal_Z_interpolators[impurity_species]
-    interpolated_values = np.power(10, interpolator((np.log10(average_electron_temp), np.log10(average_electron_density))))
-
-    atomic_number = atomic_data.datasets[impurity_species].atomic_number
-
-    interpolated_values = np.minimum(interpolated_values, atomic_number)
-    interpolated_values = np.maximum(interpolated_values, 0)
-    return interpolated_values  # type:ignore[no-any-return]
+    return impurity_concentration.groupby("dim_species").map(calc_mean_z)