Merge remote-tracking branch 'origin/main' into add-vloop-bc

docs/configuration.rst

@@ -56,18 +56,19 @@ The following inputs are valid for **time-varying-scalar** parameters:
 
 Examples:
 
-1. Define a time-dependent :math:`Z_{eff}` with piecewise linear interpolation, from :math:`t=2` to :math:`t=15`.
-:math:`Z_{eff}` drops from 2.8 to 1.5, and then stays flat.
+1. Define a time-dependent total current :math:`Ip_{tot}` with piecewise linear interpolation,
+from :math:`t=10` to :math:`t=100`. :math:`Ip_{tot}` rises from 2 to 15, and then stays flat
+due to constant extrapolation beyond the last time value.
 
 .. code-block:: python
 
-  Zeff = ({2: 2.8, 5: 2.0, 8: 1.5, 15: 1.5}, 'PIECEWISE_LINEAR')
+  Ip_tot = ({10: 2.0, 100: 15.0}, 'PIECEWISE_LINEAR')
 
 or more simply, taking advantage of the default.
 
 .. code-block:: python
 
-  Zeff = {2: 2.8, 5: 2.0, 8: 1.5, 15: 1.5}
+    Ip_tot = {10: 2.0, 100: 15.0}
 
 2. Define a time dependent internal boundary condition for ion temperature, ``Tiped``, with stepwise changes,
 starting at :math:`1~keV`` at :math:`t=2s`, transitioning to :math:`3~keV`` at :math:`t=8s`, and back down
@@ -189,24 +190,108 @@ runtime_params
 plasma_composition
 ^^^^^^^^^^^^^^^^^^
 
-Defines the distribution of ion species. Currently restricted to a single main ion, a single impurity and a flat :math:`Z_{eff}`.
+Defines the distribution of ion species.  The keys and their meanings are as follows:
 
-``Ai`` (float = 2.5)
-  Mass of main ion in amu units. For multiple-isotope plasmas, make an effective average.
+``main_ion`` (str or dict = ``{'D': 0.5, 'T': 0.5}``)
+  Specifies the main ion species.
 
-``Zi`` (float = 1.0):
-  Charge of main ion in units of electron charge.
+  *   If a string, it represents a single ion species (e.g., ``'D'`` for deuterium, ``'T'`` for tritium, ``'H'`` for hydrogen). See below for the full list of supported ions.
+  *   If a dict, it represents a mixture of ion species with given fractions. By `mixture`, we mean
+      key value pairs of ion symbols and fractional concentrations, which must sum to 1 within a tolerance of 1e-6.
+      The effective mass and charge of the mixture is the weighted average of the species masses and charges.
+      The fractions can be time-dependent, i.e. are **time-varying-scalar**. The ion mixture API thus
+      supports features such as time varying isotope ratios.
 
-``Zimp`` (float = 10.0), **time-varying-scalar**
-  Impurity charge state.
+``impurity`` (str or dict = ``'Ne'``), **time-varying-scalar**
+  Specifies the impurity species, following the same syntax as ``main_ion``. A single effective impurity species
+  is currently supported, although multiple impurities can still be defined as a mixture.
 
-``Aimp`` (float = 20.18), **time-varying-scalar**
-  Impurity mass in amu units.
-
-``Zeff`` (float = 1.0), **time-varying-scalar**
+``Zeff`` (float = 1.0), **time-varying-array**
   Plasma effective charge, defined as :math:`Z_{eff}=\sum_i Z_i^2 \hat{n}_i`, where :math:`\hat{n}_i` is
-  the normalized ion density :math:`n_i/n_e`. For a given :math:`Z_{eff}` and :math:`Z_{imp}`, a consistent :math:`\hat{n}_i` is calculated,
-  with the appropriate degree of main ion dilution.
+  the normalized ion density :math:`n_i/n_e`. For a given :math:`Z_{eff}` and impurity charge states,
+  a consistent :math:`\hat{n}_i` is calculated, with the appropriate degree of main ion dilution.
+
+``Zi_override`` (float, optional = None), **time-varying-scalar**
+  An optional override for the main ion's charge (Z) or average charge of an ion mixture.
+  If provided, this value will be used instead of the Z calculated from the ``main_ion`` specification.
+
+``Ai_override`` (float, optional = None), **time-varying-scalar**
+  An optional override for the main ion's mass (A) in amu units or average mass of an IonMixture.
+  If provided, this value will be used instead of the A calculated from the ``main_ion`` specification.
+
+``Zimp_override`` (float, optional), **time-varying-scalar**
+  As ``Zi_override``, but for the impurity ion. If provided, this value will be used instead of the Z calculated
+  from the ``impurity`` specification.
+
+``Aimp_override`` (float, optional), **time-varying-scalar**
+  As ``Ai_override``, but for the impurity ion. If provided, this value will be used instead of the A calculated
+  from the ``impurity`` specification.
+
+The average charge state of each ion in each mixture is determined by `Mavrin polynomials <https://doi.org/10.1080/10420150.2018.1462361>`_,
+which are fitted to atomic data, and in the temperature ranges of interest in the tokamak core,
+are well approximated as 1D functions of electron temperature. All ions with atomic numbers below
+Carbon are assumed to be fully ionized.
+
+Examples
+--------
+
+We remind that for all cases below, the impurity density is solely constrained by
+the input ``Zeff`` value and the impurity charge state, presently assumed to be fully ionized.
+Imminent development will support temperature-dependent impurity average charge states,
+
+* Pure deuterium plasma:
+
+  .. code-block:: python
+
+    'plasma_composition': {
+        'main_ion': 'D',
+        'impurity': 'Ne',  # Neon
+        'Zeff': 1.5,
+    }
+
+* 50-50 DT ion mixture:
+
+  .. code-block:: python
+
+    'plasma_composition': {
+        'main_ion': {'D': 0.5, 'T': 0.5},
+        'impurity': 'Be',  # Beryllium
+        'Zeff': 1.8,
+    }
+
+* Time-varying DT ion mixture:
+
+  .. code-block:: python
+
+    'plasma_composition': {
+      'main_ion': {
+        'D': {0.0: 0.1, 5.0: 0.9},  # D fraction from 0.1 to 0.9
+        'T': {0.0: 0.9, 5.0: 0.1},  # T fraction from 0.9 to 0.1
+      },
+      'impurity': 'W',  # Tungsten
+      'Zeff': 2.0,
+    }
+
+Allowed ion symbols
+-------------------
+
+The following ion symbols are recognized for ``main_ion`` and ``impurity`` input fields.
+
+  *   H  (Hydrogen)
+  *   D  (Deuterium)
+  *   T  (Tritium)
+  *   He3 (Helium-3)
+  *   He4 (Helium-4)
+  *   Li (Lithium)
+  *   Be (Beryllium)
+  *   C (Carbon)
+  *   N (Nitrogen)
+  *   O (Oxygen)
+  *   Ne (Neon)
+  *   Ar (Argon)
+  *   Kr (Krypton)
+  *   Xe (Xenon)
+  *   W (Tungsten)
 
 Profile conditions
 ^^^^^^^^^^^^^^^^^^
@@ -1240,9 +1325,9 @@ The configuration file is also available in ``torax/examples/iterhybrid_rampup.p
   CONFIG = {
       'runtime_params': {
           'plasma_composition': {
-              'Ai': 2.5,
+              'main_ion': {'D': 0.5, 'T': 0.5},
+              'impurity': 'Ne',
               'Zeff': 1.6,
-              'Zimp': 10,
           },
           'profile_conditions': {
               'Ip_tot': {0: 3, 80: 10.5},

docs/equation_summary.rst

@@ -26,6 +26,9 @@ TORAX solves coupled 1D PDEs in normalized toroidal flux coordinates,
       g_0q_i^{\mathrm{conv}}T_i\right] + Q_i
     \end{multline}
 
+  If multiple main ion species are present (e.g., a D-T mix), then :math:`n_i` represents the sum of all
+  main ions, and ion attributes like charge and mass are averaged values for the mixture, weighted by fractional abundance.
+
   Electron heat transport, governing the evolution of the electron temperature :math:`T_e`.
 
   .. math::

docs/model_integration.rst

@@ -312,14 +312,6 @@ the |torax.sim.Sim|_ object.
 
 .. code-block:: python
 
-    # in sim.py. Copied here for reference, no need to modify this.
-
-    def build_sim_object(
-        ...
-        transport_model_builder: transport_model_lib.TransportModelBuilder,
-        ...
-    ) -> Sim:
-
     # in your TORAX configuration or run file .py
 
     my_custom_transport_builder = MyCustomTransportModelBuilder()
@@ -334,7 +326,7 @@ the |torax.sim.Sim|_ object.
     my_custom_transport_builder.runtime_params.bar = 4.0
 
     # Build the Sim object.
-    sim_object = sim_lib.build_sim_object(
+    sim_object = sim_lib.Sim.create(
         ...,
         transport_model_builder=my_custom_transport_builder,
         ...
@@ -348,7 +340,7 @@ As of 7 June 2024, you cannot instantiate and configure a custom transport model
 via the config dictionary. You may still configure the other components of your
 TORAX simulation via the config dict and use other functions in
 |torax.config.build_sim|_ to convert those to the objects you can pass into
-``build_sim_object()``. We are working on making this easier, but reach out
+``Sim.create()``. We are working on making this easier, but reach out
 if this is something you need.
 
 

docs/physics_models.rst

@@ -45,16 +45,29 @@ Generalization to geometry data beyond CHEASE is also planned.
 Plasma composition, initial and prescribed conditions
 =====================================================
 
-Presently, TORAX only accommodates a single main ion species,configured with its
-atomic mass number (:math:`A_i`) and charge state (:math:`Z_i`). The plasma effective
-charge, :math:`Z_\textit{eff}`, is assumed to be radially flat and is also
-user-configurable. A single impurity with charge state :math:`Z_\textit{imp}` is
-specified to accommodate :math:`Z_\textit{eff} > 1`. The main ion density dilution
-is then calculated as follows:
+Presently, TORAX accommodates a single main ion species and single impurity species,
+which can be comprised of time-dependent mixtures of ions with fractional abundances
+summing to 1. This is useful for example for simulating isotope mixes. Based on the
+ion symbols and fractional abundances, the average mass of each species is determined.
+The average charge state of each ion in each mixture is determined by `Mavrin polynomials <https://doi.org/10.1080/10420150.2018.1462361>`_,
+which are fitted to atomic data, and in the temperature ranges of interest in the tokamak core,
+are well approximated as 1D functions of electron temperature. All ions with atomic numbers below
+Carbon are assumed to be fully ionized.
+
+The impurity and main ion densities are constrained by the plasma effective
+charge, :math:`Z_\mathrm{eff}`, which is a user-provided 2D array in both time and space,
+as well as quasineutrality.
+
+:math:`n_i`, and :math:`n_{imp}`, are solved from the
+following system of equations, where :math:`Z_\mathrm{eff}` and the electron density are
+known, and :math:`Z_\mathrm{imp}` is the average impurity charge of the impurity mixture,
+with the average charge state for each ion determined from the Mavrin polynomials.
 
 .. math::
 
-  n_i=(Z_\textit{imp}-Z_\textit{eff})/(Z_\textit{imp}-1)n_e
+  n_\mathrm{i}Z_\mathrm{i}^2 + n_\mathrm{imp}Z_\mathrm{imp}^2 = n_\mathrm{e}Z_\mathrm{eff}
+
+  n_\mathrm{i}Z_\mathrm{i} + n_\mathrm{imp}Z_\mathrm{imp} = n_\mathrm{e}
 
 Initial conditions for the evolving profiles :math:`T_i`, :math:`T_e`, :math:`n_e`,
 and :math:`\psi` are user-configurable. For :math:`T_{i,e}`, both the initial core

torax/__init__.py

@@ -22,6 +22,7 @@
 from torax.config.build_sim import build_sim_from_config
 from torax.config.config_loader import import_module
 from torax.interpolated_param import InterpolatedVarSingleAxis
+from torax.interpolated_param import InterpolationMode
 from torax.output import ToraxSimOutputs
 from torax.sim import Sim
 from torax.state import SimError