Add an output for density in LinearTransformDensity. (RobotLocomotion…

…#15189)

* Add an output for density in LinearTransformDensity.

bindings/pydrake/systems/primitives_py.cc

@@ -356,6 +356,14 @@ PYBIND11_MODULE(primitives, m) {
         .def("get_input_port_b", &LinearTransformDensity<T>::get_input_port_b,
             py_rvp::reference_internal,
             doc.LinearTransformDensity.get_input_port_b.doc)
+        .def("get_output_port_w_out",
+            &LinearTransformDensity<T>::get_output_port_w_out,
+            py_rvp::reference_internal,
+            doc.LinearTransformDensity.get_output_port_w_out.doc)
+        .def("get_output_port_w_out_density",
+            &LinearTransformDensity<T>::get_output_port_w_out_density,
+            py_rvp::reference_internal,
+            doc.LinearTransformDensity.get_output_port_w_out_density.doc)
         .def("get_distribution", &LinearTransformDensity<T>::get_distribution,
             doc.LinearTransformDensity.get_distribution.doc)
         .def("FixConstantA", &LinearTransformDensity<T>::FixConstantA,

bindings/pydrake/systems/test/primitives_test.py

@@ -310,6 +310,9 @@ def test_linear_transform_density(self, T):
 
         dut.CalcDensity(context=context)
 
+        self.assertEqual(dut.get_output_port_w_out().size(), 3)
+        self.assertEqual(dut.get_output_port_w_out_density().size(), 1)
+
     def test_vector_pass_through(self):
         model_value = BasicVector([1., 2, 3])
         system = PassThrough(vector_size=model_value.size())

systems/primitives/linear_transform_density.cc

@@ -22,8 +22,15 @@ LinearTransformDensity<T>::LinearTransformDensity(
   b_port_id_ =
       this->DeclareInputPort("b", kVectorValued, output_size_).get_index();
 
-  this->DeclareVectorOutputPort("w_out", BasicVector<T>(output_size_),
-                                &LinearTransformDensity<T>::CalcOutput);
+  w_out_port_id_ =
+      this->DeclareVectorOutputPort("w_out", BasicVector<T>(output_size_),
+                                    &LinearTransformDensity<T>::CalcOutput)
+          .get_index();
+  w_out_density_port_id_ =
+      this->DeclareVectorOutputPort(
+              "w_out_density", BasicVector<T>(1),
+              &LinearTransformDensity<T>::CalcOutputDensity)
+          .get_index();
 }
 
 template <typename T>
@@ -49,6 +56,12 @@ void LinearTransformDensity<T>::CalcOutput(const Context<T>& context,
   }
 }
 
+template <typename T>
+void LinearTransformDensity<T>::CalcOutputDensity(
+    const Context<T>& context, BasicVector<T>* w_out_density) const {
+  w_out_density->get_mutable_value()(0) = CalcDensity(context);
+}
+
 template <typename T>
 Eigen::Map<const MatrixX<T>> LinearTransformDensity<T>::GetA(
     const Context<T>& context) const {

systems/primitives/linear_transform_density.h

@@ -21,6 +21,7 @@ namespace systems {
  * - b (vector, optional)
  * output_ports:
  * - w_out
+ * - w_out_density
  * @endsystem
  *
  * The `b` port can remain disconnected, in which case it defaults to zero.
@@ -76,6 +77,14 @@ class LinearTransformDensity final : public LeafSystem<T> {
     return this->get_input_port(b_port_id_);
   }
 
+  const OutputPort<T>& get_output_port_w_out() const {
+    return this->get_output_port(w_out_port_id_);
+  }
+
+  const OutputPort<T>& get_output_port_w_out_density() const {
+    return this->get_output_port(w_out_density_port_id_);
+  }
+
   /** Gets the random distribution type. */
   RandomDistribution get_distribution() const { return distribution_; }
 
@@ -114,6 +123,9 @@ class LinearTransformDensity final : public LeafSystem<T> {
  private:
   void CalcOutput(const Context<T>& context, BasicVector<T>* w_out) const;
 
+  void CalcOutputDensity(const Context<T>& context,
+                         BasicVector<T>* w_out_density) const;
+
   Eigen::Map<const MatrixX<T>> GetA(const Context<T>& context) const;
 
   const RandomDistribution distribution_;
@@ -122,6 +134,8 @@ class LinearTransformDensity final : public LeafSystem<T> {
   InputPortIndex w_in_port_id_;
   InputPortIndex A_port_id_;
   InputPortIndex b_port_id_;
+  OutputPortIndex w_out_port_id_;
+  OutputPortIndex w_out_density_port_id_;
 };
 
 // Exclude symbolic::Expression from the scalartype conversion of

systems/primitives/test/linear_transform_density_test.cc

@@ -19,15 +19,16 @@ void TestConstructor() {
   LinearTransformDensity<T> dut(RandomDistribution::kUniform,
                                 2 /* input_size */, 3 /* output_size */);
   EXPECT_EQ(dut.num_input_ports(), 3);
-  EXPECT_EQ(dut.num_output_ports(), 1);
+  EXPECT_EQ(dut.num_output_ports(), 2);
   EXPECT_EQ(dut.get_distribution(), RandomDistribution::kUniform);
 
   EXPECT_EQ(dut.get_input_port_w_in().size(), 2);
   EXPECT_EQ(dut.get_input_port_w_in().get_random_type(),
             RandomDistribution::kUniform);
   EXPECT_EQ(dut.get_input_port_A().size(), 6);
   EXPECT_EQ(dut.get_input_port_b().size(), 3);
-  EXPECT_EQ(dut.get_output_port().size(), 3);
+  EXPECT_EQ(dut.get_output_port_w_out().size(), 3);
+  EXPECT_EQ(dut.get_output_port_w_out_density().size(), 1);
 }
 
 GTEST_TEST(LinearTransformDensityTest, Constructor) {
@@ -45,8 +46,10 @@ GTEST_TEST(LinearTransformDensityTest, ToAutoDiff) {
               dut->get_input_port_A().size());
     EXPECT_EQ(converted.get_input_port_b().size(),
               dut->get_input_port_b().size());
-    EXPECT_EQ(converted.get_output_port().size(),
-              dut->get_output_port().size());
+    EXPECT_EQ(converted.get_output_port_w_out().size(),
+              dut->get_output_port_w_out().size());
+    EXPECT_EQ(converted.get_output_port_w_out_density().size(),
+              dut->get_output_port_w_out_density().size());
   }));
 }
 
@@ -56,14 +59,16 @@ void TestCalcOutput() {
   auto context = dut.CreateDefaultContext();
 
   Eigen::Matrix<T, 3, 2> A;
-  // clang-format on
-  A << 1., 2., 3., 4., 5., 6.;
   // clang-format off
+  A << 1., 2.,
+       3., 4.,
+       5., 6.;
+  // clang-format on
   const FixedInputPortValue& port_A_val = dut.FixConstantA(context.get(), A);
   EXPECT_EQ(port_A_val.template get_vector_value<T>().size(), 6);
   // Confirm that A is stored in column-major order.
   const auto& port_A_vector =
-    port_A_val.template get_vector_value<T>().get_value();
+      port_A_val.template get_vector_value<T>().get_value();
   EXPECT_EQ(port_A_vector(0), 1.);
   EXPECT_EQ(port_A_vector(1), 3.);
   EXPECT_EQ(port_A_vector(2), 5.);
@@ -74,7 +79,7 @@ void TestCalcOutput() {
   dut.get_input_port_w_in().FixValue(context.get(), w_in);
 
   // Test when port b is not connected.
-  const auto w_out_no_b = dut.get_output_port().Eval(*context);
+  const auto w_out_no_b = dut.get_output_port_w_out().Eval(*context);
   const Vector3<T> w_out_no_b_expected = A * w_in;
   EXPECT_TRUE(CompareMatrices(w_out_no_b, w_out_no_b_expected));
 
@@ -83,7 +88,7 @@ void TestCalcOutput() {
   const auto& port_b_val = dut.FixConstantB(context.get(), b);
   EXPECT_EQ(port_b_val.template get_vector_value<T>().size(), 3);
 
-  const auto w_out = dut.get_output_port().Eval(*context);
+  const auto w_out = dut.get_output_port_w_out().Eval(*context);
   const Vector3<T> w_out_expected = A * w_in + b;
   EXPECT_TRUE(CompareMatrices(w_out, w_out_expected, 10 * kEps));
 }
@@ -93,6 +98,42 @@ GTEST_TEST(LinearTransformDensityTest, CalcOutput) {
   TestCalcOutput<AutoDiffXd>();
 }
 
+template <typename T>
+void TestCalcOutputDensity() {
+  LinearTransformDensity<T> dut(RandomDistribution::kUniform, 2, 2);
+  auto context = dut.CreateDefaultContext();
+
+  Eigen::Matrix<T, 2, 2> A;
+  // clang-format off
+  A << 1., 2.,
+       3., 4.;
+  // clang-format on
+  dut.FixConstantA(context.get(), A);
+  Vector2<T> w_in(2, 3);
+  dut.get_input_port_w_in().FixValue(context.get(), w_in);
+
+  // Test when port b is not connected.
+  const auto w_out_density_no_b =
+      dut.get_output_port_w_out_density().Eval(*context);
+  const T w_out_density_no_b_expected = dut.CalcDensity(*context);
+  EXPECT_TRUE(CompareMatrices(w_out_density_no_b,
+                              Vector1<T>(w_out_density_no_b_expected)));
+
+  // Test with port b connected.
+  Vector2<T> b(-1, -2);
+  dut.FixConstantB(context.get(), b);
+
+  const auto w_out_density = dut.get_output_port_w_out_density().Eval(*context);
+  const T w_out_density_expected = dut.CalcDensity(*context);
+  EXPECT_TRUE(
+      CompareMatrices(w_out_density, Vector1<T>(w_out_density_expected)));
+}
+
+GTEST_TEST(LinearTransformDensityTest, TestCalcOutputDensity) {
+  TestCalcOutputDensity<double>();
+  TestCalcOutputDensity<AutoDiffXd>();
+}
+
 template <typename T>
 void CheckDensity(RandomDistribution distribution,
                   const std::vector<Eigen::Vector2d>& nonzero_prob_inputs,
@@ -163,7 +204,7 @@ void CheckGaussianDensity() {
   Vector2<T> b(2, 3);
   dut.FixConstantB(context.get(), b);
   T density = dut.CalcDensity(*context);
-  auto w_out = dut.get_output_port().Eval(*context);
+  auto w_out = dut.get_output_port_w_out().Eval(*context);
   Eigen::LDLT<Matrix2<T>> ldlt_solver;
   Matrix2<T> cov = A * A.transpose();
   ldlt_solver.compute(cov);
@@ -326,7 +367,7 @@ void CheckDensityGradient(RandomDistribution distribution,
     }
   }
   Eigen::Vector2d w_out_val =
-      math::autoDiffToValueMatrix(dut.get_output_port().Eval(*context));
+      math::autoDiffToValueMatrix(dut.get_output_port_w_out().Eval(*context));
   Eigen::Vector2d b_val = math::autoDiffToValueMatrix(b);
   const double A_det_abs = std::abs(A_val.determinant());
   for (int i = 0; i < 2; ++i) {