Changes to start simplifying sim.py

Also fixes that we were using old geo and state for the final step

PiperOrigin-RevId: 706967398

torax/config/base.py

@@ -22,16 +22,17 @@
 from typing import Any, Generic, TypeVar
 
 import chex
-from torax import geometry
 from torax import interpolated_param
 from torax.config import config_args
+from torax.geometry import geometry
 
 DynamicT = TypeVar('DynamicT')
 ProviderT = TypeVar('ProviderT', bound='RuntimeParametersProvider')
 
 
 class GridType(enum.Enum):
   """Describes where interpolated values are defined on."""
+
   CELL = enum.auto()
   FACE = enum.auto()
 
@@ -76,7 +77,8 @@ def get_provider_kwargs(
 
     Args:
       torax_mesh: Required if any of the interpolated variables are both
-      temporally and radially interpolated.
+        temporally and radially interpolated.
+
     Returns:
       A dict of kwargs to be passed to the provider constructor.
     """
@@ -161,6 +163,7 @@ class RuntimeParametersProvider(Generic[DynamicT], metaclass=abc.ABCMeta):
   - any constructed interpolated variables which will already be spatially
   interpolated and could vary in time.
   """
+
   runtime_params_config: RuntimeParametersConfig
 
   def get_dynamic_params_kwargs(
@@ -177,6 +180,7 @@ def get_dynamic_params_kwargs(
 
     Args:
       t: The time to interpolate the dynamic parameters at.
+
     Returns:
       A dict of kwargs to be passed to the dynamic params constructor.
     """

torax/config/build_sim.py

@@ -18,11 +18,11 @@
 import copy
 from typing import Any
 
-from torax import geometry
-from torax import geometry_provider
 from torax import sim as sim_lib
 from torax.config import config_args
 from torax.config import runtime_params as runtime_params_lib
+from torax.geometry import geometry
+from torax.geometry import geometry_provider
 from torax.pedestal_model import pedestal_model as pedestal_model_lib
 from torax.pedestal_model import set_tped_nped
 from torax.sources import formula_config
@@ -45,6 +45,7 @@
 # pylint: disable=g-import-not-at-top
 try:
   from torax.transport_model import qualikiz_transport_model
+
   _QUALIKIZ_TRANSPORT_MODEL_AVAILABLE = True
 except ImportError:
   _QUALIKIZ_TRANSPORT_MODEL_AVAILABLE = False

torax/config/numerics.py

@@ -18,9 +18,9 @@
 
 import chex
 from torax import array_typing
-from torax import geometry
 from torax import interpolated_param
 from torax.config import base
+from torax.geometry import geometry
 from typing_extensions import override
 
 

torax/config/plasma_composition.py

@@ -21,10 +21,10 @@
 
 import chex
 from torax import array_typing
-from torax import geometry
 from torax import interpolated_param
 from torax.config import base
 from torax.config import config_args
+from torax.geometry import geometry
 
 
 # pylint: disable=invalid-name

torax/config/profile_conditions.py

@@ -21,10 +21,10 @@
 
 import chex
 from torax import array_typing
-from torax import geometry
 from torax import interpolated_param
 from torax.config import base
 from torax.config import config_args
+from torax.geometry import geometry
 from typing_extensions import override
 
 
@@ -84,7 +84,7 @@ class ProfileConditions(
   ne_bound_right: interpolated_param.TimeInterpolatedInput | None = None
   ne_bound_right_is_fGW: bool = False
   ne_bound_right_is_absolute: bool = False
-# Internal boundary condition (pedestal)
+  # Internal boundary condition (pedestal)
   # Do not set internal boundary condition if this is False
   set_pedestal: interpolated_param.TimeInterpolatedInput = True
   # current profiles (broad "Ohmic" + localized "external" currents)

torax/config/runtime_params.py

@@ -19,11 +19,11 @@
 import dataclasses
 
 import chex
-from torax import geometry
 from torax.config import base
 from torax.config import numerics as numerics_lib
 from torax.config import plasma_composition as plasma_composition_lib
 from torax.config import profile_conditions as profile_conditions_lib
+from torax.geometry import geometry
 from typing_extensions import override
 
 
@@ -92,6 +92,7 @@ def build_dynamic_params(
 @chex.dataclass
 class DynamicGeneralRuntimeParams:
   """General runtime input parameters for the `torax` module."""
+
   plasma_composition: plasma_composition_lib.DynamicPlasmaComposition
   profile_conditions: profile_conditions_lib.DynamicProfileConditions
   numerics: numerics_lib.DynamicNumerics

torax/config/runtime_params_slice.py

@@ -41,11 +41,11 @@
 import dataclasses
 
 import chex
-from torax import geometry
 from torax.config import numerics
 from torax.config import plasma_composition
 from torax.config import profile_conditions
 from torax.config import runtime_params as general_runtime_params_lib
+from torax.geometry import geometry
 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
@@ -148,7 +148,9 @@ def _build_dynamic_sources(
 ) -> dict[str, sources_params.DynamicRuntimeParams]:
   """Builds a dict of DynamicSourceConfigSlice based on the input config."""
   return {
-      source_name: input_source_config.build_dynamic_params(t,)
+      source_name: input_source_config.build_dynamic_params(
+          t,
+      )
       for source_name, input_source_config in sources.items()
   }
 
@@ -262,15 +264,11 @@ def runtime_params_provider(
 
   def _construct_providers(self):
     """Construct the providers that will give us the dynamic params."""
-    self._runtime_params_provider = (
-        self._runtime_params.make_provider(
-            self._torax_mesh
-        )
+    self._runtime_params_provider = self._runtime_params.make_provider(
+        self._torax_mesh
     )
     self._transport_runtime_params_provider = (
-        self._transport_runtime_params.make_provider(
-            self._torax_mesh
-        )
+        self._transport_runtime_params.make_provider(self._torax_mesh)
     )
     self._sources_providers = {
         key: source.make_provider(self._torax_mesh)

torax/config/tests/build_sim.py

@@ -17,11 +17,11 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import numpy as np
-from torax import geometry
-from torax import geometry_provider
 from torax.config import build_sim
 from torax.config import runtime_params as runtime_params_lib
 from torax.config import runtime_params_slice
+from torax.geometry import geometry
+from torax.geometry import geometry_provider
 from torax.pedestal_model import set_tped_nped
 from torax.sources import formula_config
 from torax.sources import formulas

torax/config/tests/numerics.py

@@ -16,9 +16,9 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
-from torax import geometry
 from torax import interpolated_param
 from torax.config import numerics
+from torax.geometry import geometry
 
 
 class NumericsTest(parameterized.TestCase):

torax/config/tests/plasma_composition.py

@@ -17,9 +17,9 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import numpy as np
-from torax import geometry
 from torax import interpolated_param
 from torax.config import plasma_composition
+from torax.geometry import geometry
 
 
 class PlasmaCompositionTest(parameterized.TestCase):

torax/config/tests/profile_conditions.py

@@ -17,10 +17,10 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import numpy as np
-from torax import geometry
 from torax import interpolated_param
 from torax.config import config_args
 from torax.config import profile_conditions
+from torax.geometry import geometry
 import xarray as xr
 
 

torax/config/tests/runtime_params.py

@@ -18,10 +18,10 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
-from torax import geometry
 from torax.config import config_args
 from torax.config import profile_conditions as profile_conditions_lib
 from torax.config import runtime_params as general_runtime_params
+from torax.geometry import geometry
 
 
 # pylint: disable=invalid-name
@@ -94,7 +94,9 @@ def test_recursive_replace(self):
     self.assertEqual(result.a4.b4.c1, 19)
     self.assertEqual(result.a4.b4.c2, 20)
 
-  def test_runtime_params_raises_for_invalid_temp_boundary_conditions(self,):
+  def test_runtime_params_raises_for_invalid_temp_boundary_conditions(
+      self,
+  ):
     """Tests that runtime params validate boundary conditions."""
     with self.assertRaises(ValueError):
       general_runtime_params.GeneralRuntimeParams(
@@ -104,12 +106,23 @@ def test_runtime_params_raises_for_invalid_temp_boundary_conditions(self,):
       )
 
   @parameterized.parameters(
-      ({0.0: {0.0: 12.0, 1.0: 2.0}}, None,),  # Ti includes 1.0.
-      ({0.0: {0.0: 12.0, 1.0: 2.0}}, 1.0,),  # Both provided.
-      ({0.0: {0.0: 12.0, 0.95: 2.0}}, 1.0,)  # Ti_bound_right provided.
+      (
+          {0.0: {0.0: 12.0, 1.0: 2.0}},
+          None,
+      ),  # Ti includes 1.0.
+      (
+          {0.0: {0.0: 12.0, 1.0: 2.0}},
+          1.0,
+      ),  # Both provided.
+      (
+          {0.0: {0.0: 12.0, 0.95: 2.0}},
+          1.0,
+      ),  # Ti_bound_right provided.
   )
   def test_runtime_params_constructs_with_valid_profile_conditions(
-      self, Ti, Ti_bound_right,
+      self,
+      Ti,
+      Ti_bound_right,
   ):
     """Tests that runtime params validate boundary conditions."""
     general_runtime_params.GeneralRuntimeParams(

torax/config/tests/runtime_params_slice.py

@@ -20,10 +20,10 @@
 from absl.testing import parameterized
 import jax
 import numpy as np
-from torax import geometry
 from torax.config import profile_conditions as profile_conditions_lib
 from torax.config import runtime_params as general_runtime_params
 from torax.config import runtime_params_slice as runtime_params_slice_lib
+from torax.geometry import geometry
 from torax.pedestal_model import set_tped_nped
 from torax.sources import electron_density_sources
 from torax.sources import formula_config
@@ -74,11 +74,15 @@ def test_time_dependent_provider_is_time_dependent(self):
         stepper=stepper_params_lib.RuntimeParams(),
         torax_mesh=self._geo.torax_mesh,
     )
-    dynamic_runtime_params_slice = provider(t=1.0,)
+    dynamic_runtime_params_slice = provider(
+        t=1.0,
+    )
     np.testing.assert_allclose(
         dynamic_runtime_params_slice.profile_conditions.Ti_bound_right, 2.5
     )
-    dynamic_runtime_params_slice = provider(t=2.0,)
+    dynamic_runtime_params_slice = provider(
+        t=2.0,
+    )
     np.testing.assert_allclose(
         dynamic_runtime_params_slice.profile_conditions.Ti_bound_right, 3.0
     )
@@ -422,9 +426,18 @@ def test_profile_conditions_set_electron_temperature_and_boundary_condition(
     )
 
   @parameterized.product(
-      ne_bound_right=[None, 1.0,],
-      ne_bound_right_is_fGW=[True, False,],
-      ne_is_fGW=[True, False,],
+      ne_bound_right=[
+          None,
+          1.0,
+      ],
+      ne_bound_right_is_fGW=[
+          True,
+          False,
+      ],
+      ne_is_fGW=[
+          True,
+          False,
+      ],
   )
   def test_profile_conditions_set_electron_density_and_boundary_condition(
       self,
@@ -630,5 +643,6 @@ def test_static_runtime_params_slice_hash_different_for_different_params(
     )
     self.assertNotEqual(hash(static_slice1), hash(static_slice2))
 
+
 if __name__ == '__main__':
   absltest.main()