Added tests

Author: nathanaelbosch

tests/test_iwp.py

@@ -60,3 +60,10 @@ def test_preconditioned_system_matrices(dt, iwp):
         precond @ precond_state_trans_mat @ precond_inv, non_precond_state_trans_mat
     )
     assert jnp.allclose(precond @ precond_proc_noice_chol, non_precond_proc_noice_chol)
+
+def test_projection_matrices(iwp):
+    P = iwp.make_projection_matrix(0)
+    assert isinstance(P, jnp.ndarray)
+    d, q = iwp.wiener_process_dimension, iwp.num_derivatives
+    assert P.shape == (d, q+1)
+    assert (P == 1).sum() == d

tornado/iwp.py

@@ -123,7 +123,7 @@ def non_preconditioned_discretize(self, dt):
 
     def make_projection_matrix(self, derivative_to_project_onto):
         """Creates a projection matrix kron(I_d, e_p)"""
-        d, q = self.num_derivatives, self.wiener_process_dimension
+        d, q = self.wiener_process_dimension, self.num_derivatives
         I_d = jnp.eye(d)
         e_p = jnp.eye(1, q + 1, derivative_to_project_onto)
         return jnp.kron(I_d, e_p)