sleipnirgroup-jormungandr 0.6.7.dev22__tar.gz → 0.6.7.dev24__tar.gz

{sleipnirgroup_jormungandr-0.6.7.dev22 → sleipnirgroup_jormungandr-0.6.7.dev24}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sleipnirgroup-jormungandr
-Version: 0.6.7.dev22
+Version: 0.6.7.dev24
 Summary: Reverse mode autodiff library and NLP solver DSL
 License-Expression: BSD-3-Clause
 License-File: LICENSE.txt

{sleipnirgroup_jormungandr-0.6.7.dev22 → sleipnirgroup_jormungandr-0.6.7.dev24}/include/sleipnir/optimization/solver/util/bounds.hpp

@@ -48,27 +48,35 @@ struct Bounds {
 /// @param inequality_constraints Variables representing the left-hand side of
 ///     cᵢ(decision_variables) ≥ 0.
 /// @param A_i The Jacobian of inequality_constraints wrt decision_variables,
-///     evaluated anywhere, in *row-major* storage; in practice, since we
-///     typically store Jacobians column-major, the user of this function must
-///     perform a transpose.
+///     evaluated anywhere.
 template <typename Scalar>
-Bounds<Scalar> get_bounds(
-    std::span<Variable<Scalar>> decision_variables,
-    std::span<Variable<Scalar>> inequality_constraints,
-    const Eigen::SparseMatrix<Scalar, Eigen::RowMajor>& A_i) {
-  // TODO: A blocked, out-of-place transpose should be much faster than
-  // traversing row major on a column major matrix unless we have few linear
-  // constraints (using a heuristic to choose between this and staying column
-  // major based on the number of constraints would be an easy performance
-  // improvement.)
-
+Bounds<Scalar> get_bounds(std::span<Variable<Scalar>> decision_variables,
+                          std::span<Variable<Scalar>> inequality_constraints,
+                          const Eigen::SparseMatrix<Scalar>& A_i) {
   // NB: Casting to long is unspecified if the size of decision_variable.size()
   // is greater than the max long value, but then we wouldn't be able to fill
   // A_i correctly anyway.
   slp_assert(static_cast<Eigen::Index>(decision_variables.size()) ==
-             A_i.innerSize());
+             A_i.cols());
   slp_assert(static_cast<Eigen::Index>(inequality_constraints.size()) ==
-             A_i.outerSize());
+             A_i.rows());
+
+  // For each row of A_i, count its structural nonzeros and record the column
+  // and value of its most recently visited nonzero. If the count is 1, that's
+  // the row's only nonzero; otherwise, the recorded column and value are
+  // meaningless, but those rows aren't bound constraints, so they're never
+  // read.
+  gch::small_vector<int> row_num_nonzeros(A_i.rows(), 0);
+  gch::small_vector<int> row_nonzero_col(A_i.rows());
+  gch::small_vector<Scalar> row_nonzero_value(A_i.rows());
+  for (int col = 0; col < A_i.outerSize(); ++col) {
+    for (typename Eigen::SparseMatrix<Scalar>::InnerIterator it{A_i, col}; it;
+         ++it) {
+      ++row_num_nonzeros[it.row()];
+      row_nonzero_col[it.row()] = it.col();
+      row_nonzero_value[it.row()] = it.value();
+    }
+  }
 
   // Maps each decision variable's index to the indices of its upper and lower
   // bounds if they exist, or NO_BOUND if they do not; used only for bookkeeping
@@ -102,10 +110,7 @@ Bounds<Scalar> get_bounds(
         ExpressionType::LINEAR) {
       continue;
     }
-    const Eigen::SparseVector<Scalar>& row_A_i =
-        A_i.innerVector(constraint_index);
-    const auto non_zeros = row_A_i.nonZeros();
-    if (non_zeros != 1) {
+    if (row_num_nonzeros[constraint_index] != 1) {
       // Constraint is not a bound because either:
       //
       // 1. It's in more than one variable
@@ -122,10 +127,9 @@ Bounds<Scalar> get_bounds(
     // and a ≠ 0. The gradient of c is then aeᵢ (where eᵢ denotes the i-th basis
     // element), and c(0) = b. If c(x) ≥ 0, then since either a < 0 or a > 0, we
     // have either x ≤ -b/a or x ≥ -b/a, respectively. ∎
-    typename Eigen::SparseVector<Scalar>::InnerIterator row_iter(row_A_i);
-    const auto constraint_coefficient =
-        row_iter
-            .value();  // The nonzero value of the j-th constraint's gradient.
+    //
+    // The nonzero value of the j-th constraint's gradient.
+    const auto constraint_coefficient = row_nonzero_value[constraint_index];
 
     // The Jacobian's sparsity pattern is structural, so the single stored value
     // can be zero if the constraint's terms cancel (e.g., x − x + 1 ≥ 0). The
@@ -134,7 +138,7 @@ Bounds<Scalar> get_bounds(
       continue;
     }
 
-    const auto decision_variable_index = row_iter.index();
+    const auto decision_variable_index = row_nonzero_col[constraint_index];
     const auto decision_variable_value =
         decision_variables[decision_variable_index].value();
     Scalar constraint_constant;

{sleipnirgroup_jormungandr-0.6.7.dev22 → sleipnirgroup_jormungandr-0.6.7.dev24}/python/src/sleipnir/__init__.py

@@ -2,4 +2,4 @@
 
 from ._sleipnir import *
 
-__version__ = "0.6.7.dev22"
+__version__ = "0.6.7.dev24"

{sleipnirgroup_jormungandr-0.6.7.dev22 → sleipnirgroup_jormungandr-0.6.7.dev24}/python/test/optimization/cart_pole_ocp_test.py

@@ -6,7 +6,6 @@ from cart_pole_util import (
     cart_pole_dynamics_double,
     cart_pole_dynamics_variable,
 )
-from rk4 import rk4
 from sleipnir.autodiff import ExpressionType, VariableMatrix
 from sleipnir.optimization import (
     OCP,
@@ -86,28 +85,37 @@ def test_cart_pole_ocp():
     assert X.value(2, 0) == pytest.approx(x_initial[2, 0], abs=1e-8)
     assert X.value(3, 0) == pytest.approx(x_initial[3, 0], abs=1e-8)
 
-    # FIXME: Replay diverges
-    if 0:
-        # Verify solution
-        x = np.zeros((4, 1))
-        u = np.zeros((1, 1))
-        for k in range(N):
-            # Cart position constraints
-            assert X[0, k] >= 0.0
-            assert X[0, k] <= d_max
-
-            # Input constraints
-            assert U[0, k] >= -u_max
-            assert U[0, k] <= u_max
-
-            # Verify state
-            assert X.value(0, k) == pytest.approx(x[0, 0], abs=1e-2)
-            assert X.value(1, k) == pytest.approx(x[1, 0], abs=1e-2)
-            assert X.value(2, k) == pytest.approx(x[2, 0], abs=1e-2)
-            assert X.value(3, k) == pytest.approx(x[3, 0], abs=1e-2)
-
-            # Project state forward
-            x = rk4(cart_pole_dynamics_double, x, u, dt)
+    # Verify solution
+    for k in range(N):
+        # Cart position constraints
+        assert X.value(0, k) >= 0.0
+        assert X.value(0, k) <= d_max
+
+        # Input constraints
+        assert U.value(0, k) >= -u_max
+        assert U.value(0, k) <= u_max
+
+        # Dynamics constraints
+        #
+        # Direct collocation constrains the system dynamics at the midpoint of a
+        # cubic Hermite spline through each pair of adjacent states.
+        f = cart_pole_dynamics_double
+        h = problem.dt().value(0, k)
+        x_begin = X[:, k : k + 1].value()
+        x_end = X[:, k + 1 : k + 2].value()
+        u_begin = U[:, k : k + 1].value()
+        u_end = U[:, k + 1 : k + 2].value()
+
+        xdot_begin = f(x_begin, u_begin)
+        xdot_end = f(x_end, u_end)
+        xdot_c = -3.0 / (2.0 * h) * (x_begin - x_end) - 0.25 * (xdot_begin + xdot_end)
+
+        x_c = 0.5 * (x_begin + x_end) + h / 8.0 * (xdot_begin - xdot_end)
+        u_c = 0.5 * (u_begin + u_end)
+
+        expected_xdot_c = f(x_c, u_c)
+        for row in range(xdot_c.shape[0]):
+            assert xdot_c[row, 0] == pytest.approx(expected_xdot_c[row, 0], abs=1e-8)
 
     # Verify final state
     assert X.value(0, N) == pytest.approx(x_final[0, 0], abs=1e-8)
@@ -121,17 +129,20 @@ def test_cart_pole_ocp():
             "Time (s),Cart position (m),Pole angle (rad),Cart velocity (m/s),Pole angular velocity (rad/s)\n"
         )
 
+        time = 0.0
         for k in range(N + 1):
             f.write(
-                f"{k * dt},{X.value(0, k)},{X.value(1, k)},{X.value(2, k)},{X.value(3, k)}\n"
+                f"{time},{X.value(0, k)},{X.value(1, k)},{X.value(2, k)},{X.value(3, k)}\n"
             )
 
+            time += problem.dt().value(0, k)
+
     # Log inputs for offline viewing
     with open("Cart-pole inputs.csv", "w") as f:
         f.write("Time (s),Cart force (N)\n")
 
+        time = 0.0
         for k in range(N + 1):
-            if k < N:
-                f.write(f"{k * dt},{U.value(0, k)}\n")
-            else:
-                f.write(f"{k * dt},0.0\n")
+            f.write(f"{time},{U.value(0, k)}\n")
+
+            time += problem.dt().value(0, k)

{sleipnirgroup_jormungandr-0.6.7.dev22 → sleipnirgroup_jormungandr-0.6.7.dev24}/python/test/optimization/differential_drive_ocp_test.py

@@ -68,29 +68,24 @@ def test_differential_drive_ocp():
     assert X.value(3, 0) == pytest.approx(x_initial[3, 0], abs=1e-8)
     assert X.value(4, 0) == pytest.approx(x_initial[4, 0], abs=1e-8)
 
-    # FIXME: Replay diverges
-    if 0:
-        # Verify solution
-        x = np.zeros((5, 1))
-        u = np.zeros((2, 1))
-        for k in range(N):
-            u = U[:, k : k + 1].value()
-
-            # Input constraints
-            assert U[0, k].value() >= -u_max[0]
-            assert U[0, k].value() <= u_max[0]
-            assert U[1, k].value() >= -u_max[1]
-            assert U[1, k].value() <= u_max[1]
-
-            # Verify state
-            assert X.value(0, k) == pytest.approx(x[0, 0], abs=1e-8)
-            assert X.value(1, k) == pytest.approx(x[1, 0], abs=1e-8)
-            assert X.value(2, k) == pytest.approx(x[2, 0], abs=1e-8)
-            assert X.value(3, k) == pytest.approx(x[3, 0], abs=1e-8)
-            assert X.value(4, k) == pytest.approx(x[4, 0], abs=1e-8)
-
-            # Project state forward
-            x = rk4(differential_drive_dynamics_double, x, u, problem.dt().value(0, k))
+    # Verify solution
+    for k in range(N):
+        # Input constraints
+        assert U.value(0, k) >= -u_max[0, 0]
+        assert U.value(0, k) <= u_max[0, 0]
+        assert U.value(1, k) >= -u_max[1, 0]
+        assert U.value(1, k) <= u_max[1, 0]
+
+        # Dynamics constraints
+        expected_x_k1 = rk4(
+            differential_drive_dynamics_double,
+            X[:, k : k + 1].value(),
+            U[:, k : k + 1].value(),
+            problem.dt().value(0, k),
+        )
+        actual_x_k1 = X[:, k + 1 : k + 2].value()
+        for row in range(actual_x_k1.shape[0]):
+            assert actual_x_k1[row, 0] == pytest.approx(expected_x_k1[row, 0], abs=1e-8)
 
     # Verify final state
     assert X.value(0, N) == pytest.approx(x_final[0, 0], abs=1e-8)