Split sim.py into two and move some bits closer to the objects they r…

…elate to

PiperOrigin-RevId: 715742159

docs/model_integration.rst

@@ -312,14 +312,6 @@ the |torax.sim.Sim|_ object.
 
 .. code-block:: python
 
-    # in sim.py. Copied here for reference, no need to modify this.
-
-    def build_sim_object(
-        ...
-        transport_model_builder: transport_model_lib.TransportModelBuilder,
-        ...
-    ) -> Sim:
-
     # in your TORAX configuration or run file .py
 
     my_custom_transport_builder = MyCustomTransportModelBuilder()
@@ -334,7 +326,7 @@ the |torax.sim.Sim|_ object.
     my_custom_transport_builder.runtime_params.bar = 4.0
 
     # Build the Sim object.
-    sim_object = sim_lib.build_sim_object(
+    sim_object = sim_lib.Sim.create(
         ...,
         transport_model_builder=my_custom_transport_builder,
         ...
@@ -348,7 +340,7 @@ As of 7 June 2024, you cannot instantiate and configure a custom transport model
 via the config dictionary. You may still configure the other components of your
 TORAX simulation via the config dict and use other functions in
 |torax.config.build_sim|_ to convert those to the objects you can pass into
-``build_sim_object()``. We are working on making this easier, but reach out
+``Sim.create()``. We are working on making this easier, but reach out
 if this is something you need.
 
 

torax/config/build_sim.py

@@ -270,7 +270,7 @@ def build_sim_from_config(
   else:
     file_restart = None
 
-  return sim_lib.build_sim_object(
+  return sim_lib.Sim.create(
       runtime_params=runtime_params,
       geometry_provider=geo_provider,
       source_models_builder=build_sources_builder_from_config(

torax/config/runtime_params_slice.py

@@ -46,6 +46,7 @@
 from torax.config import profile_conditions
 from torax.config import runtime_params as general_runtime_params_lib
 from torax.geometry import geometry
+from torax.geometry import geometry_provider as geometry_provider_lib
 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
@@ -345,6 +346,23 @@ def __call__(
     )
 
 
+def get_consistent_dynamic_runtime_params_slice_and_geometry(
+    *,
+    t: chex.Numeric,
+    dynamic_runtime_params_slice_provider: DynamicRuntimeParamsSliceProvider,
+    geometry_provider: geometry_provider_lib.GeometryProvider,
+) -> tuple[DynamicRuntimeParamsSlice, geometry.Geometry]:
+  """Returns the dynamic runtime params and geometry for a given time."""
+  geo = geometry_provider(t)
+  dynamic_runtime_params_slice = dynamic_runtime_params_slice_provider(
+      t=t,
+  )
+  dynamic_runtime_params_slice, geo = make_ip_consistent(
+      dynamic_runtime_params_slice, geo
+  )
+  return dynamic_runtime_params_slice, geo
+
+
 def make_ip_consistent(
     dynamic_runtime_params_slice: DynamicRuntimeParamsSlice,
     geo: geometry.Geometry,

torax/config/tests/build_sim.py

@@ -291,12 +291,12 @@ def test_chease_geometry_updates_Ip(self):
             torax_mesh=geo_provider.torax_mesh,
         )
     )
-    geo = geo_provider(t=0)
-    dynamic_runtime_params_slice = runtime_params_provider(
-        t=0,
-    )
-    dynamic_slice, geo = runtime_params_slice.make_ip_consistent(
-        dynamic_runtime_params_slice, geo
+    dynamic_slice, geo = (
+        runtime_params_slice.get_consistent_dynamic_runtime_params_slice_and_geometry(
+            t=0,
+            dynamic_runtime_params_slice_provider=runtime_params_provider,
+            geometry_provider=geo_provider,
+        )
     )
     self.assertIsInstance(geo, geometry.StandardGeometry)
     self.assertIsNotNone(dynamic_slice)