Fix matrix Jacobians and SciPy optimizer behavior

.gitignore

@@ -14,3 +14,7 @@ docs/tutorial/
 docs/howto/
 docs/examples/
 docs/sg_execution_times.rst
+
+# IDE files and folders
+.idea/
+.vscode/

docs/examples_src/time_switch.py

@@ -111,7 +111,7 @@ class MinimumTime(co.OptimizationProblem):
 
     class Options:
         # exact_hessian = False
-        __implementation__ = co.implementations.ScipyCG
+        __implementation__ = co.implementations.ScipySLSQP
 
 
 """

src/condor/backends/casadi/operators.py

@@ -175,11 +175,6 @@ def max(x, axis=None):
     return casadi.mmax(x)
 
 
-unsupported_jacobian_message = (
-    "jacobian of matrix expression wrt matrix variable not yet supported"
-)
-
-
 def jacobian(of, wrt):
     """jacobian of expression `of` with respect to symbols `wrt`"""
     """
@@ -203,16 +198,20 @@ def jacobian(of, wrt):
        jac(0.)
     """
     if of.size and wrt.size:
-        transpose_in = False
-        if isinstance(wrt, backend.symbol_class) and wrt.op() == casadi.OP_TRANSPOSE:
-            transpose_in = True
+        transpose_in = (
+            isinstance(wrt, backend.symbol_class) and wrt.op() == casadi.OP_TRANSPOSE
+        )
+        if transpose_in:
+            input_shape = wrt.shape
             wrt = wrt.dep()
 
-        if transpose_in and np.all(np.array(of.shape + wrt.shape) > 1):
-            raise NotImplementedError(unsupported_jacobian_message)
-
         jac = casadi.jacobian(of, wrt)
 
+        if transpose_in and len(input_shape) == 2:
+            # Match derivative columns to matrix inputs encoded as transposed symbols.
+            input_order = np.arange(wrt.numel()).reshape(input_shape).T.ravel()
+            jac = jac[:, input_order]
+
         return jac
 
     else:

src/condor/implementations/iterative.py

@@ -435,16 +435,30 @@ class ScipyMinimizeBase(OptimizationProblem):
 
     Options
     --------
+    exact_hessian : bool
+        whether to use an exact Hessian when supported; ignored by CG and SLSQP
     **options
         keyword options are passed directly to scipy.minimize's options keyword argument
     """
 
+    supports_bounds = True
+    supports_constraints = True
+
     def construct(
         self,
         model,
+        iter_callback=None,
+        init_callback=None,
+        exact_hessian=True,
         **options,
     ):
-        super().construct(model, **options)
+        """Configure SciPy options, consuming backend-independent Hessian settings."""
+        super().construct(
+            model,
+            iter_callback=iter_callback,
+            init_callback=init_callback,
+            **options,
+        )
         self.f_func = self.objective_func
         self.f_jac_func = expression_to_operator(
             [self.x, self.p],
@@ -475,6 +489,22 @@ def run_optimizer(self, model_instance):
 
         scipy_constraints = self.prepare_constraints(extra_args)
 
+        bounds = None
+        has_finite_bounds = np.any(np.isfinite(self.lbx)) or np.any(
+            np.isfinite(self.ubx)
+        )
+        has_finite_constraints = np.any(np.isfinite(self.lbg)) or np.any(
+            np.isfinite(self.ubg)
+        )
+        if has_finite_constraints and not self.supports_constraints:
+            msg = f"{self.method_string} does not support constraints"
+            raise ValueError(msg)
+        if has_finite_bounds:
+            if not self.supports_bounds:
+                msg = f"{self.method_string} does not support variable bounds"
+                raise ValueError(msg)
+            bounds = np.column_stack((self.lbx, self.ubx))
+
         if self.init_callback is not None:
             self.init_callback(
                 model_instance.parameter,
@@ -488,7 +518,7 @@ def run_optimizer(self, model_instance):
             method=self.method_string,
             args=extra_args,
             constraints=scipy_constraints,
-            bounds=np.vstack([self.lbx, self.ubx]).T,
+            bounds=bounds,
             # tol = 1E-9,
             # options=dict(disp=True),
             options=self.options,
@@ -505,6 +535,8 @@ def run_optimizer(self, model_instance):
 
 class ScipyCG(ScipyMinimizeBase):
     method_string = "CG"
+    supports_bounds = False
+    supports_constraints = False
 
 
 class ScipySLSQP(ScipyMinimizeBase):

src/condor/implementations/sgm_trajectory.py

@@ -446,7 +446,7 @@ def generate_sgm_jacobian(self, jacobian_of):
             dg_dt = jacobian(e_expr, model.t)
             dg_dp = jacobian(e_expr, self.p)
 
-            dte_dx = dg_dx / (dg_dx @ state_equation_func.expr)
+            dte_dx = dg_dx / (dg_dx @ state_equation_func.expr + dg_dt)
             dte_dp = -dg_dp / (dg_dx @ state_equation_func.expr + dg_dt)
 
             dh_dx = jacobian(h_expr.expr, self.x)

tests/test_operators.py

@@ -270,7 +270,6 @@ class TestJacobian(co.ExplicitSystem):
     ops.jacobian(TestJacobian.output.flatten(), TestJacobian.input.flatten())
 
 
-@pytest.mark.skip(reason="Casadi backend doesn't support matrix/matrix jacobian yet")
 def test_jacobian():
     A = rng.random((3, 3))  # noqa: N806
 

tests/test_optimization.py

@@ -86,6 +86,31 @@ def iter_callback(self, i, variable, objective, constraint):
     assert callback.parameter is not None
 
 
+@pytest.mark.parametrize("restriction", ["bounds", "constraints"])
+def test_scipy_cg_rejects_unsupported_restrictions(restriction):
+    if restriction == "bounds":
+
+        class Opt(co.OptimizationProblem):
+            x = variable(lower_bound=0)
+            objective = x**2
+
+            class Options:
+                __implementation__ = co.implementations.ScipyCG
+
+    else:
+
+        class Opt(co.OptimizationProblem):
+            x = variable()
+            objective = x**2
+            constraint(x <= 1)
+
+            class Options:
+                __implementation__ = co.implementations.ScipyCG
+
+    with pytest.raises(ValueError, match="CG does not support"):
+        Opt()
+
+
 def test_callback_scipy_no_instance():
     class Opt(co.OptimizationProblem):
         p = parameter()

tests/test_trajectory_analysis.py

@@ -1,6 +1,6 @@
 import numpy as np
 import pytest
-from scipy import linalg
+from scipy import linalg, signal
 
 import condor as co
 from condor.backend import operators as ops
@@ -107,84 +107,54 @@ class Options:
     np.testing.assert_allclose(k, lqr_sol.K, rtol=1e-4)
 
 
-@pytest.mark.skip(reason="Need to fix LTI function")
 def test_sp_lqr():
     # sampled LQR
     dblint_a = np.array([[0, 1], [0, 0]])
     dblint_b = np.array([[0], [1]])
     dt = 0.5
-
-    DblIntSampled = co.LTI(  # noqa: N806
-        a=dblint_a, b=dblint_b, name="DblIntSampled", dt=dt
-    )
+    ad, bd = signal.cont2discrete((dblint_a, dblint_b, None, None), dt)[:2]
+    q = np.eye(2)
+    r = np.eye(1)
+
+    class DblIntSampled(co.ODESystem):
+        x = state(shape=2)
+        accumulated_cost = state()
+        K = parameter(shape=(1, 2))
+        u = -K @ x
+        dynamic_output.u = u
+        dot[x] = np.zeros(2)
+        dot[accumulated_cost] = 0.0
+
+    class SampledStep(DblIntSampled.Event):
+        at_time = slice(dt, None, dt)
+        update[accumulated_cost] = accumulated_cost + (x.T @ q @ x + u.T @ r @ u) / 2
+        update[x] = (ad - bd @ K) @ x
 
     class DblIntSampledLQR(DblIntSampled.TrajectoryAnalysis):
         initial[x] = [1.0, 0.1]
-        # initial[u] = -k@initial[x]
-        q = np.eye(2)
-        r = np.eye(1)
-        tf = 32.0  # 12 iters, 21 calls 1E-8 jac
-        # tf = 16. # 9 iters, 20 calls, 1E-7
-        cost = trajectory_output(integrand=(x.T @ q @ x + u.T @ r @ u) / 2)
-
-        class Casadi(co.Options):
-            adjoint_adaptive_max_step_size = False
-            state_max_step_size = dt / 8
-            adjoint_max_step_size = dt / 8
+        initial[accumulated_cost] = 0.0
+        tf = 31.9
+        cost = trajectory_output(terminal_term=accumulated_cost)
 
     class SampledOptLQR(co.OptimizationProblem):
-        k = variable(shape=DblIntSampledLQR.k.shape)
+        k = variable(shape=DblIntSampledLQR.K.shape)
         sim = DblIntSampledLQR(k)
         objective = sim.cost
 
-        class Casadi(co.Options):
+        class Options:
             exact_hessian = False
+            __implementation__ = co.implementations.ScipyCG
 
-    # sim = DblIntSampledLQR([1.00842737, 0.05634044])
-
-    sim = DblIntSampledLQR([0.0, 0.0])
-    sim.implementation.callback.jac_callback(sim.implementation.callback.p, [])
-
+    SampledOptLQR.set_initial(k=[0.7, 1.3])
     lqr_sol_samp = SampledOptLQR()
 
-    # sampled_sim = DblIntSampledLQR([0., 0.])
-    # sampled_sim.implementation.callback.jac_callback([0., 0.,], [0.])
-
-    q = DblIntSampledLQR.q
-    r = DblIntSampledLQR.r
-    a = dblint_a
-    b = dblint_b
-
-    ad, bd = signal.cont2discrete((a, b, None, None), dt)[:2]
-    s = linalg.solve_discrete_are(
-        ad,
-        bd,
-        q,
-        r,
-    )
+    s = linalg.solve_discrete_are(ad, bd, q, r)
     k = linalg.solve(bd.T @ s @ bd + r, bd.T @ s @ ad)
-
-    # sim = DblIntSampledLQR([1.00842737, 0.05634044])
-    sim = DblIntSampledLQR(k)
-
-    sim.implementation.callback.jac_callback(sim.implementation.callback.p, [])
-    LTI_plot(sim)
-    plt.show()
-
-    # sim = DblIntSampledLQR([0., 0.])
-
-    # sampled_sim = DblIntSampledLQR([0., 0.])
-    # sampled_sim.implementation.callback.jac_callback([0., 0.,], [0.])
-
     sampled_sim = DblIntSampledLQR(k)
-    jac_cb = sampled_sim.implementation.callback.jac_callback
-    jac_cb(k, [0.0])
 
     assert lqr_sol_samp._stats.success
-    print(lqr_sol_samp._stats)
-    print(lqr_sol_samp.objective < sampled_sim.cost)
-    print(lqr_sol_samp.objective, sampled_sim.cost)
-    print("      ARE sol:", k, "\niterative sol:", lqr_sol_samp.k)
+    np.testing.assert_allclose(lqr_sol_samp.objective, sampled_sim.cost, rtol=1e-5)
+    np.testing.assert_allclose(lqr_sol_samp.k, k, rtol=1e-3)
 
 
 def test_time_switched():
@@ -239,7 +209,7 @@ class MinimumTime(co.OptimizationProblem):
 
         class Options:
             exact_hessian = False
-            __implementation__ = co.implementations.ScipyCG
+            __implementation__ = co.implementations.ScipySLSQP
 
     MinimumTime.set_initial(t1=2.163165480675697, t2=4.361971866705403)
     opt = MinimumTime()