Make PedestalModel a StaticDataclass

PiperOrigin-RevId: 822642397

Author: goodfeli

torax/_src/pedestal_model/no_pedestal.py

@@ -12,13 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """A pedestal model for when there is no pedestal."""
+import dataclasses
 from jax import numpy as jnp
 from torax._src import state
 from torax._src.config import runtime_params_slice
 from torax._src.geometry import geometry
 from torax._src.pedestal_model import pedestal_model
 
 
+@dataclasses.dataclass(frozen=True, eq=False)
 class NoPedestal(pedestal_model.PedestalModel):
   """A pedestal model for when there is no pedestal.
 
@@ -29,10 +31,6 @@ class NoPedestal(pedestal_model.PedestalModel):
   for the jax cond to work.
   """
 
-  def __init__(self):
-    super().__init__()
-    self._frozen = True
-
   def _call_implementation(
       self,
       runtime_params: runtime_params_slice.RuntimeParams,
@@ -46,9 +44,3 @@ def _call_implementation(
         n_e_ped=0.0,
         rho_norm_ped_top_idx=geo.torax_mesh.nx,
     )
-
-  def __hash__(self):
-    return hash('NoPedestal')
-
-  def __eq__(self, other) -> bool:
-    return isinstance(other, NoPedestal)

torax/_src/pedestal_model/pedestal_model.py

@@ -23,6 +23,7 @@
 import jax.numpy as jnp
 from torax._src import array_typing
 from torax._src import state
+from torax._src import static_dataclass
 from torax._src.config import runtime_params_slice
 from torax._src.geometry import geometry
 
@@ -47,38 +48,16 @@ class PedestalModelOutput:
   n_e_ped: array_typing.FloatScalar
 
 
-class PedestalModel(abc.ABC):
-  """Calculates temperature and density of the pedestal.
-
-  Subclass responsbilities:
-  - Must set _frozen = True at the end of the subclass __init__ method to
-    activate immutability.
-  """
-
-  def __setattr__(self, attr, value):
-    # pylint: disable=g-doc-args
-    # pylint: disable=g-doc-return-or-yield
-    """Override __setattr__ to make the class (sort of) immutable.
-
-    Note that you can still do obj.field.subfield = x, so it is not true
-    immutability, but this to helps to avoid some careless errors.
-    """
-    if getattr(self, "_frozen", False):
-      raise AttributeError("PedestalModels are immutable.")
-    return super().__setattr__(attr, value)
+@dataclasses.dataclass(frozen=True, eq=False)
+class PedestalModel(static_dataclass.StaticDataclass, abc.ABC):
+  """Calculates temperature and density of the pedestal."""
 
   def __call__(
       self,
       runtime_params: runtime_params_slice.RuntimeParams,
       geo: geometry.Geometry,
       core_profiles: state.CoreProfiles,
   ) -> PedestalModelOutput:
-    if not getattr(self, "_frozen", False):
-      raise RuntimeError(
-          f"Subclass implementation {type(self)} forgot to "
-          "freeze at the end of __init__."
-      )
-
     return jax.lax.cond(
         runtime_params.pedestal.set_pedestal,
         lambda: self._call_implementation(runtime_params, geo, core_profiles),
@@ -103,16 +82,3 @@ def _call_implementation(
       core_profiles: state.CoreProfiles,
   ) -> PedestalModelOutput:
     """Calculate the pedestal values."""
-
-  @abc.abstractmethod
-  def __hash__(self) -> int:
-    """Hash function for the pedestal model.
-
-    Needed for jax.jit caching to work.
-    """
-    ...
-
-  @abc.abstractmethod
-  def __eq__(self, other) -> bool:
-    """Equality function for the pedestal model."""
-    ...

torax/_src/pedestal_model/set_pped_tpedratio_nped.py

@@ -40,15 +40,12 @@ class RuntimeParams(runtime_params_lib.RuntimeParams):
   n_e_ped_is_fGW: array_typing.BoolScalar
 
 
+@dataclasses.dataclass(frozen=True, eq=False)
 class SetPressureTemperatureRatioAndDensityPedestalModel(
     pedestal_model.PedestalModel
 ):
   """Pedestal model with specification of pressure, temp ratio, and density."""
 
-  def __init__(self):
-    super().__init__()
-    self._frozen = True
-
   @override
   def _call_implementation(
       self,
@@ -113,9 +110,3 @@ def _call_implementation(
         rho_norm_ped_top=runtime_params.pedestal.rho_norm_ped_top,
         rho_norm_ped_top_idx=ped_idx,
     )
-
-  def __hash__(self) -> int:
-    return hash('SetPressureTemperatureRatioAndDensityPedestalModel')
-
-  def __eq__(self, other) -> bool:
-    return isinstance(other, SetPressureTemperatureRatioAndDensityPedestalModel)

torax/_src/pedestal_model/set_tped_nped.py

@@ -38,15 +38,10 @@ class RuntimeParams(runtime_params_lib.RuntimeParams):
   n_e_ped_is_fGW: array_typing.BoolScalar
 
 
+@dataclasses.dataclass(frozen=True, eq=False)
 class SetTemperatureDensityPedestalModel(pedestal_model.PedestalModel):
   """A basic version of the pedestal model that uses direct specification."""
 
-  def __init__(
-      self,
-  ):
-    super().__init__()
-    self._frozen = True
-
   @override
   def _call_implementation(
       self,
@@ -77,9 +72,3 @@ def _call_implementation(
             geo.rho_norm - pedestal_params.rho_norm_ped_top
         ).argmin(),
     )
-
-  def __hash__(self) -> int:
-    return hash('SetTemperatureDensityPedestalModel')
-
-  def __eq__(self, other) -> bool:
-    return isinstance(other, SetTemperatureDensityPedestalModel)