Add type hints to linalg matrix inputs

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 name = "thermox"
 version = "0.0.1"
-description = "OU Processes and Linear Algebra with JAX"
+description = "Exact OU processes with JAX"
 readme = "README.md"
 requires-python =">=3.9"
 license = {text = "Apache-2.0"}

thermox/linalg.py

@@ -1,12 +1,12 @@
-import jax
 import jax.numpy as jnp
-from thermox.sampler import sample, sample_identity_diffusion
 from jax.lax import fori_loop
-from jax import Array
+from jax import Array, random
+from thermox.sampler import sample, sample_identity_diffusion
+from thermox.utils import ProcessedDriftMatrix
 
 
 def solve(
-    A,
+    A: Array | ProcessedDriftMatrix,
     b,
     num_samples: int = 10000,
     dt: float = 1.0,
@@ -34,15 +34,15 @@ def solve(
         Approximate solution, x, of the linear system.
     """
     if key is None:
-        key = jax.random.PRNGKey(0)
+        key = random.PRNGKey(0)
     ts = jnp.arange(burnin, burnin + num_samples) * dt
     x0 = jnp.zeros_like(b)
     samples = sample_identity_diffusion(key, ts, x0, A, jnp.linalg.solve(A, b))
     return jnp.mean(samples, axis=0)
 
 
 def inv(
-    A,
+    A: Array,
     num_samples: int = 10000,
     dt: float = 1.0,
     burnin: int = 0,
@@ -65,16 +65,16 @@ def inv(
         Approximate inverse of A.
     """
     if key is None:
-        key = jax.random.PRNGKey(0)
+        key = random.PRNGKey(0)
     ts = jnp.arange(burnin, burnin + num_samples) * dt
     b = jnp.zeros(A.shape[0])
     x0 = jnp.zeros_like(b)
-    samples = sample(key, ts, x0, A, b, 2 * jnp.eye(A.shape[0]))
+    samples = sample_identity_diffusion(key, ts, x0, A, b, 2 * jnp.eye(A.shape[0]))
     return jnp.cov(samples.T)
 
 
 def expnegm(
-    A,
+    A: Array,
     num_samples: int = 10000,
     dt: float = 1.0,
     burnin: int = 0,
@@ -100,7 +100,7 @@ def expnegm(
         Approximate negative matrix exponential, exp(-A).
     """
     if key is None:
-        key = jax.random.PRNGKey(0)
+        key = random.PRNGKey(0)
 
     A_shifted = (A + alpha * jnp.eye(A.shape[0])) / dt
     B = A_shifted + A_shifted.T
@@ -113,7 +113,7 @@ def expnegm(
 
 
 def expm(
-    A,
+    A: Array,
     num_samples: int = 10000,
     dt: float = 1.0,
     burnin: int = 0,