Snapshot optimized solver and audit diagnostics before numerical experiments

libmeanfield/impl/fem/reference_tables.cpp

@@ -1,6 +1,8 @@
 module;
 
+#include <algorithm>
 #include <array>
+#include <limits>
 #include <cmath>
 #include <functional>
 #include <map>
@@ -11,12 +13,54 @@ module;
 #include <vector>
 
 #include <mfem.hpp>
+#include <mpi.h>
 
 module mean_field;
 import :fem.reference_tables;
 
 namespace mean_field::fem {
     namespace {
+        // Three one-dimensional contractions, with x fastest throughout.
+        // Input/output must not alias either scratch array.
+        void tensor_values(
+            int dx, int dy, int dz, int q, const double *bx, const double *by, const double *bz,
+            const double *input, double *output, double *sx, double *sxy, bool transpose
+        ) {
+            if (!transpose) {
+                for (int k = 0; k < dz; ++k) for (int j = 0; j < dy; ++j) for (int x = 0; x < q; ++x) {
+                    double v = 0.0;
+                    for (int i = 0; i < dx; ++i) v += bx[x + q*i] * input[i + dx*(j + dy*k)];
+                    sx[x + q*(j + dy*k)] = v;
+                }
+                for (int k = 0; k < dz; ++k) for (int y = 0; y < q; ++y) for (int x = 0; x < q; ++x) {
+                    double v = 0.0;
+                    for (int j = 0; j < dy; ++j) v += by[y + q*j] * sx[x + q*(j + dy*k)];
+                    sxy[x + q*(y + q*k)] = v;
+                }
+                for (int z = 0; z < q; ++z) for (int y = 0; y < q; ++y) for (int x = 0; x < q; ++x) {
+                    double v = 0.0;
+                    for (int k = 0; k < dz; ++k) v += bz[z + q*k] * sxy[x + q*(y + q*k)];
+                    output[x + q*(y + q*z)] = v;
+                }
+            } else {
+                for (int k = 0; k < dz; ++k) for (int y = 0; y < q; ++y) for (int x = 0; x < q; ++x) {
+                    double v = 0.0;
+                    for (int z = 0; z < q; ++z) v += bz[z + q*k] * input[x + q*(y + q*z)];
+                    sxy[x + q*(y + q*k)] = v;
+                }
+                for (int k = 0; k < dz; ++k) for (int j = 0; j < dy; ++j) for (int x = 0; x < q; ++x) {
+                    double v = 0.0;
+                    for (int y = 0; y < q; ++y) v += by[y + q*j] * sxy[x + q*(y + q*k)];
+                    sx[x + q*(j + dy*k)] = v;
+                }
+                for (int k = 0; k < dz; ++k) for (int j = 0; j < dy; ++j) for (int i = 0; i < dx; ++i) {
+                    double v = 0.0;
+                    for (int x = 0; x < q; ++x) v += bx[x + q*i] * sx[x + q*(j + dy*k)];
+                    output[i + dx*(j + dy*k)] = v;
+                }
+            }
+        }
+
         struct ReferenceTableKey {
             const mfem::FiniteElement *element;
             std::vector<std::array<double, 4>> points;
@@ -112,6 +156,7 @@ namespace mean_field::fem {
                 element.CalcDShape(point, gradient);
             }
         }
+        PrepareTensorEvaluation3D(element, rule);
     }
 
     const mfem::DenseMatrix &ScalarReferenceTable::GetValues() const {
@@ -130,6 +175,207 @@ namespace mean_field::fem {
         return m_dimension;
     }
 
+    void ScalarReferenceTable::PrepareTensorEvaluation3D(
+        const mfem::FiniteElement &element, const mfem::IntegrationRule &rule
+    ) {
+        // Nodal H1 and L2 hexes share the same verified tensor contract.
+        // Other elements and non-Cartesian rules keep the existing dense path.
+        const auto *hex = dynamic_cast<const mfem::NodalTensorFiniteElement *>(&element);
+        if (dynamic_cast<const mfem::H1_HexahedronElement *>(&element) == nullptr &&
+            dynamic_cast<const mfem::L2_HexahedronElement *>(&element) == nullptr) return;
+        if (hex == nullptr || element.GetMapType() != mfem::FiniteElement::VALUE ||
+            element.GetDerivType() != mfem::FiniteElement::GRAD || rule.GetNPoints() == 0) return;
+        const int d = element.GetOrder() + 1;
+        const int q = static_cast<int>(std::lround(std::cbrt(rule.GetNPoints())));
+        if (d * d * d != element.GetDof() || q * q * q != rule.GetNPoints()) return;
+        for (int z = 0; z < q; ++z) {
+            for (int y = 0; y < q; ++y) {
+                for (int x = 0; x < q; ++x) {
+                    const auto &point = rule.IntPoint(x + q * (y + q * z));
+                    if (point.x != rule.IntPoint(x).x ||
+                        point.y != rule.IntPoint(q * y).y ||
+                        point.z != rule.IntPoint(q * q * z).z) return;
+                }
+            }
+        }
+        const auto &map = hex->GetDofMap();
+        if (map.Size() != 0 && map.Size() != element.GetDof()) return;
+        m_tensorDofMap.SetSize(element.GetDof());
+        std::vector<bool> seen(element.GetDof(), false);
+        for (int i = 0; i < element.GetDof(); ++i) {
+            const int native = map.Size() == 0 ? i : map[i];
+            if (native < 0 || native >= element.GetDof() || seen[native]) return;
+            seen[native] = true;
+            m_tensorDofMap[i] = native;
+        }
+        mfem::Vector b(d), g(d);
+        for (int axis = 0; axis < 3; ++axis) {
+            m_tensorValues[axis].SetSize(q, d);
+            m_tensorGradients[axis].SetSize(q, d);
+            for (int i = 0; i < q; ++i) {
+                const double coordinate = axis == 0 ? rule.IntPoint(i).x :
+                    axis == 1 ? rule.IntPoint(q * i).y : rule.IntPoint(q * q * i).z;
+                hex->GetBasis1D().Eval(coordinate, b, g);
+                for (int j = 0; j < d; ++j) {
+                    m_tensorValues[axis](i, j) = b(j);
+                    m_tensorGradients[axis](i, j) = g(j);
+                }
+            }
+        }
+        // Verify the native DOF map and factorization once per shared table.
+        // This also guards against an unexpected MFEM basis convention.
+        const auto agrees = [](double actual, double expected) {
+            return std::isfinite(actual) && std::abs(actual - expected) <=
+                64.0 * std::numeric_limits<double>::epsilon() * std::max(1.0, std::abs(expected));
+        };
+        for (int z = 0; z < q; ++z) for (int y = 0; y < q; ++y) for (int x = 0; x < q; ++x) {
+            const int point = x + q * (y + q * z);
+            for (int k = 0; k < d; ++k) for (int j = 0; j < d; ++j) for (int i = 0; i < d; ++i) {
+                const int native = m_tensorDofMap[i + d * (j + d * k)];
+                const double bx = m_tensorValues[0](x, i), by = m_tensorValues[1](y, j),
+                             bz = m_tensorValues[2](z, k);
+                if (!agrees(bx * by * bz, m_values(point, native)) ||
+                    !agrees(m_tensorGradients[0](x, i) * by * bz, m_gradients[point](native, 0)) ||
+                    !agrees(bx * m_tensorGradients[1](y, j) * bz, m_gradients[point](native, 1)) ||
+                    !agrees(bx * by * m_tensorGradients[2](z, k), m_gradients[point](native, 2))) return;
+            }
+        }
+        m_tensorDofsPerAxis = d;
+        m_tensorPointsPerAxis = q;
+    }
+
+    bool ScalarReferenceTable::HasTensorEvaluation3D() const noexcept {
+        return m_tensorPointsPerAxis != 0;
+    }
+
+    void ScalarReferenceTable::EvaluateTensor3D(
+        const mfem::DenseMatrix &elementDofs, mfem::DenseMatrix &valuesAndGradients,
+        mfem::Vector &workspace
+    ) const {
+        MFEM_VERIFY(HasTensorEvaluation3D() && elementDofs.Height() == GetDofCount(),
+                    "Tensor evaluation requires a compatible verified hexahedral reference table.");
+        const int d = m_tensorDofsPerAxis, q = m_tensorPointsPerAxis;
+        const int nx = q * d * d, nxy = q * q * d, nq = q * q * q;
+        workspace.SetSize(2 * nx + 3 * nxy);
+        double *xValue = workspace.HostWrite(), *xGradient = xValue + nx;
+        double *xyValue = xGradient + nx, *xyDx = xyValue + nxy, *xyDy = xyDx + nxy;
+        valuesAndGradients.SetSize(nq, 4 * elementDofs.Width());
+        const double *u = elementDofs.Data();
+        double *out = valuesAndGradients.Data();
+        const double *bx = m_tensorValues[0].Data(), *by = m_tensorValues[1].Data(),
+                     *bz = m_tensorValues[2].Data(), *gx = m_tensorGradients[0].Data(),
+                     *gy = m_tensorGradients[1].Data(), *gz = m_tensorGradients[2].Data();
+        const int ndof = GetDofCount();
+        for (int c = 0; c < elementDofs.Width(); ++c) {
+            // Contract x, retaining the value and x derivative.
+            for (int k = 0; k < d; ++k) for (int j = 0; j < d; ++j) for (int x = 0; x < q; ++x) {
+                double value = 0.0, dx = 0.0;
+                for (int i = 0; i < d; ++i) {
+                    const double coefficient = u[m_tensorDofMap[i + d * (j + d * k)] + c * ndof];
+                    value += bx[x + q * i] * coefficient;
+                    dx += gx[x + q * i] * coefficient;
+                }
+                const int index = x + q * (j + d * k);
+                xValue[index] = value; xGradient[index] = dx;
+            }
+            // Contract y; propagate dx and form dy.
+            for (int k = 0; k < d; ++k) for (int y = 0; y < q; ++y) for (int x = 0; x < q; ++x) {
+                double value = 0.0, dx = 0.0, dy = 0.0;
+                for (int j = 0; j < d; ++j) {
+                    const int index = x + q * (j + d * k);
+                    value += by[y + q * j] * xValue[index];
+                    dx += by[y + q * j] * xGradient[index];
+                    dy += gy[y + q * j] * xValue[index];
+                }
+                const int index = x + q * (y + q * k);
+                xyValue[index] = value; xyDx[index] = dx; xyDy[index] = dy;
+            }
+            // Contract z and write all four fields in quadrature ordering.
+            for (int z = 0; z < q; ++z) for (int y = 0; y < q; ++y) for (int x = 0; x < q; ++x) {
+                double value = 0.0, dx = 0.0, dy = 0.0, dz = 0.0;
+                for (int k = 0; k < d; ++k) {
+                    const int index = x + q * (y + q * k);
+                    value += bz[z + q * k] * xyValue[index];
+                    dx += bz[z + q * k] * xyDx[index];
+                    dy += bz[z + q * k] * xyDy[index];
+                    dz += gz[z + q * k] * xyValue[index];
+                }
+                const int point = x + q * (y + q * z);
+                out[point + nq * (4 * c)] = value;
+                out[point + nq * (4 * c + 1)] = dx;
+                out[point + nq * (4 * c + 2)] = dy;
+                out[point + nq * (4 * c + 3)] = dz;
+            }
+        }
+    }
+
+    void ScalarReferenceTable::ApplyGradientTransposeTensor3D(
+        const mfem::DenseMatrix &quadratureDuals, mfem::DenseMatrix &elementDuals,
+        mfem::Vector &workspace
+    ) const {
+        MFEM_VERIFY(HasTensorEvaluation3D() && quadratureDuals.Height() == GetPointCount() &&
+                    quadratureDuals.Width() % 3 == 0, "Tensor gradient transpose has incompatible duals.");
+        const int d = m_tensorDofsPerAxis, q = m_tensorPointsPerAxis, n = GetDofCount();
+        workspace.SetSize(n + q*d*d + q*q*d);
+        double *lex = workspace.HostWrite(), *sx = lex + n, *sxy = sx + q*d*d;
+        elementDuals.SetSize(n, quadratureDuals.Width()/3);
+        elementDuals = 0.0;
+        for (int c = 0; c < elementDuals.Width(); ++c) for (int a = 0; a < 3; ++a) {
+            tensor_values(d, d, d, q,
+                (a == 0 ? m_tensorGradients[0] : m_tensorValues[0]).Data(),
+                (a == 1 ? m_tensorGradients[1] : m_tensorValues[1]).Data(),
+                (a == 2 ? m_tensorGradients[2] : m_tensorValues[2]).Data(),
+                quadratureDuals.GetColumn(3*c+a), lex, sx, sxy, true);
+            for (int i = 0; i < n; ++i) elementDuals(m_tensorDofMap[i], c) += lex[i];
+        }
+    }
+
+    void ScalarReferenceTable::EvaluateValuesTensor3D(
+        const mfem::Vector &elementDofs, mfem::Vector &values, mfem::Vector &workspace
+    ) const {
+        MFEM_VERIFY(HasTensorEvaluation3D() && elementDofs.Size() == GetDofCount(),
+                    "Tensor values require compatible element DOFs.");
+        const int d = m_tensorDofsPerAxis, q = m_tensorPointsPerAxis, n = GetDofCount();
+        workspace.SetSize(n + q*d*d + q*q*d);
+        double *lex = workspace.HostWrite(), *sx = lex + n, *sxy = sx + q*d*d;
+        const double *u = elementDofs.HostRead();
+        for (int i = 0; i < n; ++i) lex[i] = u[m_tensorDofMap[i]];
+        values.SetSize(GetPointCount());
+        tensor_values(d, d, d, q, m_tensorValues[0].Data(), m_tensorValues[1].Data(),
+                      m_tensorValues[2].Data(), lex, values.HostWrite(), sx, sxy, false);
+    }
+
+    void ScalarReferenceTable::ApplyTransposeTensor3D(
+        const mfem::Vector &quadratureDuals, mfem::Vector &elementDuals, mfem::Vector &workspace
+    ) const {
+        MFEM_VERIFY(HasTensorEvaluation3D() && quadratureDuals.Size() == GetPointCount(),
+                    "Tensor transpose requires compatible quadrature duals.");
+        const int d = m_tensorDofsPerAxis, q = m_tensorPointsPerAxis, n = GetDofCount();
+        workspace.SetSize(n + q*d*d + q*q*d);
+        double *lex = workspace.HostWrite(), *sx = lex + n, *sxy = sx + q*d*d;
+        tensor_values(d, d, d, q, m_tensorValues[0].Data(), m_tensorValues[1].Data(),
+                      m_tensorValues[2].Data(), quadratureDuals.HostRead(), lex, sx, sxy, true);
+        elementDuals.SetSize(n);
+        double *out = elementDuals.HostWrite();
+        for (int i = 0; i < n; ++i) out[m_tensorDofMap[i]] = lex[i];
+    }
+
+    void ScalarReferenceTable::ApplyTransposeTensor3D(
+        const mfem::DenseMatrix &quadratureDuals, mfem::DenseMatrix &elementDuals, mfem::Vector &workspace
+    ) const {
+        MFEM_VERIFY(HasTensorEvaluation3D() && quadratureDuals.Height() == GetPointCount(),
+                    "Tensor transpose requires compatible quadrature duals.");
+        const int d = m_tensorDofsPerAxis, q = m_tensorPointsPerAxis, n = GetDofCount();
+        workspace.SetSize(n + q*d*d + q*q*d);
+        double *lex = workspace.HostWrite(), *sx = lex + n, *sxy = sx + q*d*d;
+        elementDuals.SetSize(n, quadratureDuals.Width());
+        for (int c = 0; c < quadratureDuals.Width(); ++c) {
+            tensor_values(d, d, d, q, m_tensorValues[0].Data(), m_tensorValues[1].Data(),
+                          m_tensorValues[2].Data(), quadratureDuals.GetColumn(c), lex, sx, sxy, true);
+            for (int i = 0; i < n; ++i) elementDuals(m_tensorDofMap[i], c) = lex[i];
+        }
+    }
+
     VectorReferenceTable::VectorReferenceTable(
         const mfem::FiniteElement &element,
         const mfem::IntegrationRule &rule
@@ -142,6 +388,122 @@ namespace mean_field::fem {
             values.SetSize(m_dof_count, m_dimension);
             element.CalcVShape(rule.IntPoint(q), values);
         }
+        PrepareTensorEvaluation3D(element, rule);
+    }
+
+    void VectorReferenceTable::PrepareTensorEvaluation3D(
+        const mfem::FiniteElement &element, const mfem::IntegrationRule &rule
+    ) {
+        const auto *rt = dynamic_cast<const mfem::RT_HexahedronElement *>(&element);
+        if (rt == nullptr || element.GetMapType() != mfem::FiniteElement::H_DIV || rule.GetNPoints() == 0) return;
+        const int q = static_cast<int>(std::lround(std::cbrt(rule.GetNPoints())));
+        if (q*q*q != rule.GetNPoints()) return;
+        for (int z = 0; z < q; ++z) for (int y = 0; y < q; ++y) for (int x = 0; x < q; ++x) {
+            const auto &point = rule.IntPoint(x + q*(y + q*z));
+            if (point.x != rule.IntPoint(x).x || point.y != rule.IntPoint(q*y).y ||
+                point.z != rule.IntPoint(q*q*z).z) return;
+        }
+        int d = 2;
+        while (3*d*(d-1)*(d-1) < m_dof_count) ++d;
+        if (3*d*(d-1)*(d-1) != m_dof_count) return;
+        const auto &map = rt->GetDofMap();
+        if (map.Size() != m_dof_count) return;
+        m_tensorDofMap = map;
+        std::vector<bool> seen(m_dof_count, false);
+        for (int i = 0; i < m_dof_count; ++i) {
+            const int native = map[i] < 0 ? -1-map[i] : map[i];
+            if (native < 0 || native >= m_dof_count || seen[native]) return;
+            seen[native] = true;
+        }
+        const int block = m_dof_count / 3;
+        const auto basis = [&](int point, int lex, int component) {
+            const int encoded = map[lex], native = encoded < 0 ? -1-encoded : encoded;
+            return (encoded < 0 ? -1.0 : 1.0) * m_values[point](native, component);
+        };
+        // Extract owned 1D factors from the full basis. This supports nodal and
+        // integrated open bases without retaining DofToQuad pointers to rules.
+        for (int c = 0; c < 3; ++c) {
+            const int dims[3]{c == 0 ? d : d-1, c == 1 ? d : d-1, c == 2 ? d : d-1};
+            int anchor = 0;
+            for (int point = 1; point < q*q*q; ++point)
+                if (std::abs(basis(point, c*block, c)) > std::abs(basis(anchor, c*block, c))) anchor = point;
+            const double base = basis(anchor, c*block, c);
+            if (!std::isfinite(base) || base == 0.0) return;
+            const int coordinates[3]{anchor % q, (anchor/q) % q, anchor/(q*q)};
+            const int strides[3]{1, dims[0], dims[0]*dims[1]};
+            for (int axis = 0; axis < 3; ++axis) {
+                auto &b = m_tensorValues[c][axis];
+                b.SetSize(q, dims[axis]);
+                for (int i = 0; i < dims[axis]; ++i) for (int t = 0; t < q; ++t) {
+                    int xyz[3]{coordinates[0], coordinates[1], coordinates[2]};
+                    xyz[axis] = t;
+                    b(t, i) = basis(xyz[0] + q*(xyz[1] + q*xyz[2]), c*block + strides[axis]*i, c) /
+                              (axis == 0 ? 1.0 : base);
+                }
+            }
+            // Check every native basis entry, including zero components.
+            for (int k = 0; k < dims[2]; ++k) for (int j = 0; j < dims[1]; ++j)
+                for (int i = 0; i < dims[0]; ++i) {
+                    const int lex = c*block + i + dims[0]*(j + dims[1]*k);
+                    for (int z = 0; z < q; ++z) for (int y = 0; y < q; ++y) for (int x = 0; x < q; ++x) {
+                        const int point = x + q*(y + q*z);
+                        const double product = m_tensorValues[c][0](x,i) * m_tensorValues[c][1](y,j) *
+                                               m_tensorValues[c][2](z,k);
+                        for (int component = 0; component < 3; ++component) {
+                            const double expected = basis(point, lex, component);
+                            const double actual = component == c ? product : 0.0;
+                            if (!std::isfinite(actual) || std::abs(actual - expected) >
+                                128.0 * std::numeric_limits<double>::epsilon() * std::max(1.0, std::abs(expected))) return;
+                        }
+                    }
+                }
+        }
+        m_tensorClosedDofs = d;
+        m_tensorPointsPerAxis = q;
+    }
+
+    bool VectorReferenceTable::HasTensorEvaluation3D() const noexcept {
+        return m_tensorPointsPerAxis != 0;
+    }
+
+    void VectorReferenceTable::EvaluateTensor3D(
+        const mfem::Vector &elementDofs, mfem::DenseMatrix &values, mfem::Vector &workspace
+    ) const {
+        MFEM_VERIFY(HasTensorEvaluation3D() && elementDofs.Size() == m_dof_count,
+                    "RT tensor evaluation requires compatible element DOFs.");
+        const int d = m_tensorClosedDofs, q = m_tensorPointsPerAxis, block = m_dof_count/3;
+        workspace.SetSize(block + q*d*d + q*q*d);
+        double *lex = workspace.HostWrite(), *sx = lex + block, *sxy = sx + q*d*d;
+        values.SetSize(q*q*q, 3);
+        for (int c = 0; c < 3; ++c) {
+            for (int i = 0; i < block; ++i) {
+                const int encoded = m_tensorDofMap[c*block+i];
+                lex[i] = encoded < 0 ? -elementDofs(-1-encoded) : elementDofs(encoded);
+            }
+            tensor_values(c == 0 ? d : d-1, c == 1 ? d : d-1, c == 2 ? d : d-1, q,
+                          m_tensorValues[c][0].Data(), m_tensorValues[c][1].Data(), m_tensorValues[c][2].Data(),
+                          lex, values.GetColumn(c), sx, sxy, false);
+        }
+    }
+
+    void VectorReferenceTable::ApplyTransposeTensor3D(
+        const mfem::DenseMatrix &quadratureDuals, mfem::Vector &elementDuals, mfem::Vector &workspace
+    ) const {
+        MFEM_VERIFY(HasTensorEvaluation3D() && quadratureDuals.Height() == GetPointCount() && quadratureDuals.Width() == 3,
+                    "RT tensor transpose requires three quadrature dual components.");
+        const int d = m_tensorClosedDofs, q = m_tensorPointsPerAxis, block = m_dof_count/3;
+        workspace.SetSize(block + q*d*d + q*q*d);
+        double *lex = workspace.HostWrite(), *sx = lex + block, *sxy = sx + q*d*d;
+        elementDuals.SetSize(m_dof_count);
+        for (int c = 0; c < 3; ++c) {
+            tensor_values(c == 0 ? d : d-1, c == 1 ? d : d-1, c == 2 ? d : d-1, q,
+                          m_tensorValues[c][0].Data(), m_tensorValues[c][1].Data(), m_tensorValues[c][2].Data(),
+                          quadratureDuals.GetColumn(c), lex, sx, sxy, true);
+            for (int i = 0; i < block; ++i) {
+                const int encoded = m_tensorDofMap[c*block+i];
+                if (encoded < 0) elementDuals(-1-encoded) = -lex[i]; else elementDuals(encoded) = lex[i];
+            }
+        }
     }
 
     const mfem::DenseMatrix &VectorReferenceTable::GetValues(const int point) const {
@@ -156,4 +518,132 @@ namespace mean_field::fem {
     int VectorReferenceTable::GetDimension() const {
         return m_dimension;
     }
+
+    struct PreparedGeometryInspector::Storage {
+        struct Element {
+            std::shared_ptr<const ScalarReferenceTable> basis;
+            mfem::DenseMatrix inverseJacobians; // rows: points; columns: row-major entries
+        };
+        const mfem::ParFiniteElementSpace *space{nullptr};
+        const mfem::Mesh *mesh{nullptr};
+        long spaceSequence{-1}, meshSequence{-1}, nodesSequence{-1};
+        std::unique_ptr<ReferenceTableCache> tables;
+        std::vector<Element> elements;
+        std::unique_ptr<mfem::ParGridFunction> field;
+        mfem::Array<int> dofs;
+        mfem::Vector elementDofs, workspace;
+        mfem::DenseMatrix dofMatrix, samples, gradient;
+        std::uint64_t builds{0};
+
+        void Clear() {
+            elements.clear();
+            tables.reset();
+            field.reset();
+            space = nullptr;
+            mesh = nullptr;
+        }
+    };
+
+    PreparedGeometryInspector::PreparedGeometryInspector() : m_storage(std::make_unique<Storage>()) {}
+    PreparedGeometryInspector::~PreparedGeometryInspector() = default;
+    PreparedGeometryInspector::PreparedGeometryInspector(PreparedGeometryInspector &&) noexcept = default;
+    PreparedGeometryInspector &PreparedGeometryInspector::operator=(PreparedGeometryInspector &&) noexcept = default;
+    void PreparedGeometryInspector::ClearCache() const { m_storage->Clear(); }
+    int PreparedGeometryInspector::GetTensorElementCount() const noexcept {
+        int count = 0;
+        for (const auto &element : m_storage->elements)
+            count += element.basis && element.basis->HasTensorEvaluation3D();
+        return count;
+    }
+    std::uint64_t PreparedGeometryInspector::GetGeometryBuildCount() const noexcept { return m_storage->builds; }
+
+    double PreparedGeometryInspector::Inspect(
+        mfem::ParFiniteElementSpace &space, const mfem::Vector &displacement, bool reference
+    ) const {
+        MFEM_VERIFY(displacement.Size() == space.GetTrueVSize(), "Geometry inspection received incompatible DOFs.");
+        auto &s = *m_storage;
+        auto *mesh = space.GetMesh();
+        const int dimension = mesh->SpaceDimension();
+        MFEM_VERIFY(space.GetVDim() == dimension, "Geometry inspection requires a spatial displacement vector.");
+        if (s.space != &space || s.mesh != mesh || s.spaceSequence != space.GetSequence() ||
+            s.meshSequence != mesh->GetSequence() || s.nodesSequence != mesh->GetNodesSequence()) {
+            s.Clear();
+            s.space = &space; s.mesh = mesh;
+            s.spaceSequence = space.GetSequence(); s.meshSequence = mesh->GetSequence();
+            s.nodesSequence = mesh->GetNodesSequence();
+            s.field = std::make_unique<mfem::ParGridFunction>(&space);
+            s.tables = std::make_unique<ReferenceTableCache>();
+            s.elements.resize(mesh->GetNE());
+        }
+        s.field->SetFromTrueDofs(displacement);
+        double minimum = std::numeric_limits<double>::infinity();
+        int finite = 1;
+        auto accumulate = [&](double determinant) {
+            if (!std::isfinite(determinant)) finite = 0;
+            else minimum = std::min(minimum, determinant);
+        };
+        for (int e = 0; e < mesh->GetNE(); ++e) {
+            auto &entry = s.elements[e];
+            auto &transformation = *mesh->GetElementTransformation(e);
+            const auto &element = *space.GetFE(e);
+            // Preserve the existing d*p geometry inspection order, including
+            // vacuum elements and points beyond the physics quadrature.
+            const int order = std::max(element.GetOrder() + 2, 2 * dimension * element.GetOrder());
+            const auto &rule = mfem::IntRules.Get(transformation.GetGeometryType(), order);
+            if (!reference && dimension == 3 && element.GetDim() == 3 &&
+                element.GetMapType() == mfem::FiniteElement::VALUE &&
+                dynamic_cast<const mfem::H1_HexahedronElement *>(&element)) {
+                if (!entry.basis) {
+                    entry.basis = s.tables->GetScalarTable(element, rule);
+                    if (entry.basis->HasTensorEvaluation3D()) {
+                        entry.inverseJacobians.SetSize(rule.GetNPoints(), 9);
+                        for (int q = 0; q < rule.GetNPoints(); ++q) {
+                            transformation.SetIntPoint(&rule.IntPoint(q));
+                            const auto &inverse = transformation.InverseJacobian();
+                            for (int a = 0; a < 3; ++a) for (int b = 0; b < 3; ++b)
+                                entry.inverseJacobians(q, a*3+b) = inverse(a,b);
+                        }
+                        ++s.builds;
+                    }
+                }
+            }
+            if (!reference && entry.basis && entry.basis->HasTensorEvaluation3D()) {
+                auto *dofTransformation = space.GetElementVDofs(e, s.dofs);
+                s.field->GetSubVector(s.dofs, s.elementDofs);
+                if (dofTransformation) dofTransformation->InvTransformPrimal(s.elementDofs);
+                const int nd = element.GetDof();
+                s.dofMatrix.SetSize(nd, 3);
+                // GetElementVDofs groups local entries by component for either
+                // global ordering, as used by GridFunction::GetVectorGradientHat.
+                for (int c = 0; c < 3; ++c) for (int i = 0; i < nd; ++i)
+                    s.dofMatrix(i,c) = s.elementDofs(i+c*nd);
+                entry.basis->EvaluateTensor3D(s.dofMatrix, s.samples, s.workspace);
+                s.gradient.SetSize(3);
+                for (int q = 0; q < rule.GetNPoints(); ++q) {
+                    for (int c = 0; c < 3; ++c) for (int b = 0; b < 3; ++b) {
+                        double value = c == b ? 1.0 : 0.0;
+                        for (int a = 0; a < 3; ++a)
+                            value += s.samples(q, 4*c+a+1) * entry.inverseJacobians(q, a*3+b);
+                        s.gradient(c,b) = value;
+                    }
+                    accumulate(s.gradient.Det());
+                }
+            } else {
+                for (int q = 0; q < rule.GetNPoints(); ++q) {
+                    transformation.SetIntPoint(&rule.IntPoint(q));
+                    s.field->GetVectorGradient(transformation, s.gradient);
+                    for (int c = 0; c < dimension; ++c) s.gradient(c,c) += 1.0;
+                    accumulate(s.gradient.Det());
+                }
+            }
+        }
+        double globalMinimum = 0.0;
+        int globalFinite = 0;
+        const int minimumStatus = MPI_Allreduce(&minimum, &globalMinimum, 1, MPI_DOUBLE, MPI_MIN, space.GetComm());
+        const int finiteStatus = MPI_Allreduce(&finite, &globalFinite, 1, MPI_INT, MPI_MIN, space.GetComm());
+        if (minimumStatus != MPI_SUCCESS || finiteStatus != MPI_SUCCESS)
+            throw std::runtime_error("Could not reduce the mapped-geometry inspection.");
+        return globalFinite ? globalMinimum : std::numeric_limits<double>::quiet_NaN();
+    }
+
 } // namespace mean_field::fem

libmeanfield/impl/mapping/compactification/kelvin.cpp

@@ -242,6 +242,23 @@ namespace mean_field::mapping::compactification {
         return MappingStatus::valid;
     }
 
+    MappingStatus KelvinCompactification::EvaluateJacobianVariationFactors(
+        const double coordinate,
+        const mfem::Vector &coordinateGradient,
+        double &scale,
+        mfem::Vector &scaleGradient
+    ) const {
+        if (coordinateGradient.Size() <= 0) return MappingStatus::invalid_dimension;
+        if (!vector_is_finite(coordinateGradient)) return MappingStatus::non_finite_input;
+        RadialFactors factors;
+        const MappingStatus status = ComputeRadialFactors(coordinate, factors);
+        if (status != MappingStatus::valid) return status;
+        scale = factors.scale;
+        scaleGradient = coordinateGradient;
+        scaleGradient *= factors.scale_derivative;
+        return vector_is_finite(scaleGradient) ? MappingStatus::valid : MappingStatus::non_finite_result;
+    }
+
     std::string_view KelvinCompactification::GetName() const noexcept {
         return "KelvinCompactification";
     }
@@ -257,4 +274,4 @@ namespace mean_field::mapping::compactification {
     double KelvinCompactification::GetCoordinateTolerance() const noexcept {
         return m_options.coordinate_tolerance;
     }
-} // namespace mean_field::mapping::compactification
+} // namespace mean_field::mapping::compactification

libmeanfield/impl/mapping/domain_mapper.cpp

@@ -1,6 +1,9 @@
 module;
 
 #include <cmath>
+#include <array>
+#include <map>
+#include <tuple>
 #include <memory>
 #include <mfem.hpp>
 #include <stdexcept>
@@ -32,6 +35,65 @@ namespace {
 } // namespace
 
 namespace mean_field::mapping {
+    struct DomainMapper::GeometryCache {
+        bool enabled{true};
+        const mfem::Mesh *mesh{nullptr};
+        long sequence{-1}, nodesSequence{-1};
+        std::map<std::tuple<int, double, double, double>, FixedMeshGeometry> points;
+        std::uint64_t hits{0}, misses{0};
+    };
+
+    DomainMapper::~DomainMapper() = default;
+    DomainMapper::DomainMapper(DomainMapper &&) noexcept = default;
+    DomainMapper &DomainMapper::operator=(DomainMapper &&) noexcept = default;
+
+    void DomainMapper::SetGeometryCacheEnabled(bool enabled) const { m_geometryCache->enabled = enabled; }
+    bool DomainMapper::IsGeometryCacheEnabled() const noexcept { return m_geometryCache->enabled; }
+    void DomainMapper::ClearGeometryCache() const {
+        auto &cache = *m_geometryCache;
+        cache.points.clear();
+        cache.mesh = nullptr;
+        cache.sequence = cache.nodesSequence = -1;
+        cache.hits = cache.misses = 0;
+    }
+    DomainMapper::GeometryCacheStatistics DomainMapper::GetGeometryCacheStatistics() const {
+        const auto &cache = *m_geometryCache;
+        return {.hits = cache.hits, .misses = cache.misses, .points = cache.points.size()};
+    }
+
+    const DomainMapper::FixedMeshGeometry *DomainMapper::GetFixedMeshGeometry(
+        mfem::ElementTransformation &transformation, const mfem::IntegrationPoint &point
+    ) const {
+        transformation.SetIntPoint(&point);
+        auto &cache = *m_geometryCache;
+        const auto *mesh = transformation.mesh;
+        if (!cache.enabled || mesh == nullptr || transformation.ElementType != mfem::ElementTransformation::ELEMENT ||
+            transformation.ElementNo < 0 || transformation.ElementNo >= mesh->GetNE() ||
+            transformation.GetSpaceDim() != m_options.dimension) return nullptr;
+        const std::array<double, 3> coordinates{
+            point.x, m_options.dimension > 1 ? point.y : 0.0, m_options.dimension > 2 ? point.z : 0.0};
+        for (double coordinate : coordinates) if (!std::isfinite(coordinate)) return nullptr;
+        if (cache.mesh != mesh || cache.sequence != mesh->GetSequence() || cache.nodesSequence != mesh->GetNodesSequence()) {
+            cache.points.clear();
+            cache.mesh = mesh;
+            cache.sequence = mesh->GetSequence();
+            cache.nodesSequence = mesh->GetNodesSequence();
+        }
+        const auto key = std::tuple{transformation.ElementNo, coordinates[0], coordinates[1], coordinates[2]};
+        const auto found = cache.points.find(key);
+        if (found != cache.points.end()) { ++cache.hits; return &found->second; }
+        FixedMeshGeometry geometry;
+        transformation.Transform(point, geometry.position);
+        geometry.jacobian = transformation.Jacobian();
+        geometry.inverseJacobian = transformation.InverseJacobian();
+        geometry.weight = transformation.Weight();
+        ++cache.misses;
+        if (!vector_is_finite(geometry.position) || !matrix_is_finite(geometry.jacobian) ||
+            !matrix_is_finite(geometry.inverseJacobian) || !std::isfinite(geometry.weight) || geometry.weight <= 0.0)
+            return nullptr;
+        return &cache.points.emplace(key, std::move(geometry)).first->second;
+    }
+
     ElementCompactificationData::ElementCompactificationData(
         const mfem::FiniteElement &element,
         const mfem::Vector &dofs
@@ -195,7 +257,8 @@ namespace mean_field::mapping {
         const utils::DomainMapperOptions options,
         std::unique_ptr<const compactification::ExteriorDomainMap> exterior_map
     )
-        : m_options(options),
+        : m_geometryCache(std::make_unique<GeometryCache>()),
+          m_options(options),
           m_exterior_map(std::move(exterior_map)) {
         if (m_options.dimension <= 0)
             throw std::invalid_argument("The domain-mapping dimension must be positive.");
@@ -588,6 +651,16 @@ namespace mean_field::mapping {
         const mfem::IntegrationPoint &integration_point,
         Workspace &workspace,
         MappingPointContext &context
+    ) const {
+        const auto *meshGeometry = GetFixedMeshGeometry(transformation, integration_point);
+        return EvaluatePointImpl(element_data, transformation, integration_point, workspace, context, meshGeometry);
+    }
+
+    MappingStatus DomainMapper::EvaluatePointImpl(
+        const ElementMappingData &element_data, mfem::ElementTransformation &transformation,
+        const mfem::IntegrationPoint &integration_point, Workspace &workspace,
+        MappingPointContext &context, const FixedMeshGeometry *meshGeometry,
+        const ReferenceDisplacementSample *displacementSample
     ) const {
         ValidateElementData(element_data);
 
@@ -602,16 +675,28 @@ namespace mean_field::mapping {
                 "element data."
             );
 
-        transformation.SetIntPoint(&integration_point);
-
         context.reference_position.SetSize(m_options.dimension);
-        transformation.Transform(integration_point, context.reference_position);
+        if (meshGeometry) context.reference_position = meshGeometry->position;
+        else transformation.Transform(integration_point, context.reference_position);
 
         // Get the displacement field value and its Jacobian at the integration point. Note these are in the workspace to avoid repeated allocations.
-        EvaluateField(
-            element_data.displacement, transformation, integration_point, workspace, workspace.m_field_value,
-            workspace.m_field_jacobian, nullptr
-        );
+        if (displacementSample) {
+            if (element_data.displacement.GetElement().GetMapType() != mfem::FiniteElement::VALUE ||
+                displacementSample->value.Size() != m_options.dimension ||
+                displacementSample->gradient.Height() != m_options.dimension ||
+                displacementSample->gradient.Width() != m_options.dimension)
+                return MappingStatus::invalid_dimension;
+            workspace.m_field_value = displacementSample->value;
+            workspace.m_field_jacobian.SetSize(m_options.dimension);
+            mfem::Mult(displacementSample->gradient,
+                       meshGeometry ? meshGeometry->inverseJacobian : transformation.InverseJacobian(),
+                       workspace.m_field_jacobian);
+        } else {
+            EvaluateField(
+                element_data.displacement, transformation, integration_point, workspace, workspace.m_field_value,
+                workspace.m_field_jacobian, meshGeometry ? &meshGeometry->inverseJacobian : nullptr
+            );
+        }
 
         if (!vector_is_finite(context.reference_position) || !vector_is_finite(workspace.m_field_value) ||
             !matrix_is_finite(workspace.m_field_jacobian)) {
@@ -647,7 +732,7 @@ namespace mean_field::mapping {
             // the gradient of s as we move along the mesh coordinates. All of this is stashes within workspace.m_compactification_point.
             const MappingStatus coordinate_status = EvaluateCompactificationCoordinate(
                 element_data.compactification, transformation, integration_point, workspace,
-                workspace.m_compactification_point, nullptr
+                workspace.m_compactification_point, meshGeometry ? &meshGeometry->inverseJacobian : nullptr
             );
 
             if (coordinate_status != MappingStatus::valid)
@@ -701,21 +786,25 @@ namespace mean_field::mapping {
         mfem::ElementTransformation &transformation,
         const mfem::IntegrationPoint &integration_point,
         Workspace &workspace,
-        VolumeMappingContext &context
+        VolumeMappingContext &context,
+        const ReferenceDisplacementSample *displacementSample
     ) const {
+        const auto *meshGeometry = GetFixedMeshGeometry(transformation, integration_point);
         const MappingStatus point_status =
-            EvaluatePoint(element_data, transformation, integration_point, workspace, context.mapping);
+            EvaluatePointImpl(element_data, transformation, integration_point, workspace, context.mapping, meshGeometry, displacementSample);
         if (point_status != MappingStatus::valid)
             return point_status;
 
-        mfem::Mult(context.mapping.mapping_jacobian, transformation.Jacobian(), workspace.m_full_element_jacobian);
+        mfem::Mult(context.mapping.mapping_jacobian,
+                   meshGeometry ? meshGeometry->jacobian : transformation.Jacobian(), workspace.m_full_element_jacobian);
 
         context.quadrature.J_inv.SetSize(m_options.dimension, m_options.dimension);
         mfem::CalcInverse(workspace.m_full_element_jacobian, context.quadrature.J_inv);
 
         context.quadrature.detJ = context.mapping.mapping_determinant;
         context.quadrature.weight =
-            integration_point.weight * transformation.Weight() * context.mapping.mapping_determinant;
+            integration_point.weight * (meshGeometry ? meshGeometry->weight : transformation.Weight()) *
+            context.mapping.mapping_determinant;
 
         if (!matrix_is_finite(context.quadrature.J_inv) || !std::isfinite(context.quadrature.weight))
             return MappingStatus::non_finite_result;
@@ -846,7 +935,8 @@ namespace mean_field::mapping {
             );
 
         if (inverse_mesh_jacobian == nullptr) {
-            transformation.SetIntPoint(&integration_point);
+            if (const auto *geometry = GetFixedMeshGeometry(transformation, integration_point))
+                inverse_mesh_jacobian = &geometry->inverseJacobian;
         }
 
         EvaluateField(

libmeanfield/impl/operators/interior_displacement_contractions.hpp

@@ -0,0 +1,162 @@
+#pragma once
+
+#include <mfem.hpp>
+
+namespace {
+    // U and G are n x 3, component/derivative columns respectively. Keep the
+    // nine reductions together instead of invoking a general GEMM per point.
+    // The returned 3 x 3 matrix is column-major: A(c,j) = sum_i U(i,c) G(i,j).
+    inline void displacement_gradient_3d(
+        const mfem::DenseMatrix &direction, const mfem::DenseMatrix &gradients, double *a
+    ) {
+        const int n = direction.Height();
+        const double *u = direction.Data();
+        const double *g = gradients.Data();
+        for (int j = 0; j < 3; ++j) {
+            double x = 0.0, y = 0.0, z = 0.0;
+            const double *gj = g + j * n;
+            for (int i = 0; i < n; ++i) {
+                const double v = gj[i];
+                x += u[i] * v;
+                y += u[i + n] * v;
+                z += u[i + 2 * n] * v;
+            }
+            a[3 * j] = x; a[3 * j + 1] = y; a[3 * j + 2] = z;
+        }
+    }
+
+    inline void vector_basis_value_3d(
+        const mfem::DenseMatrix &basis, const mfem::Vector &coefficients, double *value
+    ) {
+        const int n = basis.Height();
+        const double *b = basis.Data();
+        const double *v = coefficients.HostRead();
+        double x = 0.0, y = 0.0, z = 0.0;
+        for (int i = 0; i < n; ++i) {
+            x += b[i] * v[i]; y += b[i + n] * v[i]; z += b[i + 2 * n] * v[i];
+        }
+        value[0] = x; value[1] = y; value[2] = z;
+    }
+
+    // Matrix cache layout: nq x 9, with row-major matrix entries in columns.
+    inline void cached_matrix_vector_3d(
+        const mfem::DenseMatrix &cache, const int point, const double *v, double *out
+    ) {
+        const int nq = cache.Height();
+        const double *m = cache.Data() + point;
+        for (int r = 0; r < 3; ++r)
+            out[r] = m[(3 * r) * nq] * v[0] + m[(3 * r + 1) * nq] * v[1] +
+                     m[(3 * r + 2) * nq] * v[2];
+    }
+
+    inline double displacement_trace_3d(
+        const mfem::DenseMatrix &direction, const mfem::DenseMatrix &gradients,
+        const mfem::DenseMatrix &inverseJacobians, const int point
+    ) {
+        double a[9];
+        displacement_gradient_3d(direction, gradients, a);
+        const double *k = inverseJacobians.Data() + point;
+        const int nq = inverseJacobians.Height();
+        double trace = 0.0;
+        for (int j = 0; j < 3; ++j)
+            for (int c = 0; c < 3; ++c) trace += k[(3 * j + c) * nq] * a[c + 3 * j];
+        return trace;
+    }
+
+    // (du/dxi) v, without constructing du/dxi. v is fixed at this point.
+    inline void displacement_directional_value_3d(
+        const mfem::DenseMatrix &direction, const mfem::DenseMatrix &gradients,
+        const double *v, double *out
+    ) {
+        const int n = direction.Height();
+        const double *u = direction.Data();
+        const double *g = gradients.Data();
+        double x = 0.0, y = 0.0, z = 0.0;
+        for (int i = 0; i < n; ++i) {
+            const double d = g[i] * v[0] + g[i + n] * v[1] + g[i + 2 * n] * v[2];
+            x += u[i] * d; y += u[i + n] * d; z += u[i + 2 * n] * d;
+        }
+        out[0] = x; out[1] = y; out[2] = z;
+    }
+
+    inline void pressure_displacement_action_3d(
+        const mfem::DenseMatrix &direction, const mfem::DenseMatrix &physicalGradients,
+        const double pressureWeight, mfem::Vector &action
+    ) {
+        // H = du/dx; d(G dx) = (tr(H) G - G H) dx.
+        double h[9];
+        displacement_gradient_3d(direction, physicalGradients, h);
+        const double trace = h[0] + h[4] + h[8];
+        const int n = direction.Height();
+        const double *g = physicalGradients.Data();
+        double *r = action.HostReadWrite();
+        for (int c = 0; c < 3; ++c) {
+            const double a = pressureWeight * (h[3 * c] - (c == 0 ? trace : 0.0));
+            const double b = pressureWeight * (h[3 * c + 1] - (c == 1 ? trace : 0.0));
+            const double d = pressureWeight * (h[3 * c + 2] - (c == 2 ? trace : 0.0));
+            for (int i = 0; i < n; ++i) r[i + c * n] += a * g[i] + b * g[i + n] + d * g[i + 2 * n];
+        }
+    }
+
+    // Element VDofs are component-blocked, including for a byVDIM global space;
+    // this is the ordering consumed by ElementDisplacementDataFromElementVDofs.
+    // For an ordinary interior VALUE field, dx = du and
+    // d(log det J_element) = tr(J_element^-1 * du/dxi).
+    inline void add_interior_functional_row(
+        const mfem::DenseMatrix &gradients,
+        const mfem::DenseMatrix &values,
+        const int point,
+        const mfem::DenseMatrix &inverseElementJacobian,
+        const double determinantCoefficient,
+        const double *positionCoefficient,
+        mfem::Vector &row
+    ) {
+        const int n = gradients.Height();
+        const int dimension = gradients.Width();
+        const int nq = values.Height();
+        const double *g = gradients.Data();
+        const double *b = values.Data() + point;
+        double *r = row.HostReadWrite();
+        for (int component = 0; component < dimension; ++component) {
+            for (int i = 0; i < n; ++i) {
+                double coefficient = positionCoefficient[component] * b[i * nq];
+                for (int derivative = 0; derivative < dimension; ++derivative) {
+                    coefficient += determinantCoefficient * inverseElementJacobian(derivative, component) *
+                                   g[i + derivative * n];
+                }
+                r[i + component * n] += coefficient;
+            }
+        }
+    }
+
+    inline double apply_interior_functional(
+        const mfem::DenseMatrix &direction,
+        const mfem::DenseMatrix &gradients,
+        const mfem::DenseMatrix &values,
+        const int point,
+        const double *inverseElementJacobian,
+        const double determinantCoefficient,
+        const double *positionCoefficient
+    ) {
+        const int n = direction.Height();
+        const int dimension = direction.Width();
+        const int nq = values.Height();
+        const double *u = direction.Data();
+        const double *g = gradients.Data();
+        const double *b = values.Data() + point;
+        double result = 0.0;
+        for (int component = 0; component < dimension; ++component) {
+            double value = 0.0;
+            for (int i = 0; i < n; ++i) value += u[i + component * n] * b[i * nq];
+            result += positionCoefficient[component] * value;
+            for (int derivative = 0; derivative < dimension; ++derivative) {
+                double gradient = 0.0;
+                for (int i = 0; i < n; ++i)
+                    gradient += u[i + component * n] * g[i + derivative * n];
+                result += determinantCoefficient * inverseElementJacobian[derivative + component * dimension] *
+                          gradient;
+            }
+        }
+        return result;
+    }
+}

libmeanfield/impl/operators/prepared_angular_momentum.cpp

@@ -11,6 +11,7 @@ module;
 
 #include <mfem.hpp>
 #include <mpi.h>
+#include "interior_displacement_contractions.hpp"
 
 module mean_field;
 
@@ -228,6 +229,7 @@ namespace mean_field::operators {
         const bool updateAngularVelocity = !m_isPrepared || dependencies.rotation != m_preparedDependencies.rotation ||
                                            angularVelocity != m_angularVelocity;
 
+        if (refreshGeometry || refreshDensity) m_jacobianRowsReady = false;
         m_isPrepared                     = false;
         if (synchronize_non_finite_failure(
                 !std::isfinite(angularVelocity), m_fem.mesh->GetComm(), "angular-velocity validity"
@@ -466,8 +468,99 @@ namespace mean_field::operators {
         ++m_residualApplicationCount;
     }
 
+    bool PreparedAngularMomentumOperator::PrepareJacobianRows() const {
+        VerifyPrepared();
+        if (m_jacobianRowsReady) return m_jacobianRowsSupported;
+        m_jacobianRowsSupported = false;
+        for (const ElementPAData &data : m_elements) {
+            const auto &element = *m_fem.displacementFes->GetFE(data.elementId);
+            if (element.GetMapType() != mfem::FiniteElement::VALUE ||
+                element.GetDerivType() != mfem::FiniteElement::GRAD ||
+                m_domainMapper.IsCompactifiedElement(*m_fem.mesh->GetElementTransformation(data.elementId))) {
+                m_jacobianRowsReady = true;
+                return false;
+            }
+        }
+        m_densityJacobianLocal.SetSize(m_fem.densityFes->GetVSize());
+        m_displacementJacobianLocal.SetSize(m_fem.displacementFes->GetVSize());
+        m_densityJacobianLocal = 0.0;
+        m_displacementJacobianLocal = 0.0;
+        mfem::Vector densityRow, displacementRow, weightedRadius;
+        mapping::VolumeMappingContext mappingContext;
+        const auto &axis = m_constraint.specification().axis();
+        const auto &center = m_constraint.specification().center();
+        for (const ElementPAData &data : m_elements) {
+            const int nq = data.integrationRule->GetNPoints();
+            weightedRadius.SetSize(nq);
+            for (int q = 0; q < nq; ++q)
+                weightedRadius(q) = data.quadratureWeights(q) * data.cylindricalRadiusSquared(q);
+            densityRow.SetSize(data.densityDofs.Size());
+            data.densityBasis->GetValues().MultTranspose(weightedRadius, densityRow);
+            displacementRow.SetSize(data.displacementDofs.Size());
+            displacementRow = 0.0;
+            const auto basis = m_fem.GetReferenceTables().GetScalarTable(
+                *m_fem.displacementFes->GetFE(data.elementId), *data.integrationRule);
+            for (int q = 0; q < nq; ++q) {
+                data.mappingContexts.Load(q, mappingContext);
+                double relative[3], positionCoefficient[3], axialPosition = 0.0;
+                for (int c = 0; c < 3; ++c) {
+                    relative[c] = mappingContext.mapping.physical_position(c) - center[c];
+                    axialPosition += axis[c] * relative[c];
+                }
+                for (int c = 0; c < 3; ++c)
+                    positionCoefficient[c] = 2.0 * data.density(q) * data.quadratureWeights(q) *
+                                             (relative[c] - axis[c] * axialPosition);
+                add_interior_functional_row(
+                    basis->GetGradients(q), basis->GetValues(), q, mappingContext.quadrature.J_inv,
+                    data.density(q) * weightedRadius(q), positionCoefficient, displacementRow);
+            }
+            if (data.densityDofTransformation) data.densityDofTransformation->TransformDual(densityRow);
+            if (data.displacementDofTransformation) data.displacementDofTransformation->TransformDual(displacementRow);
+            m_densityJacobianLocal.AddElementVector(data.densityDofs, densityRow);
+            m_displacementJacobianLocal.AddElementVector(data.displacementDofs, displacementRow);
+        }
+        verify_finite_vector(m_densityJacobianLocal, "Prepared angular-momentum density row is non-finite.");
+        verify_finite_vector(m_displacementJacobianLocal, "Prepared angular-momentum displacement row is non-finite.");
+        m_jacobianRowsSupported = true;
+        m_jacobianRowsReady = true;
+        ++m_jacobianRowPreparationCount;
+        return true;
+    }
+
+    std::uint64_t PreparedAngularMomentumOperator::GetJacobianRowPreparationCount() const noexcept {
+        return m_jacobianRowPreparationCount;
+    }
+
+    double PreparedAngularMomentumOperator::EvaluateDensityMomentActionLocal(const mfem::Vector &direction) const {
+        if (!PrepareJacobianRows()) return EvaluateDensityMomentActionLocalReference(direction);
+        true_to_local(*m_fem.densityFes, direction, m_densityVariationLocal);
+        return m_densityJacobianLocal * m_densityVariationLocal;
+    }
+
+    double PreparedAngularMomentumOperator::EvaluateDisplacementMomentActionLocal(const mfem::Vector &direction) const {
+        if (!PrepareJacobianRows()) return EvaluateDisplacementMomentActionLocalReference(direction);
+        true_to_local(*m_fem.displacementFes, direction, m_displacementVariationLocal);
+        return m_displacementJacobianLocal * m_displacementVariationLocal;
+    }
+
+    void PreparedAngularMomentumOperator::ApplyCompleteJacobianActionReference(
+        const mfem::Vector &densityVariation, const mfem::Vector &displacementVariation,
+        const double angularVelocityVariation, mfem::Vector &action
+    ) const {
+        VerifyPrepared();
+        verify_finite_vector(densityVariation, "Angular-momentum density direction is non-finite.");
+        verify_finite_vector(displacementVariation, "Angular-momentum displacement direction is non-finite.");
+        MFEM_VERIFY(std::isfinite(angularVelocityVariation), "Angular-velocity direction is non-finite.");
+        m_gravityContext.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);
+        m_gravityContext.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);
+        action.SetSize(1);
+        action(0) = m_angularVelocity * GlobalSum(EvaluateDensityMomentActionLocalReference(m_densityVariationTrue) +
+                       EvaluateDisplacementMomentActionLocalReference(m_displacementVariationTrue)) +
+                    m_momentOfInertia * angularVelocityVariation;
+    }
+
     double
-    PreparedAngularMomentumOperator::EvaluateDensityMomentActionLocal(const mfem::Vector &densityVariation) const {
+    PreparedAngularMomentumOperator::EvaluateDensityMomentActionLocalReference(const mfem::Vector &densityVariation) const {
         MFEM_VERIFY(
             densityVariation.Size() == m_fem.densityFes->GetTrueVSize(),
             "Angular-momentum density action has the wrong true-vector size."
@@ -490,7 +583,7 @@ namespace mean_field::operators {
         return localAction;
     }
 
-    double PreparedAngularMomentumOperator::EvaluateDisplacementMomentActionLocal(
+    double PreparedAngularMomentumOperator::EvaluateDisplacementMomentActionLocalReference(
         const mfem::Vector &displacementVariation
     ) const {
         MFEM_VERIFY(

libmeanfield/impl/operators/prepared_barotropic_closure.cpp

@@ -1,5 +1,7 @@
 module;
 
+#include "interior_displacement_contractions.hpp"
+
 #include <array>
 #include <cmath>
 #include <cstdint>
@@ -730,10 +732,66 @@ namespace mean_field::operators {
     }
 
     void PreparedBarotropicClosureOperator::Mult(
+        const mfem::Vector &densityVariation, const mfem::Vector &enthalpyVariation,
+        const mfem::Vector &displacementVariation, mfem::Vector &action
+    ) const {
+        MultImpl(densityVariation, enthalpyVariation, displacementVariation, action, true, true);
+    }
+
+    void PreparedBarotropicClosureOperator::MultReference(
+        const mfem::Vector &densityVariation, const mfem::Vector &enthalpyVariation,
+        const mfem::Vector &displacementVariation, mfem::Vector &action
+    ) const {
+        MultImpl(densityVariation, enthalpyVariation, displacementVariation, action, false, false);
+    }
+
+    void PreparedBarotropicClosureOperator::MultCachedReference(
+        const mfem::Vector &densityVariation, const mfem::Vector &enthalpyVariation,
+        const mfem::Vector &displacementVariation, mfem::Vector &action
+    ) const {
+        MultImpl(densityVariation, enthalpyVariation, displacementVariation, action, true, false);
+    }
+
+    int PreparedBarotropicClosureOperator::GetTensorElementCount() const noexcept {
+        int count = 0;
+        for (const auto &data : m_elements)
+            count += data.densityBasis->HasTensorEvaluation3D() &&
+                data.enthalpyBasis->HasTensorEvaluation3D() && data.displacementBasis->HasTensorEvaluation3D();
+        return count;
+    }
+
+    void PreparedBarotropicClosureOperator::ApplyEnthalpyJacobianAction(
+        const mfem::Vector &enthalpyVariation, mfem::Vector &action, const bool useTensorPath
+    ) const {
+        VerifyPrepared();
+        MFEM_VERIFY(enthalpyVariation.Size() == m_enthalpyMap.reduced_size(),
+                    "The supported enthalpy variation has the wrong size.");
+        validate_finite_vector(enthalpyVariation, "The enthalpy variation contains a non-finite value.");
+        m_enthalpyMap.scatter(enthalpyVariation, m_enthalpyVariationTrue);
+        ApplyThermodynamicActionFull(nullptr, m_enthalpyVariationTrue, m_fullThermodynamicAction, useTensorPath);
+        action.SetSize(m_densityMap.reduced_size());
+        m_densityMap.gather(m_fullThermodynamicAction, action);
+    }
+
+    void PreparedBarotropicClosureOperator::ApplyDisplacementJacobianAction(
+        const mfem::Vector &displacementVariation, mfem::Vector &action, const bool useTensorPath
+    ) const {
+        VerifyPrepared();
+        MFEM_VERIFY(displacementVariation.Size() == m_displacementMap.reduced_size(),
+                    "The supported displacement variation has the wrong size.");
+        validate_finite_vector(displacementVariation, "The displacement variation contains a non-finite value.");
+        m_displacementMap.scatter(displacementVariation, m_displacementVariationTrue);
+        ApplyDisplacementActionFull(m_displacementVariationTrue, m_fullDisplacementAction, true, useTensorPath);
+        action.SetSize(m_densityMap.reduced_size());
+        m_densityMap.gather(m_fullDisplacementAction, action);
+    }
+
+    void PreparedBarotropicClosureOperator::MultImpl(
         const mfem::Vector &densityVariation,
         const mfem::Vector &enthalpyVariation,
         const mfem::Vector &displacementVariation,
-        mfem::Vector &action
+        mfem::Vector &action,
+        const bool useFastPath, const bool useTensorPath
     ) const {
         VerifyPrepared();
 
@@ -758,9 +816,9 @@ namespace mean_field::operators {
         m_enthalpyMap.scatter(enthalpyVariation, m_enthalpyVariationTrue);
         m_displacementMap.scatter(displacementVariation, m_displacementVariationTrue);
 
-        ApplyThermodynamicActionFull(m_densityVariationTrue, m_enthalpyVariationTrue, m_fullThermodynamicAction);
+        ApplyThermodynamicActionFull(&m_densityVariationTrue, m_enthalpyVariationTrue, m_fullThermodynamicAction, useTensorPath);
 
-        ApplyDisplacementActionFull(m_displacementVariationTrue, m_fullDisplacementAction);
+        ApplyDisplacementActionFull(m_displacementVariationTrue, m_fullDisplacementAction, useFastPath, useTensorPath);
 
         MFEM_VERIFY(
             m_fullThermodynamicAction.Size() == m_densityMap.full_size() &&
@@ -797,12 +855,13 @@ namespace mean_field::operators {
     }
 
     void PreparedBarotropicClosureOperator::ApplyThermodynamicActionFull(
-        const mfem::Vector &densityVariationTrue,
+        const mfem::Vector *densityVariationTrue,
         const mfem::Vector &enthalpyVariationTrue,
-        mfem::Vector &actionTrue
+        mfem::Vector &actionTrue, const bool useTensorPath
     ) const {
         MFEM_VERIFY(
-            densityVariationTrue.Size() == m_densityMap.full_size(), "The full density variation has the wrong size."
+            densityVariationTrue == nullptr || densityVariationTrue->Size() == m_densityMap.full_size(),
+            "The full density variation has the wrong size."
         );
         MFEM_VERIFY(
             enthalpyVariationTrue.Size() == m_enthalpyMap.full_size(), "The full enthalpy variation has the wrong size."
@@ -811,7 +870,7 @@ namespace mean_field::operators {
         mfem::Vector densityVariationLocal;
         mfem::Vector enthalpyVariationLocal;
 
-        true_to_local(*m_fem.densityFes, densityVariationTrue, densityVariationLocal);
+        if (densityVariationTrue) true_to_local(*m_fem.densityFes, *densityVariationTrue, densityVariationLocal);
         true_to_local(*m_fem.enthalpyFes, enthalpyVariationTrue, enthalpyVariationLocal);
 
         mfem::Vector localAction(m_fem.densityFes->GetVSize());
@@ -825,10 +884,10 @@ namespace mean_field::operators {
         mfem::Vector elementAction;
 
         for (const ElementPAData &data : m_elements) {
-            densityVariationLocal.GetSubVector(data.densityDofs, elementDensityVariation);
+            if (densityVariationTrue) densityVariationLocal.GetSubVector(data.densityDofs, elementDensityVariation);
             enthalpyVariationLocal.GetSubVector(data.enthalpyDofs, elementEnthalpyVariation);
 
-            if (data.densityDofTransformation != nullptr) {
+            if (densityVariationTrue && data.densityDofTransformation != nullptr) {
                 data.densityDofTransformation->InvTransformPrimal(elementDensityVariation);
             }
             if (data.enthalpyDofTransformation != nullptr) {
@@ -839,17 +898,25 @@ namespace mean_field::operators {
             quadratureEnthalpyVariation.SetSize(data.quadratureWeights.Size());
             quadratureAction.SetSize(data.quadratureWeights.Size());
 
-            data.densityBasis->GetValues().Mult(elementDensityVariation, quadratureDensityVariation);
-            data.enthalpyBasis->GetValues().Mult(elementEnthalpyVariation, quadratureEnthalpyVariation);
+            if (densityVariationTrue) {
+                if (useTensorPath && data.densityBasis->HasTensorEvaluation3D())
+                    data.densityBasis->EvaluateValuesTensor3D(elementDensityVariation, quadratureDensityVariation, m_tensorWorkspace);
+                else data.densityBasis->GetValues().Mult(elementDensityVariation, quadratureDensityVariation);
+            }
+            if (useTensorPath && data.enthalpyBasis->HasTensorEvaluation3D())
+                data.enthalpyBasis->EvaluateValuesTensor3D(elementEnthalpyVariation, quadratureEnthalpyVariation, m_tensorWorkspace);
+            else data.enthalpyBasis->GetValues().Mult(elementEnthalpyVariation, quadratureEnthalpyVariation);
 
             for (int quadraturePoint = 0; quadraturePoint < quadratureAction.Size(); ++quadraturePoint) {
                 quadratureAction(quadraturePoint) =
-                    data.quadratureWeights(quadraturePoint) * quadratureDensityVariation(quadraturePoint) -
+                    (densityVariationTrue ? data.quadratureWeights(quadraturePoint) * quadratureDensityVariation(quadraturePoint) : 0.0) -
                     data.weightedEnthalpyDerivative(quadraturePoint) * quadratureEnthalpyVariation(quadraturePoint);
             }
 
             elementAction.SetSize(data.densityDofs.Size());
-            data.densityBasis->GetValues().MultTranspose(quadratureAction, elementAction);
+            if (useTensorPath && data.densityBasis->HasTensorEvaluation3D())
+                data.densityBasis->ApplyTransposeTensor3D(quadratureAction, elementAction, m_tensorWorkspace);
+            else data.densityBasis->GetValues().MultTranspose(quadratureAction, elementAction);
 
             if (data.densityDofTransformation != nullptr) {
                 data.densityDofTransformation->TransformDual(elementAction);
@@ -863,7 +930,8 @@ namespace mean_field::operators {
 
     void PreparedBarotropicClosureOperator::ApplyDisplacementActionFull(
         const mfem::Vector &displacementVariationTrue,
-        mfem::Vector &actionTrue
+        mfem::Vector &actionTrue,
+        const bool useFastPath, const bool useTensorPath
     ) const {
         MFEM_VERIFY(
             displacementVariationTrue.Size() == m_displacementMap.full_size(),
@@ -905,8 +973,29 @@ namespace mean_field::operators {
 
             m_referenceDisplacementJacobian.SetSize(dimension, dimension);
             m_quadratureDisplacementAction.SetSize(integrationRule.GetNPoints());
+            const bool tensor = useTensorPath && data.displacementBasis->HasTensorEvaluation3D();
+            if (tensor) data.displacementBasis->EvaluateTensor3D(directionDofs, m_tensorDisplacement, m_tensorWorkspace);
 
             for (int quadraturePoint = 0; quadraturePoint < integrationRule.GetNPoints(); ++quadraturePoint) {
+                if (tensor) {
+                    double trace = 0.0;
+                    for (int row = 0; row < 3; ++row) for (int column = 0; column < 3; ++column)
+                        trace += data.inverseElementJacobians(quadraturePoint, row * 3 + column) *
+                            m_tensorDisplacement(quadraturePoint, 4 * column + row + 1);
+                    m_quadratureDisplacementAction(quadraturePoint) = data.weightedResidual(quadraturePoint) * trace;
+                    MFEM_VERIFY(std::isfinite(m_quadratureDisplacementAction(quadraturePoint)),
+                                "Prepared barotropic closure displacement action encountered a non-finite quadrature value.");
+                    continue;
+                }
+                if (useFastPath && dimension == 3) {
+                    const double trace = displacement_trace_3d(
+                        directionDofs, data.displacementBasis->GetGradients(quadraturePoint),
+                        data.inverseElementJacobians, quadraturePoint);
+                    m_quadratureDisplacementAction(quadraturePoint) = data.weightedResidual(quadraturePoint) * trace;
+                    MFEM_VERIFY(std::isfinite(m_quadratureDisplacementAction(quadraturePoint)),
+                                "Prepared barotropic closure displacement action encountered a non-finite quadrature value.");
+                    continue;
+                }
                 mfem::MultAtB(
                     directionDofs, data.displacementBasis->GetGradients(quadraturePoint),
                     m_referenceDisplacementJacobian
@@ -930,7 +1019,9 @@ namespace mean_field::operators {
             }
 
             m_elementDisplacementAction.SetSize(data.densityDofs.Size());
-            data.densityBasis->GetValues().MultTranspose(m_quadratureDisplacementAction, m_elementDisplacementAction);
+            if (useTensorPath && data.densityBasis->HasTensorEvaluation3D())
+                data.densityBasis->ApplyTransposeTensor3D(m_quadratureDisplacementAction, m_elementDisplacementAction, m_tensorWorkspace);
+            else data.densityBasis->GetValues().MultTranspose(m_quadratureDisplacementAction, m_elementDisplacementAction);
 
             if (data.densityDofTransformation != nullptr) {
                 data.densityDofTransformation->TransformDual(m_elementDisplacementAction);

libmeanfield/impl/operators/prepared_gravity_displacement_force.cpp

@@ -1,5 +1,7 @@
 module;
 
+#include "interior_displacement_contractions.hpp"
+
 #include <cmath>
 #include <expected>
 #include <optional>
@@ -184,7 +186,7 @@ namespace mean_field::operators {
     std::expected<
         void,
         kernels::GravityDisplacementForceRejection>
-    PreparedGravityDisplacementForceOperator::TryPrepareElementData() {
+    PreparedGravityDisplacementForceOperator::TryPrepareElementData(const bool assembleResidual) {
         m_elements.clear();
         m_elements.reserve(m_fem.mesh->GetNE());
 
@@ -206,8 +208,14 @@ namespace mean_field::operators {
         mfem::Vector elementCompactification;
         mfem::Vector densityShape;
         mfem::Vector baseGravityReferenceValue;
+        mfem::Vector baseGravityCellValue;
         mfem::DenseMatrix gravityGradientShape;
 
+        if (assembleResidual) {
+            m_localAction.SetSize(m_fem.displacementFes->GetVSize());
+            m_localAction = 0.0;
+        }
+        mfem::Vector residualComponent;
         const int dimension = m_domainMapper.GetDimension();
 
         for (int elementId = 0; elementId < m_fem.mesh->GetNE(); ++elementId) {
@@ -262,6 +270,11 @@ namespace mean_field::operators {
                 data.meshPiolaJacobians.SetSize(data.integrationRule->GetNPoints(), dimension * dimension);
             }
 
+            data.fastApplication = dimension == 3 && m_fem.displacementFes->GetOrdering() == mfem::Ordering::byNODES &&
+                displacementElement.GetMapType() == mfem::FiniteElement::VALUE &&
+                displacementElement.GetDerivType() == mfem::FiniteElement::GRAD &&
+                data.gravityReferenceTable != nullptr;
+
             const mapping::ElementDisplacementData displacementData =
                 mapping::ElementDisplacementDataFromElementVDofs(displacementElement, elementBaseDisplacement);
             const mapping::ElementCompactificationData compactificationData(
@@ -280,6 +293,7 @@ namespace mean_field::operators {
             densityShape.SetSize(densityElement.GetDof());
             gravityGradientShape.SetSize(gravityGradientElement.GetDof(), dimension);
             baseGravityReferenceValue.SetSize(dimension);
+            if (assembleResidual) m_quadratureForce.SetSize(quadraturePointCount, dimension);
 
             for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
                 const mfem::IntegrationPoint &integrationPoint = data.integrationRule->IntPoint(quadraturePoint);
@@ -297,14 +311,26 @@ namespace mean_field::operators {
                 for (int dof = 0; dof < densityElement.GetDof(); ++dof) {
                     densityShape(dof) = data.densityReferenceTable->GetValues()(quadraturePoint, dof);
                 }
-                gravityGradientElement.CalcVShape(*transformation, gravityGradientShape);
-                gravityGradientShape.MultTranspose(elementBaseGravityGradient, baseGravityReferenceValue);
+                const auto *meshGeometry = m_domainMapper.GetFixedMeshGeometry(*transformation, integrationPoint);
+                const double meshWeight = meshGeometry ? meshGeometry->weight : transformation->Weight();
+                if (assembleResidual && meshGeometry && data.gravityReferenceTable) {
+                    baseGravityCellValue.SetSize(dimension);
+                    data.gravityReferenceTable->GetValues(quadraturePoint).MultTranspose(
+                        elementBaseGravityGradient, baseGravityCellValue);
+                    meshGeometry->jacobian.Mult(baseGravityCellValue, baseGravityReferenceValue);
+                    baseGravityReferenceValue /= meshWeight;
+                } else {
+                    gravityGradientElement.CalcVShape(*transformation, gravityGradientShape);
+                    gravityGradientShape.MultTranspose(elementBaseGravityGradient, baseGravityReferenceValue);
+                }
                 data.baseDensityValues(quadraturePoint)      = elementBaseDensity * densityShape;
-                data.referenceWeights(quadraturePoint)       = integrationPoint.weight * transformation->Weight();
+                data.referenceWeights(quadraturePoint)       = integrationPoint.weight * meshWeight;
 
-                const mfem::DenseMatrix &inverseMeshJacobian = transformation->InverseJacobian();
-                const mfem::DenseMatrix &meshJacobian        = transformation->Jacobian();
-                const double inverseMeshWeight               = 1.0 / transformation->Weight();
+                const mfem::DenseMatrix &inverseMeshJacobian =
+                    meshGeometry ? meshGeometry->inverseJacobian : transformation->InverseJacobian();
+                const mfem::DenseMatrix &meshJacobian =
+                    meshGeometry ? meshGeometry->jacobian : transformation->Jacobian();
+                const double inverseMeshWeight = 1.0 / meshWeight;
                 for (int row = 0; row < dimension; ++row) {
                     data.baseGravityReferenceValues(quadraturePoint, row) = baseGravityReferenceValue(row);
                     for (int column = 0; column < dimension; ++column) {
@@ -335,8 +361,40 @@ namespace mean_field::operators {
                         }
                     }
                 }
+
+                if (assembleResidual) {
+                    for (int row = 0; row < dimension; ++row) {
+                        double mapped = 0.0;
+                        for (int column = 0; column < dimension; ++column)
+                            mapped += mappingContext.mapping.mapping_jacobian(row, column) *
+                                      baseGravityReferenceValue(column);
+                        m_quadratureForce(quadraturePoint, row) =
+                            data.referenceWeights(quadraturePoint) * data.baseDensityValues(quadraturePoint) * mapped;
+                    }
+                }
+            }
+
+            if (assembleResidual) {
+                const int ndof = displacementElement.GetDof();
+                MFEM_VERIFY(data.displacementDofs.Size() == ndof * dimension,
+                            "Prepared force residual has an invalid displacement element size.");
+                m_elementAction.SetSize(data.displacementDofs.Size());
+                residualComponent.SetSize(ndof);
+                for (int component = 0; component < dimension; ++component) {
+                    data.displacementReferenceTable->GetValues().MultTranspose(
+                        m_quadratureForce.Data() + component * quadraturePointCount, residualComponent.GetData());
+                    for (int dof = 0; dof < ndof; ++dof)
+                        m_elementAction(vector_dof_index(m_fem.displacementFes->GetOrdering(),
+                                                         dof, component, ndof, dimension)) = residualComponent(dof);
+                }
+                if (!vector_is_finite(m_elementAction)) return std::unexpected(non_finite_rejection());
+                if (data.displacementDofTransformation != nullptr)
+                    data.displacementDofTransformation->TransformDual(m_elementAction);
+                m_localAction.AddElementVector(data.displacementDofs, m_elementAction);
             }
         }
+        if (assembleResidual && !vector_is_finite(m_localAction))
+            return std::unexpected(non_finite_rejection());
         return {};
     }
 
@@ -352,6 +410,17 @@ namespace mean_field::operators {
         PreparedGravityDisplacementForceReport,
         kernels::GravityDisplacementForceRejection>
     PreparedGravityDisplacementForceOperator::TryPrepare() {
+        return TryPrepareImpl(true);
+    }
+
+    PreparedGravityDisplacementForceReport PreparedGravityDisplacementForceOperator::PrepareReference() {
+        auto result = TryPrepareImpl(false);
+        if (!result.has_value()) throw_rejection(result.error());
+        return std::move(result).value();
+    }
+
+    std::expected<PreparedGravityDisplacementForceReport, kernels::GravityDisplacementForceRejection>
+    PreparedGravityDisplacementForceOperator::TryPrepareImpl(const bool useFastPreparation) {
         MFEM_VERIFY(
             m_gravityContext.IsPrepared(), "PreparedGravityDisplacementForceOperator requires the shared "
                                            "gravity linearization context to be prepared first."
@@ -359,20 +428,24 @@ namespace mean_field::operators {
 
         const context::gravity_field::GravityFieldRevisions &requestedRevisions = m_gravityContext.GetRevisions();
 
-        if (m_isPrepared && relevant_revisions_match(requestedRevisions, m_preparedRevisions)) {
+        if (m_isPrepared && m_preparedWithFastResidual == useFastPreparation &&
+            relevant_revisions_match(requestedRevisions, m_preparedRevisions)) {
             return {};
         }
 
         m_isPrepared             = false;
 
-        /*
-         * Build the reusable element plan before assembling the residual.
-         * This pass stops at the first invalid mapped quadrature point, so a
-         * rejected line-search candidate need not traverse the full stateless
-         * residual kernel.  Synchronize before proceeding so every rank takes
-         * the same branch.
-         */
-        const auto elementResult = TryPrepareElementData();
+        const bool inputsFinite = vector_is_finite(m_gravityContext.GetDensityTrue()) &&
+            vector_is_finite(m_gravityContext.GetGravityGradientTrue()) &&
+            vector_is_finite(m_gravityContext.GetGeometryContext().GetDisplacementTrue());
+        const auto inputStatus = synchronize_rejection(
+            inputsFinite ? std::optional<Rejection>{} : std::optional<Rejection>{non_finite_rejection()},
+            m_fem.mesh->GetComm());
+        if (!inputStatus.has_value()) return std::unexpected(inputStatus.error());
+
+        // The production path assembles local residual contributions while
+        // preparing the geometry. Synchronize validity before the true-DOF sum.
+        const auto elementResult = TryPrepareElementData(useFastPreparation);
         const std::optional<Rejection> localElementRejection =
             elementResult.has_value() ? std::optional<Rejection>{} : std::optional<Rejection>{elementResult.error()};
         auto synchronizedElement = synchronize_rejection(localElementRejection, m_fem.mesh->GetComm());
@@ -380,12 +453,14 @@ namespace mean_field::operators {
             return std::unexpected(synchronizedElement.error());
         }
 
-        auto residualResult = kernels::try_apply_gravity_displacement_force_residual(
-            m_fem, m_domainMapper, m_gravityContext.GetDensityTrue(), m_gravityContext.GetGravityGradientTrue(),
-            m_gravityContext.GetGeometryContext().GetDisplacementTrue(), m_actionTrue
-        );
-        if (!residualResult.has_value()) {
-            return std::unexpected(residualResult.error());
+        if (useFastPreparation) {
+            local_to_true(*m_fem.displacementFes, m_localAction, m_actionTrue);
+        } else {
+            auto residualResult = kernels::try_apply_gravity_displacement_force_residual(
+                m_fem, m_domainMapper, m_gravityContext.GetDensityTrue(), m_gravityContext.GetGravityGradientTrue(),
+                m_gravityContext.GetGeometryContext().GetDisplacementTrue(), m_actionTrue
+            );
+            if (!residualResult.has_value()) return std::unexpected(residualResult.error());
         }
         m_cachedResidual.SetSize(m_gravityContext.GetDisplacementMap().reduced_size());
         m_gravityContext.GetDisplacementMap().gather(m_actionTrue, m_cachedResidual);
@@ -397,6 +472,7 @@ namespace mean_field::operators {
         if (!synchronized.has_value()) {
             return std::unexpected(synchronized.error());
         }
+        m_preparedWithFastResidual = useFastPreparation;
         m_preparedRevisions = requestedRevisions;
         ++m_residualPreparationCount;
         m_isPrepared = true;
@@ -471,7 +547,8 @@ namespace mean_field::operators {
         const mfem::Vector &densityVariationTrue,
         const mfem::Vector &displacementVariationTrue,
         const mfem::Vector &gravityGradientVariationTrue,
-        mfem::Vector &actionTrue
+        mfem::Vector &actionTrue,
+        const bool useFastPath, const bool useTensorPath
     ) const {
         true_to_local(*m_fem.densityFes, densityVariationTrue, m_densityVariationLocal);
         true_to_local(*m_fem.gravityFluxFes, gravityGradientVariationTrue, m_gravityGradientVariationLocal);
@@ -513,6 +590,60 @@ namespace mean_field::operators {
             const mfem::DenseMatrix &directionDofs = directionData.GetDofMatrix();
             const int scalarDisplacementDofCount   = displacementElement.GetDof();
 
+            if (useFastPath && data.fastApplication) {
+                const int nq = data.integrationRule->GetNPoints();
+                m_quadratureDensityVariation.SetSize(nq);
+                data.densityReferenceTable->GetValues().Mult(m_elementDensityVariation, m_quadratureDensityVariation);
+                m_quadratureForce.SetSize(nq, 3);
+                double *force = m_quadratureForce.Data();
+                const bool tensorDisplacement = useTensorPath && data.displacementReferenceTable->HasTensorEvaluation3D();
+                const bool tensorGravity = useTensorPath && data.gravityReferenceTable->HasTensorEvaluation3D();
+                if (tensorDisplacement)
+                    data.displacementReferenceTable->EvaluateTensor3D(directionDofs, m_tensorDisplacement, m_tensorWorkspace);
+                if (tensorGravity)
+                    data.gravityReferenceTable->EvaluateTensor3D(m_elementGravityGradientVariation, m_tensorGravity, m_tensorWorkspace);
+                for (int q = 0; q < nq; ++q) {
+                    double base[3], cellVariation[3], meshVariation[3], material[3], mapped[3];
+                    double referenceDirection[3], geometryVariation[3];
+                    for (int c = 0; c < 3; ++c) base[c] = data.baseGravityReferenceValues(q, c);
+                    if (tensorGravity) {
+                        for (int c = 0; c < 3; ++c) cellVariation[c] = m_tensorGravity(q, c);
+                    } else vector_basis_value_3d(data.gravityReferenceTable->GetValues(q),
+                                                m_elementGravityGradientVariation, cellVariation);
+                    cached_matrix_vector_3d(data.meshPiolaJacobians, q, cellVariation, meshVariation);
+                    const double rho = data.baseDensityValues(q);
+                    const double drho = m_quadratureDensityVariation(q);
+                    for (int c = 0; c < 3; ++c) material[c] = drho * base[c] + rho * meshVariation[c];
+                    cached_matrix_vector_3d(data.mappingJacobians, q, material, mapped);
+                    cached_matrix_vector_3d(data.inverseMeshJacobians, q, base, referenceDirection);
+                    if (tensorDisplacement) {
+                        for (int c = 0; c < 3; ++c) {
+                            geometryVariation[c] = 0.0;
+                            for (int axis = 0; axis < 3; ++axis)
+                                geometryVariation[c] += m_tensorDisplacement(q, 4*c + 1 + axis) * referenceDirection[axis];
+                        }
+                    } else displacement_directional_value_3d(directionDofs,
+                        data.displacementReferenceTable->GetGradients(q), referenceDirection, geometryVariation);
+                    for (int c = 0; c < 3; ++c)
+                        force[q + c * nq] = data.referenceWeights(q) * (mapped[c] + rho * geometryVariation[c]);
+                }
+                m_elementAction.SetSize(data.displacementDofs.Size());
+                double *element = m_elementAction.HostWrite();
+                if (tensorDisplacement) {
+                    data.displacementReferenceTable->ApplyTransposeTensor3D(
+                        m_quadratureForce, m_tensorElementDual, m_tensorWorkspace);
+                    for (int i = 0; i < m_elementAction.Size(); ++i) element[i] = m_tensorElementDual.Data()[i];
+                } else {
+                    for (int c = 0; c < 3; ++c)
+                        data.displacementReferenceTable->GetValues().MultTranspose(
+                            force + c * nq, element + c * scalarDisplacementDofCount);
+                }
+                if (data.displacementDofTransformation)
+                    data.displacementDofTransformation->TransformDual(m_elementAction);
+                m_localAction.AddElementVector(data.displacementDofs, m_elementAction);
+                continue;
+            }
+
             m_densityShape.SetSize(densityElement.GetDof());
             m_displacementShape.SetSize(scalarDisplacementDofCount);
             m_gravityGradientShape.SetSize(gravityGradientElement.GetDof(), dimension);
@@ -628,6 +759,52 @@ namespace mean_field::operators {
         ++m_completeJacobianStatistics.applications;
     }
 
+    void PreparedGravityDisplacementForceOperator::ApplyCompleteJacobianActionReference(
+        const mfem::Vector &densityVariation,
+        const mfem::Vector &displacementVariation,
+        const mfem::Vector &gravityGradientVariation,
+        mfem::Vector &action
+    ) const {
+        VerifyPrepared();
+
+        m_densityVariationTrue.SetSize(m_gravityContext.GetDensityMap().full_size());
+        m_gravityGradientVariationTrue.SetSize(m_gravityContext.GetGravityGradientMap().full_size());
+        m_displacementVariationTrue.SetSize(m_gravityContext.GetDisplacementMap().full_size());
+        m_gravityContext.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);
+        m_gravityContext.GetGravityGradientMap().scatter(gravityGradientVariation, m_gravityGradientVariationTrue);
+        m_gravityContext.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);
+
+        ApplyPreparedCompleteJacobianActionTrue(
+            m_densityVariationTrue, m_displacementVariationTrue, m_gravityGradientVariationTrue, m_actionTrue, false
+        );
+        action.SetSize(m_gravityContext.GetDisplacementMap().reduced_size());
+        m_gravityContext.GetDisplacementMap().gather(m_actionTrue, action);
+
+    }
+
+    void PreparedGravityDisplacementForceOperator::ApplyCompleteJacobianActionCachedReference(
+        const mfem::Vector &densityVariation,
+        const mfem::Vector &displacementVariation,
+        const mfem::Vector &gravityGradientVariation,
+        mfem::Vector &action
+    ) const {
+        VerifyPrepared();
+
+        m_densityVariationTrue.SetSize(m_gravityContext.GetDensityMap().full_size());
+        m_gravityGradientVariationTrue.SetSize(m_gravityContext.GetGravityGradientMap().full_size());
+        m_displacementVariationTrue.SetSize(m_gravityContext.GetDisplacementMap().full_size());
+        m_gravityContext.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);
+        m_gravityContext.GetGravityGradientMap().scatter(gravityGradientVariation, m_gravityGradientVariationTrue);
+        m_gravityContext.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);
+
+        ApplyPreparedCompleteJacobianActionTrue(
+            m_densityVariationTrue, m_displacementVariationTrue, m_gravityGradientVariationTrue, m_actionTrue, true, false
+        );
+        action.SetSize(m_gravityContext.GetDisplacementMap().reduced_size());
+        m_gravityContext.GetDisplacementMap().gather(m_actionTrue, action);
+
+    }
+
     bool PreparedGravityDisplacementForceOperator::IsPrepared() const noexcept {
         if (!m_isPrepared || !m_gravityContext.IsPrepared()) {
             return false;

libmeanfield/impl/operators/prepared_gravity_source.cpp

@@ -345,6 +345,11 @@ namespace mean_field::operators {
         }
     }
 
+    void PreparedMappedGravitySourceOperator::PrepareReference(const mfem::Vector &displacement) {
+        auto result = TryPrepareImpl(displacement, PreparationMode::linearization, false);
+        if (!result.has_value()) throwGravitySourcePreparationRejection(result.error());
+    }
+
     void PreparedMappedGravitySourceOperator::PreparePrimal(const mfem::Vector &displacement) {
         MEAN_FIELD_PROFILE_SCOPE_WARMUP("PreparedMappedGravitySourceOperator::Prepare primal", 0);
         auto result = TryPrepareImpl(displacement, PreparationMode::primal);
@@ -366,7 +371,7 @@ namespace mean_field::operators {
 
     GravitySourcePreparationResult PreparedMappedGravitySourceOperator::TryPrepareImpl(
         const mfem::Vector &displacement,
-        const PreparationMode mode
+        const PreparationMode mode, const bool useCachedMeshWeight
     ) {
         MFEM_VERIFY(
             displacement.Size() == m_displacement_map.reduced_size(),
@@ -513,7 +518,12 @@ namespace mean_field::operators {
 
                 transformation.SetIntPoint(&integration_point);
 
-                const double quadrature_value = integration_point.weight * transformation.Weight() * coefficient_value;
+                // EvaluateVolume already cached the computational mesh geometry.
+                // Resetting IntPoint above invalidates MFEM's local weight cache.
+                const auto *fixed = useCachedMeshWeight ?
+                    m_domain_mapper.GetFixedMeshGeometry(transformation, integration_point) : nullptr;
+                const double meshWeight = fixed ? fixed->weight : transformation.Weight();
+                const double quadrature_value = integration_point.weight * meshWeight * coefficient_value;
 
                 if (!std::isfinite(quadrature_value) || quadrature_value <= 0.0) {
                     localNonFiniteQuadrature = true;
@@ -608,9 +618,27 @@ namespace mean_field::operators {
     }
 
     void PreparedMappedGravitySourceOperator::MultDisplacementVariationTrue(
-        const mfem::Vector &densityTrue,
-        const mfem::Vector &displacementVariationTrue,
+        const mfem::Vector &densityTrue, const mfem::Vector &displacementVariationTrue,
         mfem::Vector &actionVariationTrue
+    ) const {
+        MultDisplacementVariationTrueImpl(densityTrue, displacementVariationTrue, actionVariationTrue, true);
+    }
+    void PreparedMappedGravitySourceOperator::MultDisplacementVariationTrueReference(
+        const mfem::Vector &densityTrue, const mfem::Vector &displacementVariationTrue,
+        mfem::Vector &actionVariationTrue
+    ) const {
+        MultDisplacementVariationTrueImpl(densityTrue, displacementVariationTrue, actionVariationTrue, false);
+    }
+    int PreparedMappedGravitySourceOperator::GetTensorVariationElementCount() const noexcept {
+        int count = 0;
+        for (const auto &data : m_elements)
+            count += data.displacement_reference && data.displacement_reference->HasTensorEvaluation3D();
+        return count;
+    }
+
+    void PreparedMappedGravitySourceOperator::MultDisplacementVariationTrueImpl(
+        const mfem::Vector &densityTrue, const mfem::Vector &displacementVariationTrue,
+        mfem::Vector &actionVariationTrue, const bool useTensorPath
     ) const {
         MFEM_VERIFY(
             m_is_prepared,
@@ -667,10 +695,15 @@ namespace mean_field::operators {
 
             m_reference_displacement_jacobian.SetSize(dimension, dimension);
             m_quadrature_variation_action.SetSize(data.integration_rule->GetNPoints());
-            data.GetDensityBasis().Mult(m_element_density, m_quadrature_variation_action);
+            if (useTensorPath && data.density_reference && data.density_reference->HasTensorEvaluation3D())
+                data.density_reference->EvaluateValuesTensor3D(m_element_density, m_quadrature_variation_action, m_tensorWorkspace);
+            else data.GetDensityBasis().Mult(m_element_density, m_quadrature_variation_action);
+            const bool tensor = useTensorPath && data.displacement_reference &&
+                                data.displacement_reference->HasTensorEvaluation3D();
+            if (tensor) data.displacement_reference->EvaluateTensor3D(direction_dofs, m_tensorSamples, m_tensorWorkspace);
 
             for (int quadrature_point = 0; quadrature_point < data.integration_rule->GetNPoints(); ++quadrature_point) {
-                mfem::MultAtB(
+                if (!tensor) mfem::MultAtB(
                     direction_dofs, data.displacement_reference->GetGradients(quadrature_point),
                     m_reference_displacement_jacobian
                 );
@@ -680,7 +713,8 @@ namespace mean_field::operators {
                     for (int column = 0; column < dimension; ++column) {
                         logarithmic_jacobian_variation +=
                             data.inverse_element_jacobians(quadrature_point, row * dimension + column) *
-                            m_reference_displacement_jacobian(column, row);
+                            (tensor ? m_tensorSamples(quadrature_point, 4*column+row+1) :
+                                      m_reference_displacement_jacobian(column, row));
                     }
                 }
 
@@ -693,7 +727,9 @@ namespace mean_field::operators {
             }
 
             m_element_variation_action.SetSize(data.potential_dofs.Size());
-            data.GetPotentialBasis().MultTranspose(m_quadrature_variation_action, m_element_variation_action);
+            if (useTensorPath && data.potential_reference && data.potential_reference->HasTensorEvaluation3D())
+                data.potential_reference->ApplyTransposeTensor3D(m_quadrature_variation_action, m_element_variation_action, m_tensorWorkspace);
+            else data.GetPotentialBasis().MultTranspose(m_quadrature_variation_action, m_element_variation_action);
 
             if (data.potential_dof_transformation != nullptr) {
                 data.potential_dof_transformation->TransformDual(m_element_variation_action);

libmeanfield/impl/operators/prepared_hdiv_mass.cpp

@@ -9,6 +9,7 @@ module;
 #include <optional>
 #include <stdexcept>
 #include <string>
+#include <vector>
 
 #include <mpi.h>
 
@@ -16,6 +17,7 @@ module mean_field;
 
 import :fem.reference_tables;
 import :operators.prepared_hdiv_mass;
+import :mapping.kelvin;
 
 namespace {
     using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;
@@ -304,6 +306,124 @@ namespace {
         }
     }
 
+    // Keep MFEM's PA assembly, restrictions, signed RT ordering, and diagonal.
+    // Only the CPU element contraction changes. Scratch is sized to the actual
+    // rule, and the two domain forms visit only their marked elements.
+    class TensorHDivMassIntegrator final : public mfem::VectorFEMassIntegrator {
+    public:
+        TensorHDivMassIntegrator(mfem::MatrixCoefficient &coefficient, const mfem::Array<int> &marker)
+            : mfem::VectorFEMassIntegrator(coefficient), m_marker(marker) {}
+
+        void AssemblePA(const mfem::FiniteElementSpace &space) override {
+            mfem::VectorFEMassIntegrator::AssemblePA(space);
+            m_activeElements.clear();
+            m_supported = !mfem::Device::Allows(~mfem::Backend::CPU) &&
+                dim == 3 && trial_fetype == mfem::FiniteElement::DIV &&
+                test_fetype == trial_fetype && dofs1D == dofs1Dtest;
+            if (!m_supported) return;
+            for (int e = 0; e < space.GetNE(); ++e) {
+                const int attribute = space.GetMesh()->GetAttribute(e);
+                if (attribute > 0 && attribute <= m_marker.Size() && m_marker[attribute - 1])
+                    m_activeElements.push_back(e);
+            }
+        }
+
+        [[nodiscard]] int FastElementCount() const noexcept {
+            return m_supported ? static_cast<int>(m_activeElements.size()) : 0;
+        }
+        [[nodiscard]] bool UseFastApplication() const noexcept { return m_useFast; }
+        void SetFastApplication(bool value) const noexcept { m_useFast = value; }
+
+        void AddMultPA(const mfem::Vector &x, mfem::Vector &y) const override {
+            if (!m_useFast || !m_supported) {
+                mfem::VectorFEMassIntegrator::AddMultPA(x, y);
+                return;
+            }
+            // The current RT2 / nine-point-per-axis rule is specialized; other
+            // supported CPU orders use exactly the same algorithm at runtime.
+            if (dofs1D == 4 && quad1D == 9) Apply<4, 9>(x, y);
+            else Apply<0, 0>(x, y);
+        }
+
+    private:
+        template <int D, int Q>
+        void Apply(const mfem::Vector &xVector, mfem::Vector &yVector) const {
+            const int d = D == 0 ? dofs1D : D, q = Q == 0 ? quad1D : Q;
+            const int nq = q*q*q, nd = 3*d*(d-1)*(d-1), entries = symmetric ? 6 : 9;
+            MFEM_VERIFY(xVector.Size() == ne*nd && yVector.Size() == ne*nd,
+                        "H(div) PA element vectors have incompatible sizes.");
+            m_workspace.SetSize(3*nq + q*d*d + q*q*d);
+            double *values = m_workspace.HostWrite(), *sx = values + 3*nq, *sxy = sx + q*d*d;
+            const double *x = xVector.HostRead(), *op = pa_data.HostRead();
+            double *y = yVector.HostReadWrite();
+            const double *bc = mapsC->B.HostRead(), *bo = mapsO->B.HostRead();
+            for (const int e : m_activeElements) {
+                int offset = e*nd;
+                for (int c = 0; c < 3; ++c) {
+                    const int dx = c == 0 ? d : d-1, dy = c == 1 ? d : d-1, dz = c == 2 ? d : d-1;
+                    const double *bx = c == 0 ? bc : bo, *by = c == 1 ? bc : bo, *bz = c == 2 ? bc : bo;
+                    for (int k = 0; k < dz; ++k) for (int j = 0; j < dy; ++j) for (int a = 0; a < q; ++a) {
+                        double v = 0.0;
+                        for (int i = 0; i < dx; ++i) v += bx[a + q*i]*x[offset + i + dx*(j + dy*k)];
+                        sx[a + q*(j + dy*k)] = v;
+                    }
+                    for (int k = 0; k < dz; ++k) for (int b = 0; b < q; ++b) for (int a = 0; a < q; ++a) {
+                        double v = 0.0;
+                        for (int j = 0; j < dy; ++j) v += by[b + q*j]*sx[a + q*(j + dy*k)];
+                        sxy[a + q*(b + q*k)] = v;
+                    }
+                    for (int z = 0; z < q; ++z) for (int b = 0; b < q; ++b) for (int a = 0; a < q; ++a) {
+                        double v = 0.0;
+                        for (int k = 0; k < dz; ++k) v += bz[z + q*k]*sxy[a + q*(b + q*k)];
+                        values[a + q*(b + q*z) + nq*c] = v;
+                    }
+                    offset += dx*dy*dz;
+                }
+                const double *metric = op + e*nq*entries;
+                for (int p = 0; p < nq; ++p) {
+                    const double u = values[p], v = values[p+nq], w = values[p+2*nq];
+                    const double a00 = metric[p], a01 = metric[p+nq], a02 = metric[p+2*nq];
+                    const double a10 = symmetric ? a01 : metric[p+3*nq];
+                    const double a11 = metric[p+(symmetric ? 3 : 4)*nq];
+                    const double a12 = metric[p+(symmetric ? 4 : 5)*nq];
+                    const double a20 = symmetric ? a02 : metric[p+6*nq];
+                    const double a21 = symmetric ? a12 : metric[p+7*nq];
+                    const double a22 = metric[p+(symmetric ? 5 : 8)*nq];
+                    values[p] = a00*u + a01*v + a02*w;
+                    values[p+nq] = a10*u + a11*v + a12*w;
+                    values[p+2*nq] = a20*u + a21*v + a22*w;
+                }
+                offset = e*nd;
+                for (int c = 0; c < 3; ++c) {
+                    const int dx = c == 0 ? d : d-1, dy = c == 1 ? d : d-1, dz = c == 2 ? d : d-1;
+                    const double *bx = c == 0 ? bc : bo, *by = c == 1 ? bc : bo, *bz = c == 2 ? bc : bo;
+                    for (int k = 0; k < dz; ++k) for (int b = 0; b < q; ++b) for (int a = 0; a < q; ++a) {
+                        double v = 0.0;
+                        for (int z = 0; z < q; ++z) v += bz[z + q*k]*values[a + q*(b + q*z) + nq*c];
+                        sxy[a + q*(b + q*k)] = v;
+                    }
+                    for (int k = 0; k < dz; ++k) for (int j = 0; j < dy; ++j) for (int a = 0; a < q; ++a) {
+                        double v = 0.0;
+                        for (int b = 0; b < q; ++b) v += by[b + q*j]*sxy[a + q*(b + q*k)];
+                        sx[a + q*(j + dy*k)] = v;
+                    }
+                    for (int k = 0; k < dz; ++k) for (int j = 0; j < dy; ++j) for (int i = 0; i < dx; ++i) {
+                        double v = 0.0;
+                        for (int a = 0; a < q; ++a) v += bx[a + q*i]*sx[a + q*(j + dy*k)];
+                        y[offset + i + dx*(j + dy*k)] += v;
+                    }
+                    offset += dx*dy*dz;
+                }
+            }
+        }
+
+        mfem::Array<int> m_marker;
+        std::vector<int> m_activeElements;
+        mutable mfem::Vector m_workspace;
+        bool m_supported{false};
+        mutable bool m_useFast{true};
+    };
+
     class FrozenMappedHDivMassCoefficient final : public mfem::MatrixCoefficient {
     public:
         FrozenMappedHDivMassCoefficient(
@@ -520,6 +640,11 @@ namespace mean_field::operators {
         mfem::Vector elementDisplacement;
         mfem::Vector elementCompactification;
         mapping::VolumeMappingContext mappingContext;
+        mfem::DenseMatrix inverseMeshJacobian(3), compactificationGradients;
+        mfem::Vector coordinateGradient(3), scaleGradient(3);
+        const auto *kelvin = dynamic_cast<const mapping::compactification::KelvinCompactification *>(
+            &m_domain_mapper.GetExteriorMap()
+        );
 
         for (int elementId = 0; elementId < m_fem.mesh->GetNE(); ++elementId) {
             m_variationElements.emplace_back();
@@ -566,6 +691,24 @@ namespace mean_field::operators {
                 data.integrationRule->GetNPoints(), frozen_mapping_width(m_domain_mapper.GetDimension())
             );
 
+            const bool compactified = m_domain_mapper.IsCompactifiedElement(*transformation);
+            std::shared_ptr<const fem::ScalarReferenceTable> compactificationTable;
+            if (dimension == 3 && data.gravityReferenceTable != nullptr &&
+                displacementElement.GetMapType() == mfem::FiniteElement::VALUE &&
+                displacementElement.GetDerivType() == mfem::FiniteElement::GRAD &&
+                displacementElement.GetDim() == 3 &&
+                (!compactified || (kelvin != nullptr &&
+                    compactificationElement.GetDerivType() == mfem::FiniteElement::GRAD))) {
+                data.displacementReferenceTable =
+                    m_fem.GetReferenceTables().GetScalarTable(displacementElement, *data.integrationRule);
+                data.tangentFactors.SetSize(12, data.integrationRule->GetNPoints());
+                if (compactified) {
+                    compactificationTable = m_fem.GetReferenceTables().GetScalarTable(
+                        compactificationElement, *data.integrationRule
+                    );
+                }
+            }
+
             const mapping::ElementDisplacementData displacementData =
                 mapping::ElementDisplacementDataFromElementVDofs(displacementElement, data.baseDisplacement);
             const mapping::ElementCompactificationData compactificationData(
@@ -584,11 +727,42 @@ namespace mean_field::operators {
                     return status;
                 }
                 freeze_mapping_context(mappingContext, quadraturePoint, data.frozenMappingData);
+                if (data.displacementReferenceTable != nullptr) {
+                    // Use the same fixed-mesh inverse as EvaluateVolumeVariation.
+                    mfem::Mult(mappingContext.quadrature.J_inv, mappingContext.mapping.mapping_jacobian,
+                               inverseMeshJacobian);
+                    double scale = 1.0;
+                    scaleGradient = 0.0;
+                    if (compactified) {
+                        const auto &values = compactificationTable->GetValues();
+                        double coordinate = 0.0;
+                        for (int dof = 0; dof < data.compactification.Size(); ++dof)
+                            coordinate += values(quadraturePoint, dof) * data.compactification(dof);
+                        compactificationGradients.SetSize(compactificationElement.GetDof(), dimension);
+                        mfem::Mult(compactificationTable->GetGradients(quadraturePoint), inverseMeshJacobian,
+                                   compactificationGradients);
+                        compactificationGradients.MultTranspose(data.compactification, coordinateGradient);
+                        const auto factorStatus = kelvin->EvaluateJacobianVariationFactors(
+                            coordinate, coordinateGradient, scale, scaleGradient
+                        );
+                        if (factorStatus != mapping::MappingStatus::valid) return factorStatus;
+                    }
+                    double *factors = data.tangentFactors.GetColumn(quadraturePoint);
+                    for (int row = 0; row < 3; ++row) {
+                        for (int column = 0; column < 3; ++column) {
+                            const double value = scale * inverseMeshJacobian(row, column);
+                            if (!std::isfinite(value)) return mapping::MappingStatus::non_finite_result;
+                            factors[3 * row + column] = value;
+                        }
+                        factors[9 + row] = scaleGradient(row);
+                    }
+                }
                 if (data.gravityReferenceTable != nullptr) {
                     // CalcVShape_RT = reference_shape * J_mesh^T / Weight.
                     // Cache only this small factor, never the mapped basis.
-                    const double meshWeight                = transformation->Weight();
-                    const mfem::DenseMatrix &meshJacobian  = transformation->Jacobian();
+                    const auto *meshGeometry = m_domain_mapper.GetFixedMeshGeometry(*transformation, integrationPoint);
+                    const double meshWeight = meshGeometry ? meshGeometry->weight : transformation->Weight();
+                    const mfem::DenseMatrix &meshJacobian = meshGeometry ? meshGeometry->jacobian : transformation->Jacobian();
                     const double inverseMeshWeight         = 1.0 / meshWeight;
                     data.referenceWeights(quadraturePoint) = integrationPoint.weight * meshWeight;
                     for (int row = 0; row < dimension; ++row) {
@@ -693,8 +867,11 @@ namespace mean_field::operators {
         m_stellar_mass_form->SetAssemblyLevel(mfem::AssemblyLevel::PARTIAL);
         m_vacuum_mass_form->SetAssemblyLevel(mfem::AssemblyLevel::PARTIAL);
 
-        auto stellar_integrator = std::make_unique<mfem::VectorFEMassIntegrator>(*m_stellar_mass_coefficient);
-        auto vacuum_integrator  = std::make_unique<mfem::VectorFEMassIntegrator>(*m_vacuum_mass_coefficient);
+        auto stellar_integrator = std::make_unique<TensorHDivMassIntegrator>(*m_stellar_mass_coefficient, m_stellar_marker);
+        auto vacuum_integrator  = std::make_unique<TensorHDivMassIntegrator>(*m_vacuum_mass_coefficient, m_vacuum_marker);
+        auto *stellarIntegratorView = stellar_integrator.get();
+        auto *vacuumIntegratorView = vacuum_integrator.get();
+        m_fastMassElements = 0;
 
         m_fem.quadratureFactory->configure_gravity_hdiv_mass(
             *stellar_integrator, quadrature::QuadratureRole::discretization, stellar_element, stellar_transformation,
@@ -731,6 +908,7 @@ namespace mean_field::operators {
             return preparationResult;
         }
 
+        m_fastMassElements = stellarIntegratorView->FastElementCount() + vacuumIntegratorView->FastElementCount();
         m_has_variation_data = mode == PreparationMode::linearization;
         m_is_prepared        = true;
         ++m_preparation_count;
@@ -791,10 +969,71 @@ namespace mean_field::operators {
         }
     }
 
+    void PreparedMappedHDivMassOperator::MultReference(
+        const mfem::Vector &gravityGradient, mfem::Vector &action
+    ) const {
+        MFEM_VERIFY(m_is_prepared && m_stellar_mass_form && m_vacuum_mass_form,
+                    "The H(div) mass reference requires a prepared operator.");
+        const auto *stellar = dynamic_cast<const TensorHDivMassIntegrator *>((*m_stellar_mass_form->GetDBFI())[0]);
+        const auto *vacuum = dynamic_cast<const TensorHDivMassIntegrator *>((*m_vacuum_mass_form->GetDBFI())[0]);
+        MFEM_VERIFY(stellar && vacuum, "The H(div) mass forms have unexpected integrators.");
+        struct RestoreMode {
+            const TensorHDivMassIntegrator *integrator;
+            bool previous;
+            ~RestoreMode() { integrator->SetFastApplication(previous); }
+        } restoreStellar{stellar, stellar->UseFastApplication()}, restoreVacuum{vacuum, vacuum->UseFastApplication()};
+        stellar->SetFastApplication(false);
+        vacuum->SetFastApplication(false);
+        Mult(gravityGradient, action);
+    }
+
     void PreparedMappedHDivMassOperator::MultDisplacementVariationTrue(
         const mfem::Vector &gravityGradientTrue,
         const mfem::Vector &displacementVariationTrue,
         mfem::Vector &actionVariationTrue
+    ) const {
+        MEAN_FIELD_PROFILE_SCOPE("PreparedMappedHDivMassOperator::MultDisplacementVariationTrue");
+        MultDisplacementVariationTrueImpl(gravityGradientTrue, displacementVariationTrue, actionVariationTrue, true, true);
+    }
+
+    void PreparedMappedHDivMassOperator::MultDisplacementVariationTrueReference(
+        const mfem::Vector &gravityGradientTrue,
+        const mfem::Vector &displacementVariationTrue,
+        mfem::Vector &actionVariationTrue
+    ) const {
+        MEAN_FIELD_PROFILE_SCOPE("PreparedMappedHDivMassOperator::MultDisplacementVariationTrueReference");
+        MultDisplacementVariationTrueImpl(gravityGradientTrue, displacementVariationTrue, actionVariationTrue, false, false);
+    }
+
+    void PreparedMappedHDivMassOperator::MultDisplacementVariationTrueCachedReference(
+        const mfem::Vector &gravityGradientTrue, const mfem::Vector &displacementVariationTrue,
+        mfem::Vector &actionVariationTrue
+    ) const {
+        MultDisplacementVariationTrueImpl(gravityGradientTrue, displacementVariationTrue, actionVariationTrue, true, false);
+    }
+
+    int PreparedMappedHDivMassOperator::GetTensorVariationElementCount() const noexcept {
+        if (!m_is_prepared || !m_has_variation_data) return 0;
+        int count = 0;
+        for (const auto &data : m_variationElements)
+            if (data.displacementReferenceTable && data.displacementReferenceTable->HasTensorEvaluation3D() &&
+                data.gravityReferenceTable && data.gravityReferenceTable->HasTensorEvaluation3D()) ++count;
+        return count;
+    }
+
+    int PreparedMappedHDivMassOperator::GetFastVariationElementCount() const noexcept {
+        if (!m_is_prepared || !m_has_variation_data) return 0;
+        int count = 0;
+        for (const auto &data : m_variationElements)
+            if (data.displacementReferenceTable != nullptr) ++count;
+        return count;
+    }
+
+    void PreparedMappedHDivMassOperator::MultDisplacementVariationTrueImpl(
+        const mfem::Vector &gravityGradientTrue,
+        const mfem::Vector &displacementVariationTrue,
+        mfem::Vector &actionVariationTrue,
+        const bool useFastPath, const bool useTensorPath
     ) const {
         MFEM_VERIFY(
             m_is_prepared, "PreparedMappedHDivMassOperator must be prepared before applying a displacement variation."
@@ -837,6 +1076,11 @@ namespace mean_field::operators {
                 data.displacementDofTransformation->InvTransformPrimal(m_elementDisplacementVariation);
             }
 
+            if (useFastPath && data.displacementReferenceTable != nullptr) {
+                AddFastElementVariation(data, useTensorPath);
+                continue;
+            }
+
             const mfem::FiniteElement &gravityGradientElement  = *m_fem.gravityFluxFes->GetFE(data.elementId);
             const mfem::FiniteElement &displacementElement     = *m_fem.displacementFes->GetFE(data.elementId);
             const mfem::FiniteElement &compactificationElement = *m_fem.compactificationFes->GetFE(data.elementId);
@@ -924,6 +1168,111 @@ namespace mean_field::operators {
         local_to_true(*m_fem.gravityFluxFes, m_localVariationAction, actionVariationTrue);
     }
 
+    void PreparedMappedHDivMassOperator::AddFastElementVariation(
+        const ElementVariationData &data, const bool useTensorPath
+    ) const {
+        // GetElementVDofs always returns component blocks, independently of the
+        // global FES ordering. Signed DOFs and primal transforms were applied by
+        // the caller; preserve the dual transform and scatter below.
+        const int dofs = data.displacementReferenceTable->GetDofCount();
+        MFEM_VERIFY(m_elementDisplacementVariation.Size() == 3 * dofs,
+                    "Prepared H(div) tangent has an incompatible displacement layout.");
+        const auto &values = data.displacementReferenceTable->GetValues();
+        const bool compactified = data.frozenMappingData(0, 47) != 0.0;
+        m_elementVariationAction.SetSize(m_elementGravityGradient.Size());
+        m_elementVariationAction = 0.0;
+        m_gravityReferenceCellValue.SetSize(3);
+        m_referenceCellDual.SetSize(3);
+        const bool tensorDisplacement = useTensorPath && data.displacementReferenceTable->HasTensorEvaluation3D();
+        const bool tensorGravity = useTensorPath && data.gravityReferenceTable->HasTensorEvaluation3D();
+        if (tensorDisplacement) {
+            const mfem::DenseMatrix direction(m_elementDisplacementVariation.GetData(), dofs, 3);
+            data.displacementReferenceTable->EvaluateTensor3D(direction, m_tensorDisplacement, m_tensorWorkspace);
+        }
+        if (tensorGravity) {
+            data.gravityReferenceTable->EvaluateTensor3D(m_elementGravityGradient, m_tensorGravity, m_tensorWorkspace);
+            m_tensorDual.SetSize(data.integrationRule->GetNPoints(), 3);
+        }
+
+        for (int q = 0; q < data.integrationRule->GetNPoints(); ++q) {
+            const auto &gradients = data.displacementReferenceTable->GetGradients(q);
+            const double *factors = data.tangentFactors.GetColumn(q);
+            double deltaJ[3][3]{};
+            for (int component = 0; component < 3; ++component) {
+                const double *coefficients = m_elementDisplacementVariation.GetData() + component * dofs;
+                double value = 0.0;
+                double referenceGradient[3]{};
+                if (tensorDisplacement) {
+                    value = compactified ? m_tensorDisplacement(q, 4*component) : 0.0;
+                    for (int axis = 0; axis < 3; ++axis)
+                        referenceGradient[axis] = m_tensorDisplacement(q, 4*component + 1 + axis);
+                } else {
+                    for (int dof = 0; dof < dofs; ++dof) {
+                        const double coefficient = coefficients[dof];
+                        if (compactified) value += coefficient * values(q, dof);
+                        for (int axis = 0; axis < 3; ++axis)
+                            referenceGradient[axis] += coefficient * gradients(dof, axis);
+                    }
+                }
+                for (int column = 0; column < 3; ++column) {
+                    deltaJ[component][column] = value * factors[9 + column];
+                    for (int axis = 0; axis < 3; ++axis)
+                        deltaJ[component][column] += referenceGradient[axis] * factors[3 * axis + column];
+                    MFEM_VERIFY(std::isfinite(deltaJ[component][column]),
+                                "Prepared H(div) tangent produced a non-finite mapping derivative.");
+                }
+            }
+
+            const auto &referenceShape = data.gravityReferenceTable->GetValues(q);
+            if (tensorGravity) {
+                for (int c = 0; c < 3; ++c) m_gravityReferenceCellValue(c) = m_tensorGravity(q, c);
+            } else referenceShape.MultTranspose(m_elementGravityGradient, m_gravityReferenceCellValue);
+            double piola[3][3], jacobian[3][3];
+            double gravity[3]{}, mappedGravity[3]{}, mappedVariation[3]{};
+            double trace = 0.0;
+            for (int row = 0; row < 3; ++row) {
+                for (int column = 0; column < 3; ++column) {
+                    piola[row][column] = data.meshPiolaJacobians(q, 3 * row + column);
+                    jacobian[row][column] = data.frozenMappingData(q, 18 + 3 * row + column);
+                    gravity[row] += piola[row][column] * m_gravityReferenceCellValue(column);
+                    trace += data.frozenMappingData(q, 27 + 3 * row + column) * deltaJ[column][row];
+                }
+            }
+            for (int row = 0; row < 3; ++row) {
+                for (int column = 0; column < 3; ++column) {
+                    mappedGravity[row] += jacobian[row][column] * gravity[column];
+                    mappedVariation[row] += deltaJ[row][column] * gravity[column];
+                }
+            }
+            // M = J^T J / det(J). Apply delta M directly to the three-vector:
+            // [delta J^T (Jg) + J^T(delta J g - tr(J^-1 delta J) Jg)] / det(J).
+            const double inverseDeterminant = 1.0 / data.frozenMappingData(q, 45);
+            double massAction[3]{};
+            for (int column = 0; column < 3; ++column) {
+                for (int row = 0; row < 3; ++row) {
+                    massAction[column] += deltaJ[row][column] * mappedGravity[row] +
+                        jacobian[row][column] * (mappedVariation[row] - trace * mappedGravity[row]);
+                }
+                massAction[column] *= inverseDeterminant;
+                MFEM_VERIFY(std::isfinite(massAction[column]),
+                            "Prepared H(div) tangent produced a non-finite mass action.");
+            }
+            for (int column = 0; column < 3; ++column) {
+                double dual = 0.0;
+                for (int row = 0; row < 3; ++row) dual += piola[row][column] * massAction[row];
+                m_referenceCellDual(column) = dual;
+            }
+            if (tensorGravity) {
+                for (int c = 0; c < 3; ++c) m_tensorDual(q, c) = data.referenceWeights(q) * m_referenceCellDual(c);
+            } else referenceShape.AddMult(m_referenceCellDual, m_elementVariationAction, data.referenceWeights(q));
+        }
+        if (tensorGravity)
+            data.gravityReferenceTable->ApplyTransposeTensor3D(m_tensorDual, m_elementVariationAction, m_tensorWorkspace);
+        if (data.gravityGradientDofTransformation != nullptr)
+            data.gravityGradientDofTransformation->TransformDual(m_elementVariationAction);
+        m_localVariationAction.AddElementVector(data.gravityGradientDofs, m_elementVariationAction);
+    }
+
     void PreparedMappedHDivMassOperator::AssembleDiagonal(mfem::Vector &diagonal) const {
         mfem::Vector true_diagonal;
         AssembleTrueDiagonal(true_diagonal);

libmeanfield/impl/operators/prepared_hydrostatic_equilibrium.cpp

@@ -11,6 +11,7 @@ module;
 #include <utility>
 
 #include <mfem.hpp>
+#include "interior_displacement_contractions.hpp"
 
 module mean_field;
 
@@ -384,7 +385,30 @@ namespace mean_field::operators {
         const context::hydrostatic::HydrostaticEquilibriumDependencies &dependencies,
         const physics::RigidRotation &rotation
     ) {
-        const bool wasPrepared = m_isPrepared;
+        return TryPrepareImpl(state, dependencies, rotation, true);
+    }
+    HydrostaticEquilibriumPreparationResult PreparedHydrostaticEquilibriumOperator::TryPrepareReference(
+        const context::hydrostatic::HydrostaticEquilibriumStateView &state,
+        const context::hydrostatic::HydrostaticEquilibriumDependencies &dependencies,
+        const physics::RigidRotation &rotation
+    ) {
+        return TryPrepareImpl(state, dependencies, rotation, false);
+    }
+    PreparedHydrostaticEquilibriumReport PreparedHydrostaticEquilibriumOperator::PrepareReference(
+        const context::hydrostatic::HydrostaticEquilibriumStateView &state,
+        const context::hydrostatic::HydrostaticEquilibriumDependencies &dependencies,
+        const physics::RigidRotation &rotation
+    ) {
+        auto result = TryPrepareReference(state, dependencies, rotation);
+        if (!result.has_value()) throwHydrostaticEquilibriumPreparationRejection(result.error());
+        return std::move(result).value();
+    }
+    HydrostaticEquilibriumPreparationResult PreparedHydrostaticEquilibriumOperator::TryPrepareImpl(
+        const context::hydrostatic::HydrostaticEquilibriumStateView &state,
+        const context::hydrostatic::HydrostaticEquilibriumDependencies &dependencies,
+        const physics::RigidRotation &rotation, const bool useTensorPreparation
+    ) {
+        const bool wasPrepared = m_isPrepared && m_preparedWithTensorPreparation == useTensorPreparation;
         const bool rotationObjectChanged =
             !wasPrepared || !m_context.IsPrepared() || dependencies.rotation != m_context.GetDependencies().rotation;
 
@@ -409,7 +433,7 @@ namespace mean_field::operators {
         }
 
         if (report.contextReport.preparedGeometryState || !wasPrepared) {
-            const auto mappingFailure = synchronize_mapping_failure(PrepareGeometry(), m_fem.mesh->GetComm());
+            const auto mappingFailure = synchronize_mapping_failure(PrepareGeometry(useTensorPreparation), m_fem.mesh->GetComm());
             if (mappingFailure.has_value()) {
                 const auto reason = *mappingFailure == mapping::MappingStatus::non_positive_determinant
                                         ? HydrostaticEquilibriumPreparationRejectionReason::inverted_geometry
@@ -419,7 +443,7 @@ namespace mean_field::operators {
                 );
             }
 
-            if (synchronize_non_finite_failure(PrepareAlgebraicJacobianBlocks(), m_fem.mesh->GetComm())) {
+            if (synchronize_non_finite_failure(PrepareAlgebraicJacobianBlocks(useTensorPreparation), m_fem.mesh->GetComm())) {
                 return std::unexpected(
                     HydrostaticEquilibriumPreparationRejection{
                         .reason        = HydrostaticEquilibriumPreparationRejectionReason::non_finite_geometry,
@@ -441,7 +465,7 @@ namespace mean_field::operators {
         }
 
         if (report.contextReport.preparedBaseState || !wasPrepared) {
-            if (synchronize_non_finite_failure(PrepareBaseState(), m_fem.mesh->GetComm())) {
+            if (synchronize_non_finite_failure(PrepareBaseState(useTensorPreparation), m_fem.mesh->GetComm())) {
                 return std::unexpected(
                     HydrostaticEquilibriumPreparationRejection{
                         .reason = HydrostaticEquilibriumPreparationRejectionReason::non_finite_residual
@@ -480,6 +504,7 @@ namespace mean_field::operators {
             ++m_residualPreparationCount;
         }
 
+        m_preparedWithTensorPreparation = useTensorPreparation;
         m_isPrepared = true;
         return report;
     }
@@ -529,10 +554,18 @@ namespace mean_field::operators {
             data.enthalpyReferenceTable = referenceTables.GetScalarTable(enthalpyElement, *data.integrationRule);
             data.gravityPotentialReferenceTable =
                 referenceTables.GetScalarTable(gravityPotentialElement, *data.integrationRule);
+            const auto &displacementElement = *m_fem.displacementFes->GetFE(elementId);
+            if (m_fem.mesh->Dimension() == 3 &&
+                displacementElement.GetMapType() == mfem::FiniteElement::VALUE &&
+                displacementElement.GetDerivType() == mfem::FiniteElement::GRAD &&
+                !m_domainMapper.IsCompactifiedElement(*transformation)) {
+                data.displacementReferenceTable =
+                    referenceTables.GetScalarTable(displacementElement, *data.integrationRule);
+            }
         }
     }
 
-    std::optional<mapping::MappingStatus> PreparedHydrostaticEquilibriumOperator::PrepareGeometry() {
+    std::optional<mapping::MappingStatus> PreparedHydrostaticEquilibriumOperator::PrepareGeometry(const bool useTensorPreparation) {
         mfem::Vector displacementLocal;
 
         true_to_local(*m_fem.displacementFes, m_context.GetDisplacementTrue(), displacementLocal);
@@ -544,6 +577,8 @@ namespace mean_field::operators {
 
         mfem::Vector elementDisplacement;
         mfem::Vector elementCompactification;
+        mfem::Vector sampleValue(3), tensorWorkspace;
+        mfem::DenseMatrix sampleGradient(3), tensorSamples;
 
         for (ElementPAData &data : m_elements) {
             mfem::ElementTransformation *transformation = m_fem.mesh->GetElementTransformation(data.elementId);
@@ -590,14 +625,27 @@ namespace mean_field::operators {
             data.quadratureWeights.SetSize(quadraturePointCount);
 
             data.baseMappingContexts.SetSize(quadraturePointCount, m_fem.mesh->Dimension());
+            if (data.displacementReferenceTable != nullptr)
+                data.displacementInverseJacobians.SetSize(9, quadraturePointCount);
+
+            const bool tensor = useTensorPreparation && data.displacementReferenceTable &&
+                                data.displacementReferenceTable->HasTensorEvaluation3D();
+            if (tensor) data.displacementReferenceTable->EvaluateTensor3D(
+                data.baseDisplacementData->GetDofMatrix(), tensorSamples, tensorWorkspace);
 
             for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
                 const mfem::IntegrationPoint &integrationPoint = data.integrationRule->IntPoint(quadraturePoint);
 
                 transformation->SetIntPoint(&integrationPoint);
 
+                if (tensor) for (int c = 0; c < 3; ++c) {
+                    sampleValue(c) = tensorSamples(quadraturePoint, 4*c);
+                    for (int a = 0; a < 3; ++a)
+                        sampleGradient(c,a) = tensorSamples(quadraturePoint, 4*c+a+1);
+                }
+                const mapping::DomainMapper::ReferenceDisplacementSample sample{sampleValue, sampleGradient};
                 const mapping::MappingStatus mappingStatus = m_domainMapper.EvaluateVolume(
-                    mappingData, *transformation, integrationPoint, workspace, mappingContext
+                    mappingData, *transformation, integrationPoint, workspace, mappingContext, tensor ? &sample : nullptr
                 );
 
                 if (mappingStatus != mapping::MappingStatus::valid) {
@@ -615,6 +663,16 @@ namespace mean_field::operators {
 
                 data.baseMappingContexts.Store(quadraturePoint, mappingContext);
                 data.quadratureWeights(quadraturePoint) = quadratureWeight;
+                if (data.displacementInverseJacobians.Width() != 0) {
+                    if (mappingContext.mapping.compactified) {
+                        data.displacementInverseJacobians.SetSize(0, 0);
+                    } else {
+                        double *inverse = data.displacementInverseJacobians.Data() + 9 * quadraturePoint;
+                        for (int c = 0; c < 3; ++c)
+                            for (int r = 0; r < 3; ++r)
+                                inverse[r + 3 * c] = mappingContext.quadrature.J_inv(r, c);
+                    }
+                }
 
                 for (int component = 0; component < m_fem.mesh->Dimension(); ++component) {
                     const double position = mappingContext.mapping.physical_position(component);
@@ -631,7 +689,8 @@ namespace mean_field::operators {
         return std::nullopt;
     }
 
-    bool PreparedHydrostaticEquilibriumOperator::PrepareAlgebraicJacobianBlocks() {
+    bool PreparedHydrostaticEquilibriumOperator::PrepareAlgebraicJacobianBlocks(const bool useTensorPreparation) {
+        mfem::Vector dual, column, workspace;
         for (ElementPAData &data : m_elements) {
             const mfem::DenseMatrix &enthalpyBasis         = data.GetEnthalpyBasis();
             const mfem::DenseMatrix &gravityPotentialBasis = data.GetGravityPotentialBasis();
@@ -658,7 +717,21 @@ namespace mean_field::operators {
             data.gravityPotentialJacobian  = 0.0;
             data.bernoulliConstantJacobian = 0.0;
 
-            for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
+            if (useTensorPreparation && data.enthalpyReferenceTable->HasTensorEvaluation3D()) {
+                dual.SetSize(quadraturePointCount);
+                auto assemble = [&](const mfem::DenseMatrix &basis, mfem::DenseMatrix &matrix) {
+                    for (int j = 0; j < basis.Width(); ++j) {
+                        for (int q = 0; q < quadraturePointCount; ++q)
+                            dual(q) = data.quadratureWeights(q)*basis(q,j);
+                        data.enthalpyReferenceTable->ApplyTransposeTensor3D(dual, column, workspace);
+                        for (int i = 0; i < enthalpyDofCount; ++i) matrix(i,j) = column(i);
+                    }
+                };
+                assemble(enthalpyBasis, data.enthalpyJacobian);
+                assemble(gravityPotentialBasis, data.gravityPotentialJacobian);
+                for (int q = 0; q < quadraturePointCount; ++q) dual(q) = -data.quadratureWeights(q);
+                data.enthalpyReferenceTable->ApplyTransposeTensor3D(dual, data.bernoulliConstantJacobian, workspace);
+            } else for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
                 const double quadratureWeight = data.quadratureWeights(quadraturePoint);
 
                 for (int testDof = 0; testDof < enthalpyDofCount; ++testDof) {
@@ -738,7 +811,7 @@ namespace mean_field::operators {
         return false;
     }
 
-    bool PreparedHydrostaticEquilibriumOperator::PrepareBaseState() {
+    bool PreparedHydrostaticEquilibriumOperator::PrepareBaseState(const bool useTensorPreparation) {
         mfem::Vector enthalpyLocal;
         mfem::Vector gravityPotentialLocal;
 
@@ -750,6 +823,7 @@ namespace mean_field::operators {
         mfem::Vector elementGravityPotential;
         mfem::Vector quadratureEnthalpy;
         mfem::Vector quadratureGravityPotential;
+        mfem::Vector workspace;
 
         for (ElementPAData &data : m_elements) {
             enthalpyLocal.GetSubVector(data.enthalpyDofs, elementEnthalpy);
@@ -770,9 +844,12 @@ namespace mean_field::operators {
 
             quadratureGravityPotential.SetSize(quadraturePointCount);
 
-            data.GetEnthalpyBasis().Mult(elementEnthalpy, quadratureEnthalpy);
-
-            data.GetGravityPotentialBasis().Mult(elementGravityPotential, quadratureGravityPotential);
+            if (useTensorPreparation && data.enthalpyReferenceTable->HasTensorEvaluation3D())
+                data.enthalpyReferenceTable->EvaluateValuesTensor3D(elementEnthalpy, quadratureEnthalpy, workspace);
+            else data.GetEnthalpyBasis().Mult(elementEnthalpy, quadratureEnthalpy);
+            if (useTensorPreparation && data.gravityPotentialReferenceTable->HasTensorEvaluation3D())
+                data.gravityPotentialReferenceTable->EvaluateValuesTensor3D(elementGravityPotential, quadratureGravityPotential, workspace);
+            else data.GetGravityPotentialBasis().Mult(elementGravityPotential, quadratureGravityPotential);
 
             MFEM_VERIFY(
                 data.rotationPotential.Size() == quadraturePointCount, "Prepared hydrostatic base state has stale "
@@ -1156,6 +1233,38 @@ namespace mean_field::operators {
     void PreparedHydrostaticEquilibriumOperator::ApplyDisplacementJacobianAction(
         const mfem::Vector &displacementVariation,
         mfem::Vector &action
+    ) const {
+        ApplyDisplacementJacobianActionImpl(displacementVariation, action, true, true);
+    }
+
+    void PreparedHydrostaticEquilibriumOperator::ApplyDisplacementJacobianActionReference(
+        const mfem::Vector &displacementVariation, mfem::Vector &action
+    ) const {
+        ApplyDisplacementJacobianActionImpl(displacementVariation, action, false, false);
+    }
+
+    void PreparedHydrostaticEquilibriumOperator::ApplyDisplacementJacobianActionCachedReference(
+        const mfem::Vector &displacementVariation, mfem::Vector &action
+    ) const {
+        ApplyDisplacementJacobianActionImpl(displacementVariation, action, true, false);
+    }
+
+    int PreparedHydrostaticEquilibriumOperator::GetTensorDisplacementElementCount() const noexcept {
+        int count = 0;
+        for (const auto &data : m_elements)
+            count += data.displacementInverseJacobians.Width() != 0 && data.displacementReferenceTable != nullptr &&
+                     data.displacementReferenceTable->HasTensorEvaluation3D();
+        return count;
+    }
+
+    int PreparedHydrostaticEquilibriumOperator::GetFastDisplacementElementCount() const noexcept {
+        int count = 0;
+        for (const auto &data : m_elements) count += data.displacementInverseJacobians.Width() != 0;
+        return count;
+    }
+
+    void PreparedHydrostaticEquilibriumOperator::ApplyDisplacementJacobianActionImpl(
+        const mfem::Vector &displacementVariation, mfem::Vector &action, const bool useFastPath, const bool useTensorPath
     ) const {
         VerifyPrepared();
 
@@ -1226,8 +1335,42 @@ namespace mean_field::operators {
             );
 
             weightedQuadratureVariation.SetSize(quadraturePointCount);
+            const bool tensor = useFastPath && useTensorPath && data.displacementInverseJacobians.Width() != 0 &&
+                data.displacementReferenceTable->HasTensorEvaluation3D();
+            if (tensor) data.displacementReferenceTable->EvaluateTensor3D(
+                directionData.GetDofMatrix(), m_tensorDisplacementValuesAndGradients, m_tensorDisplacementWorkspace);
 
             for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
+                if (tensor) {
+                    const double *inverse = data.displacementInverseJacobians.Data() + 9 * quadraturePoint;
+                    double trace = 0.0, position = 0.0;
+                    for (int c = 0; c < 3; ++c) {
+                        position += data.rotationGradient(quadraturePoint, c) *
+                            m_tensorDisplacementValuesAndGradients(quadraturePoint, 4 * c);
+                        for (int derivative = 0; derivative < 3; ++derivative)
+                            trace += inverse[derivative + 3 * c] *
+                                m_tensorDisplacementValuesAndGradients(quadraturePoint, 4 * c + 1 + derivative);
+                    }
+                    const double value = data.quadratureWeights(quadraturePoint) *
+                        (data.hydrostaticImbalance(quadraturePoint) * trace - position);
+                    MFEM_VERIFY(std::isfinite(value), "Prepared hydrostatic tensor action is non-finite.");
+                    weightedQuadratureVariation(quadraturePoint) = value;
+                    continue;
+                }
+                if (useFastPath && data.displacementInverseJacobians.Width() != 0) {
+                    const double weight = data.quadratureWeights(quadraturePoint);
+                    double positionCoefficient[3];
+                    for (int c = 0; c < 3; ++c)
+                        positionCoefficient[c] = -weight * data.rotationGradient(quadraturePoint, c);
+                    const double value = apply_interior_functional(
+                        directionData.GetDofMatrix(), data.displacementReferenceTable->GetGradients(quadraturePoint),
+                        data.displacementReferenceTable->GetValues(), quadraturePoint,
+                        data.displacementInverseJacobians.Data() + 9 * quadraturePoint,
+                        weight * data.hydrostaticImbalance(quadraturePoint), positionCoefficient);
+                    MFEM_VERIFY(std::isfinite(value), "Prepared hydrostatic displacement action is non-finite.");
+                    weightedQuadratureVariation(quadraturePoint) = value;
+                    continue;
+                }
                 const mfem::IntegrationPoint &integrationPoint = data.integrationRule->IntPoint(quadraturePoint);
                 data.baseMappingContexts.Load(quadraturePoint, mappingContext);
 
@@ -1279,7 +1422,7 @@ namespace mean_field::operators {
         action.SetSize(m_context.GetEnthalpyMap().reduced_size());
         m_context.GetEnthalpyMap().gather(m_fullEnthalpyAction, action);
 
-        ++m_displacementJacobianStatistics.applications;
+        if (useFastPath && useTensorPath) ++m_displacementJacobianStatistics.applications;
     }
 
     void PreparedHydrostaticEquilibriumOperator::ApplyCompleteJacobianAction(

libmeanfield/impl/operators/prepared_mass_normalization.cpp

@@ -9,6 +9,7 @@ module;
 #include <utility>
 
 #include <mpi.h>
+#include "interior_displacement_contractions.hpp"
 
 module mean_field;
 
@@ -321,6 +322,7 @@ namespace mean_field::operators {
         const bool updateTargetMass = !m_isPrepared || dependencies.targetMass != m_preparedDependencies.targetMass ||
                                       state.targetMass != m_targetMass;
 
+        if (refreshGeometry || refreshDensity) m_jacobianRowsReady = false;
         m_isPrepared                = false;
 
         PreparedMassNormalizationReport report;
@@ -603,7 +605,83 @@ namespace mean_field::operators {
         ++m_residualApplicationCount;
     }
 
-    double PreparedMassNormalizationOperator::EvaluateDensityActionLocal(const mfem::Vector &densityVariation) const {
+    bool PreparedMassNormalizationOperator::PrepareJacobianRows() const {
+        VerifyPrepared();
+        if (m_jacobianRowsReady) return m_jacobianRowsSupported;
+        m_jacobianRowsSupported = false;
+        const int dimension = m_fem.mesh->Dimension();
+        for (const ElementPAData &data : m_elements) {
+            const auto &element = *m_fem.displacementFes->GetFE(data.elementId);
+            if (dimension < 1 || dimension > 3 || element.GetMapType() != mfem::FiniteElement::VALUE ||
+                element.GetDerivType() != mfem::FiniteElement::GRAD ||
+                m_domainMapper.IsCompactifiedElement(*m_fem.mesh->GetElementTransformation(data.elementId))) {
+                m_jacobianRowsReady = true;
+                return false;
+            }
+        }
+        m_densityJacobianLocal.SetSize(m_fem.densityFes->GetVSize());
+        m_displacementJacobianLocal.SetSize(m_fem.displacementFes->GetVSize());
+        m_densityJacobianLocal = 0.0;
+        m_displacementJacobianLocal = 0.0;
+        mfem::Vector densityRow, displacementRow;
+        mapping::VolumeMappingContext mappingContext;
+        constexpr double zeroPosition[3]{0.0, 0.0, 0.0};
+        for (const ElementPAData &data : m_elements) {
+            densityRow.SetSize(data.densityDofs.Size());
+            data.densityBasis->GetValues().MultTranspose(data.quadratureWeights, densityRow);
+            displacementRow.SetSize(data.displacementDofs.Size());
+            displacementRow = 0.0;
+            const auto basis = m_fem.GetReferenceTables().GetScalarTable(
+                *m_fem.displacementFes->GetFE(data.elementId), *data.integrationRule);
+            for (int q = 0; q < data.integrationRule->GetNPoints(); ++q) {
+                data.mappingContexts.Load(q, mappingContext);
+                add_interior_functional_row(
+                    basis->GetGradients(q), basis->GetValues(), q, mappingContext.quadrature.J_inv,
+                    data.density(q) * data.quadratureWeights(q), zeroPosition, displacementRow);
+            }
+            if (data.densityDofTransformation) data.densityDofTransformation->TransformDual(densityRow);
+            if (data.displacementDofTransformation) data.displacementDofTransformation->TransformDual(displacementRow);
+            m_densityJacobianLocal.AddElementVector(data.densityDofs, densityRow);
+            m_displacementJacobianLocal.AddElementVector(data.displacementDofs, displacementRow);
+        }
+        validate_finite_vector(m_densityJacobianLocal, "Prepared mass density row is non-finite.");
+        validate_finite_vector(m_displacementJacobianLocal, "Prepared mass displacement row is non-finite.");
+        m_jacobianRowsSupported = true;
+        m_jacobianRowsReady = true;
+        ++m_jacobianRowPreparationCount;
+        return true;
+    }
+
+    std::uint64_t PreparedMassNormalizationOperator::GetJacobianRowPreparationCount() const noexcept {
+        return m_jacobianRowPreparationCount;
+    }
+
+    double PreparedMassNormalizationOperator::EvaluateDensityActionLocal(const mfem::Vector &direction) const {
+        if (!PrepareJacobianRows()) return EvaluateDensityActionLocalReference(direction);
+        true_to_local(*m_fem.densityFes, direction, m_densityVariationLocal);
+        return m_densityJacobianLocal * m_densityVariationLocal;
+    }
+
+    double PreparedMassNormalizationOperator::EvaluateDisplacementActionLocal(const mfem::Vector &direction) const {
+        if (!PrepareJacobianRows()) return EvaluateDisplacementActionLocalReference(direction);
+        true_to_local(*m_fem.displacementFes, direction, m_displacementVariationLocal);
+        return m_displacementJacobianLocal * m_displacementVariationLocal;
+    }
+
+    void PreparedMassNormalizationOperator::ApplyCompleteJacobianActionReference(
+        const mfem::Vector &densityVariation, const mfem::Vector &displacementVariation, mfem::Vector &action
+    ) const {
+        VerifyPrepared();
+        validate_finite_vector(densityVariation, "Mass-normalization density direction is non-finite.");
+        validate_finite_vector(displacementVariation, "Mass-normalization displacement direction is non-finite.");
+        m_gravityContext.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);
+        m_gravityContext.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);
+        action.SetSize(1);
+        action(0) = GlobalSum(EvaluateDensityActionLocalReference(m_densityVariationTrue) +
+                              EvaluateDisplacementActionLocalReference(m_displacementVariationTrue));
+    }
+
+    double PreparedMassNormalizationOperator::EvaluateDensityActionLocalReference(const mfem::Vector &densityVariation) const {
         MFEM_VERIFY(
             densityVariation.Size() == m_fem.densityFes->GetTrueVSize(),
             "Mass-normalization density action received a vector with the "
@@ -635,7 +713,7 @@ namespace mean_field::operators {
         return localAction;
     }
 
-    double PreparedMassNormalizationOperator::EvaluateDisplacementActionLocal(
+    double PreparedMassNormalizationOperator::EvaluateDisplacementActionLocalReference(
         const mfem::Vector &displacementVariation
     ) const {
         MFEM_VERIFY(
@@ -898,8 +976,18 @@ namespace mean_field::operators {
         VerifyPrepared();
         MFEM_VERIFY(std::isfinite(residualDual), "Mass-normalization transpose action received a non-finite dual.");
 
-        AssembleDensityTransposeAction(residualDual, densityDual);
-        AssembleDisplacementTransposeAction(residualDual, displacementDual);
+        if (PrepareJacobianRows()) {
+            mfem::Vector trueDual;
+            local_to_true(*m_fem.densityFes, m_densityJacobianLocal, trueDual);
+            densityDual = m_gravityContext.GetDensityMap().gather(trueDual);
+            densityDual *= residualDual;
+            local_to_true(*m_fem.displacementFes, m_displacementJacobianLocal, trueDual);
+            displacementDual = m_gravityContext.GetDisplacementMap().gather(trueDual);
+            displacementDual *= residualDual;
+        } else {
+            AssembleDensityTransposeAction(residualDual, densityDual);
+            AssembleDisplacementTransposeAction(residualDual, displacementDual);
+        }
         ++m_actionStatistics.transposeApplications;
     }
 

libmeanfield/impl/operators/prepared_pressure_force.cpp

@@ -1,5 +1,8 @@
 module;
 
+#include "interior_displacement_contractions.hpp"
+
+#include <algorithm>
 #include <array>
 #include <cmath>
 #include <cstdint>
@@ -406,13 +409,42 @@ namespace mean_field::operators {
     PreparedPressureForceOperator::TryPrepare(
         const context::pressure_force::PressureForceStateView &state,
         const context::pressure_force::PressureForceDependencies &dependencies
+    ) {
+        return TryPrepareImpl(state, dependencies, true);
+    }
+
+    PreparedPressureForceReport PreparedPressureForceOperator::PrepareReference(
+        const context::pressure_force::PressureForceStateView &state,
+        const context::pressure_force::PressureForceDependencies &dependencies
+    ) {
+        auto result = TryPrepareImpl(state, dependencies, false, false);
+        if (!result.has_value()) throw_rejection(result.error());
+        return std::move(result).value();
+    }
+
+    PreparedPressureForceReport PreparedPressureForceOperator::PrepareCachedReference(
+        const context::pressure_force::PressureForceStateView &state,
+        const context::pressure_force::PressureForceDependencies &dependencies
+    ) {
+        auto result = TryPrepareImpl(state, dependencies, true, false);
+        if (!result.has_value()) throw_rejection(result.error());
+        return std::move(result).value();
+    }
+
+    std::expected<PreparedPressureForceReport, PressureForcePreparationRejection>
+    PreparedPressureForceOperator::TryPrepareImpl(
+        const context::pressure_force::PressureForceStateView &state,
+        const context::pressure_force::PressureForceDependencies &dependencies,
+        const bool useFastPreparation, const bool useTensorPreparation
     ) {
         const bool wasPrepared = m_isPrepared;
+        const bool implementationChanged = m_preparedWithFastMaterialAssembly != useFastPreparation ||
+            m_preparedWithTensorPreparation != useTensorPreparation;
         PreparedPressureForceReport report;
 
         report.contextReport = m_context.Prepare(state, dependencies);
 
-        if (!report.contextReport.DidAnyWork() && wasPrepared) {
+        if (!report.contextReport.DidAnyWork() && wasPrepared && !implementationChanged) {
             return report;
         }
 
@@ -433,16 +465,16 @@ namespace mean_field::operators {
             PrepareStaticPlan();
         }
 
-        if (report.contextReport.preparedGeometryState || !wasPrepared) {
-            const auto globalGeometryFailure = synchronize_rejection(PrepareGeometry(), m_fem.enthalpyFes->GetComm());
+        if (report.contextReport.preparedGeometryState || !wasPrepared || implementationChanged) {
+            const auto globalGeometryFailure = synchronize_rejection(PrepareGeometry(useTensorPreparation), m_fem.enthalpyFes->GetComm());
             if (globalGeometryFailure.has_value()) {
                 return std::unexpected(*globalGeometryFailure);
             }
         }
 
         std::optional<PressureForcePreparationRejection> localMaterialFailure;
-        if (report.contextReport.preparedMaterialState || !wasPrepared) {
-            localMaterialFailure = PrepareMaterialState();
+        if (report.contextReport.preparedMaterialState || !wasPrepared || implementationChanged) {
+            localMaterialFailure = PrepareMaterialState(useFastPreparation, useTensorPreparation);
         }
 
         const auto globalMaterialFailure = synchronize_rejection(localMaterialFailure, m_fem.enthalpyFes->GetComm());
@@ -450,7 +482,7 @@ namespace mean_field::operators {
             return std::unexpected(*globalMaterialFailure);
         }
 
-        if (report.contextReport.preparedMaterialState || !wasPrepared) {
+        if (report.contextReport.preparedMaterialState || !wasPrepared || implementationChanged) {
             FinalizeDisplacementJacobianPreparation();
 
             const auto globalAssemblyFailure =
@@ -482,6 +514,8 @@ namespace mean_field::operators {
             "supported displacement size."
         );
 
+        m_preparedWithTensorPreparation = useTensorPreparation;
+        m_preparedWithFastMaterialAssembly = useFastPreparation;
         m_isPrepared = true;
 
         return report;
@@ -564,7 +598,7 @@ namespace mean_field::operators {
         }
     }
 
-    std::optional<PressureForcePreparationRejection> PreparedPressureForceOperator::PrepareGeometry() {
+    std::optional<PressureForcePreparationRejection> PreparedPressureForceOperator::PrepareGeometry(const bool useTensorPreparation) {
         mfem::Vector displacementLocal;
 
         true_to_local(*m_fem.displacementFes, m_baseDisplacementTrue, displacementLocal);
@@ -574,6 +608,8 @@ namespace mean_field::operators {
 
         mfem::Vector elementDisplacement;
         mfem::Vector elementCompactification;
+        mfem::Vector sampleValue(3), tensorWorkspace;
+        mfem::DenseMatrix sampleGradient(3), tensorSamples;
 
         for (ElementPAData &data : m_elements) {
             MFEM_VERIFY(data.integrationRule != nullptr, "Prepared pressure-force geometry has no integration rule.");
@@ -626,13 +662,24 @@ namespace mean_field::operators {
 
             data.physicalTestGradients.resize(quadraturePointCount);
 
+            const bool tensor = useTensorPreparation && data.displacementReferenceTable &&
+                                data.displacementReferenceTable->HasTensorEvaluation3D();
+            if (tensor) data.displacementReferenceTable->EvaluateTensor3D(
+                data.baseDisplacementData->GetDofMatrix(), tensorSamples, tensorWorkspace);
+
             for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
                 const mfem::IntegrationPoint &integrationPoint = data.integrationRule->IntPoint(quadraturePoint);
 
                 transformation->SetIntPoint(&integrationPoint);
 
+                if (tensor) for (int c = 0; c < 3; ++c) {
+                    sampleValue(c) = tensorSamples(quadraturePoint, 4*c);
+                    for (int a = 0; a < 3; ++a)
+                        sampleGradient(c,a) = tensorSamples(quadraturePoint, 4*c+a+1);
+                }
+                const mapping::DomainMapper::ReferenceDisplacementSample sample{sampleValue, sampleGradient};
                 const mapping::MappingStatus mappingStatus = m_domainMapper.EvaluateVolume(
-                    mappingData, *transformation, integrationPoint, workspace, mappingContext
+                    mappingData, *transformation, integrationPoint, workspace, mappingContext, tensor ? &sample : nullptr
                 );
 
                 if (mappingStatus != mapping::MappingStatus::valid) {
@@ -675,7 +722,9 @@ namespace mean_field::operators {
         return std::nullopt;
     }
 
-    std::optional<PressureForcePreparationRejection> PreparedPressureForceOperator::PrepareMaterialState() {
+    std::optional<PressureForcePreparationRejection> PreparedPressureForceOperator::PrepareMaterialState(
+        const bool useFastPreparation, const bool useTensorPreparation
+    ) {
         mfem::Vector enthalpyLocal;
 
         true_to_local(*m_fem.enthalpyFes, m_baseEnthalpyTrue, enthalpyLocal);
@@ -687,6 +736,8 @@ namespace mean_field::operators {
 
         const mfem::Ordering::Type displacementOrdering = m_fem.displacementFes->GetOrdering();
         std::optional<PressureForcePreparationRejection> materialFailure;
+        mfem::DenseMatrix weightedGradients;
+        mfem::DenseMatrix tensorFactors, tensorInverse;
 
         for (ElementPAData &data : m_elements) {
             enthalpyLocal.GetSubVector(data.enthalpyDofs, elementEnthalpy);
@@ -717,7 +768,9 @@ namespace mean_field::operators {
 
             quadratureEnthalpy.SetSize(quadraturePointCount);
 
-            enthalpyBasis.Mult(elementEnthalpy, quadratureEnthalpy);
+            if (useTensorPreparation && data.enthalpyReferenceTable->HasTensorEvaluation3D())
+                data.enthalpyReferenceTable->EvaluateValuesTensor3D(elementEnthalpy, quadratureEnthalpy, m_tensorWorkspace);
+            else enthalpyBasis.Mult(elementEnthalpy, quadratureEnthalpy);
 
             data.pressure.SetSize(quadraturePointCount);
 
@@ -796,14 +849,72 @@ namespace mean_field::operators {
 
                         data.elementResidual(vectorDof) -= residualContribution;
 
-                        for (int enthalpyDof = 0; enthalpyDof < enthalpyDofCount; ++enthalpyDof) {
-                            const double jacobianContribution =
-                                pressureDerivative * weightedTestGradient * enthalpyBasis(quadraturePoint, enthalpyDof);
-                            if (!std::isfinite(jacobianContribution)) {
-                                materialFailure = non_finite_rejection();
-                                continue;
+                        if (!useFastPreparation) {
+                            for (int enthalpyDof = 0; enthalpyDof < enthalpyDofCount; ++enthalpyDof) {
+                                const double jacobianContribution = pressureDerivative * weightedTestGradient *
+                                    enthalpyBasis(quadraturePoint, enthalpyDof);
+                                if (!std::isfinite(jacobianContribution)) {
+                                    materialFailure = non_finite_rejection();
+                                    continue;
+                                }
+                                data.enthalpyJacobian(vectorDof, enthalpyDof) -= jacobianContribution;
+                            }
+                        }
+                    }
+                }
+            }
+
+            if (useFastPreparation && useTensorPreparation &&
+                data.displacementReferenceTable->HasTensorEvaluation3D() && !materialFailure.has_value()) {
+                // Build each enthalpy column by a tensor gradient transpose.
+                // These are the same -w P'(h) grad(test) B_h entries.
+                tensorFactors.SetSize(quadraturePointCount, 9);
+                m_tensorDuals.SetSize(quadraturePointCount, 9);
+                for (int q = 0; q < quadraturePointCount; ++q) {
+                    data.baseMappingContexts.LoadInverseJacobian(q, tensorInverse);
+                    const double w = -data.pressureDerivative(q)*data.quadratureWeights(q);
+                    for (int c = 0; c < 3; ++c) for (int a = 0; a < 3; ++a)
+                        tensorFactors(q, 3*c+a) = w*tensorInverse(a,c);
+                }
+                for (int col = 0; col < enthalpyDofCount; ++col) {
+                    for (int a = 0; a < 9; ++a) for (int q = 0; q < quadraturePointCount; ++q)
+                        m_tensorDuals(q,a) = tensorFactors(q,a)*enthalpyBasis(q,col);
+                    data.displacementReferenceTable->ApplyGradientTransposeTensor3D(
+                        m_tensorDuals, m_tensorElement, m_tensorWorkspace);
+                    for (int c = 0; c < 3; ++c) for (int i = 0; i < scalarDisplacementDofCount; ++i)
+                        data.enthalpyJacobian(i+c*scalarDisplacementDofCount,col) = m_tensorElement(i,c);
+                }
+            } else if (useFastPreparation && !materialFailure.has_value()) {
+                // J_h = -sum_q w_q P'(h_q) grad(test)_q B_h(q).
+                // Pack once, then accumulate small contiguous output tiles.
+                // The old loop is retained above as the independent reference.
+                weightedGradients.SetSize(displacementDofCount, quadraturePointCount);
+                for (int q = 0; q < quadraturePointCount; ++q) {
+                    for (int c = 0; c < dimension; ++c) {
+                        for (int i = 0; i < scalarDisplacementDofCount; ++i) {
+                            const int row = vector_dof_index(
+                                displacementOrdering, i, c, scalarDisplacementDofCount, dimension);
+                            weightedGradients(row, q) = -data.pressureDerivative(q) *
+                                (data.quadratureWeights(q) * data.physicalTestGradients[q](i, c));
+                        }
+                    }
+                }
+                const double *g = weightedGradients.Data();
+                const double *b = enthalpyBasis.Data();
+                double *j = data.enthalpyJacobian.Data();
+                constexpr int rowBlock = 32, columnBlock = 8;
+                for (int colBegin = 0; colBegin < enthalpyDofCount; colBegin += columnBlock) {
+                    const int colEnd = std::min(colBegin + columnBlock, enthalpyDofCount);
+                    for (int rowBegin = 0; rowBegin < displacementDofCount; rowBegin += rowBlock) {
+                        const int rowEnd = std::min(rowBegin + rowBlock, displacementDofCount);
+                        for (int q = 0; q < quadraturePointCount; ++q) {
+                            const double *gradient = g + q * displacementDofCount;
+                            for (int col = colBegin; col < colEnd; ++col) {
+                                const double basis = b[q + col * quadraturePointCount];
+                                double *column = j + col * displacementDofCount;
+                                for (int row = rowBegin; row < rowEnd; ++row)
+                                    column[row] += gradient[row] * basis;
                             }
-                            data.enthalpyJacobian(vectorDof, enthalpyDof) -= jacobianContribution;
                         }
                     }
                 }
@@ -938,8 +1049,32 @@ namespace mean_field::operators {
     }
 
     void PreparedPressureForceOperator::ApplyDisplacementJacobianAction(
-        const mfem::Vector &displacementVariation,
-        mfem::Vector &action
+        const mfem::Vector &displacementVariation, mfem::Vector &action
+    ) const {
+        ApplyDisplacementJacobianActionImpl(displacementVariation, action, true);
+    }
+
+    void PreparedPressureForceOperator::ApplyDisplacementJacobianActionReference(
+        const mfem::Vector &displacementVariation, mfem::Vector &action
+    ) const {
+        ApplyDisplacementJacobianActionImpl(displacementVariation, action, false, false);
+    }
+
+    void PreparedPressureForceOperator::ApplyDisplacementJacobianActionCachedReference(
+        const mfem::Vector &displacementVariation, mfem::Vector &action
+    ) const {
+        ApplyDisplacementJacobianActionImpl(displacementVariation, action, true, false);
+    }
+
+    int PreparedPressureForceOperator::GetTensorElementCount() const noexcept {
+        int count = 0;
+        for (const auto &data : m_elements)
+            count += data.displacementReferenceTable->HasTensorEvaluation3D();
+        return count;
+    }
+
+    void PreparedPressureForceOperator::ApplyDisplacementJacobianActionImpl(
+        const mfem::Vector &displacementVariation, mfem::Vector &action, const bool useFastPath, const bool useTensorPath
     ) const {
         VerifyPrepared();
 
@@ -1017,7 +1152,39 @@ namespace mean_field::operators {
             matrixTemporary.SetSize(dimension, dimension);
             physicalTestGradientVariation.SetSize(scalarDisplacementDofCount, dimension);
 
-            for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
+            const bool tensor = useTensorPath && data.displacementReferenceTable->HasTensorEvaluation3D();
+            if (tensor) {
+                data.displacementReferenceTable->EvaluateTensor3D(directionDofs, m_tensorSamples, m_tensorWorkspace);
+                m_tensorDuals.SetSize(quadraturePointCount, 9);
+                for (int q = 0; q < quadraturePointCount; ++q) {
+                    data.baseMappingContexts.LoadInverseJacobian(q, inverseElementJacobian);
+                    double gradient[3][3]{};
+                    for (int c = 0; c < 3; ++c) for (int j = 0; j < 3; ++j)
+                        for (int a = 0; a < 3; ++a)
+                            gradient[c][j] += m_tensorSamples(q, 4*c+a+1)*inverseElementJacobian(a,j);
+                    const double trace = gradient[0][0]+gradient[1][1]+gradient[2][2];
+                    const double weight = data.pressure(q)*data.quadratureWeights(q);
+                    // d(-P grad(test) dx): grad_ref(test) J^-1 (du/dx - tr(du/dx) I).
+                    for (int c = 0; c < 3; ++c) for (int a = 0; a < 3; ++a) {
+                        double value = 0.0;
+                        for (int j = 0; j < 3; ++j)
+                            value += inverseElementJacobian(a,j)*(gradient[j][c]-(j==c ? trace : 0.0));
+                        m_tensorDuals(q, 3*c+a) = weight*value;
+                    }
+                }
+                data.displacementReferenceTable->ApplyGradientTransposeTensor3D(
+                    m_tensorDuals, m_tensorElement, m_tensorWorkspace);
+                for (int c = 0; c < 3; ++c) for (int i = 0; i < scalarDisplacementDofCount; ++i)
+                    elementAction(i+c*scalarDisplacementDofCount) = m_tensorElement(i,c);
+            }
+
+            for (int quadraturePoint = 0; !tensor && quadraturePoint < quadraturePointCount; ++quadraturePoint) {
+                if (useFastPath && dimension == 3) {
+                    pressure_displacement_action_3d(
+                        directionDofs, data.physicalTestGradients[quadraturePoint],
+                        data.pressure(quadraturePoint) * data.quadratureWeights(quadraturePoint), elementAction);
+                    continue;
+                }
                 const mfem::DenseMatrix &referenceTestGradient =
                     data.displacementReferenceTable->GetGradients(quadraturePoint);
 
@@ -1076,7 +1243,7 @@ namespace mean_field::operators {
 
         m_displacementMap.gather(m_fullDisplacementAction, action);
 
-        ++m_displacementJacobianStatistics.applications;
+        if (useFastPath) ++m_displacementJacobianStatistics.applications;
     }
 
     void PreparedPressureForceOperator::ApplyCompleteJacobianAction(

libmeanfield/impl/operators/prepared_rotation_displacement_force.cpp

@@ -1,5 +1,7 @@
 module;
 
+#include "interior_displacement_contractions.hpp"
+
 #include <array>
 #include <cmath>
 #include <expected>
@@ -14,6 +16,7 @@ module mean_field;
 
 import :operators.kernels.rotational_displacement_force;
 import :operators.prepared_rotational_displacement_force;
+import :fem.reference_tables;
 
 namespace {
     using DomainSchema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;
@@ -194,7 +197,7 @@ namespace mean_field::operators {
     std::expected<
         void,
         kernels::RotationalDisplacementForceRejection>
-    PreparedRotationalDisplacementForceOperator::TryPrepareElementData() {
+    PreparedRotationalDisplacementForceOperator::TryPrepareElementData(const bool assembleResidual) {
         MFEM_VERIFY(m_rotation.has_value(), "Prepared rotational force has no frozen rotation state.");
 
         m_elements.clear();
@@ -214,7 +217,20 @@ namespace mean_field::operators {
         mfem::Vector densityShape;
         mfem::Vector potentialGradient;
 
+        if (assembleResidual) {
+            m_localAction.SetSize(m_fem.displacementFes->GetVSize());
+            m_localAction = 0.0;
+        }
+        mfem::Vector residualComponent;
         const int dimension = m_domainMapper.GetDimension();
+        // Rigid rotation has a position-independent acceleration derivative.
+        m_centrifugalJacobian.SetSize(dimension, dimension);
+        mfem::Vector unit(dimension), derivative(dimension);
+        for (int c = 0; c < dimension; ++c) {
+            unit = 0.0; unit(c) = 1.0;
+            m_rotation->potential_gradient_directional_derivative(unit, derivative);
+            for (int r = 0; r < dimension; ++r) m_centrifugalJacobian(r, c) = -derivative(r);
+        }
 
         for (int elementId = 0; elementId < m_fem.mesh->GetNE(); ++elementId) {
             mfem::ElementTransformation *transformation = m_fem.mesh->GetElementTransformation(elementId);
@@ -249,6 +265,11 @@ namespace mean_field::operators {
             const mfem::FiniteElement &displacementElement     = *m_fem.displacementFes->GetFE(elementId);
             const mfem::FiniteElement &compactificationElement = *m_fem.compactificationFes->GetFE(elementId);
             data.integrationRule                               = &get_rotation_force_rule(m_fem, *transformation);
+            data.densityReferenceTable = m_fem.GetReferenceTables().GetScalarTable(densityElement, *data.integrationRule);
+            data.displacementReferenceTable = m_fem.GetReferenceTables().GetScalarTable(displacementElement, *data.integrationRule);
+            data.fastApplication = dimension == 3 && m_fem.displacementFes->GetOrdering() == mfem::Ordering::byNODES &&
+                displacementElement.GetMapType() == mfem::FiniteElement::VALUE &&
+                displacementElement.GetDerivType() == mfem::FiniteElement::GRAD;
 
             const mapping::ElementDisplacementData displacementData =
                 mapping::ElementDisplacementDataFromElementVDofs(displacementElement, elementBaseDisplacement);
@@ -266,6 +287,7 @@ namespace mean_field::operators {
             data.quadratureWeights.SetSize(quadraturePointCount);
             densityShape.SetSize(densityElement.GetDof());
             potentialGradient.SetSize(dimension);
+            if (assembleResidual) m_quadratureForce.SetSize(quadraturePointCount, dimension);
 
             for (int quadraturePoint = 0; quadraturePoint < quadraturePointCount; ++quadraturePoint) {
                 const mfem::IntegrationPoint &integrationPoint = data.integrationRule->IntPoint(quadraturePoint);
@@ -304,8 +326,36 @@ namespace mean_field::operators {
                         }
                     }
                 }
+
+                if (assembleResidual) {
+                    const double weightedDensity =
+                        data.quadratureWeights(quadraturePoint) * data.baseDensityValues(quadraturePoint);
+                    for (int component = 0; component < dimension; ++component)
+                        m_quadratureForce(quadraturePoint, component) = -weightedDensity * potentialGradient(component);
+                }
+            }
+
+            if (assembleResidual) {
+                const int ndof = displacementElement.GetDof();
+                MFEM_VERIFY(data.displacementDofs.Size() == ndof * dimension,
+                            "Prepared force residual has an invalid displacement element size.");
+                m_elementAction.SetSize(data.displacementDofs.Size());
+                residualComponent.SetSize(ndof);
+                for (int component = 0; component < dimension; ++component) {
+                    data.displacementReferenceTable->GetValues().MultTranspose(
+                        m_quadratureForce.Data() + component * quadraturePointCount, residualComponent.GetData());
+                    for (int dof = 0; dof < ndof; ++dof)
+                        m_elementAction(vector_dof_index(m_fem.displacementFes->GetOrdering(),
+                                                         dof, component, ndof, dimension)) = residualComponent(dof);
+                }
+                if (!vector_is_finite(m_elementAction)) return std::unexpected(non_finite_rejection());
+                if (data.displacementDofTransformation != nullptr)
+                    data.displacementDofTransformation->TransformDual(m_elementAction);
+                m_localAction.AddElementVector(data.displacementDofs, m_elementAction);
             }
         }
+        if (assembleResidual && !vector_is_finite(m_localAction))
+            return std::unexpected(non_finite_rejection());
         return {};
     }
 
@@ -328,15 +378,36 @@ namespace mean_field::operators {
         const context::rotational_displacement_force::RotationalDisplacementForceStateView &state,
         const context::rotational_displacement_force::RotationalDisplacementForceDependencies &dependencies,
         const physics::RigidRotation &rotation
+    ) {
+        return TryPrepareImpl(state, dependencies, rotation, true);
+    }
+
+    PreparedRotationalDisplacementForceReport PreparedRotationalDisplacementForceOperator::PrepareReference(
+        const context::rotational_displacement_force::RotationalDisplacementForceStateView &state,
+        const context::rotational_displacement_force::RotationalDisplacementForceDependencies &dependencies,
+        const physics::RigidRotation &rotation
+    ) {
+        auto result = TryPrepareImpl(state, dependencies, rotation, false);
+        if (!result.has_value()) throw_rejection(result.error());
+        return std::move(result).value();
+    }
+
+    std::expected<PreparedRotationalDisplacementForceReport, kernels::RotationalDisplacementForceRejection>
+    PreparedRotationalDisplacementForceOperator::TryPrepareImpl(
+        const context::rotational_displacement_force::RotationalDisplacementForceStateView &state,
+        const context::rotational_displacement_force::RotationalDisplacementForceDependencies &dependencies,
+        const physics::RigidRotation &rotation,
+        const bool useFastPreparation
     ) {
         const bool wasPrepared = m_isPrepared;
+        const bool implementationChanged = m_preparedWithFastResidual != useFastPreparation;
         const bool rotationChanged =
             !m_context.IsPrepared() || dependencies.rotation != m_context.GetDependencies().rotation;
 
         PreparedRotationalDisplacementForceReport report;
         report.contextReport = m_context.Prepare(state, dependencies);
 
-        if (!report.contextReport.DidAnyWork() && wasPrepared) {
+        if (!report.contextReport.DidAnyWork() && wasPrepared && !implementationChanged) {
             return report;
         }
 
@@ -352,29 +423,38 @@ namespace mean_field::operators {
                                     "rotation state."
         );
 
-        if (report.contextReport.preparedBaseState || !wasPrepared) {
-            auto residualResult = kernels::try_apply_rotational_displacement_force_residual(
-                m_fem, m_domainMapper, *m_rotation, m_context.GetBaseDensityTrue(), m_context.GetDisplacementTrue(),
-                m_actionTrue
-            );
-            if (!residualResult.has_value()) {
-                return std::unexpected(residualResult.error());
+        if (report.contextReport.preparedBaseState || !wasPrepared || implementationChanged) {
+            const bool inputsFinite = vector_is_finite(m_context.GetBaseDensityTrue()) &&
+                                      vector_is_finite(m_context.GetDisplacementTrue());
+            const auto inputStatus = synchronize_rejection(
+                inputsFinite ? std::optional<Rejection>{} : std::optional<Rejection>{non_finite_rejection()},
+                m_fem.mesh->GetComm());
+            if (!inputStatus.has_value()) return std::unexpected(inputStatus.error());
+
+            // Retain the original two-pass path for an independent residual
+            // comparison. The production path evaluates the geometry once.
+            if (!useFastPreparation) {
+                auto residualResult = kernels::try_apply_rotational_displacement_force_residual(
+                    m_fem, m_domainMapper, *m_rotation, m_context.GetBaseDensityTrue(), m_context.GetDisplacementTrue(),
+                    m_actionTrue);
+                if (!residualResult.has_value()) return std::unexpected(residualResult.error());
             }
+            const auto elementResult = TryPrepareElementData(useFastPreparation);
+            const std::optional<Rejection> localRejection = elementResult.has_value()
+                ? std::optional<Rejection>{} : std::optional<Rejection>{elementResult.error()};
+            const auto elementStatus = synchronize_rejection(localRejection, m_fem.mesh->GetComm());
+            if (!elementStatus.has_value()) return std::unexpected(elementStatus.error());
+
+            if (useFastPreparation) local_to_true(*m_fem.displacementFes, m_localAction, m_actionTrue);
             m_cachedResidual.SetSize(m_context.GetDisplacementMap().reduced_size());
             m_context.GetDisplacementMap().gather(m_actionTrue, m_cachedResidual);
+            const auto residualStatus = synchronize_rejection(
+                vector_is_finite(m_cachedResidual) ? std::optional<Rejection>{}
+                    : std::optional<Rejection>{non_finite_rejection()},
+                m_fem.mesh->GetComm());
+            if (!residualStatus.has_value()) return std::unexpected(residualStatus.error());
 
-            const auto elementResult                = TryPrepareElementData();
-            std::optional<Rejection> localRejection = elementResult.has_value()
-                                                          ? std::optional<Rejection>{}
-                                                          : std::optional<Rejection>{elementResult.error()};
-            if (!vector_is_finite(m_cachedResidual)) {
-                localRejection = non_finite_rejection();
-            }
-            auto synchronized = synchronize_rejection(localRejection, m_fem.mesh->GetComm());
-            if (!synchronized.has_value()) {
-                return std::unexpected(synchronized.error());
-            }
-
+            m_preparedWithFastResidual = useFastPreparation;
             ++m_residualPreparationCount;
             report.preparedResidual = true;
         }
@@ -437,7 +517,8 @@ namespace mean_field::operators {
     void PreparedRotationalDisplacementForceOperator::ApplyPreparedCompleteJacobianActionTrue(
         const mfem::Vector &densityVariationTrue,
         const mfem::Vector &displacementVariationTrue,
-        mfem::Vector &actionTrue
+        mfem::Vector &actionTrue,
+        const bool useFastPath, const bool useTensorPath
     ) const {
         true_to_local(*m_fem.densityFes, densityVariationTrue, m_densityVariationLocal);
         true_to_local(*m_fem.displacementFes, displacementVariationTrue, m_displacementVariationLocal);
@@ -466,6 +547,61 @@ namespace mean_field::operators {
             const mfem::DenseMatrix &directionDofs = directionData.GetDofMatrix();
             const int scalarDisplacementDofCount   = displacementElement.GetDof();
 
+            if (useFastPath && data.fastApplication) {
+                const int nq = data.integrationRule->GetNPoints();
+                m_quadratureDensityVariation.SetSize(nq);
+                if (useTensorPath && data.densityReferenceTable->HasTensorEvaluation3D())
+                    data.densityReferenceTable->EvaluateValuesTensor3D(m_elementDensityVariation, m_quadratureDensityVariation, m_tensorWorkspace);
+                else data.densityReferenceTable->GetValues().Mult(m_elementDensityVariation, m_quadratureDensityVariation);
+                m_quadraturePositionVariation.SetSize(nq, 3);
+                m_quadratureForce.SetSize(nq, 3);
+                double *position = m_quadraturePositionVariation.Data();
+                double *force = m_quadratureForce.Data();
+                const double *u = directionDofs.Data();
+                const auto &values = data.displacementReferenceTable->GetValues();
+                const bool tensor = useTensorPath && data.displacementReferenceTable->HasTensorEvaluation3D();
+                if (tensor) {
+                    data.displacementReferenceTable->EvaluateTensor3D(directionDofs, m_tensorDisplacement, m_tensorWorkspace);
+                    for (int c = 0; c < 3; ++c) for (int q = 0; q < nq; ++q)
+                        position[q + c * nq] = m_tensorDisplacement(q, 4 * c);
+                } else {
+                    for (int c = 0; c < 3; ++c)
+                        values.Mult(u + c * scalarDisplacementDofCount, position + c * nq);
+                }
+                for (int q = 0; q < nq; ++q) {
+                    double trace = 0.0;
+                    if (tensor) {
+                        for (int row = 0; row < 3; ++row) for (int column = 0; column < 3; ++column)
+                            trace += data.inverseElementJacobians(q, row * 3 + column) *
+                                m_tensorDisplacement(q, 4 * column + row + 1);
+                    } else trace = displacement_trace_3d(directionDofs,
+                        data.displacementReferenceTable->GetGradients(q), data.inverseElementJacobians, q);
+                    const double rho = data.baseDensityValues(q);
+                    const double densityWeight = m_quadratureDensityVariation(q) + rho * trace;
+                    for (int c = 0; c < 3; ++c) {
+                        const double da = m_centrifugalJacobian(c, 0) * position[q] +
+                            m_centrifugalJacobian(c, 1) * position[q + nq] +
+                            m_centrifugalJacobian(c, 2) * position[q + 2 * nq];
+                        force[q + c * nq] = data.quadratureWeights(q) *
+                            (densityWeight * data.centrifugalAccelerations(q, c) + rho * da);
+                    }
+                }
+                m_elementAction.SetSize(data.displacementDofs.Size());
+                double *element = m_elementAction.HostWrite();
+                if (tensor) {
+                    data.displacementReferenceTable->ApplyTransposeTensor3D(m_quadratureForce, m_tensorElement, m_tensorWorkspace);
+                    for (int c = 0; c < 3; ++c) for (int i = 0; i < scalarDisplacementDofCount; ++i)
+                        element[i + c * scalarDisplacementDofCount] = m_tensorElement(i, c);
+                } else {
+                    for (int c = 0; c < 3; ++c)
+                        values.MultTranspose(force + c * nq, element + c * scalarDisplacementDofCount);
+                }
+                if (data.displacementDofTransformation)
+                    data.displacementDofTransformation->TransformDual(m_elementAction);
+                m_localAction.AddElementVector(data.displacementDofs, m_elementAction);
+                continue;
+            }
+
             m_densityShape.SetSize(densityElement.GetDof());
             m_displacementShape.SetSize(scalarDisplacementDofCount);
             m_referenceDisplacementDShape.SetSize(scalarDisplacementDofCount, dimension);
@@ -546,6 +682,50 @@ namespace mean_field::operators {
         ++m_completeJacobianStatistics.applications;
     }
 
+    void PreparedRotationalDisplacementForceOperator::ApplyCompleteJacobianActionReference(
+        const mfem::Vector &densityVariation,
+        const mfem::Vector &displacementVariation,
+        mfem::Vector &action
+    ) const {
+        VerifyPrepared();
+
+        m_densityVariationTrue.SetSize(m_context.GetDensityMap().full_size());
+        m_displacementVariationTrue.SetSize(m_context.GetDisplacementMap().full_size());
+        m_context.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);
+        m_context.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);
+
+        ApplyPreparedCompleteJacobianActionTrue(m_densityVariationTrue, m_displacementVariationTrue, m_actionTrue, false);
+        action.SetSize(m_context.GetDisplacementMap().reduced_size());
+        m_context.GetDisplacementMap().gather(m_actionTrue, action);
+
+    }
+
+    void PreparedRotationalDisplacementForceOperator::ApplyCompleteJacobianActionCachedReference(
+        const mfem::Vector &densityVariation,
+        const mfem::Vector &displacementVariation,
+        mfem::Vector &action
+    ) const {
+        VerifyPrepared();
+
+        m_densityVariationTrue.SetSize(m_context.GetDensityMap().full_size());
+        m_displacementVariationTrue.SetSize(m_context.GetDisplacementMap().full_size());
+        m_context.GetDensityMap().scatter(densityVariation, m_densityVariationTrue);
+        m_context.GetDisplacementMap().scatter(displacementVariation, m_displacementVariationTrue);
+
+        ApplyPreparedCompleteJacobianActionTrue(m_densityVariationTrue, m_displacementVariationTrue, m_actionTrue, true, false);
+        action.SetSize(m_context.GetDisplacementMap().reduced_size());
+        m_context.GetDisplacementMap().gather(m_actionTrue, action);
+
+    }
+
+    int PreparedRotationalDisplacementForceOperator::GetTensorElementCount() const noexcept {
+        int count = 0;
+        for (const auto &data : m_elements)
+            count += data.fastApplication && data.densityReferenceTable->HasTensorEvaluation3D() &&
+                data.displacementReferenceTable->HasTensorEvaluation3D();
+        return count;
+    }
+
     bool PreparedRotationalDisplacementForceOperator::IsPrepared() const noexcept {
         return m_isPrepared && m_rotation.has_value() && m_context.MatchesDependencies(m_preparedDependencies);
     }

libmeanfield/impl/preconditioning/gravity_field.cpp

@@ -3,6 +3,7 @@ module;
 #include <cmath>
 #include <memory>
 #include <mfem.hpp>
+#include <mpi.h>
 #include <stdexcept>
 
 module mean_field;
@@ -41,9 +42,33 @@ namespace mean_field::preconditioning {
         return assembled;
     }
 
+    const mfem::HypreParMatrix &PreparedGravityDivergenceCache::Get(const fem::FEM &f) {
+        if (!f.mesh || !f.gravityFluxFes || !f.gravityPotentialFes || !f.quadratureFactory)
+            throw std::invalid_argument("Fixed divergence caching requires a complete gravity discretization.");
+        const auto *mesh = f.mesh.get();
+        const auto *flux = f.gravityFluxFes.get(), *potential = f.gravityPotentialFes.get();
+        const int localChanged = !m_matrix || m_mesh != mesh || m_flux != flux || m_potential != potential ||
+            m_policy != f.quadratureFactory.get() || m_meshSequence != mesh->GetSequence() ||
+            m_nodesSequence != mesh->GetNodesSequence() || m_fluxSequence != flux->GetSequence() ||
+            m_potentialSequence != potential->GetSequence();
+        int changed = 0;
+        if (MPI_Allreduce(&localChanged, &changed, 1, MPI_INT, MPI_MAX, f.mesh->GetComm()) != MPI_SUCCESS)
+            throw std::runtime_error("Could not synchronize fixed divergence cache validity.");
+        if (changed) {
+            auto matrix = assembleGravityDivergenceSurrogate(f);
+            m_matrix = std::move(matrix);
+            m_mesh = mesh; m_flux = flux; m_potential = potential; m_policy = f.quadratureFactory.get();
+            m_meshSequence = mesh->GetSequence(); m_nodesSequence = mesh->GetNodesSequence();
+            m_fluxSequence = flux->GetSequence(); m_potentialSequence = potential->GetSequence();
+            ++m_builds;
+        }
+        return *m_matrix;
+    }
+
     std::unique_ptr<mfem::HypreParMatrix> assembleGravityPotentialSchurSurrogate(
         const fem::FEM &f,
-        const mfem::Vector &trueMassDiagonal
+        const mfem::Vector &trueMassDiagonal,
+        PreparedGravityDivergenceCache *divergenceCache
     ) {
         if (f.gravityFluxFes == nullptr || trueMassDiagonal.Size() != f.gravityFluxFes->GetTrueVSize()) {
             throw std::invalid_argument(
@@ -62,12 +87,14 @@ namespace mean_field::preconditioning {
             inverseMassDiagonal(index) = 1.0 / entry;
         }
 
-        std::unique_ptr<mfem::HypreParMatrix> divergence = assembleGravityDivergenceSurrogate(f);
-        std::unique_ptr<mfem::HypreParMatrix> inverseMassDivergenceTranspose(divergence->Transpose());
+        std::unique_ptr<mfem::HypreParMatrix> uncachedDivergence;
+        if (!divergenceCache) uncachedDivergence = assembleGravityDivergenceSurrogate(f);
+        const mfem::HypreParMatrix &divergence = divergenceCache ? divergenceCache->Get(f) : *uncachedDivergence;
+        std::unique_ptr<mfem::HypreParMatrix> inverseMassDivergenceTranspose(divergence.Transpose());
         inverseMassDivergenceTranspose->ScaleRows(inverseMassDiagonal);
 
         std::unique_ptr<mfem::HypreParMatrix> schur(
-            mfem::ParMult(divergence.get(), inverseMassDivergenceTranspose.get())
+            mfem::ParMult(&divergence, inverseMassDivergenceTranspose.get())
         );
         if (schur == nullptr) {
             throw std::runtime_error("MFEM did not assemble the gravity potential-Schur surrogate.");