Add the Thomas Algorithm Block Tridiagonal solver

PiperOrigin-RevId: 884493419

Author: joeljennings

torax/_src/fvm/implicit_solve_block.py

@@ -19,6 +19,7 @@
 import dataclasses
 
 import jax
+from torax._src import tridiagonal
 from torax._src.fvm import block_1d_coeffs
 from torax._src.fvm import cell_variable
 from torax._src.fvm import fvm_conversions
@@ -29,6 +30,7 @@
     static_argnames=[
         'convection_dirichlet_mode',
         'convection_neumann_mode',
+        'implicit_solver_type',
         'theta_implicit',
     ],
 )
@@ -41,6 +43,7 @@ def implicit_solve_block(
     theta_implicit: float = 1.0,
     convection_dirichlet_mode: str = 'ghost',
     convection_neumann_mode: str = 'ghost',
+    implicit_solver_type: tridiagonal.SolverType = tridiagonal.SolverType.THOMAS,
 ) -> tuple[cell_variable.CellVariable, ...]:
   # pyformat: disable  # pyformat removes line breaks needed for readability
   """Runs one time step of an implicit solver on the equation defined by `coeffs`.
@@ -65,6 +68,9 @@ def implicit_solve_block(
       `dirichlet_mode` argument.
     convection_neumann_mode: See docstring of the `convection_terms` function,
       `neumann_mode` argument.
+    implicit_solver_type: The tridiagonal solver algorithm to use for the
+      implicit linear system solve. Either SolverType.THOMAS (default) for the
+      Thomas algorithm, or SolverType.DENSE for a dense matrix solve.
 
   Returns:
     x_new: Tuple, with x_new[i] giving channel i of x at the next time step
@@ -94,7 +100,7 @@ def implicit_solve_block(
   )
 
   rhs_result = rhs_matrix.matvec(x_old_array) + rhs_vec - lhs_vec
-  x_new = lhs_matrix.solve(rhs_result)
+  x_new = lhs_matrix.solve(rhs_result, solver_type=implicit_solver_type)
 
   # Create updated CellVariable instances based on state_plus_dt which has
   # updated boundary conditions and prescribed profiles.

torax/_src/solver/predictor_corrector_method.py

@@ -105,6 +105,7 @@ def loop_body(x_new_guess):
         theta_implicit=solver_params.theta_implicit,
         convection_dirichlet_mode=(solver_params.convection_dirichlet_mode),
         convection_neumann_mode=(solver_params.convection_neumann_mode),
+        implicit_solver_type=solver_params.implicit_solver_type,
     )
 
   if solver_params.use_predictor_corrector:

torax/_src/solver/pydantic_model.py

@@ -19,6 +19,7 @@
 
 import pydantic
 from torax._src import models as models_lib
+from torax._src import tridiagonal
 from torax._src.fvm import enums
 from torax._src.solver import linear_theta_method
 from torax._src.solver import nonlinear_theta_method
@@ -45,6 +46,9 @@ class BaseSolver(torax_pydantic.BaseModelFrozen, abc.ABC):
       `neumann_mode` argument.
     use_pereverzev: Use pereverzev terms for linear solver. Is only applied in
       the nonlinear solver for the optional initial guess from the linear solver
+    implicit_solver_type: The tridiagonal solver algorithm to use for the
+      implicit linear system solve. Either SolverType.THOMAS (default) for the
+      Thomas algorithm, or SolverType.DENSE for a dense matrix solve.
     chi_pereverzev: (deliberately) large heat conductivity for Pereverzev rule.
     D_pereverzev: (deliberately) large particle diffusion for Pereverzev rule.
   """
@@ -63,6 +67,9 @@ class BaseSolver(torax_pydantic.BaseModelFrozen, abc.ABC):
       Literal['ghost', 'semi-implicit'], torax_pydantic.JAX_STATIC
   ] = 'ghost'
   use_pereverzev: Annotated[bool, torax_pydantic.JAX_STATIC] = False
+  implicit_solver_type: Annotated[
+      tridiagonal.SolverType, torax_pydantic.JAX_STATIC
+  ] = tridiagonal.SolverType.THOMAS
   chi_pereverzev: pydantic.PositiveFloat = 30.0
   D_pereverzev: pydantic.NonNegativeFloat = 15.0
 
@@ -107,6 +114,7 @@ def build_runtime_params(self) -> runtime_params.RuntimeParams:
         convection_neumann_mode=self.convection_neumann_mode,
         use_pereverzev=self.use_pereverzev,
         use_predictor_corrector=self.use_predictor_corrector,
+        implicit_solver_type=self.implicit_solver_type,
         chi_pereverzev=self.chi_pereverzev,
         D_pereverzev=self.D_pereverzev,
         n_corrector_steps=self.n_corrector_steps,
@@ -160,6 +168,7 @@ def build_runtime_params(
         convection_neumann_mode=self.convection_neumann_mode,
         use_pereverzev=self.use_pereverzev,
         use_predictor_corrector=self.use_predictor_corrector,
+        implicit_solver_type=self.implicit_solver_type,
         chi_pereverzev=self.chi_pereverzev,
         D_pereverzev=self.D_pereverzev,
         maxiter=self.n_max_iterations,
@@ -210,6 +219,7 @@ def build_runtime_params(
         convection_neumann_mode=self.convection_neumann_mode,
         use_pereverzev=self.use_pereverzev,
         use_predictor_corrector=self.use_predictor_corrector,
+        implicit_solver_type=self.implicit_solver_type,
         chi_pereverzev=self.chi_pereverzev,
         D_pereverzev=self.D_pereverzev,
         n_max_iterations=self.n_max_iterations,

torax/_src/solver/runtime_params.py

@@ -15,6 +15,7 @@
 import dataclasses
 
 import jax
+from torax._src import tridiagonal
 
 
 @jax.tree_util.register_dataclass
@@ -27,5 +28,8 @@ class RuntimeParams:
   convection_dirichlet_mode: str = dataclasses.field(metadata={'static': True})
   convection_neumann_mode: str = dataclasses.field(metadata={'static': True})
   use_pereverzev: bool = dataclasses.field(metadata={'static': True})
+  implicit_solver_type: tridiagonal.SolverType = dataclasses.field(
+      metadata={'static': True}
+  )
   chi_pereverzev: float
   D_pereverzev: float  # pylint: disable=invalid-name

torax/_src/tests/tridiagonal_test.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from absl.testing import absltest
+import jax
 import jax.numpy as jnp
 import numpy as np
 from torax._src import tridiagonal
@@ -282,7 +283,7 @@ def test_solve(self):
     x_true = jnp.array(rng.randn(4, 3), dtype=jnp.float64)
     rhs = bt.matvec(x_true)
 
-    x_solved = bt.solve(rhs)
+    x_solved = bt.solve(rhs, solver_type=tridiagonal.SolverType.THOMAS)
 
     np.testing.assert_allclose(x_solved, x_true, atol=1e-10)
 
@@ -292,7 +293,7 @@ def test_solve_recovers_rhs(self):
     rng = np.random.RandomState(55)
     rhs = jnp.array(rng.randn(3, 2), dtype=jnp.float64)
 
-    x = bt.solve(rhs)
+    x = bt.solve(rhs, solver_type=tridiagonal.SolverType.THOMAS)
     reconstructed_rhs = bt.matvec(x)
 
     np.testing.assert_allclose(reconstructed_rhs, rhs, atol=1e-10)
@@ -345,5 +346,174 @@ def test_from_tridiagonals_to_dense_matches_per_channel(self):
     np.testing.assert_allclose(dense, expected_full)
 
 
+class ThomasSolveTest(absltest.TestCase):
+  """Tests specifically targeting the Thomas algorithm for block-tridiagonal."""
+
+  def _make_nonsingular_block_tridiag(
+      self, num_blocks: int, block_size: int, seed: int = 0
+  ) -> tridiagonal.BlockTriDiagonal:
+    """Helper to create a diagonally-dominant BlockTriDiagonal."""
+    rng = np.random.RandomState(seed)
+    lower = jnp.array(
+        rng.randn(num_blocks - 1, block_size, block_size), dtype=jnp.float64
+    )
+    upper = jnp.array(
+        rng.randn(num_blocks - 1, block_size, block_size), dtype=jnp.float64
+    )
+    diag_blocks = jnp.array(
+        rng.randn(num_blocks, block_size, block_size), dtype=jnp.float64
+    )
+    diag_blocks = diag_blocks + 10.0 * jnp.eye(block_size, dtype=jnp.float64)
+    return tridiagonal.BlockTriDiagonal(
+        lower=lower, diagonal=diag_blocks, upper=upper
+    )
+
+  def test_thomas_matches_dense_solve(self):
+    """Thomas and dense solvers should produce the same result."""
+    bt = self._make_nonsingular_block_tridiag(num_blocks=5, block_size=3)
+    rng = np.random.RandomState(42)
+    rhs = jnp.array(rng.randn(5, 3), dtype=jnp.float64)
+
+    x_thomas = tridiagonal.thomas_solve(bt, rhs)
+    x_dense = tridiagonal.dense_solve(bt, rhs)
+
+    np.testing.assert_allclose(x_thomas, x_dense, atol=1e-10)
+
+  def test_thomas_small_known_system(self):
+    """Thomas algorithm on a small 2-block system with known answer."""
+    # 2x2 block system: [[D0, U0], [L0, D1]] @ x = rhs
+    # D0 = [[10, 0], [0, 10]], U0 = [[1, 0], [0, 1]]
+    # L0 = [[1, 0], [0, 1]], D1 = [[10, 0], [0, 10]]
+    # This is close to identity so the answer is close to rhs/10.
+    diag = jnp.array(
+        [[[10.0, 0.0], [0.0, 10.0]], [[10.0, 0.0], [0.0, 10.0]]],
+        dtype=jnp.float64,
+    )
+    upper = jnp.array([[[1.0, 0.0], [0.0, 1.0]]], dtype=jnp.float64)
+    lower = jnp.array([[[1.0, 0.0], [0.0, 1.0]]], dtype=jnp.float64)
+    bt = tridiagonal.BlockTriDiagonal(lower=lower, diagonal=diag, upper=upper)
+
+    x_true = jnp.array([[1.0, 2.0], [3.0, 4.0]], dtype=jnp.float64)
+    rhs = bt.matvec(x_true)
+
+    x_solved = tridiagonal.thomas_solve(bt, rhs)
+
+    np.testing.assert_allclose(x_solved, x_true, atol=1e-12)
+
+  def test_thomas_identity_blocks(self):
+    """Solving with block-identity should return the RHS itself."""
+    num_blocks = 4
+    block_size = 2
+    bt = tridiagonal.BlockTriDiagonal(
+        lower=jnp.zeros(
+            (num_blocks - 1, block_size, block_size), dtype=jnp.float64
+        ),
+        diagonal=jnp.tile(
+            jnp.eye(block_size, dtype=jnp.float64), (num_blocks, 1, 1)
+        ),
+        upper=jnp.zeros(
+            (num_blocks - 1, block_size, block_size), dtype=jnp.float64
+        ),
+    )
+    rng = np.random.RandomState(11)
+    rhs = jnp.array(rng.randn(num_blocks, block_size), dtype=jnp.float64)
+
+    x = tridiagonal.thomas_solve(bt, rhs)
+
+    np.testing.assert_allclose(x, rhs, atol=1e-14)
+
+  def test_thomas_scalar_blocks(self):
+    """Thomas algorithm with block_size=1 should match scalar tridiagonal."""
+    bt = self._make_nonsingular_block_tridiag(
+        num_blocks=6, block_size=1, seed=7
+    )
+    rng = np.random.RandomState(13)
+    rhs = jnp.array(rng.randn(6, 1), dtype=jnp.float64)
+
+    x = tridiagonal.thomas_solve(bt, rhs)
+
+    np.testing.assert_allclose(bt.matvec(x), rhs, atol=1e-12)
+
+  def test_thomas_two_blocks(self):
+    """Minimal multi-block case: 2 blocks."""
+    bt = self._make_nonsingular_block_tridiag(
+        num_blocks=2, block_size=2, seed=99
+    )
+    rng = np.random.RandomState(17)
+    x_true = jnp.array(rng.randn(2, 2), dtype=jnp.float64)
+    rhs = bt.matvec(x_true)
+
+    x_solved = tridiagonal.thomas_solve(bt, rhs)
+
+    np.testing.assert_allclose(x_solved, x_true, atol=1e-10)
+
+  def test_thomas_large_system(self):
+    """Thomas should handle larger systems accurately."""
+    bt = self._make_nonsingular_block_tridiag(
+        num_blocks=50, block_size=4, seed=22
+    )
+    rng = np.random.RandomState(33)
+    x_true = jnp.array(rng.randn(50, 4), dtype=jnp.float64)
+    rhs = bt.matvec(x_true)
+
+    x_solved = tridiagonal.thomas_solve(bt, rhs)
+
+    np.testing.assert_allclose(x_solved, x_true, atol=1e-8)
+
+  def test_solver_type_dispatch_thomas(self):
+    """solve() with SolverType.THOMAS should use thomas_solve."""
+    bt = self._make_nonsingular_block_tridiag(num_blocks=3, block_size=2)
+    rng = np.random.RandomState(44)
+    rhs = jnp.array(rng.randn(3, 2), dtype=jnp.float64)
+
+    x_via_type = bt.solve(rhs, solver_type=tridiagonal.SolverType.THOMAS)
+    x_direct = tridiagonal.thomas_solve(bt, rhs)
+
+    np.testing.assert_allclose(x_via_type, x_direct, atol=1e-14)
+
+  def test_solver_type_dispatch_dense(self):
+    """solve() with SolverType.DENSE should use dense_solve."""
+    bt = self._make_nonsingular_block_tridiag(num_blocks=3, block_size=2)
+    rng = np.random.RandomState(44)
+    rhs = jnp.array(rng.randn(3, 2), dtype=jnp.float64)
+
+    x_via_type = bt.solve(rhs, solver_type=tridiagonal.SolverType.DENSE)
+    x_direct = tridiagonal.dense_solve(bt, rhs)
+
+    np.testing.assert_allclose(x_via_type, x_direct, atol=1e-14)
+
+  def test_thomas_jit_compatible(self):
+    """thomas_solve should work under jax.jit."""
+    bt = self._make_nonsingular_block_tridiag(num_blocks=4, block_size=2)
+    rng = np.random.RandomState(66)
+    x_true = jnp.array(rng.randn(4, 2), dtype=jnp.float64)
+    rhs = bt.matvec(x_true)
+
+    jitted_solve = jax.jit(tridiagonal.thomas_solve)
+    x_solved = jitted_solve(bt, rhs)
+
+    np.testing.assert_allclose(x_solved, x_true, atol=1e-10)
+
+  def test_thomas_from_tridiagonals(self):
+    """Thomas solve on a block system built from per-channel scalar tridiagonals."""
+    ch0 = tridiagonal.TriDiagonal(
+        diagonal=jnp.array([10.0, 12.0, 14.0], dtype=jnp.float64),
+        above=jnp.array([1.0, 3.0], dtype=jnp.float64),
+        below=jnp.array([5.0, 7.0], dtype=jnp.float64),
+    )
+    ch1 = tridiagonal.TriDiagonal(
+        diagonal=jnp.array([11.0, 13.0, 15.0], dtype=jnp.float64),
+        above=jnp.array([2.0, 4.0], dtype=jnp.float64),
+        below=jnp.array([6.0, 8.0], dtype=jnp.float64),
+    )
+    bt = tridiagonal.BlockTriDiagonal.from_tridiagonals([ch0, ch1])
+    rng = np.random.RandomState(77)
+    rhs = jnp.array(rng.randn(3, 2), dtype=jnp.float64)
+
+    x = tridiagonal.thomas_solve(bt, rhs)
+
+    np.testing.assert_allclose(bt.matvec(x), rhs, atol=1e-12)
+
+
 if __name__ == '__main__':
   absltest.main()