Jax euler inverse (#24)

* use `solve(A, Bv)` instead of `inv(A) @ B` in `jax_euler_step`

* Remove dead code

* Fix euler step test code

cge_modeling/base/cge.py

@@ -891,13 +891,6 @@ def f_euler(theta_final, n_steps, progressbar=True, update_freq=1, **data):
                         current_variable_vals = current_step[: len(self.variable_names)]
                         current_parameter_vals = current_step[len(self.variable_names) :]
 
-                        # current_variable_vals = flat_current_step[
-                        #     : len(self.variable_names)
-                        # ]
-                        # current_parameter_vals = flat_current_step[
-                        #     len(self.variable_names) :
-                        # ]
-
                         current_variables = {
                             k: current_variable_vals[i] for i, k in enumerate(self.variable_names)
                         }

cge_modeling/pytensorf/compile.py

@@ -472,10 +472,8 @@ def step(**kwargs):
         step_size = (theta_final_vec - theta0_vec) / n_steps
 
         A = jax.jacobian(f_sys, 0)(x_vec, theta_vec)
-        A_inv = jax.scipy.linalg.solve(A, jnp.eye(A.shape[0]))
-
         _, Bv = jax.jvp(lambda theta: f_sys(x_vec, theta), (theta_vec,), (step_size,))
-        step = -A_inv @ Bv
+        step = -jax.scipy.linalg.solve(A, Bv)
 
         x_next_vec = x_vec + step
         theta_next_vec = theta_vec + step_size

tests/test_pytensorf.py

@@ -13,10 +13,9 @@ def test_euler_approximation_1d():
     n_steps = pt.iscalar("n_steps")
 
     A = make_jacobian(eq, [y])
-    A_inv = 1 / A
     B = make_jacobian(eq, [x])
 
-    x0_final, result = euler_approximation(A_inv, B, variables=[y], parameters=[x], n_steps=n_steps)
+    x0_final, result = euler_approximation(A, B, variables=[y], parameters=[x], n_steps=n_steps)
     f = pytensor.function([x, y, x0_final, n_steps], result)
     y_values, x_values = f(0, 0, np.array([10.0]), 10_000)
     true = -np.cos(np.array([10])) + 1