Connect MHD and sawtooth configs and models to TORAX.

1. sawtooth_model is now available in step_fn
2. mhd.sawtooth dynamic_runtime_params now available in dynamic_runtime_params_provider

Nothing yet is done with the models.

Next PRs will focus on model implementation.

PiperOrigin-RevId: 739937629

Author: jcitrin

torax/config/build_runtime_params.py

@@ -27,6 +27,8 @@
 from torax.config import runtime_params_slice
 from torax.geometry import geometry
 from torax.geometry import geometry_provider as geometry_provider_lib
+from torax.mhd import pydantic_model as mhd_pydantic_model
+from torax.mhd import runtime_params as mhd_runtime_params
 from torax.pedestal_model import pydantic_model as pedestal_pydantic_model
 from torax.sources import pydantic_model as sources_pydantic_model
 from torax.stepper import pydantic_model as stepper_pydantic_model
@@ -128,6 +130,7 @@ def __init__(
       transport: transport_model_pydantic_model.Transport | None = None,
       sources: sources_pydantic_model.Sources | None = None,
       stepper: stepper_pydantic_model.Stepper | None = None,
+      mhd: mhd_pydantic_model.MHD | None = None,
       torax_mesh: torax_pydantic.Grid1D | None = None,
   ):
     """Constructs a build_simulation_params.DynamicRuntimeParamsSliceProvider.
@@ -142,6 +145,8 @@ def __init__(
         defaults to an empty dict (i.e. no sources).
       stepper: The stepper configuration to use. If None, defaults to the
         default stepper configuration.
+      mhd: The mhd configuration to use. If None, defaults to an empty MHD
+        object.
       torax_mesh: The torax mesh to use. If the slice provider doesn't need to
         construct any rho interpolated values, this can be None, else an error
         will be raised within the constructor of the interpolated variable.
@@ -153,13 +158,15 @@ def __init__(
     sources = sources or sources_pydantic_model.Sources()
     stepper = stepper or stepper_pydantic_model.Stepper()
     pedestal = pedestal or pedestal_pydantic_model.Pedestal()
+    mhd = mhd or mhd_pydantic_model.MHD()
     torax_pydantic.set_grid(sources, torax_mesh, mode='relaxed')
     self._torax_mesh = torax_mesh
     self._sources = sources
     self._runtime_params = runtime_params
     self._transport_model = transport
     self._stepper = stepper
     self._pedestal = pedestal
+    self._mhd = mhd
 
   @property
   def sources(self) -> sources_pydantic_model.Sources:
@@ -196,6 +203,11 @@ def __call__(
         profile_conditions=dynamic_general_runtime_params.profile_conditions,
         numerics=dynamic_general_runtime_params.numerics,
         pedestal=self._pedestal.build_dynamic_params(t),
+        mhd=mhd_runtime_params.DynamicMHDParams(**{
+            mhd_model_name: mhd_model_config.build_dynamic_params(t)
+            for mhd_model_name, mhd_model_config in self._mhd
+            if mhd_model_config is not None
+        }),
     )
 
 

torax/config/runtime_params_slice.py

@@ -43,6 +43,7 @@
 from torax.config import profile_conditions
 from torax.geometry import geometry
 from torax.geometry import standard_geometry
+from torax.mhd import runtime_params as mhd_runtime_params
 from torax.pedestal_model import runtime_params as pedestal_model_params
 from torax.sources import runtime_params as sources_params
 from torax.stepper import runtime_params as stepper_params
@@ -76,10 +77,14 @@ class DynamicRuntimeParamsSlice:
   This class contains "slices" of various RuntimeParams attributes defined
   throughout TORAX:
 
-  - from the "general" runtime params
+  - from the profile_conditions runtime params
+  - from the numerics runtime params
+  - from the plasma_composition runtime params
   - from the transport model's runtime params
   - from the stepper's runtime params
   - from each of the sources' runtime params
+  - from the pedestal model's runtime params
+  - from each of the mhd models' runtime params
 
   This class packages all these together for convenience, as it simplifies many
   of the internal APIs within TORAX.
@@ -92,6 +97,7 @@ class DynamicRuntimeParamsSlice:
   numerics: numerics.DynamicNumerics
   sources: Mapping[str, sources_params.DynamicRuntimeParams]
   pedestal: pedestal_model_params.DynamicRuntimeParams
+  mhd: mhd_runtime_params.DynamicMHDParams
 
 
 @chex.dataclass(frozen=True)

torax/mhd/base.py

@@ -0,0 +1,25 @@
+# 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.
+
+"""Base classes for MHD models."""
+
+import chex
+from torax.mhd.sawtooth import sawtooth_model
+
+
+@chex.dataclass
+class MHDModels:
+  """Container for instantiated MHD model objects."""
+
+  sawtooth: sawtooth_model.SawtoothModel | None = None

torax/mhd/pydantic_model.py

@@ -0,0 +1,41 @@
+# 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.
+
+"""Pydantic config for MHD models."""
+
+import pydantic
+from torax.mhd import base
+from torax.mhd.sawtooth import pydantic_model as sawtooth_pydantic_model
+from torax.torax_pydantic import torax_pydantic
+
+
+class MHD(torax_pydantic.BaseModelFrozen):
+  """Config for MHD models.
+
+  Attributes:
+    sawtooth: Config for sawtooth models.
+  """
+
+  sawtooth: sawtooth_pydantic_model.SawtoothConfig | None = pydantic.Field(
+      default=None
+  )
+
+  def build_mhd_models(self) -> base.MHDModels:
+    """Builds and returns a container with instantiated MHD model objects."""
+
+    return base.MHDModels(
+        sawtooth=self.sawtooth.build_model()
+        if self.sawtooth is not None
+        else None,
+    )

torax/mhd/runtime_params.py

@@ -0,0 +1,25 @@
+# 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.
+
+"""Container for MHD model dynamic runtime params."""
+
+import chex
+from torax.mhd.sawtooth import runtime_params as sawtooth_runtime_params
+
+
+@chex.dataclass(frozen=True)
+class DynamicMHDParams:
+  """Container for dynamic parameters of all configured MHD models."""
+
+  sawtooth: sawtooth_runtime_params.DynamicRuntimeParams | None = None

torax/mhd/sawtooth/pydantic_model.py

@@ -45,7 +45,7 @@ class SawtoothConfig(torax_pydantic.BaseModelFrozen):
   )
   crash_step_duration: torax_pydantic.Second = 1e-3
 
-  def build_sawtooth_model(self) -> sawtooth_model.SawtoothModel:
+  def build_model(self) -> sawtooth_model.SawtoothModel:
     return sawtooth_model.SawtoothModel(
         trigger_model=self.trigger_model_config.build_trigger_model(),
         redistribution_model=self.redistribution_model_config.build_redistribution_model(),

torax/mhd/sawtooth/sawtooth_model.py

@@ -51,7 +51,7 @@ def __call__(
 
 
 class SawtoothModel:
-  """Container for sawtooth trigger and redistribution models."""
+  """Sawtooth trigger and redistribution, and carries out sawtooth step."""
 
   def __init__(
       self,

torax/mhd/sawtooth/simple_redistribution.py

@@ -39,7 +39,7 @@ def __call__(
     return core_profiles
 
 
-@chex.dataclass
+@chex.dataclass(frozen=True)
 class DynamicRuntimeParams(runtime_params.RedistributionDynamicRuntimeParams):
   # TODO(b/317360481): implement redistribution model. For now, does nothing.
   """Dynamic runtime params for simple redistribution model."""

torax/mhd/tests/pydantic_model_test.py

@@ -0,0 +1,147 @@
+# 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.
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from torax.config import build_runtime_params
+from torax.config import runtime_params as general_runtime_params
+from torax.geometry import pydantic_model as geometry_pydantic_model
+from torax.mhd import pydantic_model as mhd_pydantic_model
+from torax.mhd import runtime_params as mhd_runtime_params
+from torax.mhd.sawtooth import pydantic_model as sawtooth_pydantic_model
+from torax.mhd.sawtooth import runtime_params as sawtooth_runtime_params
+from torax.mhd.sawtooth import sawtooth_model
+from torax.stepper import pydantic_model as stepper_pydantic_model
+from torax.tests.test_lib import default_sources
+from torax.torax_pydantic import model_config
+
+
+class MHDPydanticModelTest(parameterized.TestCase):
+  """Tests for the MHD Pydantic model and dynamic params construction."""
+
+  def setUp(self):
+    super().setUp()
+    self.geo = geometry_pydantic_model.CircularConfig().build_geometry()
+    self.runtime_params = general_runtime_params.GeneralRuntimeParams()
+    self.sources = default_sources.get_default_sources()
+    self.stepper = stepper_pydantic_model.Stepper()
+
+  def test_no_mhd_config(self):
+    """Tests the case where the 'mhd' key is entirely absent."""
+    config_dict = {
+        'runtime_params': {},
+        'geometry': {'geometry_type': 'circular'},
+        'pedestal': {},
+        'sources': {},
+        'stepper': {},
+        'time_step_calculator': {},
+        'transport': {},
+    }
+    torax_config = model_config.ToraxConfig.from_dict(config_dict)
+
+    self.assertIs(torax_config.mhd, None)
+    provider = build_runtime_params.DynamicRuntimeParamsSliceProvider(
+        runtime_params=self.runtime_params,
+        sources=self.sources,
+        stepper=self.stepper,
+        torax_mesh=self.geo.torax_mesh,
+        mhd=torax_config.mhd,
+    )
+    dynamic_slice = provider(t=0.0)
+    self.assertIsInstance(
+        dynamic_slice.mhd, mhd_runtime_params.DynamicMHDParams
+    )
+    self.assertIs(dynamic_slice.mhd.sawtooth, None)
+
+  def test_empty_mhd_config(self):
+    """Tests the case where 'mhd' key exists but is an empty dict."""
+    config_dict = {
+        'runtime_params': {},
+        'geometry': {'geometry_type': 'circular'},
+        'pedestal': {},
+        'sources': {},
+        'stepper': {},
+        'time_step_calculator': {},
+        'transport': {},
+        'mhd': {},
+    }
+    torax_config = model_config.ToraxConfig.from_dict(config_dict)
+
+    self.assertIsInstance(torax_config.mhd, mhd_pydantic_model.MHD)
+    assert isinstance(torax_config.mhd, mhd_pydantic_model.MHD)
+    mhd_models = torax_config.mhd.build_mhd_models()
+    self.assertIs(mhd_models.sawtooth, None)
+    provider = build_runtime_params.DynamicRuntimeParamsSliceProvider(
+        runtime_params=self.runtime_params,
+        sources=self.sources,
+        stepper=self.stepper,
+        torax_mesh=self.geo.torax_mesh,
+        mhd=torax_config.mhd,
+    )
+    dynamic_slice = provider(t=0.0)
+    self.assertIsInstance(
+        dynamic_slice.mhd, mhd_runtime_params.DynamicMHDParams
+    )
+    self.assertIs(dynamic_slice.mhd.sawtooth, None)
+
+  def test_mhd_config_with_sawtooth(self):
+    """Tests the case with a valid sawtooth configuration."""
+    config_dict = {
+        'runtime_params': {},
+        'geometry': {'geometry_type': 'circular'},
+        'pedestal': {},
+        'sources': {},
+        'stepper': {},
+        'time_step_calculator': {},
+        'transport': {},
+        'mhd': {
+            'sawtooth': {
+                'trigger_model_config': {'trigger_model_type': 'simple'},
+                'redistribution_model_config': {
+                    'redistribution_model_type': 'simple'
+                },
+                'minimum_radius': 0.06,
+            }
+        },
+    }
+    torax_config = model_config.ToraxConfig.from_dict(config_dict)
+
+    self.assertIsInstance(torax_config.mhd, mhd_pydantic_model.MHD)
+    assert torax_config.mhd is not None
+    self.assertIsInstance(
+        torax_config.mhd.sawtooth, sawtooth_pydantic_model.SawtoothConfig
+    )
+
+    mhd_models = torax_config.mhd.build_mhd_models()
+    self.assertIn('sawtooth', mhd_models)
+    self.assertIsInstance(mhd_models['sawtooth'], sawtooth_model.SawtoothModel)
+
+    provider = build_runtime_params.DynamicRuntimeParamsSliceProvider(
+        runtime_params=self.runtime_params,
+        sources=self.sources,
+        stepper=self.stepper,
+        torax_mesh=self.geo.torax_mesh,
+        mhd=torax_config.mhd,
+    )
+    dynamic_slice = provider(t=0.0)
+    self.assertIn('sawtooth', dynamic_slice.mhd)
+    sawtooth_dynamic_params = dynamic_slice.mhd.sawtooth
+    self.assertIsInstance(
+        sawtooth_dynamic_params, sawtooth_runtime_params.DynamicRuntimeParams
+    )
+    self.assertEqual(sawtooth_dynamic_params.minimum_radius, 0.06)
+
+
+if __name__ == '__main__':
+  absltest.main()

torax/orchestration/run_simulation.py

@@ -55,11 +55,18 @@ def run_simulation(
       pedestal_model=pedestal_model,
   )
 
+  mhd_models = (
+      torax_config.mhd.build_mhd_models()
+      if torax_config.mhd is not None
+      else None
+  )
+
   step_fn = step_function.SimulationStepFn(
       stepper=stepper,
       time_step_calculator=torax_config.time_step_calculator.time_step_calculator,
       transport_model=transport_model,
       pedestal_model=pedestal_model,
+      mhd_models=mhd_models,
   )
 
   static_runtime_params_slice = (
@@ -78,6 +85,7 @@ def run_simulation(
           transport=torax_config.transport,
           sources=torax_config.sources,
           stepper=torax_config.stepper,
+          mhd=torax_config.mhd,
           torax_mesh=torax_config.geometry.build_provider.torax_mesh,
       )
   )