Delete sim.Sim which has been replaced by run_simulation

PiperOrigin-RevId: 738807490

Author: tamaranorman

Diff for: torax/sim.py

@@ -26,7 +26,6 @@
 
 import dataclasses
 import time
-from typing import Optional
 
 from absl import logging
 import jax
@@ -36,27 +35,13 @@
 from torax import post_processing
 from torax import state
 from torax.config import build_runtime_params
-from torax.config import config_args
-from torax.config import runtime_params as general_runtime_params
 from torax.config import runtime_params_slice
 from torax.core_profiles import initialization
 from torax.geometry import geometry
 from torax.geometry import geometry_provider as geometry_provider_lib
 from torax.orchestration import step_function
-from torax.pedestal_model import pedestal_model as pedestal_model_lib
-from torax.pedestal_model import pydantic_model as pedestal_pydantic_model
-from torax.sources import pydantic_model as source_pydantic_model
-from torax.sources import source_models as source_models_lib
 from torax.sources import source_profile_builders
-from torax.stepper import pydantic_model as stepper_pydantic_model
-from torax.stepper import stepper as stepper_lib
-from torax.time_step_calculator import chi_time_step_calculator
-from torax.time_step_calculator import time_step_calculator as ts
-from torax.torax_pydantic import file_restart as file_restart_pydantic_model
-from torax.transport_model import pydantic_model as transport_model_pydantic_model
-from torax.transport_model import transport_model as transport_model_lib
 import tqdm
-import typing_extensions
 import xarray as xr
 
 
@@ -102,347 +87,6 @@ def get_initial_state(
   )
 
 
-class Sim:
-  """A lightweight object holding all components of a simulation.
-
-  Use of this object is optional, it is also fine to hold these objects
-  in local variables of a script and call `run_simulation` directly.
-
-  The main purpose of the Sim object is to enable configuration via
-  constructor arguments. Components are reused in subsequent simulation runs, so
-  if a component is compiled, it will be reused for the next `Sim.run()` call
-  and will not be recompiled unless a static argument or shape changes.
-  """
-
-  def __init__(
-      self,
-      static_runtime_params_slice: runtime_params_slice.StaticRuntimeParamsSlice,
-      dynamic_runtime_params_slice_provider: build_runtime_params.DynamicRuntimeParamsSliceProvider,
-      geometry_provider: geometry_provider_lib.GeometryProvider,
-      initial_state: state.ToraxSimState,
-      step_fn: step_function.SimulationStepFn,
-      file_restart: file_restart_pydantic_model.FileRestart | None = None,
-  ):
-    self._static_runtime_params_slice = static_runtime_params_slice
-    self._dynamic_runtime_params_slice_provider = (
-        dynamic_runtime_params_slice_provider
-    )
-    self._geometry_provider = geometry_provider
-    self._initial_state = initial_state
-    self._step_fn = step_fn
-    self._file_restart = file_restart
-
-  @property
-  def file_restart(self) -> file_restart_pydantic_model.FileRestart | None:
-    return self._file_restart
-
-  @property
-  def time_step_calculator(self) -> ts.TimeStepCalculator:
-    return self._step_fn.time_step_calculator
-
-  @property
-  def initial_state(self) -> state.ToraxSimState:
-    return self._initial_state
-
-  @property
-  def geometry_provider(self) -> geometry_provider_lib.GeometryProvider:
-    return self._geometry_provider
-
-  @property
-  def dynamic_runtime_params_slice_provider(
-      self,
-  ) -> build_runtime_params.DynamicRuntimeParamsSliceProvider:
-    return self._dynamic_runtime_params_slice_provider
-
-  @property
-  def static_runtime_params_slice(
-      self,
-  ) -> runtime_params_slice.StaticRuntimeParamsSlice:
-    return self._static_runtime_params_slice
-
-  @property
-  def step_fn(self) -> step_function.SimulationStepFn:
-    return self._step_fn
-
-  @property
-  def stepper(self) -> stepper_lib.Stepper:
-    return self._step_fn.stepper
-
-  @property
-  def transport_model(self) -> transport_model_lib.TransportModel:
-    return self.stepper.transport_model
-
-  @property
-  def pedestal_model(self) -> pedestal_model_lib.PedestalModel:
-    return self.stepper.pedestal_model
-
-  @property
-  def source_models(self) -> source_models_lib.SourceModels:
-    return self.stepper.source_models
-
-  def update_base_components(
-      self,
-      *,
-      allow_recompilation: bool = False,
-      static_runtime_params_slice: (
-          runtime_params_slice.StaticRuntimeParamsSlice | None
-      ) = None,
-      dynamic_runtime_params_slice_provider: (
-          build_runtime_params.DynamicRuntimeParamsSliceProvider | None
-      ) = None,
-      geometry_provider: geometry_provider_lib.GeometryProvider | None = None,
-  ):
-    """Updates the Sim object with components that have already been updated.
-
-    Currently this only supports updating the geometry provider and the dynamic
-    runtime params slice provider, both of which can be updated without
-    recompilation.
-
-    Args:
-      allow_recompilation: Whether recompilation is allowed. If True, the static
-        runtime params slice can be updated. NOTE: recompilaton may still occur
-        if the mesh is updated or if the shapes returned in the dynamic runtime
-        params slice provider change even if this is False.
-      static_runtime_params_slice: The new static runtime params slice. If None,
-        the existing one is kept.
-      dynamic_runtime_params_slice_provider: The new dynamic runtime params
-        slice provider. This should already have been updated with modifications
-        to the various components. If None, the existing one is kept.
-      geometry_provider: The new geometry provider. If None, the existing one is
-        kept.
-
-    Raises:
-      ValueError: If the Sim object has a file restart or if the geometry
-        provider has a different mesh than the existing one.
-    """
-    if self._file_restart is not None:
-      # TODO(b/384767453): Add support for updating a Sim object with a file
-      # restart.
-      raise ValueError('Cannot update a Sim object with a file restart.')
-    if not allow_recompilation and static_runtime_params_slice is not None:
-      raise ValueError(
-          'Cannot update a Sim object with a static runtime params slice if '
-          'recompilation is not allowed.'
-      )
-
-    if static_runtime_params_slice is not None:
-      assert isinstance(  # Avoid pytype error.
-          self._static_runtime_params_slice,
-          runtime_params_slice.StaticRuntimeParamsSlice,
-      )
-      self._static_runtime_params_slice.validate_new(
-          static_runtime_params_slice
-      )
-      self._static_runtime_params_slice = static_runtime_params_slice
-    if dynamic_runtime_params_slice_provider is not None:
-      assert isinstance(  # Avoid pytype error.
-          self._dynamic_runtime_params_slice_provider,
-          build_runtime_params.DynamicRuntimeParamsSliceProvider,
-      )
-      self._dynamic_runtime_params_slice_provider.validate_new(
-          dynamic_runtime_params_slice_provider
-      )
-      self._dynamic_runtime_params_slice_provider = (
-          dynamic_runtime_params_slice_provider
-      )
-    if geometry_provider is not None:
-      if geometry_provider.torax_mesh != self._geometry_provider.torax_mesh:
-        raise ValueError(
-            'Cannot update a Sim object with a geometry provider with a '
-            'different mesh.'
-        )
-      self._geometry_provider = geometry_provider
-
-    dynamic_runtime_params_slice_for_init, geo_for_init = (
-        build_runtime_params.get_consistent_dynamic_runtime_params_slice_and_geometry(
-            t=self._dynamic_runtime_params_slice_provider._runtime_params.numerics.t_initial,  # pylint: disable=protected-access
-            dynamic_runtime_params_slice_provider=self._dynamic_runtime_params_slice_provider,
-            geometry_provider=self._geometry_provider,
-        )
-    )
-    self._initial_state = get_initial_state(
-        static_runtime_params_slice=self._static_runtime_params_slice,
-        dynamic_runtime_params_slice=dynamic_runtime_params_slice_for_init,
-        geo=geo_for_init,
-        step_fn=self._step_fn,
-    )
-
-  def run(
-      self,
-      log_timestep_info: bool = False,
-  ) -> output.ToraxSimOutputs:
-    """Runs the transport simulation over a prescribed time interval.
-
-    See `run_simulation` for details.
-
-    Args:
-      log_timestep_info: See `run_simulation()`.
-
-    Returns:
-      Tuple of all ToraxSimStates, one per time step and an additional one at
-      the beginning for the starting state.
-    """
-    return _run_simulation(
-        static_runtime_params_slice=self.static_runtime_params_slice,
-        dynamic_runtime_params_slice_provider=self.dynamic_runtime_params_slice_provider,
-        geometry_provider=self.geometry_provider,
-        initial_state=self.initial_state,
-        step_fn=self.step_fn,
-        log_timestep_info=log_timestep_info,
-    )
-
-  @classmethod
-  def create(
-      cls,
-      *,
-      runtime_params: general_runtime_params.GeneralRuntimeParams,
-      geometry_provider: geometry_provider_lib.GeometryProvider,
-      stepper: stepper_pydantic_model.Stepper,
-      transport_model: transport_model_pydantic_model.Transport,
-      sources: source_pydantic_model.Sources,
-      pedestal: pedestal_pydantic_model.Pedestal,
-      time_step_calculator: Optional[ts.TimeStepCalculator] = None,
-      file_restart: file_restart_pydantic_model.FileRestart | None = None,
-  ) -> typing_extensions.Self:
-    """Builds a Sim object from the input runtime params and sim components.
-
-    Args:
-      runtime_params: The input runtime params used throughout the simulation
-        run.
-      geometry_provider: The geometry used throughout the simulation run.
-      stepper: The stepper config that can be used to build the stepper.
-      transport_model: The transport model config that can be used to build the
-        transport model.
-      sources: Builds the sources.
-      pedestal: The pedestal config that can be used to build the pedestal.
-      time_step_calculator: The time_step_calculator, if built, otherwise a
-        ChiTimeStepCalculator will be built by default.
-      file_restart: If provided we will reconstruct the initial state from the
-        provided file at the given time step. This state from the file will only
-        be used for constructing the initial state (as well as the config) and
-        for all subsequent steps, the evolved state and runtime parameters from
-        config are used.
-
-    Returns:
-      sim: The built Sim instance.
-    """
-    pedestal_model = pedestal.build_pedestal_model()
-
-    # TODO(b/385788907): Document all changes that lead to recompilations.
-    static_runtime_params_slice = (
-        build_runtime_params.build_static_runtime_params_slice(
-            runtime_params=runtime_params,
-            sources=sources,
-            torax_mesh=geometry_provider.torax_mesh,
-            stepper=stepper,
-        )
-    )
-    dynamic_runtime_params_slice_provider = (
-        build_runtime_params.DynamicRuntimeParamsSliceProvider(
-            runtime_params=runtime_params,
-            transport=transport_model,
-            sources=sources,
-            stepper=stepper,
-            torax_mesh=geometry_provider.torax_mesh,
-            pedestal=pedestal,
-        )
-    )
-    source_models = source_models_lib.SourceModels(
-        sources=sources.source_model_config
-    )
-    transport_model = transport_model.build_transport_model()
-    stepper_model = stepper.build_stepper_model(
-        transport_model=transport_model,
-        source_models=source_models,
-        pedestal_model=pedestal_model,
-    )
-
-    if time_step_calculator is None:
-      time_step_calculator = chi_time_step_calculator.ChiTimeStepCalculator()
-
-    # Build dynamic_runtime_params_slice at t_initial for initial conditions.
-    dynamic_runtime_params_slice_for_init, geo_for_init = (
-        build_runtime_params.get_consistent_dynamic_runtime_params_slice_and_geometry(
-            t=runtime_params.numerics.t_initial,
-            dynamic_runtime_params_slice_provider=dynamic_runtime_params_slice_provider,
-            geometry_provider=geometry_provider,
-        )
-    )
-    if file_restart is not None and file_restart.do_restart:
-      data_tree = output.load_state_file(file_restart.filename)
-      # Find the closest time in the given dataset.
-      data_tree = data_tree.sel(time=file_restart.time, method='nearest')
-      t_restart = data_tree.time.item()
-      core_profiles_dataset = data_tree.children[output.CORE_PROFILES].dataset
-      # Remap coordinates in saved file to be consistent with expectations of
-      # how config_args parses xarrays.
-      core_profiles_dataset = core_profiles_dataset.rename(
-          {output.RHO_CELL_NORM: config_args.RHO_NORM}
-      )
-      core_profiles_dataset = core_profiles_dataset.squeeze()
-      if t_restart != runtime_params.numerics.t_initial:
-        logging.warning(
-            'Requested restart time %f not exactly available in state file %s.'
-            ' Restarting from closest available time %f instead.',
-            file_restart.time,
-            file_restart.filename,
-            t_restart,
-        )
-      # Override some of dynamic runtime params slice from t=t_initial.
-      dynamic_runtime_params_slice_for_init, geo_for_init = (
-          _override_initial_runtime_params_from_file(
-              dynamic_runtime_params_slice_for_init,
-              geo_for_init,
-              t_restart,
-              core_profiles_dataset,
-          )
-      )
-      post_processed_dataset = data_tree.children[
-          output.POST_PROCESSED_OUTPUTS
-      ].dataset
-      post_processed_dataset = post_processed_dataset.rename(
-          {output.RHO_CELL_NORM: config_args.RHO_NORM}
-      )
-      post_processed_dataset = post_processed_dataset.squeeze()
-      post_processed_outputs = (
-          _override_initial_state_post_processed_outputs_from_file(
-              geo_for_init,
-              post_processed_dataset,
-          )
-      )
-
-    step_fn = step_function.SimulationStepFn(
-        stepper=stepper_model,
-        time_step_calculator=time_step_calculator,
-        transport_model=transport_model,
-        pedestal_model=pedestal_model,
-    )
-
-    initial_state = get_initial_state(
-        static_runtime_params_slice=static_runtime_params_slice,
-        dynamic_runtime_params_slice=dynamic_runtime_params_slice_for_init,
-        geo=geo_for_init,
-        step_fn=step_fn,
-    )
-
-    # If we are restarting from a file, we need to override the initial state
-    # post processed outputs such that cumulative outputs remain correct.
-    if file_restart is not None and file_restart.do_restart:
-      initial_state = dataclasses.replace(
-          initial_state,
-          post_processed_outputs=post_processed_outputs,  # pylint: disable=undefined-variable
-      )
-
-    return cls(
-        static_runtime_params_slice=static_runtime_params_slice,
-        dynamic_runtime_params_slice_provider=dynamic_runtime_params_slice_provider,
-        geometry_provider=geometry_provider,
-        initial_state=initial_state,
-        step_fn=step_fn,
-        file_restart=file_restart,
-    )
-
-
 def _override_initial_runtime_params_from_file(
     dynamic_runtime_params_slice: runtime_params_slice.DynamicRuntimeParamsSlice,
     geo: geometry.Geometry,