Add experimental compile to linear solver.

torax/_src/solver/linear_theta_method.py

@@ -13,8 +13,11 @@
 # limitations under the License.
 
 """The LinearThetaMethod solver class."""
+import functools
+
 import jax
 from torax._src import state
+from torax._src import xnp
 from torax._src.config import runtime_params_slice
 from torax._src.core_profiles import convertors
 from torax._src.fvm import calc_coeffs
@@ -28,6 +31,14 @@
 class LinearThetaMethod(solver_lib.Solver):
   """Time step update using theta method, linearized on coefficients at t."""
 
+  @functools.partial(
+      xnp.jit,
+      static_argnames=[
+          'self',
+          'static_runtime_params_slice',
+          'evolving_names',
+      ],
+  )
   def _x_new(
       self,
       dt: jax.Array,

torax/_src/solver/solver.py

@@ -29,6 +29,7 @@
 from torax._src.sources import source_models as source_models_lib
 from torax._src.sources import source_profiles
 from torax._src.transport_model import transport_model as transport_model_lib
+import typing_extensions
 
 
 class Solver(abc.ABC):
@@ -65,6 +66,24 @@ def __init__(
     self.neoclassical_models = neoclassical_models
     self.static_runtime_params_slice = static_runtime_params_slice
 
+  def __hash__(self) -> int:
+    return hash((
+        self.static_runtime_params_slice,
+        self.transport_model,
+        self.source_models,
+        self.pedestal_model,
+    ))
+
+  def __eq__(self, other: typing_extensions.Self) -> bool:
+    return (
+        self.static_runtime_params_slice == other.static_runtime_params_slice
+        and self.static_runtime_params_slice
+        == other.static_runtime_params_slice
+        and self.transport_model == other.transport_model
+        and self.source_models == other.source_models
+        and self.pedestal_model == other.pedestal_model
+    )
+
   @functools.cached_property
   def evolving_names(self) -> tuple[str, ...]:
     """The names of core_profiles variables that are evolved by the solver."""

torax/tests/sim_experimental_compile_test.py

@@ -31,10 +31,16 @@
 class SimExperimentalCompileTest(sim_test_case.SimTestCase):
 
   @parameterized.named_parameters(
+      # Using newton raphson non linear solver.
       (
           'test_iterhybrid_rampup',
           'test_iterhybrid_rampup.py',
       ),
+      # Using linear solver.
+      (
+          'test_iterhybrid_predictor_corrector',
+          'test_iterhybrid_predictor_corrector.py',
+      ),
   )
   def test_run_simulation_experimental_compile(
       self,