Added tests

Author: nathanaelbosch

tests/test_iwp.py

@@ -62,6 +62,14 @@ def test_preconditioned_system_matrices(dt, iwp):
     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
+
+
 def test_reorder_states():
     # Transition handles reordering
     iwp = tornado.iwp.IntegratedWienerTransition(

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)