Migrate SPARC configs to combined transport model.

tamaranorman · Torax team · commit 8b29202552cd · 2026-06-29T14:13:42.000-07:00
PiperOrigin-RevId: 938598306
diff --git a/torax/_src/fvm/tests/calc_coeffs_test.py b/torax/_src/fvm/tests/calc_coeffs_test.py
@@ -56,7 +56,11 @@ def test_calc_coeffs_smoke_test(
             solver=dict(
                 use_predictor_corrector=False, theta_implicit=theta_implicit
             ),
-            transport=dict(model_name='constant', chi_min=0, chi_i=1),
+            transport=dict(
+                model_name='combined',
+                transport_models=[dict(model_name='constant', chi_i=1)],
+                chi_min=0,
+            ),
             time_step_calculator=dict(),
         )
     )
diff --git a/torax/_src/fvm/tests/fvm_test.py b/torax/_src/fvm/tests/fvm_test.py
@@ -225,7 +225,11 @@ def test_nonlinear_solve_block_loss_minimum(
             solver=dict(
                 use_predictor_corrector=False, theta_implicit=theta_implicit
             ),
-            transport=dict(model_name='constant', chi_min=0, chi_i=1),
+            transport=dict(
+                model_name='combined',
+                transport_models=[dict(model_name='constant', chi_i=1)],
+                chi_min=0,
+            ),
             time_step_calculator=dict(),
         )
     )
@@ -333,7 +337,11 @@ def test_implicit_solve_block_uses_updated_boundary_conditions(self):
             pedestal=dict(),
             sources=source_config,
             solver=dict(use_predictor_corrector=False, theta_implicit=1.0),
-            transport=dict(model_name='constant', chi_min=0, chi_i=1),
+            transport=dict(
+                model_name='combined',
+                transport_models=[dict(model_name='constant', chi_i=1)],
+                chi_min=0,
+            ),
             time_step_calculator=dict(),
         )
     )
@@ -444,7 +452,11 @@ def test_theta_residual_uses_updated_boundary_conditions(self):
             pedestal=dict(),
             sources=source_config,
             solver=dict(use_predictor_corrector=False, theta_implicit=0.0),
-            transport=dict(model_name='constant', chi_min=0, chi_i=1),
+            transport=dict(
+                model_name='combined',
+                transport_models=[dict(model_name='constant', chi_i=1)],
+                chi_min=0,
+            ),
             time_step_calculator=dict(),
         )
     )
diff --git a/torax/examples/basic_config.py b/torax/examples/basic_config.py
@@ -41,7 +41,10 @@
     },
     'pedestal': {},
     'transport': {
-        'model_name': 'constant',
+        'model_name': 'combined',
+        'transport_models': [
+            {'model_name': 'constant'},
+        ],
     },
     'solver': {
         'solver_type': 'linear',
diff --git a/torax/examples/step_flattop_bgb.py b/torax/examples/step_flattop_bgb.py
@@ -133,21 +133,24 @@
         },
     },
     "transport": {
-        "model_name": "bohm-gyrobohm",
-        # BgB settings from [2] sec 3.3
-        # Tuning factor to achieve desired confinement
-        "chi_e_bohm_multiplier": bgb_multiplier,
-        "chi_i_bohm_multiplier": bgb_multiplier,
-        "chi_e_gyrobohm_multiplier": bgb_multiplier,
-        "chi_i_gyrobohm_multiplier": bgb_multiplier,
-        # Base coefficients
-        "chi_e_bohm_coeff": 0.01 * 2e-4,
-        "chi_e_gyrobohm_coeff": 50 * 5e-6,
-        "chi_i_bohm_coeff": 0.001 * 2e-4,
-        "chi_i_gyrobohm_coeff": 1.0 * 5e-6,
-        "D_face_c1": 1,
-        "D_face_c2": 0.3,
-        "V_face_coeff": -0.1,
+        "model_name": "combined",
+        "transport_models": [{
+            "model_name": "bohm-gyrobohm",
+            # BgB settings from [2] sec 3.3
+            # Tuning factor to achieve desired confinement
+            "chi_e_bohm_multiplier": bgb_multiplier,
+            "chi_i_bohm_multiplier": bgb_multiplier,
+            "chi_e_gyrobohm_multiplier": bgb_multiplier,
+            "chi_i_gyrobohm_multiplier": bgb_multiplier,
+            # Base coefficients
+            "chi_e_bohm_coeff": 0.01 * 2e-4,
+            "chi_e_gyrobohm_coeff": 50 * 5e-6,
+            "chi_i_bohm_coeff": 0.001 * 2e-4,
+            "chi_i_gyrobohm_coeff": 1.0 * 5e-6,
+            "D_face_c1": 1,
+            "D_face_c2": 0.3,
+            "V_face_coeff": -0.1,
+        }],
         # Clipping
         "chi_min": 0.15,
         "chi_max": 100.0,
diff --git a/torax/tests/test_data/compilation_benchmark.py b/torax/tests/test_data/compilation_benchmark.py
@@ -65,10 +65,12 @@
         'n_e_ped': 1e20,
     },
     'transport': {
-        'model_name': 'qlknn',
-        'qlknn_params': {
+        'model_name': 'combined',
+        'transport_models': [{
+            'model_name': 'qlknn',
             'DV_effective': False,
-        },
+        }],
+        'smoothing_width': 0.1,
     },
     'solver': {
         'solver_type': 'newton_raphson',
diff --git a/torax/tests/test_data/default_config.py b/torax/tests/test_data/default_config.py
@@ -41,7 +41,10 @@
         'ohmic': {},
     },
     'transport': {
-        'model_name': 'constant',
+        'model_name': 'combined',
+        'transport_models': [
+            {'model_name': 'constant'},
+        ],
     },
     'pedestal': {},
     'solver': {
diff --git a/torax/tests/test_data/test_all_transport_fusion_qlknn.py b/torax/tests/test_data/test_all_transport_fusion_qlknn.py
@@ -70,8 +70,12 @@
         'n_e_ped': 1.0e20,
     },
     'transport': {
-        'model_name': 'qlknn',
-        'DV_effective': False,
+        'model_name': 'combined',
+        'transport_models': [{
+            'model_name': 'qlknn',
+            'DV_effective': False,
+        }],
+        'smoothing_width': 0.1,
     },
     'solver': {
         'solver_type': 'linear',
diff --git a/torax/tests/test_data/test_bohmgyrobohm_all.py b/torax/tests/test_data/test_bohmgyrobohm_all.py
@@ -70,7 +70,10 @@
         'n_e_ped_is_fGW': True,
     },
     'transport': {
-        'model_name': 'bohm-gyrobohm',
+        'model_name': 'combined',
+        'transport_models': [{
+            'model_name': 'bohm-gyrobohm',
+        }],
     },
     'solver': {
         'solver_type': 'linear',
diff --git a/torax/tests/test_data/test_bremsstrahlung_time_dependent_Zimp.py b/torax/tests/test_data/test_bremsstrahlung_time_dependent_Zimp.py
@@ -59,9 +59,12 @@
         'set_pedestal': True,
     },
     'transport': {
-        'model_name': 'constant',
-        'chi_i': 0.5,
-        'chi_e': 0.5,
+        'model_name': 'combined',
+        'transport_models': [{
+            'model_name': 'constant',
+            'chi_i': 0.5,
+            'chi_e': 0.5,
+        }],
     },
     'solver': {
         'solver_type': 'linear',
diff --git a/torax/tests/test_data/test_chease.py b/torax/tests/test_data/test_chease.py
@@ -64,7 +64,10 @@
     },
     'pedestal': {},
     'transport': {
-        'model_name': 'constant',
+        'model_name': 'combined',
+        'transport_models': [{
+            'model_name': 'constant',
+        }],
     },
     'solver': {
         'solver_type': 'linear',
diff --git a/torax/tests/test_data/test_crank_nicolson.py b/torax/tests/test_data/test_crank_nicolson.py
@@ -47,7 +47,10 @@
     },
     'pedestal': {},
     'transport': {
-        'model_name': 'constant',
+        'model_name': 'combined',
+        'transport_models': [{
+            'model_name': 'constant',
+        }],
     },
     'solver': {
         'solver_type': 'linear',
diff --git a/torax/tests/test_data/test_fixed_dt.py b/torax/tests/test_data/test_fixed_dt.py
@@ -43,7 +43,11 @@
         'set_pedestal': True,
     },
     'transport': {
-        'model_name': 'qlknn',
+        'model_name': 'combined',
+        'transport_models': [{
+            'model_name': 'qlknn',
+        }],
+        'smoothing_width': 0.1,
     },
     'solver': {
         'solver_type': 'linear',
diff --git a/torax/tests/test_data/test_implicit_short_optimizer.py b/torax/tests/test_data/test_implicit_short_optimizer.py
@@ -41,7 +41,10 @@
     },
     'pedestal': {},
     'transport': {
-        'model_name': 'constant',
+        'model_name': 'combined',
+        'transport_models': [{
+            'model_name': 'constant',
+        }],
     },
     'solver': {
         'solver_type': 'optimizer',
diff --git a/torax/tests/test_data/test_ne_qlknn_deff_veff.py b/torax/tests/test_data/test_ne_qlknn_deff_veff.py
@@ -68,9 +68,12 @@
         'n_e_ped': 1.0e20,
     },
     'transport': {
-        'model_name': 'qlknn',
-        # qlknn params.
-        'DV_effective': True,
+        'model_name': 'combined',
+        'transport_models': [{
+            'model_name': 'qlknn',
+            'DV_effective': True,
+        }],
+        'smoothing_width': 0.1,
     },
     'solver': {
         'solver_type': 'linear',
diff --git a/torax/tests/test_data/test_ne_qlknn_defromchie.py b/torax/tests/test_data/test_ne_qlknn_defromchie.py
@@ -68,9 +68,12 @@
         'n_e_ped': 1.0e20,
     },
     'transport': {
-        'model_name': 'qlknn',
-        # qlknn params.
-        'DV_effective': False,
+        'model_name': 'combined',
+        'transport_models': [{
+            'model_name': 'qlknn',
+            'DV_effective': False,
+        }],
+        'smoothing_width': 0.1,
     },
     'solver': {
         'solver_type': 'linear',
diff --git a/torax/tests/test_data/test_particle_sources_cgm.py b/torax/tests/test_data/test_particle_sources_cgm.py
@@ -61,10 +61,12 @@
         'n_e_ped': 1.0e20,
     },
     'transport': {
-        'model_name': 'CGM',
-        # cgm params.
-        'chi_D_ratio': 8,
-        'VR_D_ratio': {0: -3.0, 1: 0.0},
+        'model_name': 'combined',
+        'transport_models': [{
+            'model_name': 'CGM',
+            'chi_D_ratio': 8,
+            'VR_D_ratio': {0: -3.0, 1: 0.0},
+        }],
     },
     'solver': {
         'solver_type': 'linear',
diff --git a/torax/tests/test_data/test_prescribed_generic_current_source.py b/torax/tests/test_data/test_prescribed_generic_current_source.py
@@ -81,12 +81,12 @@ def gaussian(r, center, width, amplitude):
     },
     'pedestal': {},
     'transport': {
-        'model_name': 'constant',
-        # constant params.
-        # diffusion coefficient in electron density equation in m^2/s
-        'D_e': 0.5,
-        # convection coefficient in electron density equation in m^2/s
-        'V_e': -0.2,
+        'model_name': 'combined',
+        'transport_models': [{
+            'model_name': 'constant',
+            'D_e': 0.5,
+            'V_e': -0.2,
+        }],
     },
     'solver': {
         'solver_type': 'linear',
diff --git a/torax/tests/test_data/test_prescribed_timedependent_ne.nc b/torax/tests/test_data/test_prescribed_timedependent_ne.nc
diff --git a/torax/tests/test_data/test_prescribed_timedependent_ne.py b/torax/tests/test_data/test_prescribed_timedependent_ne.py
@@ -65,11 +65,23 @@
         'generic_current': {},
     },
     'transport': {
-        'model_name': 'qlknn',
-        'apply_inner_patch': True,
-        'chi_i_inner': 2.0,
-        'chi_e_inner': 2.0,
-        'rho_inner': 0.3,
+        'model_name': 'combined',
+        'transport_models': [
+            {
+                'model_name': 'qlknn',
+                'rho_min': 0.3,
+            },
+            # Inner patch
+            {
+                'model_name': 'constant',
+                'chi_i': 2.0,
+                'chi_e': 2.0,
+                'D_e': 0.2,
+                'V_e': 0.0,
+                'rho_max': 0.3,
+            },
+        ],
+        'smoothing_width': 0.1,
     },
     'solver': {
         'solver_type': 'linear',
diff --git a/torax/tests/test_data/test_prescribed_transport.py b/torax/tests/test_data/test_prescribed_transport.py
@@ -35,19 +35,20 @@
 CONFIG = copy.deepcopy(default_config.CONFIG)
 
 CONFIG['transport'] = {
-    # Note: prescribed transport is currently handled by ConstantTransportModel
-    # The name of this model will change in a future release
-    'model_name': 'constant',
-    'chi_i': (
-        np.array([0.0, 5.0]),
-        x,
-        1.5 * np.stack([chi_base_t1, chi_base_t2]),
-    ),
-    'chi_e': (
-        np.array([0.0, 5.0]),
-        x,
-        np.stack([chi_base_t1, chi_base_t2]),
-    ),
-    'D_e': 1.0,
-    'V_e': -0.33,
+    'model_name': 'combined',
+    'transport_models': [{
+        'model_name': 'constant',
+        'chi_i': (
+            np.array([0.0, 5.0]),
+            x,
+            1.5 * np.stack([chi_base_t1, chi_base_t2]),
+        ),
+        'chi_e': (
+            np.array([0.0, 5.0]),
+            x,
+            np.stack([chi_base_t1, chi_base_t2]),
+        ),
+        'D_e': 1.0,
+        'V_e': -0.33,
+    }],
 }
diff --git a/torax/tests/test_data/test_psi_and_heat.py b/torax/tests/test_data/test_psi_and_heat.py
@@ -48,7 +48,11 @@
         'set_pedestal': True,
     },
     'transport': {
-        'model_name': 'qlknn',
+        'model_name': 'combined',
+        'transport_models': [{
+            'model_name': 'qlknn',
+        }],
+        'smoothing_width': 0.1,
     },
     'solver': {
         'solver_type': 'linear',
diff --git a/torax/tests/test_data/test_psi_heat_dens.py b/torax/tests/test_data/test_psi_heat_dens.py
@@ -45,12 +45,12 @@
     },
     'pedestal': {'model_name': 'set_T_ped_n_ped', 'set_pedestal': True},
     'transport': {
-        'model_name': 'constant',
-        # constant params.
-        # diffusion coefficient in electron density equation in m^2/s
-        'D_e': 0.5,
-        # convection coefficient in electron density equation in m^2/s
-        'V_e': -0.2,
+        'model_name': 'combined',
+        'transport_models': [{
+            'model_name': 'constant',
+            'D_e': 0.5,
+            'V_e': -0.2,
+        }],
     },
     'solver': {
         'solver_type': 'linear',
diff --git a/torax/tests/test_data/test_psichease_ip_chease_vloop.py b/torax/tests/test_data/test_psichease_ip_chease_vloop.py
diff --git a/torax/tests/test_data/test_psichease_ip_parameters_vloop_varying.py b/torax/tests/test_data/test_psichease_ip_parameters_vloop_varying.py
diff --git a/torax/tests/test_data/test_psichease_prescribed_johm.py b/torax/tests/test_data/test_psichease_prescribed_johm.py
diff --git a/torax/tests/test_data/test_psichease_prescribed_jtot.py b/torax/tests/test_data/test_psichease_prescribed_jtot.py
diff --git a/torax/tests/test_data/test_semiimplicit_convection.py b/torax/tests/test_data/test_semiimplicit_convection.py
diff --git a/torax/tests/test_data/test_timedependence.nc b/torax/tests/test_data/test_timedependence.nc
diff --git a/torax/tests/test_data/test_timedependence.py b/torax/tests/test_data/test_timedependence.py

Original file line number	Diff line number	Diff line change
`@@ -56,7 +56,11 @@ def test_calc_coeffs_smoke_test(`
`56`	`56`	`solver=dict(`
`57`	`57`	`use_predictor_corrector=False, theta_implicit=theta_implicit`
`58`	`58`	`),`
`59`		`- transport=dict(model_name='constant', chi_min=0, chi_i=1),`
	`59`	`+ transport=dict(`
	`60`	`+ model_name='combined',`
	`61`	`+ transport_models=[dict(model_name='constant', chi_i=1)],`
	`62`	`+ chi_min=0,`
	`63`	`+ ),`
`60`	`64`	`time_step_calculator=dict(),`
`61`	`65`	`)`
`62`	`66`	`)`