Integration tests with extended-lengyel.

torax/_src/edge/collisional_radiative_models.py

@@ -81,10 +81,11 @@ def calculate_mavrin_2017(
     ion_symbol_lookup = ion_symbol
 
   if ion_symbol_lookup not in cr_module.COEFFS.keys():
-    raise ValueError(
-        f'Invalid ion symbol: {ion_symbol}. Allowed symbols are:'
-        f' {cr_module.COEFFS.keys()}'
-    )
+    # If the ion is not supported by the edge radiation model, we assume it
+    # negligibly contributes to the edge physics (radiation or Z_eff/dilution in
+    # the divertor). This is a good assumption for heavy impurities like W,
+    # which this case covers. This behaviour is silent to avoid log spam.
+    return jnp.zeros_like(T_e)
 
   # Mavrin 2017 formulas are constructed for [eV] temperature input
   T_e_ev = T_e * 1e3

torax/_src/edge/pydantic_model.py

@@ -71,11 +71,6 @@ class ExtendedLengyelConfig(base.EdgeModelConfig):
           extended_lengyel_defaults.DIVERTOR_BROADENING_FACTOR
       )
   )
-  ratio_bpol_omp_to_bpol_avg: torax_pydantic.PositiveTimeVaryingScalar = (
-      torax_pydantic.ValidatedDefault(
-          extended_lengyel_defaults.RATIO_BPOL_OMP_TO_BPOL_AVG
-      )
-  )
   sheath_heat_transmission_factor: torax_pydantic.PositiveTimeVaryingScalar = (
       torax_pydantic.ValidatedDefault(
           extended_lengyel_defaults.SHEATH_HEAT_TRANSMISSION_FACTOR
@@ -134,16 +129,15 @@ class ExtendedLengyelConfig(base.EdgeModelConfig):
   connection_length_divertor: (
       torax_pydantic.PositiveTimeVaryingScalar | None
   ) = None
-  angle_of_incidence_target: torax_pydantic.PositiveTimeVaryingScalar = (
-      torax_pydantic.ValidatedDefault(
-          extended_lengyel_defaults.ANGLE_OF_INCIDENCE_TARGET
-      )
+  angle_of_incidence_target: torax_pydantic.PositiveTimeVaryingScalar | None = (
+      None
   )
-  toroidal_flux_expansion: torax_pydantic.PositiveTimeVaryingScalar = (
-      torax_pydantic.ValidatedDefault(
-          extended_lengyel_defaults.TOROIDAL_FLUX_EXPANSION
-      )
+  toroidal_flux_expansion: torax_pydantic.PositiveTimeVaryingScalar | None = (
+      None
   )
+  ratio_bpol_omp_to_bpol_avg: (
+      torax_pydantic.PositiveTimeVaryingScalar | None
+  ) = None
 
   # Optional input for inverse mode
   T_e_target: torax_pydantic.PositiveTimeVaryingScalar | None = None
@@ -343,7 +337,6 @@ def _get_optional_value(
         newton_raphson_tol=self.newton_raphson_tol,
         ne_tau=self.ne_tau.get_value(t),
         divertor_broadening_factor=self.divertor_broadening_factor.get_value(t),
-        ratio_bpol_omp_to_bpol_avg=self.ratio_bpol_omp_to_bpol_avg.get_value(t),
         sheath_heat_transmission_factor=self.sheath_heat_transmission_factor.get_value(
             t
         ),
@@ -367,8 +360,15 @@ def _get_optional_value(
         connection_length_divertor=_get_optional_value(
             self.connection_length_divertor, t
         ),
-        angle_of_incidence_target=self.angle_of_incidence_target.get_value(t),
-        toroidal_flux_expansion=self.toroidal_flux_expansion.get_value(t),
+        angle_of_incidence_target=_get_optional_value(
+            self.angle_of_incidence_target, t
+        ),
+        toroidal_flux_expansion=_get_optional_value(
+            self.toroidal_flux_expansion, t
+        ),
+        ratio_bpol_omp_to_bpol_avg=_get_optional_value(
+            self.ratio_bpol_omp_to_bpol_avg, t
+        ),
         seed_impurity_weights=seed_impurity_weights,
         fixed_impurity_concentrations=fixed_impurity_concentrations,
         enrichment_factor=enrichment_factor,

torax/_src/edge/tests/edge_pydantic_model_test.py

@@ -379,6 +379,9 @@ def test_run_standalone_from_pydantic_config(self):
         'fixed_impurity_concentrations': {'He': 0.01},
         'connection_length_target': 20.0,
         'connection_length_divertor': 5.0,
+        'angle_of_incidence_target': 3.0,
+        'toroidal_flux_expansion': 1.0,
+        'ratio_bpol_omp_to_bpol_avg': 4.0/3.0,
         'computation_mode': 'inverse',
         'solver_mode': 'fixed_point',
         'use_enrichment_model': False,

torax/_src/edge/tests/extended_lengyel_model_test.py

@@ -159,6 +159,9 @@ def test_call_inverse_mode(self):
         T_e_target=2.34,
         connection_length_target=20.0,
         connection_length_divertor=5.0,
+        angle_of_incidence_target=3.0,
+        ratio_bpol_omp_to_bpol_avg=4.0/3.0,
+        toroidal_flux_expansion=1.0,
         seed_impurity_weights={'N': 1.0, 'Ar': 0.05},
         fixed_impurity_concentrations={'He': 0.01},
         enrichment_factor={'N': 1.0, 'Ar': 1.0, 'He': 1.0},
@@ -301,6 +304,7 @@ def test_fixed_impurity_update_from_core(self, mock_run_standalone):
         connection_length_divertor=5.0,
         angle_of_incidence_target=1.0,
         toroidal_flux_expansion=1.0,
+        ratio_bpol_omp_to_bpol_avg=1.0,
         use_enrichment_model=False,
         diverted=True,
     )
@@ -417,6 +421,7 @@ def test_broadening_limited_logic(
         connection_length_divertor=5.0,
         angle_of_incidence_target=1.0,
         toroidal_flux_expansion=1.0,
+        ratio_bpol_omp_to_bpol_avg=1.0,
         use_enrichment_model=True,
         diverted=diverted,
         divertor_broadening_factor=_CONFIG_BROADENING,
@@ -631,6 +636,7 @@ def test_edge_model_coupling_smoke(self):
             'connection_length_divertor': 10.0,
             'toroidal_flux_expansion': 4.0,
             'angle_of_incidence_target': 3.0,
+            'ratio_bpol_omp_to_bpol_avg': 4.0 / 3.0,
             'use_enrichment_model': False,
             'diverted': True,
         },
@@ -691,6 +697,7 @@ def test_temperature_boundary_condition_updates(
             'connection_length_divertor': 10.0,
             'toroidal_flux_expansion': 4.0,
             'angle_of_incidence_target': 3.0,
+            'ratio_bpol_omp_to_bpol_avg': 4.0 / 3.0,
             # Test parameters
             'update_temperatures': update_temperatures,
             'T_i_T_e_ratio_target': ion_to_electron_ratio,
@@ -789,6 +796,7 @@ def test_inverse_mode_updates_impurities(self):
             'connection_length_divertor': 15.0,
             'toroidal_flux_expansion': 4.0,
             'angle_of_incidence_target': 3.0,
+            'ratio_bpol_omp_to_bpol_avg': 4.0 / 3.0,
             'use_enrichment_model': False,
             'diverted': True,
         },
@@ -881,6 +889,7 @@ def setUp(self):
             'connection_length_divertor': 10.0,
             'toroidal_flux_expansion': 4.0,
             'angle_of_incidence_target': 3.0,
+            'ratio_bpol_omp_to_bpol_avg': 4.0 / 3.0,
             'diverted': True,
             # 'use_enrichment_model' and 'diverted' are overridden in tests.
         },

torax/_src/geometry/fbt.py

@@ -479,8 +479,8 @@ def _get_val(key):
     connection_length_target = _get_val('Lpar_target')
     # Lpar_div -> connection_length_divertor [m]
     connection_length_divertor = _get_val('Lpar_div')
-    # alpha_target -> angle_of_incidence_target [degrees]
-    angle_of_incidence_target = _get_val('alpha_target')
+    # alpha_target [radians] -> angle_of_incidence_target [degrees]
+    angle_of_incidence_target = np.rad2deg(_get_val('alpha_target'))
     # r_OMP -> R_OMP [m]
     R_OMP = _get_val('r_OMP')
     # r_target -> R_target [m]

torax/_src/geometry/tests/fbt_test.py

@@ -205,7 +205,7 @@ def test_fbt_edge_parameters(
     # Domain 1 (Upper): [15, 25] over time for Lpar_target.
     LY['Lpar_target'] = np.array([[10.0, 20.0], [15.0, 25.0]])
     LY['Lpar_div'] = np.array([[1.0, 2.0], [1.5, 2.5]])
-    LY['alpha_target'] = np.array([[3.0, 4.0], [3.5, 4.5]])
+    LY['alpha_target'] = np.deg2rad(np.array([[3.0, 4.0], [3.5, 4.5]]))
     LY['r_OMP'] = np.array([[6.0, 7.0], [6.5, 7.5]])
     LY['r_target'] = np.array([[5.0, 6.0], [5.5, 6.5]])
     LY['Bp_OMP'] = np.array([[0.5, 0.7], [0.6, 0.8]])

torax/_src/output_tools/output.py

@@ -839,9 +839,20 @@ def _save_edge_outputs(self) -> xr.Dataset:
       xr_dict["solver_iterations"] = self._pack_into_data_array(
           "solver_iterations", numerics.iterations
       )
-      xr_dict["solver_residual"] = self._pack_into_data_array(
-          "solver_residual", numerics.residual
-      )
+      # Handle solver_residual explicitly because it is a vector
+      # (time, n_unknowns) which _pack_into_data_array doesn't support
+      # automatically.
+      if numerics.residual is not None and numerics.residual.ndim == 2:
+        xr_dict["solver_residual"] = xr.DataArray(
+            numerics.residual,
+            dims=[TIME, "solver_unknown_idx"],
+            name="solver_residual",
+        )
+      else:
+        xr_dict["solver_residual"] = self._pack_into_data_array(
+            "solver_residual", numerics.residual
+        )
+
       xr_dict["solver_error"] = self._pack_into_data_array(
           "solver_error", numerics.error
       )

torax/_src/torax_pydantic/model_config.py

@@ -274,7 +274,10 @@ def _validate_edge_core_impurity_consistency(self) -> typing_extensions.Self:
         raise ValueError(
             'Mismatch between core plasma composition impurities and edge'
             f' impurities. Core: {core_species}, Edge: {edge_species}.'
-            f' Difference: {core_species.symmetric_difference(edge_species)}'
+            f' Difference: {core_species.symmetric_difference(edge_species)}.'
+            ' Likely reason: edge.fixed_impurity_concentrations and/or'
+            ' edge.seed_impurity_weights do not match plasma_composition.'
+            'impurity.species.'
         )
 
     return self

torax/tests/sim_test.py

@@ -226,6 +226,11 @@ class SimTest(sim_test_case.SimTestCase):
           'test_iterhybrid_predictor_corrector_mavrin_n_e_ratios',
           'test_iterhybrid_predictor_corrector_mavrin_n_e_ratios.py',
       ),
+      # Predictor-corrector w/ Mavrin, n_e_ratios, forward lengyel
+      (
+          'test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel',
+          'test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.py',
+      ),
       # Predictor-corrector with Mavrin and n_e_ratios_Z_eff impurity mode.
       (
           'test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_z_eff',

torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_n_e_ratios_lengyel.py

@@ -0,0 +1,46 @@
+# Copyright 2024 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""test_iterhybrid_predictor_corrector with Mavrin radiation, n_e_ratios, x-lengyel."""
+
+import copy
+
+from torax.tests.test_data import test_iterhybrid_predictor_corrector
+
+CONFIG = copy.deepcopy(test_iterhybrid_predictor_corrector.CONFIG)
+
+assert isinstance(CONFIG['plasma_composition'], dict)
+# If present, Z_eff would be ignored in the 'n_e_ratios' impurity mode.
+del CONFIG['plasma_composition']['Z_eff']
+CONFIG['plasma_composition']['impurity'] = {
+    'impurity_mode': 'n_e_ratios',
+    'species': {'Ne': 0.01, 'W': 1e-5},
+}
+CONFIG['sources']['impurity_radiation'] = {
+    'model_name': 'mavrin_fit',
+}
+
+CONFIG['edge'] = {
+    'model_name': 'extended_lengyel',
+    'computation_mode': 'forward',
+    'fixed_impurity_concentrations': {'Ne': 0.01, 'W': 1e-5},
+    'impurity_sot': 'core',
+    'connection_length_target': 30.0,
+    'connection_length_divertor': 10.0,
+    'angle_of_incidence_target': 3.0,
+    'ratio_bpol_omp_to_bpol_avg': 4.0 / 3.0,
+    'toroidal_flux_expansion': 1.0,
+    'use_enrichment_model': True,
+    'diverted': True,
+}