Reorganize impurity_radiation_heat sink module.

Split into separate class and model_func files in separate directory. Preparation for multiple model_func implementations.

Followup PR will be the Mavrin polynomial fits for impurity radiation.

PiperOrigin-RevId: 719243188

torax/sources/impurity_radiation_heat_sink/__init__.py

@@ -0,0 +1,15 @@
+# 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.
+
+"""This package contains the Source class and model functions for impurity radiation."""

torax/sources/impurity_radiation_heat_sink.py → torax/sources/impurity_radiation_heat_sink/impurity_radiation_constant_fraction.py

@@ -13,10 +13,9 @@
 # limitations under the License.
 
 
-"""Basic impurity radiation heat sink for electron heat equation.."""
+"""Impurity radiation heat sink for electron heat equation based on constant fraction of total power density."""
 
 import dataclasses
-from typing import ClassVar
 
 import chex
 import jax
@@ -30,6 +29,8 @@
 from torax.sources import source as source_lib
 from torax.sources import source_models as source_models_lib
 
+MODEL_FUNCTION_NAME = "radially_constant_fraction_of_Pin"
+
 
 def radially_constant_fraction_of_Pin(  # pylint: disable=invalid-name
     static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
@@ -133,27 +134,3 @@ def build_dynamic_params(
 @chex.dataclass(frozen=True)
 class DynamicRuntimeParams(runtime_params_lib.DynamicRuntimeParams):
   fraction_of_total_power_density: array_typing.ScalarFloat
-
-
-@dataclasses.dataclass(kw_only=True, frozen=True, eq=True)
-class ImpurityRadiationHeatSink(source_lib.Source):
-  """Impurity radiation heat sink for electron heat equation."""
-
-  SOURCE_NAME = "impurity_radiation_heat_sink"
-  DEFAULT_MODEL_FUNCTION_NAME: ClassVar[str] = (
-      "radially_constant_fraction_of_Pin"
-  )
-  model_func: source_lib.SourceProfileFunction = (
-      radially_constant_fraction_of_Pin
-  )
-  source_models: source_models_lib.SourceModels
-
-  @property
-  def source_name(self) -> str:
-    return self.SOURCE_NAME
-
-  @property
-  def affected_core_profiles(
-      self,
-  ) -> tuple[source_lib.AffectedCoreProfile, ...]:
-    return (source_lib.AffectedCoreProfile.TEMP_EL,)

torax/sources/impurity_radiation_heat_sink/impurity_radiation_heat_sink.py

@@ -0,0 +1,48 @@
+# 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.
+
+
+"""Class for impurity radiation heat sinks.
+
+Model functions are in separate files.
+"""
+
+import dataclasses
+from typing import ClassVar
+
+from torax.sources import source as source_lib
+from torax.sources import source_models as source_models_lib
+from torax.sources.impurity_radiation_heat_sink import impurity_radiation_constant_fraction
+
+
+@dataclasses.dataclass(kw_only=True, frozen=True, eq=True)
+class ImpurityRadiationHeatSink(source_lib.Source):
+  """Impurity radiation heat sink for electron heat equation."""
+
+  SOURCE_NAME = "impurity_radiation_heat_sink"
+  DEFAULT_MODEL_FUNCTION_NAME: ClassVar[str] = (
+      impurity_radiation_constant_fraction.MODEL_FUNCTION_NAME
+  )
+  model_func: source_lib.SourceProfileFunction
+  source_models: source_models_lib.SourceModels
+
+  @property
+  def source_name(self) -> str:
+    return self.SOURCE_NAME
+
+  @property
+  def affected_core_profiles(
+      self,
+  ) -> tuple[source_lib.AffectedCoreProfile, ...]:
+    return (source_lib.AffectedCoreProfile.TEMP_EL,)

torax/sources/register_source.py

@@ -47,12 +47,13 @@ class to build, the runtime associated with that source and (optionally) the
 from torax.sources import fusion_heat_source
 from torax.sources import generic_current_source
 from torax.sources import generic_ion_el_heat_source as ion_el_heat
-from torax.sources import impurity_radiation_heat_sink
 from torax.sources import ion_cyclotron_source
 from torax.sources import ohmic_heat_source
 from torax.sources import qei_source
 from torax.sources import runtime_params
 from torax.sources import source
+from torax.sources.impurity_radiation_heat_sink import impurity_radiation_constant_fraction
+from torax.sources.impurity_radiation_heat_sink import impurity_radiation_heat_sink
 
 
 @dataclasses.dataclass(frozen=True)
@@ -66,6 +67,7 @@ class ModelFunction:
 @dataclasses.dataclass(frozen=True)
 class SupportedSource:
   """Source that can be used in TORAX and any associated model functions."""
+
   source_class: Type[source.Source]
   model_functions: dict[str, ModelFunction]
 
@@ -196,8 +198,8 @@ class SupportedSource:
         source_class=impurity_radiation_heat_sink.ImpurityRadiationHeatSink,
         model_functions={
             impurity_radiation_heat_sink.ImpurityRadiationHeatSink.DEFAULT_MODEL_FUNCTION_NAME: ModelFunction(
-                source_profile_function=impurity_radiation_heat_sink.radially_constant_fraction_of_Pin,
-                runtime_params_class=impurity_radiation_heat_sink.RuntimeParams,
+                source_profile_function=impurity_radiation_constant_fraction.radially_constant_fraction_of_Pin,
+                runtime_params_class=impurity_radiation_constant_fraction.RuntimeParams,
                 links_back=True,
             )
         },

torax/sources/tests/impurity_radiation_heat_sink.py

@@ -25,26 +25,27 @@
 from torax.config import runtime_params_slice
 from torax.geometry import circular_geometry
 from torax.sources import generic_ion_el_heat_source
-from torax.sources import (
-    impurity_radiation_heat_sink as impurity_radiation_heat_sink_lib,
-)
 from torax.sources import runtime_params as runtime_params_lib
 from torax.sources import source as source_lib
 from torax.sources import source_models as source_models_lib
+from torax.sources.impurity_radiation_heat_sink import impurity_radiation_constant_fraction
+from torax.sources.impurity_radiation_heat_sink import (
+    impurity_radiation_heat_sink as impurity_radiation_heat_sink_lib,
+)
 from torax.sources.tests import test_lib
 
 
-class ImpurityRadiationHeatSinkTest(test_lib.SourceTestCase):
-  """Tests for ImpurityRadiationHeatSink."""
+class ImpurityRadiationConstantFractionTest(test_lib.SourceTestCase):
+  """Tests impurity_radiation_constant_fraction implementation of ImpurityRadiationHeatSink."""
 
   @classmethod
   def setUpClass(cls):
     super().setUpClass(
         source_class=impurity_radiation_heat_sink_lib.ImpurityRadiationHeatSink,
-        runtime_params_class=impurity_radiation_heat_sink_lib.RuntimeParams,
+        runtime_params_class=impurity_radiation_constant_fraction.RuntimeParams,
         source_name=impurity_radiation_heat_sink_lib.ImpurityRadiationHeatSink.SOURCE_NAME,
         links_back=True,
-        model_func=impurity_radiation_heat_sink_lib.radially_constant_fraction_of_Pin
+        model_func=impurity_radiation_constant_fraction.radially_constant_fraction_of_Pin,
     )
 
   def test_source_value(self):
@@ -120,7 +121,7 @@ def test_source_value(self):
 
     assert isinstance(
         impurity_radiation_sink_dynamic_runtime_params_slice,
-        impurity_radiation_heat_sink_lib.DynamicRuntimeParams,
+        impurity_radiation_constant_fraction.DynamicRuntimeParams,
     )
     assert isinstance(
         heat_source_dynamic_runtime_params_slice,