Implement hash and eq for sawtooth models

PiperOrigin-RevId: 745079788

Author: jcitrin

Diff for: torax/mhd/sawtooth/sawtooth_model.py

@@ -37,6 +37,17 @@ def __call__(
   ) -> tuple[array_typing.ScalarBool, array_typing.ScalarFloat]:
     """Indicates if a crash is triggered and the radius of the q=1 surface."""
 
+  @abc.abstractmethod
+  def __hash__(self) -> int:
+    """Returns a hash of the trigger model.
+
+    Should be implemented to support jax.jit caching.
+    """
+
+  @abc.abstractmethod
+  def __eq__(self, other: object) -> bool:
+    """Equality method to be implemented to support jax.jit caching."""
+
 
 class RedistributionModel(abc.ABC):
   """Abstract base class for sawtooth redistribution models."""
@@ -52,6 +63,17 @@ def __call__(
   ) -> state.CoreProfiles:
     """Returns a redistributed core_profiles if sawtooth has been triggered."""
 
+  @abc.abstractmethod
+  def __hash__(self) -> int:
+    """Returns a hash of the redistribution model.
+
+    Should be implemented to support jax.jit caching.
+    """
+
+  @abc.abstractmethod
+  def __eq__(self, other) -> bool:
+    """Equality method to be implemented to support jax.jit caching."""
+
 
 class SawtoothModel:
   """Sawtooth trigger and redistribution, and carries out sawtooth step."""
@@ -112,3 +134,13 @@ def __call__(
     # modify output state with new time, dt, and core_profiles if triggered.
 
     return input_state
+
+  def __hash__(self) -> int:
+    return hash((self.trigger_model, self.redistribution_model))
+
+  def __eq__(self, other: object) -> bool:
+    return (
+        isinstance(other, SawtoothModel)
+        and self.trigger_model == other.trigger_model
+        and self.redistribution_model == other.redistribution_model
+    )

Diff for: torax/mhd/sawtooth/simple_redistribution.py

@@ -54,6 +54,12 @@ def __call__(
 
     return core_profiles
 
+  def __hash__(self) -> int:
+    return hash(self.__class__.__name__)
+
+  def __eq__(self, other: object) -> bool:
+    return isinstance(other, SimpleRedistribution)
+
 
 @chex.dataclass(frozen=True)
 class DynamicRuntimeParams(runtime_params.RedistributionDynamicRuntimeParams):

Diff for: torax/mhd/sawtooth/simple_trigger.py

@@ -103,6 +103,12 @@ def __call__(
         rho_norm_q1,
     )
 
+  def __hash__(self) -> int:
+    return hash(self.__class__.__name__)
+
+  def __eq__(self, other: object) -> bool:
+    return isinstance(other, SimpleTrigger)
+
 
 @chex.dataclass(frozen=True)
 class DynamicRuntimeParams(runtime_params.TriggerDynamicRuntimeParams):

Diff for: torax/tests/sim_test.py

@@ -405,7 +405,7 @@ class SimTest(sim_test_case.SimTestCase):
           'test_iterhybrid_newton',
           'test_iterhybrid_newton.py',
           _ALL_PROFILES,
-          5e-7,
+          1e-6,
       ),
       # Tests current and density rampup for for ITER-hybrid-like-config
       # using Newton-Raphson. Only case which reverts to coarse_tol for several