Implement an isotropic and blockdiagonal TS1

docs/examples_basic/posterior_uncertainties.py

@@ -42,7 +42,7 @@ def vf(y, *, t):  # noqa: ARG001
 # Set up a solver
 # To all users: Try replacing the fixedpoint-smoother with a filter!
 tcoeffs = taylor.odejet_padded_scan(lambda y: vf(y, t=t0), (u0,), num=3)
-init, ibm, ssm = ivpsolvers.prior_wiener_integrated(tcoeffs, ssm_fact="dense")
+init, ibm, ssm = ivpsolvers.prior_wiener_integrated(tcoeffs, ssm_fact="blockdiag")
 ts = ivpsolvers.correction_ts1(vf, ssm=ssm)
 strategy = ivpsolvers.strategy_fixedpoint(ssm=ssm)
 solver = ivpsolvers.solver_mle(strategy, prior=ibm, correction=ts, ssm=ssm)
@@ -60,13 +60,20 @@ def vf(y, *, t):  # noqa: ARG001
 u_std = ssm.stats.qoi_from_sample(std)
 
 # Plot the solution
-fig, axes = plt.subplots(nrows=2, ncols=len(tcoeffs), tight_layout=True, figsize=(8, 3))
+fig, axes = plt.subplots(
+    nrows=3,
+    ncols=len(tcoeffs),
+    sharex="col",
+    tight_layout=True,
+    figsize=(len(u_std) * 2, 5),
+)
 for i, (u_i, std_i, ax_i) in enumerate(zip(sol.u, u_std, axes.T)):
     # Set up titles and axis descriptions
     if i == 0:
         ax_i[0].set_title("State")
-        ax_i[0].set_ylabel("Predators")
-        ax_i[1].set_ylabel("Prey")
+        ax_i[0].set_ylabel("Prey")
+        ax_i[1].set_ylabel("Predators")
+        ax_i[2].set_ylabel("Std.-dev.")
     elif i == 1:
         ax_i[0].set_title(f"{i}st deriv.")
     elif i == 2:
@@ -76,7 +83,7 @@ def vf(y, *, t):  # noqa: ARG001
     else:
         ax_i[0].set_title(f"{i}th deriv.")
 
-    ax_i[1].set_xlabel("Time")
+    ax_i[-1].set_xlabel("Time")
 
     for m, std, ax in zip(u_i.T, std_i.T, ax_i):
         # Plot the mean
@@ -87,5 +94,9 @@ def vf(y, *, t):  # noqa: ARG001
         ax.fill_between(sol.t, lower, upper, alpha=0.3)
         ax.set_xlim((jnp.amin(ts), jnp.amax(ts)))
 
+    ax_i[2].semilogy(sol.t, std_i[:, 0], label="Prey")
+    ax_i[2].semilogy(sol.t, std_i[:, 1], label="Predators")
+    ax_i[2].legend(fontsize="x-small")
+
 fig.align_ylabels()
 plt.show()

probdiffeq/backend/linalg.py

@@ -67,6 +67,10 @@ def diagonal(arr, /):
     return jnp.diagonal(arr)
 
 
+def trace(arr, /):
+    return jnp.trace(arr)
+
+
 def diagonal_matrix(arr, /):
     return jnp.diag(arr)
 

probdiffeq/backend/random.py

@@ -7,5 +7,13 @@ def prng_key(*, seed):
     return jax.random.PRNGKey(seed=seed)
 
 
+def split(key, num):
+    return jax.random.split(key, num=num)
+
+
 def normal(key, /, shape):
     return jax.random.normal(key, shape=shape)
+
+
+def rademacher(key, /, shape, dtype):
+    return jax.random.rademacher(key, shape=shape, dtype=dtype)