Use subd quadrature in VolumetricElement

multibody/fem/tet_subdivision_quadrature.cc

@@ -1 +1 @@
-#include "drake/multibody/fem/tet_subdivision_quadrature.h"
+#include "drake/multibody/fem/tet_subdivision_quadrature.h"

multibody/fem/volumetric_element.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <array>
+#include <optional>
 #include <type_traits>
 #include <utility>
 
@@ -18,7 +19,7 @@ namespace internal {
  FemElement for the requirement. We define it here instead of nesting it in the
  traits class below due to #17109. */
 template <typename ConstitutiveModelType, int num_dofs,
-          int num_quadrature_points>
+          int num_quadrature_points, int num_subd_quadrature_points>
 struct VolumetricElementData {
   using T = typename ConstitutiveModelType::T;
   /* The states evaluated at nodes of the element. */
@@ -28,10 +29,16 @@ struct VolumetricElementData {
   Vector<T, num_dofs> element_a;
   /* The current locations of the quadrature points in the world frame. */
   std::array<Vector<T, 3>, num_quadrature_points> quadrature_positions;
+  /* The current locations of the quadrature points used to integrate external
+   forces in the world frame. */
+  std::optional<std::array<Vector<T, 3>, num_subd_quadrature_points>>
+      subd_quadrature_positions;
 
   using DeformationGradientData = typename ConstitutiveModelType::Data;
   std::array<DeformationGradientData, num_quadrature_points>
       deformation_gradient_data;
+  std::optional<std::array<T, num_subd_quadrature_points>>
+      subd_deformation_gradient_determinant;
   /* The elastic energy density evaluated at quadrature points. Note that this
    is energy per unit of "reference" volume. */
   std::array<T, num_quadrature_points> Psi;
@@ -44,14 +51,18 @@ struct VolumetricElementData {
 
 /* Forward declaration needed for defining the traits below. */
 template <class IsoparametricElementType, class QuadratureType,
-          class ConstitutiveModelType>
+          class ConstitutiveModelType,
+          class SubdIsoparametricElementType = IsoparametricElementType,
+          class SubdQuadratureType = QuadratureType>
 class VolumetricElement;
 
 /* The traits class for volumetric elasticity FEM element. */
 template <class IsoparametricElementType, class QuadratureType,
-          class ConstitutiveModelType>
+          class ConstitutiveModelType, class SubdIsoparametricElementType,
+          class SubdQuadratureType>
 struct FemElementTraits<VolumetricElement<
-    IsoparametricElementType, QuadratureType, ConstitutiveModelType>> {
+    IsoparametricElementType, QuadratureType, ConstitutiveModelType,
+    SubdIsoparametricElementType, SubdQuadratureType>> {
   /* Check that template parameters are of the correct types. */
   static_assert(
       is_isoparametric_element<IsoparametricElementType>::value,
@@ -91,21 +102,26 @@ struct FemElementTraits<VolumetricElement<
                 "The natural dimension of the isoparametric element must be 3 "
                 "for volumetric FEM elements. Codimensional objects are not "
                 "yet supported.");
+  // TODO(xuchenhan-tri): check consistency of isoparametric element/quadrature
+  // with their subd counterparts.
 
   using T = typename ConstitutiveModelType::T;
   using ConstitutiveModel = ConstitutiveModelType;
 
   static constexpr int num_nodes = IsoparametricElementType::num_nodes;
   static constexpr int num_quadrature_points =
       QuadratureType::num_quadrature_points;
+  static constexpr int num_subd_quadrature_points =
+      SubdQuadratureType::num_quadrature_points;
   static constexpr int natural_dimension = QuadratureType::natural_dimension;
   /* The number of degrees of freedom is equal to the spatial dimension (which
    gives the number of degrees of freedom for a single node) times the number of
    nodes. */
   static constexpr int num_dofs = 3 * num_nodes;
 
-  using Data = VolumetricElementData<ConstitutiveModelType, num_dofs,
-                                     num_quadrature_points>;
+  using Data =
+      VolumetricElementData<ConstitutiveModelType, num_dofs,
+                            num_quadrature_points, num_subd_quadrature_points>;
 };
 
 /* This class models a single 3D elasticity FEM element in which the
@@ -122,26 +138,35 @@ struct FemElementTraits<VolumetricElement<
                                 VolumetricElement. ConstitutiveModelType must be
                                 derived from ConstitutiveModel. */
 template <class IsoparametricElementType, class QuadratureType,
-          class ConstitutiveModelType>
+          class ConstitutiveModelType, class SubdIsoparametricElementType,
+          class SubdQuadratureType>
 class VolumetricElement
     : public FemElement<VolumetricElement<
-          IsoparametricElementType, QuadratureType, ConstitutiveModelType>> {
+          IsoparametricElementType, QuadratureType, ConstitutiveModelType,
+          SubdIsoparametricElementType, SubdQuadratureType>> {
  public:
   DRAKE_DEFAULT_COPY_AND_MOVE_AND_ASSIGN(VolumetricElement);
 
-  using ElementType = VolumetricElement<IsoparametricElementType,
-                                        QuadratureType, ConstitutiveModelType>;
+  using ElementType =
+      VolumetricElement<IsoparametricElementType, QuadratureType,
+                        ConstitutiveModelType, SubdIsoparametricElementType,
+                        SubdQuadratureType>;
   using IsoparametricElement = IsoparametricElementType;
   using Quadrature = QuadratureType;
+  using SubdIsoparametricElement = SubdIsoparametricElementType;
+  using SubdQuadrature = SubdQuadratureType;
 
   using Traits = FemElementTraits<ElementType>;
   using Data = typename Traits::Data;
   using T = typename Traits::T;
   static constexpr int natural_dimension = Traits::natural_dimension;
   static constexpr int kSpatialDimension = 3;
   static constexpr int num_quadrature_points = Traits::num_quadrature_points;
+  static constexpr int num_subd_quadrature_points =
+      Traits::num_subd_quadrature_points;
   static constexpr int num_dofs = Traits::num_dofs;
   static constexpr int num_nodes = Traits::num_nodes;
+  static constexpr bool use_subd = !std::is_same_v<SubdQuadrature, Quadrature>;
 
   /* Constructs a new VolumetricElement. In that process, precomputes the mass
    matrix and the gravity force acting on the element.
@@ -197,6 +222,34 @@ class VolumetricElement
       dSdX_transpose_[q] = dSdX[q].transpose();
     }
     mass_matrix_ = PrecomputeMassMatrix();
+
+    if constexpr (use_subd) {
+      // TODO(xuchenhan-tri): The following code block is duplicated in the
+      // previous block. Consider refactoring.
+      subd_quadrature_ = SubdQuadrature{};
+      subd_isoparametric_element_ =
+          SubdIsoparametricElement(subd_quadrature_.get_points());
+      /* Computes the Jacobian of the change of variable function X(ξ) at subd
+       quadrature locations. */
+      const std::array<Eigen::Matrix<T, 3, natural_dimension>,
+                       num_subd_quadrature_points>
+          subd_dXdxi = subd_isoparametric_element_.value().CalcJacobian(
+              reference_positions);
+      /* Record the quadrature point volume in reference configuration for each
+       quadrature location. */
+      std::array<T, num_subd_quadrature_points> subd_reference_volume;
+      for (int q = 0; q < num_subd_quadrature_points; ++q) {
+        /* The scale to transform quadrature weight in parent coordinates to
+         reference coordinates. */
+        T volume_scale;
+        volume_scale = subd_dXdxi[q].determinant();
+        /* Degenerate element in the initial configuration is not allowed. */
+        DRAKE_DEMAND(volume_scale > 0);
+        subd_reference_volume[q] =
+            volume_scale * subd_quadrature_.value().get_weight(q);
+      }
+      subd_reference_volume_ = std::move(subd_reference_volume);
+    }
   }
 
   /* Calculates the elastic potential energy (in joules) stored in this element
@@ -373,9 +426,9 @@ class VolumetricElement
             element_q_reshaped);
 
     std::array<Matrix3<T>, num_quadrature_points> F =
-        CalcDeformationGradient(data.element_q);
-    std::array<Matrix3<T>, num_quadrature_points> F0 =
-        CalcDeformationGradient(data.element_q0);
+        CalcDeformationGradient(data.element_q, isoparametric_element_, dxidX_);
+    std::array<Matrix3<T>, num_quadrature_points> F0 = CalcDeformationGradient(
+        data.element_q0, isoparametric_element_, dxidX_);
 
     for (int q = 0; q < num_quadrature_points; ++q) {
       data.deformation_gradient_data[q].UpdateData(F[q], F0[q]);
@@ -386,13 +439,37 @@ class VolumetricElement
       this->constitutive_model().CalcFirstPiolaStressDerivative(
           data.deformation_gradient_data[q], &(data.dPdF[q]));
     }
+    if constexpr (use_subd) {
+      data.subd_quadrature_positions =
+          subd_isoparametric_element_.template InterpolateNodalValues<3>(
+              element_q_reshaped);
+      std::array<T, num_subd_quadrature_points> subd_determinant;
+      std::array<Matrix3<T>, num_subd_quadrature_points> subd_F =
+          CalcDeformationGradient(data.element_q, subd_isoparametric_element_,
+                                  subd_dxidX_);
+      for (int q = 0; q < num_subd_quadrature_points; ++q) {
+        subd_determinant[q] = subd_F[q].determinant();
+      }
+      data.subd_deformation_gradient_determinant = std::move(subd_determinant);
+    }
     return data;
   }
 
+  // TODO(xuchenhan-tri): Refactor to reduce code duplication.
   void DoAddScaledExternalForces(const Data& data,
                                  const FemPlantData<T>& plant_data,
                                  const T& scale,
                                  EigenPtr<Vector<T, num_dofs>> result) const {
+    if constexpr (use_subd) {
+      AddScaledExternalForcesSubd(data, plant_data, scale, result);
+    } else {
+      AddScaledExternalForcesNoSubd(data, plant_data, scale, result);
+    }
+  }
+
+  void AddScaledExternalForcesNoSubd(
+      const Data& data, const FemPlantData<T>& plant_data, const T& scale,
+      EigenPtr<Vector<T, num_dofs>> result) const {
     const std::array<Vector<T, 3>, num_quadrature_points>&
         quadrature_positions = data.quadrature_positions;
     const std::array<Vector<T, num_nodes>, num_quadrature_points>& S =
@@ -420,21 +497,68 @@ class VolumetricElement
     }
   }
 
+  void AddScaledExternalForcesSubd(const Data& data,
+                                   const FemPlantData<T>& plant_data,
+                                   const T& scale,
+                                   EigenPtr<Vector<T, num_dofs>> result) const {
+    const std::array<Vector<T, 3>, num_subd_quadrature_points>&
+        quadrature_positions = data.subd_quadrature_positions;
+    const std::array<Vector<T, num_nodes>, num_subd_quadrature_points>& S =
+        subd_isoparametric_element_.GetShapeFunctions();
+    for (int q = 0; q < num_quadrature_points; ++q) {
+      const T& det = data.subd_deformation_gradient_determinant[q];
+      Vector3<T> scaled_force = Vector3<T>::Zero();
+      for (const multibody::ForceDensityField<T>* force_density :
+           plant_data.force_density_fields) {
+        DRAKE_ASSERT(force_density != nullptr);
+        const T change_of_volume =
+            force_density->density_type() ==
+                    multibody::ForceDensityType::kPerReferenceVolume
+                ? 1.0
+                : det;
+        scaled_force += scale *
+                        force_density->EvaluateAt(plant_data.plant_context,
+                                                  quadrature_positions[q]) *
+                        subd_reference_volume_[q] * change_of_volume;
+      }
+      for (int n = 0; n < num_nodes; ++n) {
+        result->template segment<3>(3 * n) += scaled_force * S[q](n);
+      }
+    }
+  }
+
   /* Calculates the deformation gradient at all quadrature points in this
    element.
    @param[in] element_q  The positions of the nodes of the element in a flat
    vector. */
   std::array<Matrix3<T>, num_quadrature_points> CalcDeformationGradient(
       const Eigen::Ref<const VectorX<T>>& element_q) const {
-    std::array<Matrix3<T>, num_quadrature_points> F;
+    return CalcDeformationGradient(element_q, isoparametric_element_, dxidX_);
+  }
+
+  /* Calculates the deformation gradient at all quadrature points in this
+   element.
+   @param[in] element_q  The positions of the nodes of the element in a flat
+   vector. */
+  template <class IsoparametricElement>
+  std::array<Matrix3<T>, IsoparametricElement::num_sample_locations>
+  CalcDeformationGradient(
+      const Eigen::Ref<const VectorX<T>>& element_q,
+      const IsoparametricElement& element,
+      const std::array<Eigen::Matrix<T, natural_dimension, 3>,
+                       IsoparametricElement::num_sample_locations>& dxidX)
+      const {
+    constexpr int num_sample_locations =
+        IsoparametricElement::num_sample_locations;
+    std::array<Matrix3<T>, num_sample_locations> F;
     const auto& element_q_reshaped =
         Eigen::Map<const Eigen::Matrix<T, 3, num_nodes>>(element_q.data(), 3,
                                                          num_nodes);
-    const std::array<typename IsoparametricElementType::JacobianMatrix,
-                     num_quadrature_points>
-        dxdxi = isoparametric_element_.CalcJacobian(element_q_reshaped);
-    for (int quad = 0; quad < num_quadrature_points; ++quad) {
-      F[quad] = dxdxi[quad] * dxidX_[quad];
+    const std::array<typename IsoparametricElement::JacobianMatrix,
+                     num_sample_locations>
+        dxdxi = element.CalcJacobian(element_q_reshaped);
+    for (int quad = 0; quad < num_sample_locations; ++quad) {
+      F[quad] = dxdxi[quad] * dxidX[quad];
     }
     return F;
   }
@@ -507,6 +631,17 @@ class VolumetricElement
   T density_;
   /* Precomputed mass matrix. */
   Eigen::Matrix<T, num_dofs, num_dofs> mass_matrix_;
+
+  /* Optional data member mirror their non-subd counterparts. These have values
+   when `use_subd` flag is on, and stores data associated with the subd
+   quadrature points used to integrate external forces more precisely. */
+  std::optional<SubdQuadratureType> subd_quadrature_;
+  std::optional<SubdIsoparametricElementType> subd_isoparametric_element_;
+  std::optional<std::array<Eigen::Matrix<T, natural_dimension, 3>,
+                           num_subd_quadrature_points>>
+      subd_dxidX_;
+  std::optional<std::array<T, num_subd_quadrature_points>>
+      subd_reference_volume_;
 };
 
 }  // namespace internal