jaxrts.helpers.cramer_solve

jaxrts.helpers.cramer_solve(A, b, N_max=4)[source]

Solve A x = b via Cramer’s rule for small, fixed system sizes, avoiding the fixed dispatch overhead of a batched jnp.linalg.solve. Falls back to jnp.linalg.solve for sizes not hand-coded below. For small matrices, e.g., 2 component HNC, this does serve as a notable speedup.

Warning

Using this function if a tradeoff of speed (for small systems) over numerical stability. We found that for the application in HNC with few components, this works reasonably well. Setting M_max to 0 allows to bypass explicit calculation, if required.

A: (…, N, N), b: (…, N). N must be static (known at trace time). Works either called directly on a batched array, or from inside a vmap-mapped function (where A, b appear unbatched, shape (N,N) & (N,)).

Parameters:

  • A (jnp.ndarray) – Shape (…, N, N) matrix to be inverted

  • b (jnp.ndarray) – Shape (…, N) array of solutions

  • N_max (int) – Maximal size of the matrix for which the hand-written inversion should be performed. If the size is bigger than N_max, fall back to jax.linalg.solve.

Returns:

x (jnp.narray) – Shape (…, N), so that \(A x = b\)