Add internal boundary conditions to ProfileConditions

Summary of changes:
- Necessitated separating adaptive_source out from internal_boundary_conditions to avoid a circular import
- Added an iterhybrid test case with an internal boundary condition
- Combines internal boundary conditions from pedestal with ones set by user

PiperOrigin-RevId: 868117952

Author: theo-brown

torax/_src/core_profiles/profile_conditions.py

@@ -22,6 +22,7 @@
 import numpy as np
 import pydantic
 from torax._src import array_typing
+from torax._src.internal_boundary_conditions import internal_boundary_conditions as internal_boundary_conditions_lib
 from torax._src.torax_pydantic import torax_pydantic
 from typing_extensions import Self
 
@@ -66,6 +67,9 @@ class RuntimeParams:
   n_e_nbar_is_fGW: bool
   n_e_right_bc: array_typing.FloatScalar
   n_e_right_bc_is_fGW: bool
+  internal_boundary_conditions: (
+      internal_boundary_conditions_lib.InternalBoundaryConditions
+  )
   current_profile_nu: float
   initial_j_is_total_current: bool = dataclasses.field(
       metadata={'static': True}
@@ -186,6 +190,11 @@ class ProfileConditions(torax_pydantic.BaseModelFrozen):
   n_e_nbar_is_fGW: bool = False
   n_e_right_bc: torax_pydantic.TimeVaryingScalar | None = None
   n_e_right_bc_is_fGW: bool = False
+  internal_boundary_conditions: (
+      internal_boundary_conditions_lib.InternalBoundaryConditionsConfig
+  ) = torax_pydantic.ValidatedDefault(
+      internal_boundary_conditions_lib.InternalBoundaryConditionsConfig()
+  )
   current_profile_nu: float = 1.0
   initial_j_is_total_current: Annotated[bool, torax_pydantic.JAX_STATIC] = False
   # TODO(b/434175938): Remove this before the V2 API release in place of
@@ -428,9 +437,18 @@ def build_runtime_params(self, t: chex.Numeric) -> RuntimeParams:
     else:
       runtime_params['n_e_right_bc_is_absolute'] = True
 
+    runtime_params['internal_boundary_conditions'] = (
+        self.internal_boundary_conditions.build_runtime_params(t)
+    )
+
     def _get_value(x):
       if isinstance(
-          x, (torax_pydantic.TimeVaryingScalar, torax_pydantic.TimeVaryingArray)
+          x,
+          (
+              torax_pydantic.TimeVaryingScalar,
+              torax_pydantic.TimeVaryingArray,
+              torax_pydantic.TimeVaryingPoints,
+          ),
       ):
         return x.get_value(t)
       else:

torax/_src/core_profiles/tests/profile_conditions_test.py

@@ -521,6 +521,32 @@ def test_multiple_validation_errors(self):
     with self.assertRaisesRegex(ValueError, '3 errors were found'):
       profile_conditions.ProfileConditions(**config_overrides)
 
+  def test_internal_boundary_conditions_set_correctly(self):
+    """Tests that internal_boundary_conditions is populated from config."""
+    config = default_configs.get_default_config_dict()
+    config['profile_conditions'] = {
+        'internal_boundary_conditions': {
+            'T_i': {
+                0.0: {0: 1.0, 1: 2.0},
+                1.0: {0: 3.0, 1: 4.0},
+            },
+        },
+    }
+    torax_config = model_config.ToraxConfig.from_dict(config)
+    runtime_params_provider = (
+        build_runtime_params.RuntimeParamsProvider.from_config(torax_config)
+    )
+
+    runtime_params = runtime_params_provider(t=0.0)
+    self.assertIsNotNone(
+        runtime_params.profile_conditions.internal_boundary_conditions
+    )
+    # Basic check to ensure the config was actually used.
+    np.testing.assert_array_equal(
+        runtime_params.profile_conditions.internal_boundary_conditions.T_i,
+        np.array([1.0, 0.0, 0.0, 2.0]),
+    )
+
 
 if __name__ == '__main__':
   absltest.main()

torax/_src/fvm/calc_coeffs.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Calculates Block1DCoeffs for a time step."""
+
 import functools
 
 import jax
@@ -26,7 +27,7 @@
 from torax._src.fvm import block_1d_coeffs
 from torax._src.fvm import cell_variable
 from torax._src.geometry import geometry
-from torax._src.internal_boundary_conditions import internal_boundary_conditions as internal_boundary_conditions_lib
+from torax._src.internal_boundary_conditions import adaptive_source
 from torax._src.pedestal_model import pedestal_model as pedestal_model_lib
 from torax._src.sources import source_profile_builders
 from torax._src.sources import source_profiles as source_profiles_lib
@@ -298,6 +299,9 @@ def _calc_coeffs_full(
   pedestal_model_output = physics_models.pedestal_model(
       runtime_params, geo, core_profiles
   )
+  internal_boundary_conditions_from_config = (
+      runtime_params.profile_conditions.internal_boundary_conditions
+  )
 
   conductivity = (
       physics_models.neoclassical_models.conductivity.calculate_conductivity(
@@ -521,27 +525,31 @@ def _calc_coeffs_full(
       * core_profiles.psi.grad()
   )
 
-  # Add internal boundary condition source terms
+  # Add internal boundary condition source terms, combining user-specified
+  # boundary conditions with pedestal model output.
+  # Note that the pedestal model will overwrite any user-specified boundary
+  # conditions, since the pedestal model is applied last.
+  combined_internal_boundary_conditions = (
+      internal_boundary_conditions_from_config.update(
+          pedestal_model_output.to_internal_boundary_conditions(geo)
+      )
+  )
   (
       source_i,
       source_e,
       source_n_e,
       source_mat_ii,
       source_mat_ee,
       source_mat_nn,
-  ) = internal_boundary_conditions_lib.apply_adaptive_source(
+  ) = adaptive_source.apply_adaptive_source(
       source_T_i=source_i,
       source_T_e=source_e,
       source_n_e=source_n_e,
       source_mat_ii=source_mat_ii,
       source_mat_ee=source_mat_ee,
       source_mat_nn=source_mat_nn,
       runtime_params=runtime_params,
-      # Pedestal contributes an internal boundary condition to the source
-      # terms at the pedestal top.
-      internal_boundary_conditions=pedestal_model_output.to_internal_boundary_conditions(
-          geo
-      ),
+      internal_boundary_conditions=combined_internal_boundary_conditions,
   )
 
   # Build arguments to solver based on which variables are evolving

torax/_src/internal_boundary_conditions/adaptive_source.py

@@ -0,0 +1,83 @@
+# Copyright 2026 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.
+
+"""Adaptive source for internal boundary conditions."""
+
+import jax.numpy as jnp
+from torax._src import array_typing
+from torax._src.config import runtime_params as runtime_params_lib
+from torax._src.internal_boundary_conditions import internal_boundary_conditions as internal_boundary_conditions_lib
+
+
+def apply_adaptive_source(
+    *,
+    source_T_i: array_typing.FloatVectorCell,
+    source_T_e: array_typing.FloatVectorCell,
+    source_n_e: array_typing.FloatVectorCell,
+    source_mat_ii: array_typing.FloatVectorCell,
+    source_mat_ee: array_typing.FloatVectorCell,
+    source_mat_nn: array_typing.FloatVectorCell,
+    runtime_params: runtime_params_lib.RuntimeParams,
+    internal_boundary_conditions: internal_boundary_conditions_lib.InternalBoundaryConditions,
+) -> tuple[
+    array_typing.FloatVectorCell,
+    array_typing.FloatVectorCell,
+    array_typing.FloatVectorCell,
+    array_typing.FloatVectorCell,
+    array_typing.FloatVectorCell,
+    array_typing.FloatVectorCell,
+]:
+  """Applies an adaptive source to the source profiles to set internal boundary conditions."""
+
+  # Ion temperature
+  source_T_i += (
+      runtime_params.numerics.adaptive_T_source_prefactor
+      * internal_boundary_conditions.T_i
+  )
+  source_mat_ii -= jnp.where(
+      internal_boundary_conditions.T_i != 0.0,
+      runtime_params.numerics.adaptive_T_source_prefactor,
+      0.0,
+  )
+
+  # Electron temperature
+  source_T_e += (
+      runtime_params.numerics.adaptive_T_source_prefactor
+      * internal_boundary_conditions.T_e
+  )
+  source_mat_ee -= jnp.where(
+      internal_boundary_conditions.T_e != 0.0,
+      runtime_params.numerics.adaptive_T_source_prefactor,
+      0.0,
+  )
+
+  # Density
+  source_n_e += (
+      runtime_params.numerics.adaptive_n_source_prefactor
+      * internal_boundary_conditions.n_e
+  )
+  source_mat_nn -= jnp.where(
+      internal_boundary_conditions.n_e != 0.0,
+      runtime_params.numerics.adaptive_n_source_prefactor,
+      0.0,
+  )
+
+  return (
+      source_T_i,
+      source_T_e,
+      source_n_e,
+      source_mat_ii,
+      source_mat_ee,
+      source_mat_nn,
+  )

torax/_src/internal_boundary_conditions/internal_boundary_conditions.py

@@ -15,14 +15,15 @@
 """Internal boundary conditions."""
 
 import dataclasses
+from typing import Annotated
 
+import chex
 import jax
 import jax.numpy as jnp
 from torax._src import array_typing
 from torax._src import jax_utils
-from torax._src.config import runtime_params as runtime_params_lib
 from torax._src.geometry import geometry
-
+from torax._src.torax_pydantic import torax_pydantic
 # pylint: disable=invalid-name
 
 
@@ -72,64 +73,24 @@ def empty(cls, geo: geometry.Geometry) -> 'InternalBoundaryConditions':
     )
 
 
-def apply_adaptive_source(
-    *,
-    source_T_i: array_typing.FloatVectorCell,
-    source_T_e: array_typing.FloatVectorCell,
-    source_n_e: array_typing.FloatVectorCell,
-    source_mat_ii: array_typing.FloatVectorCell,
-    source_mat_ee: array_typing.FloatVectorCell,
-    source_mat_nn: array_typing.FloatVectorCell,
-    runtime_params: runtime_params_lib.RuntimeParams,
-    internal_boundary_conditions: InternalBoundaryConditions,
-) -> tuple[
-    array_typing.FloatVectorCell,
-    array_typing.FloatVectorCell,
-    array_typing.FloatVectorCell,
-    array_typing.FloatVectorCell,
-    array_typing.FloatVectorCell,
-    array_typing.FloatVectorCell,
-]:
-  """Applies an adaptive source to the source profiles to set internal boundary conditions."""
-
-  # Ion temperature
-  source_T_i += (
-      runtime_params.numerics.adaptive_T_source_prefactor
-      * internal_boundary_conditions.T_i
-  )
-  source_mat_ii -= jnp.where(
-      internal_boundary_conditions.T_i != 0.0,
-      runtime_params.numerics.adaptive_T_source_prefactor,
-      0.0,
-  )
-
-  # Electron temperature
-  source_T_e += (
-      runtime_params.numerics.adaptive_T_source_prefactor
-      * internal_boundary_conditions.T_e
-  )
-  source_mat_ee -= jnp.where(
-      internal_boundary_conditions.T_e != 0.0,
-      runtime_params.numerics.adaptive_T_source_prefactor,
-      0.0,
-  )
-
-  # Density
-  source_n_e += (
-      runtime_params.numerics.adaptive_n_source_prefactor
-      * internal_boundary_conditions.n_e
-  )
-  source_mat_nn -= jnp.where(
-      internal_boundary_conditions.n_e != 0.0,
-      runtime_params.numerics.adaptive_n_source_prefactor,
-      0.0,
-  )
-
-  return (
-      source_T_i,
-      source_T_e,
-      source_n_e,
-      source_mat_ii,
-      source_mat_ee,
-      source_mat_nn,
-  )
+class InternalBoundaryConditionsConfig(torax_pydantic.BaseModelFrozen):
+  """Pydantic model for internal boundary conditions."""
+
+  # Set to zero by default, which is ignored by the adaptive source.
+  T_i: Annotated[
+      torax_pydantic.TimeVaryingPoints, torax_pydantic.JAX_STATIC
+  ] = torax_pydantic.ValidatedDefault(0.0)
+  T_e: Annotated[
+      torax_pydantic.TimeVaryingPoints, torax_pydantic.JAX_STATIC
+  ] = torax_pydantic.ValidatedDefault(0.0)
+  n_e: Annotated[
+      torax_pydantic.TimeVaryingPoints, torax_pydantic.JAX_STATIC
+  ] = torax_pydantic.ValidatedDefault(0.0)
+
+  def build_runtime_params(self, t: chex.Numeric) -> InternalBoundaryConditions:
+    """Builds the runtime params for the internal boundary conditions."""
+    kwargs = {
+        field.name: getattr(self, field.name).get_value(t)
+        for field in dataclasses.fields(InternalBoundaryConditions)
+    }
+    return InternalBoundaryConditions(**kwargs)

torax/_src/internal_boundary_conditions/tests/internal_boundary_conditions_test.py

@@ -13,14 +13,50 @@
 # limitations under the License.
 
 from absl.testing import absltest
+from absl.testing import parameterized
 import jax.numpy as jnp
 import numpy as np
+from torax._src.geometry import circular_geometry
+from torax._src.internal_boundary_conditions import adaptive_source
 from torax._src.internal_boundary_conditions import internal_boundary_conditions
+from torax._src.torax_pydantic import torax_pydantic
 
 # pylint: disable=invalid-name
 
 
-class InternalBoundaryConditionsTest(absltest.TestCase):
+class InternalBoundaryConditionsTest(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      ('initial_time', 0.0, [1.0, 0.0, 0.0, 2.0]),
+      ('intermediate_time', 0.5, [2.0, 0.0, 0.0, 3.0]),
+      ('final_time', 1.0, [3.0, 0.0, 0.0, 4.0]),
+      ('after_final_time', 1.5, [3.0, 0.0, 0.0, 4.0]),
+  )
+  def test_internal_boundary_conditions_config_build_runtime_params(
+      self, t, expected_T_i
+  ):
+    ibc_config = internal_boundary_conditions.InternalBoundaryConditionsConfig(
+        T_i={
+            0.0: {0: 1.0, 1: 2.0},
+            1.0: {0: 3.0, 1: 4.0},
+        },
+    )
+    geo = circular_geometry.CircularConfig(n_rho=4).build_geometry()
+    torax_pydantic.set_grid(ibc_config, geo.torax_mesh)
+
+    runtime_params = ibc_config.build_runtime_params(t=t)
+    np.testing.assert_array_equal(
+        runtime_params.T_i,
+        np.array(expected_T_i),
+    )
+    np.testing.assert_array_equal(
+        runtime_params.T_e,
+        np.array([0.0, 0.0, 0.0, 0.0]),
+    )
+    np.testing.assert_array_equal(
+        runtime_params.n_e,
+        np.array([0.0, 0.0, 0.0, 0.0]),
+    )
 
   def test_update(self):
     ibc1 = internal_boundary_conditions.InternalBoundaryConditions(
@@ -69,7 +105,7 @@ class MockRuntimeParams:
         source_mat_ii,
         source_mat_ee,
         source_mat_nn,
-    ) = internal_boundary_conditions.apply_adaptive_source(
+    ) = adaptive_source.apply_adaptive_source(
         source_T_i=source_T_i,
         source_T_e=source_T_e,
         source_n_e=source_n_e,