Add conversion between full and active index in DofMask.

bindings/generated_docstrings/planning.h

@@ -1800,6 +1800,18 @@ vector instead of writing to an output argument.
 Precondition:
     ``full_vec.size() == size()``.)""";
         } GetFromArray;
+        // Symbol: drake::planning::DofMask::GetFullToSelectedIndex
+        struct /* GetFullToSelectedIndex */ {
+          // Source: drake/planning/dof_mask.h
+          const char* doc =
+R"""(The inverse mapping of GetSelectedToFullIndex(). If we have q_selected
+= dof_mask.GetFromArray(q_full), the this function returns the mapping
+from q_full index to q_selected index. Namely if dof_mask[i] is true,
+namely q_full[i] is selected, then
+q_selected[dof_mask.GetFullToSelectedMap()[i]] is the same as
+q_full[i]; if dof_mask[i] is false, then
+dof_mask.GetFullToSelected()[i] is nullopt.)""";
+        } GetFullToSelectedIndex;
         // Symbol: drake::planning::DofMask::GetJoints
         struct /* GetJoints */ {
           // Source: drake/planning/dof_mask.h
@@ -1814,6 +1826,15 @@ order.
 Precondition:
     ``plant.num_positions() == size()``.)""";
         } GetJoints;
+        // Symbol: drake::planning::DofMask::GetSelectedToFullIndex
+        struct /* GetSelectedToFullIndex */ {
+          // Source: drake/planning/dof_mask.h
+          const char* doc =
+R"""(If we have q_selected = dof_mask.GetFromArray(q_full), then this
+function returns a mapping from q_selected index to q_full index, such
+that q_selected[i] is the same as
+q_full[dof_mask.GetSelectedToFull()[i]].)""";
+        } GetSelectedToFullIndex;
         // Symbol: drake::planning::DofMask::Intersect
         struct /* Intersect */ {
           // Source: drake/planning/dof_mask.h

bindings/pydrake/planning/planning_py_dof_mask.cc

@@ -48,7 +48,11 @@ void DefinePlanningDofMask(py::module m) {
         .def("Intersect", &Class::Intersect, py::arg("other"),
             cls_doc.Intersect.doc)
         .def("Subtract", &Class::Subtract, py::arg("other"),
-            cls_doc.Subtract.doc);
+            cls_doc.Subtract.doc)
+        .def("GetFullToSelectedIndex", &Class::GetFullToSelectedIndex,
+            cls_doc.GetFullToSelectedIndex.doc)
+        .def("GetSelectedToFullIndex", &Class::GetSelectedToFullIndex,
+            cls_doc.GetSelectedToFullIndex.doc);
     py::implicitly_convertible<std::vector<bool>, Class>();
   }
 }

bindings/pydrake/planning/test/dof_mask_test.py

@@ -43,3 +43,7 @@ def test_dof_set(self):
         self.assertEqual(mut.DofMask([True, False]).count(), 1)
         self.assertTrue(mut.DofMask([True, False])[0])
         self.assertFalse(mut.DofMask([True, False])[1])
+        self.assertEqual(
+            mut.DofMask([True, False]).GetFullToSelectedIndex(), [0, None])
+        self.assertEqual(
+            mut.DofMask([False, True]).GetSelectedToFullIndex(), [1])

planning/dof_mask.cc

@@ -203,5 +203,30 @@ bool DofMask::operator==(const DofMask& o) const {
   return result;
 }
 
+std::vector<int> DofMask::GetSelectedToFullIndex() const {
+  std::vector<int> ret(count_);
+  int count = 0;
+  for (int i = 0; i < this->size(); ++i) {
+    if (data_[i]) {
+      ret[count++] = i;
+    }
+  }
+  return ret;
+}
+
+std::vector<std::optional<int>> DofMask::GetFullToSelectedIndex() const {
+  std::vector<std::optional<int>> ret(this->size(), std::nullopt);
+  int count = 0;
+  for (int i = 0; i < this->size(); ++i) {
+    if (data_[i]) {
+      ret[i] = count++;
+      if (count >= count_) {
+        break;
+      }
+    }
+  }
+  return ret;
+}
+
 }  // namespace planning
 }  // namespace drake

planning/dof_mask.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <initializer_list>
+#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
@@ -234,6 +235,20 @@ class DofMask {
   @pre `output.size() == size()`. */
   void SetInArray(const Eigen::Ref<const Eigen::VectorXd>& vec,
                   drake::EigenPtr<Eigen::VectorXd> output) const;
+
+  /** If we have q_selected = dof_mask.GetFromArray(q_full), then this function
+   returns a mapping from q_selected index to q_full index, such that
+   q_selected[i] is the same as q_full[dof_mask.GetSelectedToFull()[i]]. */
+  [[nodiscard]] std::vector<int> GetSelectedToFullIndex() const;
+
+  /** The inverse mapping of GetSelectedToFullIndex().
+   If we have q_selected = dof_mask.GetFromArray(q_full), the this function
+   returns the mapping from q_full index to q_selected index. Namely
+   if dof_mask[i] is true, namely q_full[i] is selected,  then
+   q_selected[dof_mask.GetFullToSelectedMap()[i]] is the same as q_full[i]; if
+   dof_mask[i] is false, then dof_mask.GetFullToSelected()[i] is nullopt.
+   */
+  [[nodiscard]] std::vector<std::optional<int>> GetFullToSelectedIndex() const;
   //@}
 
  private:

planning/test/dof_mask_test.cc

@@ -1,6 +1,7 @@
 #include "drake/planning/dof_mask.h"
 
 #include <algorithm>
+#include <optional>
 #include <string>
 #include <utility>
 #include <vector>
@@ -295,6 +296,25 @@ GTEST_TEST(DofMaskTest, SetInArray) {
   EXPECT_EQ(full_new, full_new_expected);
 }
 
+GTEST_TEST(DofMaskTest, GetSelectedToFullIndex) {
+  const DofMask dofs1{true, false, false, true, false};
+  EXPECT_EQ(dofs1.GetSelectedToFullIndex(), std::vector<int>({0, 3}));
+
+  const DofMask dofs2{false, false, false, false};
+  EXPECT_TRUE(dofs2.GetSelectedToFullIndex().empty());
+}
+
+GTEST_TEST(DofMaskTest, GetFullToSelectedIndex) {
+  const DofMask dofs1{true, false, false, true, false, true, false};
+  EXPECT_EQ(dofs1.GetFullToSelectedIndex(),
+            (std::vector<std::optional<int>>{0, std::nullopt, std::nullopt, 1,
+                                             std::nullopt, 2, std::nullopt}));
+
+  const DofMask dofs2{false, false, false, false};
+  EXPECT_EQ(dofs2.GetFullToSelectedIndex(),
+            std::vector<std::optional<int>>(4, std::nullopt));
+}
+
 }  // namespace
 }  // namespace planning
 }  // namespace drake