[symbolic] Support half-angle formulas in sin/cos substitution. (Robo…

…tLocomotion#17282)

* [symbolic] Support half-angle formulas in sin/cos substitution.

bindings/pydrake/symbolic_py.cc

@@ -513,11 +513,25 @@ PYBIND11_MODULE(symbolic, m) {
           const Expression& e) { return Substitute(M, var, e); },
       py::arg("m"), py::arg("var"), py::arg("e"), doc.Substitute.doc_3args);
 
+  {
+    using Enum = SinCosSubstitutionType;
+    constexpr auto& enum_doc = doc.SinCosSubstitutionType;
+    py::enum_<Enum> enum_py(m, "SinCosSubstitutionType", enum_doc.doc);
+    enum_py  // BR
+        .value("kAngle", Enum::kAngle, enum_doc.kAngle.doc)
+        .value("kHalfAnglePreferSin", Enum::kHalfAnglePreferSin,
+            enum_doc.kHalfAnglePreferSin.doc)
+        .value("kHalfAnglePreferCos", Enum::kHalfAnglePreferCos,
+            enum_doc.kHalfAnglePreferCos.doc);
+  }
+
   py::class_<SinCos>(m, "SinCos", doc.SinCos.doc)
-      .def(py::init<const Variable&, const Variable&>(), py::arg("s"),
-          py::arg("c"), doc.SinCos.ctor.doc)
+      .def(py::init<Variable, Variable, SinCosSubstitutionType>(), py::arg("s"),
+          py::arg("c"), py::arg("type") = SinCosSubstitutionType::kAngle,
+          doc.SinCos.ctor.doc)
       .def_readwrite("s", &SinCos::s, doc.SinCos.s.doc)
-      .def_readwrite("c", &SinCos::c, doc.SinCos.c.doc);
+      .def_readwrite("c", &SinCos::c, doc.SinCos.c.doc)
+      .def_readwrite("type", &SinCos::type, doc.SinCos.type.doc);
 
   m.def(
       "Substitute",

bindings/pydrake/test/symbolic_test.py

@@ -1928,7 +1928,7 @@ def test_basics(self):
 
 
 class TestSinCosSubstitution(unittest.TestCase):
-    def basic_test(self):
+    def test_basics(self):
         x = sym.Variable("x")
         y = sym.Variable("y")
         sx = sym.Variable("sx")
@@ -1937,8 +1937,47 @@ def basic_test(self):
         cy = sym.Variable("cy")
         subs = {x: sym.SinCos(s=sx, c=cx), y: sym.SinCos(s=sy, c=cy)}
 
-        self.assert_equal(sym.Substitute(e=np.sin(x + y), subs=subs),
-                          sx * cy + cx * sy)
+        self.assertTrue(
+            sym.Substitute(e=np.sin(x + y),
+                           subs=subs).EqualTo(sx * cy + cx * sy))
         m = np.array([np.sin(x), np.cos(y), 1, np.sin(2*x)])
-        me = np.array([sx, cy, 1, 2*sx*cx])
-        np.testing.assert_equal(sym.Substitute(m=m, subs=subs), me)
+        me = np.array([
+            sym.Expression(sx),
+            sym.Expression(cy),
+            sym.Expression(1), 2 * sx * cx
+        ])
+        msubs = sym.Substitute(m=m, subs=subs)
+        self.assertEqual(msubs.shape, (4, 1))
+        for i in range(4):
+            self.assertTrue(msubs[i, 0].EqualTo(me[i]))
+
+    def test_halfangle(self):
+        x = sym.Variable("x")
+        y = sym.Variable("y")
+        sx = sym.Variable("sx")
+        sy = sym.Variable("sy")
+        cx = sym.Variable("cx")
+        cy = sym.Variable("cy")
+        subs = {
+            x:
+            sym.SinCos(s=sx,
+                       c=cx,
+                       type=sym.SinCosSubstitutionType.kHalfAnglePreferSin),
+            y:
+            sym.SinCos(s=sy,
+                       c=cy,
+                       type=sym.SinCosSubstitutionType.kHalfAnglePreferCos)
+        }
+
+        self.assertTrue(
+            sym.Substitute(e=np.sin(x), subs=subs).EqualTo(2 * sx * cx))
+        self.assertTrue(
+            sym.Substitute(e=np.cos(x), subs=subs).EqualTo(1 - 2 * sx * sx))
+        self.assertTrue(
+            sym.Substitute(e=np.cos(y), subs=subs).EqualTo(2 * cy * cy - 1))
+        m = np.array([np.sin(0.5*x), np.cos(0.5*y), 1, np.cos(x)])
+        me = np.array([sx, cy, 1, 1 - 2 * sx**2])
+        msubs = sym.Substitute(m=m, subs=subs)
+        self.assertEqual(msubs.shape, (4, 1))
+        for i in range(4):
+            self.assertTrue(msubs[i, 0].EqualTo(me[i]))

common/symbolic_trigonometric_polynomial.cc

@@ -3,6 +3,7 @@
 #include <map>
 #include <optional>
 #include <stdexcept>
+#include <string>
 #include <utility>
 
 #define DRAKE_COMMON_SYMBOLIC_DETAIL_HEADER
@@ -25,9 +26,9 @@ class SinCosVisitor {
  public:
   DRAKE_NO_COPY_NO_MOVE_NO_ASSIGN(SinCosVisitor)
 
-  explicit SinCosVisitor(SinCosSubstitution s) : subs_{std::move(s)} {};
+  explicit SinCosVisitor(SinCosSubstitution s)
+      : subs_{std::move(s)} {};
 
-  // Generates latex expression for the expression @p e.
   [[nodiscard]] Expression Substitute(
       const Expression& e,
       std::optional<bool> needs_substitution = std::nullopt) const {
@@ -54,9 +55,18 @@ class SinCosVisitor {
     if (iter == subs_.end()) {
       return e;
     } else if (status == kInsideSin) {
-      return iter->second.s;
+      return iter->second.type == SinCosSubstitutionType::kAngle
+                 ? iter->second.s
+                 : 2 * iter->second.s * iter->second.c;
     } else if (status == kInsideCos) {
-      return iter->second.c;
+      switch (iter->second.type) {
+        case SinCosSubstitutionType::kAngle:
+          return iter->second.c;
+        case SinCosSubstitutionType::kHalfAnglePreferSin:
+          return 1 - 2 * iter->second.s * iter->second.s;
+        case SinCosSubstitutionType::kHalfAnglePreferCos:
+          return 2 * iter->second.c * iter->second.c - 1;
+      }
     }
     return e;
   }
@@ -112,16 +122,36 @@ class SinCosVisitor {
     const std::map<Expression, Expression>& base_to_exponent_map{
         get_base_to_exponent_map_in_multiplication(e)};
     if (status == kInsideSin) {
-      double c_int;
+      std::string msg = fmt::format(
+          "Got sin({}), but we only support sin(c*x) where c is an integer (c "
+          "= +/- 0.5 is also supported if performing a half-angle "
+          "substitution), and x is a variable to be substituted.",
+          e.to_string());
       if (base_to_exponent_map.size() != 1 ||
-          base_to_exponent_map.begin()->second != 1.0 ||
-          modf(c, &c_int) != 0.0) {
-        throw std::runtime_error(
-            fmt::format("Got sin({}), but we only support sin(c*x) where c is "
-                        "an integer, and x is a variable to be substituted.",
-                        e.to_string()));
+          base_to_exponent_map.begin()->second != 1.0) {
+        throw std::runtime_error(msg);
       }
       Expression x = base_to_exponent_map.begin()->first;
+      if (!is_variable(x)) {
+        throw std::runtime_error(msg);
+      }
+      auto iter = subs_.find(get_variable(x));
+      if (iter == subs_.end()) {
+        throw std::runtime_error(msg);
+      }
+      if (iter->second.type != SinCosSubstitutionType::kAngle) {
+        // Handle special case of sin(±0.5*x).
+        if (c == 0.5) {
+          return iter->second.s;
+        } else if (c == -0.5) {
+          return -iter->second.s;
+        }
+        // TODO(russt): Handle the cases where c is an odd multiple of 0.5.
+      }
+      double c_int;
+      if (modf(c, &c_int) != 0.0) {
+        throw std::runtime_error(msg);
+      }
       if (c_int == 0.0) {
         return 0.0;  // Don't expect to get here, though.
       } else if (c_int == 1.0) {
@@ -138,16 +168,34 @@ class SinCosVisitor {
       }
       throw std::runtime_error("Got unexpected 0 < c_int < 1");
     } else if (status == kInsideCos) {
-      double c_int;
+      std::string msg = fmt::format(
+          "Got cos({}), but we only support cos(c*x) where c is an integer (c "
+          "= +/- 0.5 is also supported if performing a half-angle "
+          "substitution), and x is a variable to be substituted.",
+          e.to_string());
       if (base_to_exponent_map.size() != 1 ||
-          base_to_exponent_map.begin()->second != 1.0 ||
-          modf(c, &c_int) != 0.0) {
-        throw std::runtime_error(
-            fmt::format("Got cos({}), but we only support cos(c*x) where c is "
-                        "an integer, and x is a variable to be substituted.",
-                        e.to_string()));
+          base_to_exponent_map.begin()->second != 1.0) {
+        throw std::runtime_error(msg);
       }
       Expression x = base_to_exponent_map.begin()->first;
+      if (!is_variable(x)) {
+        throw std::runtime_error(msg);
+      }
+      auto iter = subs_.find(get_variable(x));
+      if (iter == subs_.end()) {
+          throw std::runtime_error(msg);
+      }
+      if (iter->second.type != SinCosSubstitutionType::kAngle) {
+        // Handle special case of cos(±0.5*x).
+        if (c == 0.5 || c == -0.5) {
+          return iter->second.c;
+        }
+        // TODO(russt): Handle the cases where c is an odd multiple of 0.5.
+      }
+      double c_int;
+      if (modf(c, &c_int) != 0.0) {
+        throw std::runtime_error(msg);
+      }
       if (c_int == 0.0) {
         return 1.0;  // Don't expect to get here, though.
       } else if (c_int == 1.0) {
@@ -345,8 +393,7 @@ class SinCosVisitor {
 
 }  // namespace
 
-Expression Substitute(const Expression& e,
-                      const SinCosSubstitution& subs) {
+Expression Substitute(const Expression& e, const SinCosSubstitution& subs) {
   return SinCosVisitor(subs).Substitute(e);
 }
 

common/symbolic_trigonometric_polynomial.h

@@ -10,17 +10,33 @@
 namespace drake {
 namespace symbolic {
 
+enum class SinCosSubstitutionType {
+  /** Substitutes s <=> sin(q), c <=> cos(q). */
+  kAngle,
+  /** Substitutes s <=> sin(q/2), c <=> cos(q/2), and prefers sin when the
+   choice is ambiguous; e.g. cos(q) => 1 - 2s². */
+  kHalfAnglePreferSin,
+  /** Subsitutes s <=> sin(q/2), c <=> cos(q/2), and prefers cos when the
+   choice is ambiguous; e.g. cos(q) => 2c² - 1. */
+  kHalfAnglePreferCos,
+};
+
 /** Represents a pair of Variables corresponding to sin(q) and cos(q). */
 struct SinCos {
   DRAKE_DEFAULT_COPY_AND_MOVE_AND_ASSIGN(SinCos);
 
-  SinCos(Variable s_var, Variable c_var)
-      : s(std::move(s_var)), c(std::move(c_var)) {}
+  SinCos(Variable _s, Variable _c,
+         SinCosSubstitutionType _type = SinCosSubstitutionType::kAngle)
+      : s(std::move(_s)), c(std::move(_c)), type(_type) {}
 
   /** sin variable. */
   Variable s{};
   /** cos variable. */
   Variable c{};
+
+  /** Allows a user to specify non-default substitutions, such as using
+   half-angle formulas. */
+  SinCosSubstitutionType type{};
 };
 
 using SinCosSubstitution = std::unordered_map<Variable, SinCos>;
@@ -38,13 +54,20 @@ using SinCosSubstitution = std::unordered_map<Variable, SinCos>;
  @endverbatim
  will result in the expression `x * (sx*cy + cx*sy)`.
 
+ @param half_angle If true, then the same workflow replaces instances of
+ sin(q/2) and cos(q/2) in `e` will be replaced with `s`, and `c`.
+ @default false.
+
+ The half-angle representation is more natural in many analysis computations
+ for robots, for instance:
+ https://underactuated.csail.mit.edu/lyapunov.html#trig_quadratic
+
  @throws std::exception if a trigonometric function is not a trigonometric
  polynomial in `q` or if the `e` requires a trigonometric expansion that not
  supported yet.
  @pydrake_mkdoc_identifier{sincos}
 */
-Expression Substitute(const Expression& e,
-                      const SinCosSubstitution& subs);
+Expression Substitute(const Expression& e, const SinCosSubstitution& subs);
 
 /** Matrix version of sin/cos substitution.
  @pydrake_mkdoc_identifier{sincos_matrix} */
@@ -58,8 +81,9 @@ Substitute(const Eigen::MatrixBase<Derived>& m,
                 "Substitute only accepts a matrix of symbolic::Expression.");
   // Note that the return type is written out explicitly to help gcc 5 (on
   // ubuntu).
-  return m.unaryExpr(
-      [&subs](const Expression& e) { return Substitute(e, subs); });
+  return m.unaryExpr([&subs](const Expression& e) {
+    return Substitute(e, subs);
+  });
 }
 
 }  // namespace symbolic

common/test/symbolic_trigonometric_polynomial_test.cc

@@ -137,6 +137,12 @@ GTEST_TEST(SymbolicLatex, BasicTest) {
   DRAKE_EXPECT_THROWS_MESSAGE(Substitute(cos(if_then_else(x > y, x, y)), subs),
                               ".*'status == kNotSinCos' failed.*");
 
+  // Limited half-angle support.
+  DRAKE_EXPECT_THROWS_MESSAGE(Substitute(sin(0.5*x), subs),
+                              ".*only support sin.* where c is an integer.*");
+  DRAKE_EXPECT_THROWS_MESSAGE(Substitute(cos(0.5*x), subs),
+                              ".*only support cos.* where c is an integer.*");
+
   // Matrix<Expression>
   Eigen::Matrix<Expression, 2, 2> m;
   m << x, sin(x), sin(y), sin(x + y);
@@ -147,6 +153,39 @@ GTEST_TEST(SymbolicLatex, BasicTest) {
   EXPECT_PRED2(ExprEqual, m2(1, 1), sx * cy + sy * cx);
 }
 
+GTEST_TEST(SymbolicLatex, HalfAngleTest) {
+  Variable x{"x"}, y{"y"}, z{"z"};
+  Variable sx{"sx"}, sy{"sy"}, cx{"cx"}, cy{"cy"};
+  SinCosSubstitution subs;
+  subs.emplace(x, SinCos(sx, cx, SinCosSubstitutionType::kHalfAnglePreferSin));
+  subs.emplace(y, SinCos(sy, cy, SinCosSubstitutionType::kHalfAnglePreferCos));
+
+  EXPECT_PRED2(ExprEqual, Substitute(sin(x), subs), 2 * sx * cx);
+  EXPECT_PRED2(ExprEqual, Substitute(sin(1 * x), subs), 2 * sx * cx);
+  EXPECT_PRED2(ExprEqual, Substitute(sin(2 * x), subs),
+               4 * sx * cx * (1 - 2 * sx * sx));
+  EXPECT_PRED2(ExprEqual, Substitute(cos(x), subs), 1- 2 * sx * sx);
+  EXPECT_PRED2(ExprEqual, Substitute(sin(y), subs), 2 * sy * cy);
+  EXPECT_PRED2(ExprEqual, Substitute(cos(y), subs), 2 * cy * cy - 1);
+  EXPECT_PRED2(ExprEqual, Substitute(sin(0.5*x), subs), sx);
+  EXPECT_PRED2(ExprEqual, Substitute(cos(0.5*x), subs), cx);
+  EXPECT_PRED2(ExprEqual, Substitute(sin(-0.5*x), subs), -sx);
+  EXPECT_PRED2(ExprEqual, Substitute(cos(-0.5*x), subs), cx);
+  // Note: We don't support division yet, but calling Expand makes this work:
+  EXPECT_PRED2(ExprEqual, Substitute(sin(x / 2.0).Expand(), subs), sx);
+
+  // Matrix half-angle.
+  Eigen::Matrix<Expression, 2, 2> m;
+  m << x, sin(x), sin(y), sin(x + y);
+  const Eigen::Matrix<Expression, 2, 2> m2 = Substitute(m, subs);
+  EXPECT_PRED2(ExprEqual, m2(0, 0), x);
+  EXPECT_PRED2(ExprEqual, m2(0, 1), 2 * sx * cx);
+  EXPECT_PRED2(ExprEqual, m2(1, 0), 2 * sy * cy);
+  EXPECT_PRED2(
+      ExprEqual, m2(1, 1),
+      (2 * sx * cx) * (2 * cy * cy - 1) + (2 * sy * cy) * (1 - 2 * sx * sx));
+}
+
 }  // namespace
 }  // namespace symbolic
 }  // namespace drake