Extract base class from IonCyclotronSourceConfig for ICRH.

Rename IonCyclotronSourceConfig to ToricNNIonCyclotronSourceConfig and
introduce a base IonCyclotronSourceConfig class that holds shared ICRH
configuration (P_total, absorption_fraction, mode). This prepares for
future ICRH model variants beyond the ToricNN surrogate.

Key changes:
- IonCyclotronSourceConfig is now the base class with shared fields.
- ToricNNIonCyclotronSourceConfig inherits from it with ToricNN-specific
  fields (model_path, wall geometry, frequency, minority params).
- Update pydantic_model.py type annotation to ToricNNIonCyclotronSourceConfig.
- Update model_config.py validator to use isinstance check.
- Update tests to reference ToricNNIonCyclotronSourceConfig.

PiperOrigin-RevId: 892473531

Author: hamelphi

torax/_src/sources/ion_cyclotron_source.py

@@ -11,7 +11,7 @@
 # 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.
-"""Surrogate model for ion-cyclotron resonance heating (ICRH) model."""
+"""Ion-cyclotron resonance heating (ICRH) source models."""
 
 import dataclasses
 import functools
@@ -580,7 +580,7 @@ def _build_fast_ions(
 
 @dataclasses.dataclass(kw_only=True, frozen=True, eq=False)
 class IonCyclotronSource(source.Source):
-  """Ion cyclotron source with surrogate model."""
+  """Ion cyclotron source."""
 
   SOURCE_NAME: ClassVar[str] = 'icrh'
   AFFECTED_CORE_PROFILES: ClassVar[tuple[source.AffectedCoreProfile, ...]] = (
@@ -615,7 +615,37 @@ def _icrh_model_func_with_toric_nn(
 
 
 class IonCyclotronSourceConfig(base.SourceModelBase):
-  """Configuration for the IonCyclotronSource.
+  """Base configuration for IonCyclotronSource.
+
+  This base class contains fields common to all ICRH model implementations.
+  Subclasses implement the specific model logic (e.g., ToricNN, scaled curves)
+  and must override `model_name` with a `Literal` to serve as discriminator.
+
+  Attributes:
+    model_name: Discriminator field for Pydantic. Subclasses must override with
+      a `Literal` value.
+    P_total: Total heating power [W].
+    absorption_fraction: Fraction of absorbed power.
+    mode: Defines how the source values are computed.
+  """
+
+  model_name: Annotated[str, torax_pydantic.JAX_STATIC] = ''
+  P_total: torax_pydantic.TimeVaryingScalar = torax_pydantic.ValidatedDefault(
+      10e6
+  )
+  absorption_fraction: torax_pydantic.PositiveTimeVaryingScalar = (
+      torax_pydantic.ValidatedDefault(1.0)
+  )
+  mode: Annotated[source_runtime_params_lib.Mode, torax_pydantic.JAX_STATIC] = (
+      source_runtime_params_lib.Mode.MODEL_BASED
+  )
+
+  def build_source(self) -> IonCyclotronSource:
+    return IonCyclotronSource(model_func=self.model_func)
+
+
+class ToricNNIonCyclotronSourceConfig(IonCyclotronSourceConfig):
+  """Configuration for the IonCyclotronSource using the ToricNN model.
 
   Attributes:
     model_path: Path to JSON weights and model config of ToricNN model.
@@ -633,8 +663,6 @@ class IonCyclotronSourceConfig(base.SourceModelBase):
       plasma_composition instead of using minority_concentration parameter. The
       species can be either a main ion (if hydrogenic) or an impurity (if
       helium).
-    P_total: Total heating power [W].
-    absorption_fraction: Fraction of absorbed power.
   """
 
   model_name: Annotated[Literal['toric_nn'], torax_pydantic.JAX_STATIC] = (
@@ -651,15 +679,6 @@ class IonCyclotronSourceConfig(base.SourceModelBase):
       torax_pydantic.ValidatedDefault(0.03)
   )
   minority_species: Annotated[str | None, torax_pydantic.JAX_STATIC] = None
-  P_total: torax_pydantic.TimeVaryingScalar = torax_pydantic.ValidatedDefault(
-      10e6
-  )
-  absorption_fraction: torax_pydantic.PositiveTimeVaryingScalar = (
-      torax_pydantic.ValidatedDefault(1.0)
-  )
-  mode: Annotated[source_runtime_params_lib.Mode, torax_pydantic.JAX_STATIC] = (
-      source_runtime_params_lib.Mode.MODEL_BASED
-  )
 
   @property
   def model_func(self) -> source.SourceProfileFunction:
@@ -713,6 +732,3 @@ def build_runtime_params(
         P_total=self.P_total.get_value(t),
         absorption_fraction=self.absorption_fraction.get_value(t),
     )
-
-  def build_source(self) -> IonCyclotronSource:
-    return IonCyclotronSource(model_func=self.model_func)

torax/_src/sources/pydantic_model.py

@@ -91,7 +91,7 @@ class Sources(torax_pydantic.BaseModelFrozen):
       discriminator='model_name',
       default=None,
   )
-  icrh: ion_cyclotron_source_lib.IonCyclotronSourceConfig | None = (
+  icrh: ion_cyclotron_source_lib.ToricNNIonCyclotronSourceConfig | None = (
       pydantic.Field(
           discriminator='model_name',
           default=None,

torax/_src/sources/tests/ion_cyclotron_source_test.py

@@ -100,14 +100,14 @@ def setUp(self):
     self.dummy_input = model_input
     self.dummy_output = model_output
     super().setUp(
-        source_config_class=ion_cyclotron_source.IonCyclotronSourceConfig,
+        source_config_class=ion_cyclotron_source.ToricNNIonCyclotronSourceConfig,
         source_name=ion_cyclotron_source.IonCyclotronSource.SOURCE_NAME,
     )
     # pytype: enable=signature-mismatch
 
   def config_raises_if_model_path_does_not_exist(self):
     with self.assertRaises(FileNotFoundError):
-      ion_cyclotron_source.IonCyclotronSourceConfig.from_dict(
+      ion_cyclotron_source.ToricNNIonCyclotronSourceConfig.from_dict(
           {"model_path": "/tmp/non_existent_file.json"}
       )
 
@@ -338,21 +338,25 @@ def test_source_raises_if_minority_species_not_in_composition(self):
 
   def test_minority_concentration_warning_by_default(self):
     with self.assertLogs(level="WARNING") as cm:
-      ion_cyclotron_source.IonCyclotronSourceConfig()
+      ion_cyclotron_source.ToricNNIonCyclotronSourceConfig()
     self.assertTrue(
         any("minority_concentration is provided" in o for o in cm.output)
     )
 
   def test_minority_concentration_warning_when_explicitly_set_to_non_none(self):
     with self.assertLogs(level="WARNING") as cm:
-      ion_cyclotron_source.IonCyclotronSourceConfig(minority_concentration=0.05)
+      ion_cyclotron_source.ToricNNIonCyclotronSourceConfig(
+          minority_concentration=0.05
+      )
     self.assertTrue(
         any("minority_concentration is provided" in o for o in cm.output)
     )
 
   def test_minority_concentration_no_warning_when_none(self):
     with self.assertNoLogs(level="WARNING"):
-      ion_cyclotron_source.IonCyclotronSourceConfig(minority_concentration=None)
+      ion_cyclotron_source.ToricNNIonCyclotronSourceConfig(
+          minority_concentration=None
+      )
 
   def test_icrh_returns_fast_ion_data(self):
     config = default_configs.get_default_config_dict()

torax/_src/torax_pydantic/model_config.py

@@ -350,9 +350,10 @@ def _validate_toric_nn_he3_presence(self) -> typing_extensions.Self:
     is not set.
     """
     if (
-        self.sources.icrh is not None
-        and self.sources.icrh.model_name
-        == ion_cyclotron_source_lib.DEFAULT_MODEL_FUNCTION_NAME
+        isinstance(
+            self.sources.icrh,
+            ion_cyclotron_source_lib.ToricNNIonCyclotronSourceConfig,
+        )
         and self.sources.icrh.minority_species is not None
     ):
       he3_present = (