Introduce IonMixture dataclass.

First step in API change for plasma composition where ion attributes are set from ion names.

The IonMixture dataclass is used for bundled isotope mixtures (e.g. DT) and bundled impurities (e.g. of similar Z).

No used yet. Will be implemented in upcoming PRs.

PiperOrigin-RevId: 714119004

torax/config/plasma_composition.py

@@ -16,18 +16,78 @@
 
 from __future__ import annotations
 
+from collections.abc import Mapping
 import dataclasses
 import logging
 
 import chex
+import numpy as np
 from torax import array_typing
+from torax import constants
 from torax import interpolated_param
 from torax.config import base
 from torax.config import config_args
 from torax.geometry import geometry
 
 
 # pylint: disable=invalid-name
+@chex.dataclass(frozen=True)
+class IonMixture:
+  """Represents a mixture of ion species. The mixture can depend on time.
+
+  Main use cases:
+  1. Represent a bundled mixture of hydrogenic main ions (e.g. D and T)
+  2. Represent a bundled impurity species where the avg charge state, mass,
+    and radiation is consistent with each fractional concentration, and these
+    quantities are then averaged over the mixture to represent a single impurity
+    species in the transport equations for efficiency.
+
+  Attributes:
+    species: A dict mapping ion symbols (from ION_SYMBOLS) to their fractional
+      concentration in the mixture. The fractions must sum to 1.
+    tolerance: The tolerance used to check if the fractions sum to 1.
+  """
+
+  species: Mapping[
+      constants.ION_SYMBOLS, interpolated_param.TimeInterpolatedInput
+  ]
+  tolerance: float = 1e-6
+
+  def __post_init__(self):
+
+    if not self.species:
+      raise ValueError(self.__class__.__name__ + ' species cannot be empty.')
+
+    if not isinstance(self.species, Mapping):
+      raise ValueError('species must be a Mapping')
+
+    time_arrays = []
+    fraction_arrays = []
+
+    for value in self.species.values():
+      time_array, fraction_array, _, _ = (
+          interpolated_param.convert_input_to_xs_ys(value)
+      )
+      time_arrays.append(time_array)
+      fraction_arrays.append(fraction_array)
+
+    # Check if all time arrays are equal
+    # Note that if the TimeInterpolatedInput is a constant fraction (float) then
+    # convert_input_to_xs_ys returns a single-element array for t with value=0
+    if not all(np.array_equal(time_arrays[0], x) for x in time_arrays[1:]):
+      raise ValueError(
+          'All time indexes for '
+          + self.__class__.__name__
+          + ' fractions must be equal.'
+      )
+
+    # Check if the ion fractions sum to 1 at all times
+    fraction_sum = np.sum(fraction_arrays, axis=0)
+    if not np.allclose(fraction_sum, 1.0, rtol=self.tolerance):
+      raise ValueError(
+          'Fractional concentrations in an IonMixture must sum to 1 at all'
+          ' times.'
+      )
 
 
 @chex.dataclass

torax/config/tests/plasma_composition.py

@@ -23,7 +23,7 @@
 
 
 class PlasmaCompositionTest(parameterized.TestCase):
-  """Unit tests for the `torax.config.plasma_composition` module."""
+  """Unit tests for methods in the `torax.config.plasma_composition` module."""
 
   def test_plasma_composition_make_provider(self):
     """Checks provider construction with no issues."""
@@ -117,5 +117,45 @@ def test_interpolated_vars_are_only_constructed_once(
         self.assertIs(value, interpolated_params[field])
 
 
+class IonMixtureTest(parameterized.TestCase):
+  """Unit tests for constructing the IonMixture class."""
+
+  @parameterized.named_parameters(
+      ('valid_constant', {'D': 0.5, 'T': 0.5}, False),
+      (
+          'valid_time_dependent',
+          {'D': {0: 0.6, 1: 0.7}, 'T': {0: 0.4, 1: 0.3}},
+          False,
+      ),
+      (
+          'valid_time_dependent_with_step',
+          {'D': ({0: 0.6, 1: 0.7}, 'step'), 'T': {0: 0.4, 1: 0.3}},
+          False,
+      ),
+      ('invalid_empty', {}, True),
+      ('invalid_fractions', {'D': 0.4, 'T': 0.3}, True),
+      (
+          'invalid_time_mismatch',
+          {'D': {0: 0.5, 1: 0.6}, 'T': {0: 0.5, 2: 0.4}},
+          True,
+      ),
+      (
+          'invalid_time_dependent_fractions',
+          {'D': {0: 0.6, 1: 0.7}, 'T': {0: 0.5, 1: 0.4}},
+          True,
+      ),
+      ('valid_tolerance', {'D': 0.49999999, 'T': 0.5}, False),
+      ('invalid_tolerance', {'D': 0.4999, 'T': 0.5}, True),
+      ('invalid_not_mapping', 'D', True),
+  )
+  def test_ion_mixture_constructor(self, input_species, should_raise):
+    """Tests various cases of IonMixture construction."""
+    if should_raise:
+      with self.assertRaises(ValueError):
+        plasma_composition.IonMixture(species=input_species)
+    else:
+      plasma_composition.IonMixture(species=input_species)
+
+
 if __name__ == '__main__':
   absltest.main()