From cfe56b9bd21033418e3cf3252a6bd8d69392f336 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mag=C3=AD=20Romany=C3=A0?= Date: Thu, 5 Dec 2024 12:27:53 +0100 Subject: [PATCH 1/2] Add Differentiability MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Magí Romanyà --- src/Mandos/Core/CMakeLists.txt | 6 +- src/Mandos/Core/DiffRigidBody.cpp | 69 +++ src/Mandos/Core/DiffRigidBody.hpp | 38 ++ src/Mandos/Core/Differentiable.cpp | 124 ++++++ src/Mandos/Core/Differentiable.hpp | 49 +++ .../Core/Energies/LumpedMassInertia.cpp | 27 ++ .../Core/Energies/LumpedMassInertia.hpp | 11 +- src/Mandos/Core/Energies/RigidBodyInertia.cpp | 126 ++++++ src/Mandos/Core/Energies/RigidBodyInertia.hpp | 21 + src/Mandos/Core/Model.cpp | 61 +++ src/Mandos/Core/Model.hpp | 36 +- src/Mandos/Core/differentiable.cpp | 394 ------------------ src/Mandos/Core/differentiable.hpp | 37 -- tests/CMakeLists.txt | 1 + .../RotationMatrixDerivatives/CMakeLists.txt | 8 + .../tst_RotationMatrixDerivatives.cpp | 221 ++++++++++ 16 files changed, 793 insertions(+), 436 deletions(-) create mode 100644 src/Mandos/Core/DiffRigidBody.cpp create mode 100644 src/Mandos/Core/DiffRigidBody.hpp create mode 100644 src/Mandos/Core/Differentiable.cpp create mode 100644 src/Mandos/Core/Differentiable.hpp delete mode 100644 src/Mandos/Core/differentiable.cpp delete mode 100644 src/Mandos/Core/differentiable.hpp create mode 100644 tests/RotationMatrixDerivatives/CMakeLists.txt create mode 100644 tests/RotationMatrixDerivatives/tst_RotationMatrixDerivatives.cpp diff --git a/src/Mandos/Core/CMakeLists.txt b/src/Mandos/Core/CMakeLists.txt index bd655c70..96bdd4e5 100644 --- a/src/Mandos/Core/CMakeLists.txt +++ b/src/Mandos/Core/CMakeLists.txt @@ -33,7 +33,8 @@ set(HEADERS Collisions/ContactEvent.hpp Collisions/CollisionDetection.hpp # async_simulation_loop.hpp - # differentiable.hpp + Differentiable.hpp + DiffRigidBody.hpp # edge.hpp # fem_element.hpp # gravity.hpp @@ -79,7 +80,8 @@ set(SOURCES Collisions/CollisionDetection.cpp # async_simulation_loop.cpp # compute_tetrahedrons.cpp - # differentiable.cpp + Differentiable.cpp + DiffRigidBody.cpp # edge.cpp # fem_element.cpp # inertia_energies.cpp diff --git a/src/Mandos/Core/DiffRigidBody.cpp b/src/Mandos/Core/DiffRigidBody.cpp new file mode 100644 index 00000000..ea4a78d6 --- /dev/null +++ b/src/Mandos/Core/DiffRigidBody.cpp @@ -0,0 +1,69 @@ +#include + +#include +#include + +namespace mandos::core +{ + +void applyGlobalToLocal(const Model &model, Vec &x) +{ + auto offset = 0; + for (auto node : model.freeSimulationObjects()) { + const auto &simObjectV = model.graph()[node]; + if (const auto *handle = std::get_if>>(&simObjectV)) { + const auto &simObj = handle->simulationObject(); + const auto size = simObj.mstate.size(); + for (auto i = 0; i < static_cast(simObj.mstate.m_x.size()); ++i) { + const Vec3 theta = handle->simulationObject().mstate.m_x[static_cast(i)].segment(3, 3); + const Mat3 &J = computeGlobalToLocalAxisAngleJacobian(theta); + x.segment<3>(offset + 6 * i + 3) = x.segment<3>(offset + 6 * i + 3).transpose() * J; + } + offset += size; + } + } +} + +void applyLocalToGlobal(const Model &model, Vec &x) +{ + auto offset = 0; + for (auto node : model.freeSimulationObjects()) { + const auto &simObjectV = model.graph()[node]; + if (const auto *handle = std::get_if>>(&simObjectV)) { + const auto &simObj = handle->simulationObject(); + const auto size = simObj.mstate.size(); + for (auto i = 0; i < static_cast(simObj.mstate.m_x.size()); ++i) { + const Vec3 theta = handle->simulationObject().mstate.m_x[static_cast(i)].segment(3, 3); + const Mat3 &J = computeLocalToGlobalAxisAngleJacobian(theta); + x.segment<3>(offset + 6 * i + 3) = x.segment<3>(offset + 6 * i + 3).transpose() * J; + } + offset += size; + } + } +} + +void applyComposeAxisAngleJacobian(Scalar h, const Model &model, Vec &v) +{ + unsigned int offset = 0; + for (auto node : model.freeSimulationObjects()) { + const auto &simObjectV = model.graph()[node]; + std::visit( + [&model, &offset, &v, h](const auto &handle) { + const auto &simulationObject = model.simulationObject(handle); + using SimulationObjectT = std::remove_reference_t>; + const auto size = simulationObject.mstate.size(); + if constexpr (std::is_same>()) { + for (unsigned int i = 0; i < simulationObject.mstate.m_x.size(); ++i) { + const Vec3 theta = simulationObject.mstate.m_x[i].template segment<3>(3); + const Vec3 omega = simulationObject.mstate.m_v[i].template segment<3>(3); + const Mat3 J = rotationExpMap(h * omega); + v.segment(offset + 6 * i + 3, 3) = v.segment(offset + 6 * i + 3, 3).transpose() * J; + } + } + offset += static_cast(size); + }, + simObjectV); + } +} + +} // namespace mandos::core diff --git a/src/Mandos/Core/DiffRigidBody.hpp b/src/Mandos/Core/DiffRigidBody.hpp new file mode 100644 index 00000000..4b2f11f3 --- /dev/null +++ b/src/Mandos/Core/DiffRigidBody.hpp @@ -0,0 +1,38 @@ +#ifndef MANDOS_CORE_DIFF_RIGID_BODY_H_ +#define MANDOS_CORE_DIFF_RIGID_BODY_H_ + +#include +#include + +namespace mandos::core +{ + +/** + * @brief Transform the angular segments of a derivative vector from Global to Local Axis Angle derivatives. + * + * Global Axis Angle derivatives are done at the origin tangent plane (Lie Algebra), with the full exponential map. + * Local Axis Angle derivatives are done at the point's tangent plane, with a linearization of the exponential map. Note + * that we use left incremental rotations: exp(theta) * R. + * + * @param model The simulation model + * @param x The derivative vector (size = nDof) + */ +void applyGlobalToLocal(const Model &model, Vec &x); + +/** + * @brief Transform the angular segments of a derivative vector from Local to Global Axis Angle derivatives. + * + * Global Axis Angle derivatives are done at the origin tangent plane (Lie Algebra), with the full exponential map. + * Local Axis Angle derivatives are done at the point's tangent plane, with a linearization of the exponential map. Note + * that we use left incremental rotations: exp(theta) * R. + * + * @param model The simulation model + * @param x The derivative vector (size = nDof) + */ +void applyLocalToGlobal(const Model &model, Vec &x); + +void applyComposeAxisAngleJacobian(Scalar h, const Model &model, Vec &v); + +} // namespace mandos::core + +#endif // MANDOS_CORE_DIFF_RIGID_BODY_H_ diff --git a/src/Mandos/Core/Differentiable.cpp b/src/Mandos/Core/Differentiable.cpp new file mode 100644 index 00000000..0c2bd47f --- /dev/null +++ b/src/Mandos/Core/Differentiable.cpp @@ -0,0 +1,124 @@ +#include +#include +#include + +namespace +{ + +mandos::core::Vec approximateHzFiniteDiff(mandos::core::Model &model, + const mandos::core::Vec &grad, + const mandos::core::Vec &x0, + const mandos::core::Vec &v0, + const mandos::core::Vec &z, + mandos::core::Scalar h, + mandos::core::Scalar dx) +{ + mandos::core::Vec gradient = mandos::core::Vec::Zero(model.nDof()); + mandos::core::Vec x_current = mandos::core::Vec::Zero(model.nDof()); + mandos::core::Vec v_current = mandos::core::Vec::Zero(model.nDof()); + + // Finite diff + model.state(x_current, v_current); + model.updateState(z * dx, x0, v0, h); + model.computeEnergyAndGradient(h, gradient); + + model.setState(x_current, v_current); // Recover state + + return (gradient - grad) / dx; +} + +} // namespace + +namespace mandos::core +{ + +Vec computeLossFunctionGradientBackpropagation(Scalar h, + Model &model, + const Trajectory &trajectory, + const LossFunctionAndGradients &loss, + const Mat &dx0_dp, + const Mat &dv0_dp, + unsigned int maxIterations) +{ + const int nParameters = loss.getNParameters(); + const int nDof = trajectory.getNDof(); + const int nSteps = trajectory.getNSteps(); + + // Initialize the loss function gradients dg_dp, dg_dx and dg_dv + // ------------------------------------------------------------------------- + Vec lossGradient = loss.lossParameterPartialDerivative; + Vec lossPositionGradient = loss.lossPositionPartialDerivative[static_cast(nSteps)]; + Vec lossVelocityGradient = loss.lossVelocityPartialDerivative[static_cast(nSteps)]; + // Backward loop + // ------------------------------------------------------------------------- + const Scalar one_over_h = 1.0 / h; + for (int i = nSteps - 1; i >= 0; --i) { + // Set the correc state + const Vec &x = trajectory.positions[static_cast(i) + 1]; + const Vec &v = trajectory.velocities[static_cast(i) + 1]; + const Vec &x0 = trajectory.positions[static_cast(i)]; + const Vec &v0 = trajectory.velocities[static_cast(i)]; + + model.setState(x0, v0); + model.computeAdvection(h); + model.setState(x, v); + + // Compute 2nd order derivatives + SystemMatrix hessian(static_cast(nDof), static_cast(nDof)); + // At the moment they ignore mappings + SparseMat dgradE_dx0(static_cast(nDof), static_cast(nDof)); + SparseMat dgradE_dv0(static_cast(nDof), static_cast(nDof)); + + Vec grad = Vec::Zero(nDof); + model.computeEnergyGradientAndHessian(h, grad, hessian); + Vec grad_test = Vec::Zero(nDof); + model.computeEnergyAndGradient(h, grad_test); + + // TODO Handle gradients with respect to paramters! + const Mat dgradE_dp = Mat::Zero(nDof, nParameters); + + const Vec equation_vector = -(lossPositionGradient + one_over_h * lossVelocityGradient); + + Eigen::ConjugateGradient solver; + solver.compute(hessian); + Vec z = solver.solve(equation_vector); // Adjoint + + // Multiple backward iterations + Vec h_grad = -equation_vector; + Vec dz = Vec::Zero(nDof); + constexpr Scalar dx = 1e-8; + for (unsigned int j = 0; j < maxIterations; ++j) { + const Vec Hz = approximateHzFiniteDiff(model, grad, x0, v0, z, h, dx); + if (Hz.hasNaN()) { + throw std::runtime_error( + "Back Propagation failed to converge with the current number of backward iterations."); + } + h_grad = Hz - equation_vector; + dz = solver.solve(-h_grad); + + z += dz; + } + + // Update the loss function gradients + // ------------------------------------------------------------------------- + applyComposeAxisAngleJacobian(h, model, lossVelocityGradient); + + model.computeEnergyRetardedPositionHessian(h, dgradE_dx0); + lossPositionGradient = loss.lossPositionPartialDerivative[static_cast(i)].transpose() - + one_over_h * lossVelocityGradient.transpose() + z.transpose() * dgradE_dx0; + + model.computeEnergyRetardedVelocityHessian(h, dgradE_dv0); + lossVelocityGradient = + loss.lossVelocityPartialDerivative[static_cast(i)].transpose() + z.transpose() * dgradE_dv0; + + lossGradient += z.transpose() * dgradE_dp; + } + // Add the initial conditions term + // ------------------------------------------------------------------------- + applyLocalToGlobal(model, lossPositionGradient); + lossGradient += lossPositionGradient.transpose() * dx0_dp + lossVelocityGradient.transpose() * dv0_dp; + + return lossGradient; +} + +} // namespace mandos::core diff --git a/src/Mandos/Core/Differentiable.hpp b/src/Mandos/Core/Differentiable.hpp new file mode 100644 index 00000000..e3383e28 --- /dev/null +++ b/src/Mandos/Core/Differentiable.hpp @@ -0,0 +1,49 @@ +#ifndef MANDOS_CORE_DIFFERENTIABLE_HPP +#define MANDOS_CORE_DIFFERENTIABLE_HPP + +#include +#include + +namespace mandos::core +{ +struct MANDOS_CORE_EXPORT Trajectory { + std::vector positions; + std::vector velocities; + + int getNDof() const + { + return static_cast(positions[0].size()); + } + int getNStates() const + { + return static_cast(positions.size()); + } + int getNSteps() const + { + return getNStates() - 1; + } +}; + +struct MANDOS_CORE_EXPORT LossFunctionAndGradients { + Scalar loss; // g + std::vector lossPositionPartialDerivative; // dg_dx_i + std::vector lossVelocityPartialDerivative; // dg_dv_i + Vec lossParameterPartialDerivative; // dg_dp + + int getNParameters() const + { + return static_cast(lossParameterPartialDerivative.size()); + } +}; + +MANDOS_CORE_EXPORT Vec computeLossFunctionGradientBackpropagation(Scalar h, + Model &model, + const Trajectory &trajectory, + const LossFunctionAndGradients &loss, + const Mat &dx0_dp, + const Mat &dv0_dp, + unsigned int maxIterations = 0); + +} // namespace mandos::core + +#endif // MANDOS_CORE_DIFFERENTIABLE_HPP diff --git a/src/Mandos/Core/Energies/LumpedMassInertia.cpp b/src/Mandos/Core/Energies/LumpedMassInertia.cpp index c2089f8f..b1a1738f 100644 --- a/src/Mandos/Core/Energies/LumpedMassInertia.cpp +++ b/src/Mandos/Core/Energies/LumpedMassInertia.cpp @@ -62,4 +62,31 @@ const std::vector &LumpedMassInertia::vertexMass() const { return m_vertexMass; } + +void LumpedMassInertia::computeEnergyRetardedPositionHessian(MechanicalState & /*unused*/, + Scalar h, + unsigned int offset, + std::vector &triplets) const +{ + for (auto i{0UL}; i < m_vertexMass.size(); ++i) { + const Scalar dgradE_dx0 = -1.0 / (h * h) * m_vertexMass[i]; + triplets.emplace_back(offset + 3 * i + 0, offset + 3 * i + 0, dgradE_dx0); + triplets.emplace_back(offset + 3 * i + 1, offset + 3 * i + 1, dgradE_dx0); + triplets.emplace_back(offset + 3 * i + 2, offset + 3 * i + 2, dgradE_dx0); + } +} + +void LumpedMassInertia::computeEnergyRetardedVelocityHessian(MechanicalState & /*unused*/, + Scalar h, + unsigned int offset, + std::vector &triplets) const +{ + for (auto i{0UL}; i < m_vertexMass.size(); ++i) { + const Scalar dgradE_dv0 = -1.0 / h * m_vertexMass[i]; + triplets.emplace_back(offset + 3 * i + 0, offset + 3 * i + 0, dgradE_dv0); + triplets.emplace_back(offset + 3 * i + 1, offset + 3 * i + 1, dgradE_dv0); + triplets.emplace_back(offset + 3 * i + 2, offset + 3 * i + 2, dgradE_dv0); + } +} + } // namespace mandos::core diff --git a/src/Mandos/Core/Energies/LumpedMassInertia.hpp b/src/Mandos/Core/Energies/LumpedMassInertia.hpp index 98bde4fe..39ba2a56 100644 --- a/src/Mandos/Core/Energies/LumpedMassInertia.hpp +++ b/src/Mandos/Core/Energies/LumpedMassInertia.hpp @@ -58,6 +58,15 @@ public: */ Scalar computeEnergyGradientAndHessian(MechanicalState &mstate, Scalar h) const; + void computeEnergyRetardedPositionHessian(MechanicalState &mstate, + Scalar h, + unsigned int offset, + std::vector &triplets) const; + void computeEnergyRetardedVelocityHessian(MechanicalState &mstate, + Scalar h, + unsigned int offset, + std::vector &triplets) const; + /** * @brief Read/Write accessor to vertex mass * @@ -79,4 +88,4 @@ private: }; } // namespace mandos::core -#endif // MANDOS_ENERGIES_LUMPEDMASSINERTIA_H \ No newline at end of file +#endif // MANDOS_ENERGIES_LUMPEDMASSINERTIA_H diff --git a/src/Mandos/Core/Energies/RigidBodyInertia.cpp b/src/Mandos/Core/Energies/RigidBodyInertia.cpp index e8279883..7138ae00 100644 --- a/src/Mandos/Core/Energies/RigidBodyInertia.cpp +++ b/src/Mandos/Core/Energies/RigidBodyInertia.cpp @@ -167,6 +167,74 @@ inline mandos::core::Mat6 rigidBodyGlobalHessian(const mandos::core::Vec6 &x, return hess; } +inline mandos::core::Mat6 rigidBodyLocalRetardedPositionHessian(const mandos::core::Mat3 &R, + const mandos::core::Vec3 &theta0, + const mandos::core::Vec3 &omega0, + mandos::core::Scalar mass, + const mandos::core::Mat3 &inertiaTensor, + mandos::core::Scalar h) +{ + using mandos::core::Scalar; + const mandos::core::Scalar h2 = h * h; + + const Eigen::Matrix vLeviCivita = mandos::core::vectorizedLeviCivita(); + // const Eigen::Matrix dvecRguess_dtheta0 = + // 2 * computeVectorizedRotationMatrixDerivativeGlobal(theta0) - + // computeVectorizedRotationMatrixDerivativeGlobal(theta0 - omega0 * h); + + const Eigen::Matrix dvecRguess_dtheta0 = + mandos::core::blockDiagonalMatrix<3>(2 * mandos::core::Mat3::Identity() - + mandos::core::rotationExpMap(h * omega0).transpose()) * + mandos::core::computeVectorizedRotationMatrixDerivativeLocal(theta0).transpose(); + + const Eigen::Matrix dvecRMRguess_dtheta0 = + mandos::core::blockMatrix<3>(R * inertiaTensor) * dvecRguess_dtheta0; + + const Eigen::Matrix dvecAdtheta0 = + 0.5 * (mandos::core::transposeVectorizedMatrixRows<9, 3>(dvecRMRguess_dtheta0) - dvecRMRguess_dtheta0); + + // Hessian + mandos::core::Mat6 hess = mandos::core::Mat6::Zero(); + hess.block<3, 3>(0, 0) = -mass / h2 * mandos::core::Mat3::Identity(); + hess.block<3, 3>(3, 3) = 1.0 / h2 * vLeviCivita * dvecAdtheta0; + + return hess; +} + +inline mandos::core::Mat6 rigidBodyLocalRetardedVelocityHessian(const mandos::core::Mat3 &R, + const mandos::core::Vec3 &theta0, + const mandos::core::Vec3 &omega0, + mandos::core::Scalar mass, + const mandos::core::Mat3 &inertiaTensor, + mandos::core::Scalar h) +{ + using mandos::core::Scalar; + const mandos::core::Scalar h2 = h * h; + const Eigen::Matrix vLeviCivita = mandos::core::vectorizedLeviCivita(); + + // const Eigen::Matrix dvecRguess_domega0 = + // h * computeVectorizedRotationMatrixDerivativeGlobal(theta0 - omega0 * h); + + const mandos::core::Mat3 R0 = mandos::core::rotationExpMap(theta0); + const Eigen::Matrix dvecRguess_domega0 = + -h * mandos::core::transposeVectorizedMatrixRows<9, 3>( + mandos::core::blockDiagonalMatrix<3>(R0.transpose()) * + mandos::core::computeVectorizedRotationMatrixDerivativeLocal(h * omega0).transpose()); + + const Eigen::Matrix dvecRMR_guess_domega0 = + mandos::core::blockMatrix<3>(R * inertiaTensor) * dvecRguess_domega0; + + const Eigen::Matrix dvecAdomega0 = + 0.5 * (mandos::core::transposeVectorizedMatrixRows<9, 3>(dvecRMR_guess_domega0) - dvecRMR_guess_domega0); + + // Hessian + mandos::core::Mat6 hess = mandos::core::Mat6::Zero(); + hess.block<3, 3>(0, 0) = -mass / h * mandos::core::Mat3::Identity(); + hess.block<3, 3>(3, 3) = 1.0 / h2 * vLeviCivita * dvecAdomega0; + + return hess; +} + } // namespace namespace mandos::core @@ -176,6 +244,9 @@ void mandos::core::RigidBodyInertia::advect(const MechanicalState { m_advX.resize(mstate.m_x.size()); m_advXR.resize(mstate.m_x.size()); + m_x0 = mstate.m_x; + m_v0 = mstate.m_v; + for (auto i = 0UL; i < mstate.m_x.size(); ++i) { rigidBodyLocalAdvect(mstate.m_x[i], mstate.m_v[i], h, m_advX[i], m_advXR[i]); } @@ -298,6 +369,61 @@ Scalar RigidBodyInertia::computeEnergyGradientAndHessian(MechanicalState &mstate, + Scalar h, + unsigned int offset, + std::vector &triplets) const +{ + for (unsigned int i = 0; i < mstate.m_x.size(); ++i) { + const Mat3 R = rotationExpMap(mstate.m_x[i].segment<3>(3)); + const Mat6 hess = rigidBodyLocalRetardedPositionHessian( + R, m_x0[i].segment<3>(3), m_v0[i].segment<3>(3), m_mass[i], m_inertiaTensor0[i], h); + for (unsigned int j = 0; j < 3; ++j) { + triplets.emplace_back(offset + 6 * i + j, offset + 6 * i + j, hess(j, j)); // Diagonal translation hessian + for (unsigned int k = 0; k < 3; ++k) { + triplets.emplace_back(offset + 6 * i + j + 3, + offset + 6 * i + k + 3, + hess(3 + j, 3 + k)); // Non-diagonal rotation hessian + } + } + } +} + +void RigidBodyInertia::computeEnergyRetardedVelocityHessian(MechanicalState &mstate, + Scalar h, + unsigned int offset, + std::vector &triplets) const +{ + for (unsigned int i = 0; i < mstate.m_x.size(); ++i) { + const Mat3 R = rotationExpMap(mstate.m_x[i].segment<3>(3)); + const Mat6 hess = rigidBodyLocalRetardedVelocityHessian( + R, m_x0[i].segment<3>(3), m_v0[i].segment<3>(3), m_mass[i], m_inertiaTensor0[i], h); + for (unsigned int j = 0; j < 3; ++j) { + triplets.emplace_back(offset + 6 * i + j, offset + 6 * i + j, hess(j, j)); // Diagonal translation hessian + for (unsigned int k = 0; k < 3; ++k) { + triplets.emplace_back(offset + 6 * i + j + 3, + offset + 6 * i + k + 3, + hess(3 + j, 3 + k)); // Non-diagonal rotation hessian + } + } + } +} + +void RigidBodyInertia::computeEnergyRetardedPositionHessian(MechanicalState & /*unused*/, + Scalar /*unused*/, + unsigned int /*unused*/, + std::vector & /*unused*/) const +{ + // TODO +} +void RigidBodyInertia::computeEnergyRetardedVelocityHessian(MechanicalState & /*unused*/, + Scalar /*unused*/, + unsigned int /*unused*/, + std::vector & /*unused*/) const +{ + // TODO +} + std::vector &mandos::core::RigidBodyInertia::inertiaTensor() { return m_inertiaTensor0; diff --git a/src/Mandos/Core/Energies/RigidBodyInertia.hpp b/src/Mandos/Core/Energies/RigidBodyInertia.hpp index de369ca6..8d1badbf 100644 --- a/src/Mandos/Core/Energies/RigidBodyInertia.hpp +++ b/src/Mandos/Core/Energies/RigidBodyInertia.hpp @@ -62,6 +62,24 @@ public: Scalar computeEnergyGradientAndHessian(MechanicalState &mstate, Scalar h) const; Scalar computeEnergyGradientAndHessian(MechanicalState &mstate, Scalar h) const; + void computeEnergyRetardedPositionHessian(MechanicalState &mstate, + Scalar h, + unsigned int offset, + std::vector &triplets) const; + void computeEnergyRetardedVelocityHessian(MechanicalState &mstate, + Scalar h, + unsigned int offset, + std::vector &triplets) const; + + void computeEnergyRetardedPositionHessian(MechanicalState &mstate, + Scalar h, + unsigned int offset, + std::vector &triplets) const; + void computeEnergyRetardedVelocityHessian(MechanicalState &mstate, + Scalar h, + unsigned int offset, + std::vector &triplets) const; + /** * @brief Read/Write accessor to rigid body mass * @@ -96,6 +114,9 @@ private: std::vector m_advX; std::vector m_advXR; + + std::vector m_x0; + std::vector m_v0; }; } // namespace mandos::core diff --git a/src/Mandos/Core/Model.cpp b/src/Mandos/Core/Model.cpp index 1defac69..a83410fd 100644 --- a/src/Mandos/Core/Model.cpp +++ b/src/Mandos/Core/Model.cpp @@ -474,6 +474,11 @@ void Model::updateColliders() m_simulationObjects); } +const std::vector &Model::freeSimulationObjects() const +{ + return m_freeSimulationObjects; +} + void Model::detectCollisions() { utilities::static_for_each( @@ -633,4 +638,60 @@ void Model::commit() }, m_simulationObjects); } + +void Model::computeEnergyRetardedPositionHessian(Scalar h, SparseMat &hessian) +{ + ZoneScopedN("Model.computeEnergyRetardedPositionHessian"); + std::vector triplets; + // // We compute the retarded position hessian for each simulation object with inertia + // Only compute free objects + unsigned int offset = 0; + for (auto node : m_freeSimulationObjects) { + const auto &simObjectV = m_simulationObjectsGraph[node]; + std::visit( + [this, &offset, &triplets, h](auto &handle) { + auto &simulationObject = this->simulationObject(handle); + using SimulationObjectT = std::remove_reference_t>; + const auto size = simulationObject.mstate.size(); + if constexpr (SimulationObjectT::hasInertias) { + utilities::static_for_each( + [&simulationObject, h, offset, &triplets](auto &inertia) { + inertia.computeEnergyRetardedPositionHessian(simulationObject.mstate, h, offset, triplets); + }, + simulationObject.inertias()); + } + offset += static_cast(size); + }, + simObjectV); + } + hessian.setFromTriplets(triplets.begin(), triplets.end()); +} + +void Model::computeEnergyRetardedVelocityHessian(Scalar h, SparseMat &hessian) +{ + ZoneScopedN("Model.computeEnergyRetardedVelocityHessian"); + // // We compute the retarded position hessian for each simulation object with inertia + std::vector triplets; + // Only compute free objects + unsigned int offset = 0; + for (auto node : m_freeSimulationObjects) { + const auto &simObjectV = m_simulationObjectsGraph[node]; + std::visit( + [this, &offset, &triplets, h](auto &handle) { + auto &simulationObject = this->simulationObject(handle); + using SimulationObjectT = std::remove_reference_t>; + const auto size = simulationObject.mstate.size(); + if constexpr (SimulationObjectT::hasInertias) { + utilities::static_for_each( + [&simulationObject, h, offset, &triplets](auto &inertia) { + inertia.computeEnergyRetardedVelocityHessian(simulationObject.mstate, h, offset, triplets); + }, + simulationObject.inertias()); + } + offset += static_cast(size); + }, + simObjectV); + } + hessian.setFromTriplets(triplets.begin(), triplets.end()); +} } // namespace mandos::core diff --git a/src/Mandos/Core/Model.hpp b/src/Mandos/Core/Model.hpp index 2f3cff87..270038c2 100644 --- a/src/Mandos/Core/Model.hpp +++ b/src/Mandos/Core/Model.hpp @@ -12,11 +12,9 @@ #include "Mandos/Core/Collisions/CollisionPair.hpp" #include "Mandos/Core/Collisions/SimulationCollider.hpp" -#include "Mandos/Core/KinematicGraph.hpp" #include "Mandos/Core/Mappings/BarycentricMapping.hpp" #include "Mandos/Core/MechanicalStates/Particle3D.hpp" #include "Mandos/Core/MechanicalStates/RigidBody.hpp" -#include "Mandos/Core/MechanicalStates/RigidBodyGlobal.hpp" #include #include @@ -256,12 +254,46 @@ public: */ Scalar computeEnergyGradientAndHessian(Scalar h, Vec &gradient, SystemMatrix &hessian); + /** + * @brief Computes the derivative of the energy gradient wrt the previous positions x0 + + * [!NOTE] The only energies that depend on previous positions are the ones that depend on velocities. This includes + inertia energies and dissipative or damping potentials. + * [!NOTE] This function is needed exclusively for differentiability. It is necessary to compute how does the state + after taking one simulation step changes if we change the previous positions x0. + * [!NOTE] This function at the moment does not support mapped objects and will only work for free simulation + objects + + * @param model The simulation model + * @param h Span of time to compute the energy model + * @param hessian The derivative of the energy gradient wrt x0 + */ + void computeEnergyRetardedPositionHessian(Scalar h, SparseMat &hessian); + + /** + * @brief Computes the derivative of the energy gradient wrt the previous velocities v0 + + * [!NOTE] The only energies that depend on previous velocities are the ones that depend on acceleration. This means + that only affects inertial energies. + * [!NOTE] This function is needed exclusively for differentiability. It is necessary to compute how does the state + after taking one simulation step changes if we change the previous velocity v0. + * [!NOTE] This function at the moment does not support mapped objects and will only work for free simulation + objects + + * @param model The simulation model + * @param h Span of time to compute the energy model + * @param hessian The derivative of the energy gradient wrt v0 + */ + void computeEnergyRetardedVelocityHessian(Scalar h, SparseMat &hessian); + void updateColliders(); void detectCollisions(); void updateCollisionMappings(); KinematicGraph &graph(); + const std::vector &freeSimulationObjects() const; + const KinematicGraph &graph() const; private: diff --git a/src/Mandos/Core/differentiable.cpp b/src/Mandos/Core/differentiable.cpp deleted file mode 100644 index 51cc5115..00000000 --- a/src/Mandos/Core/differentiable.cpp +++ /dev/null @@ -1,394 +0,0 @@ -#include -#include -#include // For inverse - -#include -#include -#include -#include -#include -#include -#include -#include - -#include - -namespace mandos::core -{ - -Mat3 rotation_inertia_energy_hessian(const Mat3 &J_inertia_tensor, - const Vec3 &theta, - const Vec3 &theta0, - const Vec3 &omega0, - Scalar TimeStep) -{ - const Mat3 R = compute_rotation_matrix_rodrigues(theta); - const Mat3 R0 = compute_rotation_matrix_rodrigues(theta0); - const Mat3 R0old = compute_rotation_matrix_rodrigues(theta0 - omega0 * TimeStep); - const Mat3 R_guess = (R0 + (R0 - R0old)); // x0 + h* v0 - - const Mat3 rot_inertia = R * J_inertia_tensor * R_guess.transpose(); - const Mat3 S = (rot_inertia + rot_inertia.transpose()) / 2; // Exact hessian - // const Mat3 S = R * J_inertia_tensor * R.transpose(); // Linear approximation - const Scalar h2 = TimeStep * TimeStep; - - const Mat3 hessian = 1.0 / h2 * (S.trace() * Mat3::Identity() - S); - return hessian; -} - -inline Mat3 rotation_inertia_dgradE_dtheta(const Mat3 &J_inertia_tensor, - const Vec3 &theta, - const Vec3 &theta0, - const Vec3 &omega0, - Scalar TimeStep) -{ - const Mat3 R0 = compute_rotation_matrix_rodrigues(theta0); - const Mat3 R0old = compute_rotation_matrix_rodrigues(theta0 - TimeStep * omega0); - const Mat3 Rguess = (R0 + (R0 - R0old)); // x0 + h* v0 - - const Scalar h2 = TimeStep * TimeStep; - const Eigen::Matrix vLeviCivita = vectorized_levi_civita(); - const Eigen::Matrix dvecR_dtheta = dvecR_dtheta_global(theta); - - const Eigen::Matrix dvecRMR_guess_dtheta = - block_matrix<3, 3>(Rguess * J_inertia_tensor) * dvecR_dtheta.transpose(); - const Eigen::Matrix dvecAdtheta = - 0.5 * (dvecRMR_guess_dtheta - transpose_vectorized_matrix_N<9, 3>(dvecRMR_guess_dtheta)); - - Mat3 H = 1.0 / h2 * vLeviCivita * dvecAdtheta; - return H; -} - -inline Mat3 rotation_inertia_dgradE_dtheta0(const Mat3 &J_inertia_tensor, - const Vec3 &theta, - const Vec3 &theta0, - const Vec3 &omega0, - Scalar TimeStep) -{ - const Mat3 R = compute_rotation_matrix_rodrigues(theta); - - const Scalar h2 = TimeStep * TimeStep; - const Eigen::Matrix vLeviCivita = vectorized_levi_civita(); - const Eigen::Matrix dvecRguess_dtheta0 = - 2 * dvecR_dtheta_global(theta0) - dvecR_dtheta_global(theta0 - omega0 * TimeStep); - - const Eigen::Matrix dvecRMR_guess_dtheta0 = - block_matrix<3, 3>(R * J_inertia_tensor) * dvecRguess_dtheta0.transpose(); - ; - const Eigen::Matrix dvecAdtheta0 = - 0.5 * (transpose_vectorized_matrix_N<9, 3>(dvecRMR_guess_dtheta0) - dvecRMR_guess_dtheta0); - - Mat3 H = 1.0 / h2 * vLeviCivita * dvecAdtheta0; - return H; -} - -inline Mat3 rotation_inertia_dgradE_domega0(const Mat3 &J_inertia_tensor, - const Vec3 &theta, - const Vec3 &theta0, - const Vec3 &omega0, - Scalar TimeStep) -{ - const Mat3 R = compute_rotation_matrix_rodrigues(theta); - - const Scalar h2 = TimeStep * TimeStep; - const Eigen::Matrix vLeviCivita = vectorized_levi_civita(); - const Eigen::Matrix dvecRguess_domega0 = TimeStep * dvecR_dtheta_global(theta0 - omega0 * TimeStep); - - const Eigen::Matrix dvecRMR_guess_domega0 = - block_matrix<3, 3>(R * J_inertia_tensor) * dvecRguess_domega0.transpose(); - ; - const Eigen::Matrix dvecAdtheta0 = - 0.5 * (transpose_vectorized_matrix_N<9, 3>(dvecRMR_guess_domega0) - dvecRMR_guess_domega0); - - Mat3 H = 1.0 / h2 * vLeviCivita * dvecAdtheta0; - return H; -} - -#define DIFF_LOCAL_RB -void compute_gradient_partial_derivatives(Scalar TimeStep, - const Energies &energies, - const PhysicsState &state, - const PhysicsState &state0, - SparseMat &dgradE_dx, - SparseMat &dgradE_dx0, - SparseMat &dgradE_dv0) -{ - const unsigned int nDoF = state.get_nDoF(); - EnergyAndDerivatives f(nDoF); - std::vector dgradE_dx0_triplets; - std::vector dgradE_dv0_triplets; - - // Linear inertias - // ---------------------------------------------------------------------------------------------------- - const Scalar one_over_h2 = 1.0f / (TimeStep * TimeStep); - for (unsigned int i = 0; i < energies.inertial_energies.linear_inertias.size(); i++) { - LinearInertia e = energies.inertial_energies.linear_inertias[i]; - const Mat3 dgradE_dx = one_over_h2 * e.Mass; - const Mat3 dgradE_dx0 = -dgradE_dx; - const Mat3 dgradE_dv0 = -TimeStep * dgradE_dx; - - for (unsigned int i = 0; i < 3; i++) { - for (unsigned int j = 0; j < 3; j++) { - f.hessian_triplets.emplace_back(e.p.index + i, e.p.index + j, dgradE_dx(i, j)); - dgradE_dx0_triplets.emplace_back(e.p.index + i, e.p.index + j, dgradE_dx0(i, j)); - dgradE_dv0_triplets.emplace_back(e.p.index + i, e.p.index + j, dgradE_dv0(i, j)); - } - } - } - - // Rotational inertias - // ---------------------------------------------------------------------------------------------------- - for (unsigned int i = 0; i < energies.inertial_energies.rotational_inertias.size(); i++) { - const RotationalInertia &e = energies.inertial_energies.rotational_inertias[i]; - const Mat3 J_inertia_tensor = e.rb.J_inertia_tensor0; - const Vec3 theta = e.rb.get_axis_angle(state.x); - const Vec3 theta0 = e.rb.get_axis_angle(state0.x); - const Vec3 omega0 = e.rb.get_axis_angle(state0.v); -#ifdef DIFF_LOCAL_RB - const Mat3 dgradE_dtheta = rotation_inertia_energy_hessian(J_inertia_tensor, theta, theta0, omega0, TimeStep); -#else - const Mat3 dgradE_dtheta = rotation_inertia_dgradE_dtheta(J_inertia_tensor, theta, theta0, omega0, TimeStep); -#endif - const Mat3 dgradE_dtheta0 = rotation_inertia_dgradE_dtheta0(J_inertia_tensor, theta, theta0, omega0, TimeStep); - const Mat3 dgradE_domega0 = rotation_inertia_dgradE_domega0(J_inertia_tensor, theta, theta0, omega0, TimeStep); - - for (unsigned int i = 0; i < 3; i++) { - for (unsigned int j = 0; j < 3; j++) { - f.hessian_triplets.emplace_back(e.rb.index + 3 + i, e.rb.index + 3 + j, dgradE_dtheta(i, j)); - dgradE_dx0_triplets.emplace_back(e.rb.index + 3 + i, e.rb.index + 3 + j, dgradE_dtheta0(i, j)); - dgradE_dv0_triplets.emplace_back(e.rb.index + 3 + i, e.rb.index + 3 + j, dgradE_domega0(i, j)); - } - } - } - - // Potential energies - // --------------------------------------------------------------------------------- - - energies.potential_energies.for_each(ComputePotentialEnergyAndDerivativesVisitor(TimeStep, state, f)); - - // NOTE: IGNORING POTENTIALS WITH VELOCITY DEPENDENCE - - dgradE_dx.setFromTriplets(f.hessian_triplets.begin(), f.hessian_triplets.end()); - dgradE_dx0.setFromTriplets(dgradE_dx0_triplets.begin(), dgradE_dx0_triplets.end()); - dgradE_dv0.setFromTriplets(dgradE_dv0_triplets.begin(), dgradE_dv0_triplets.end()); -} - -void compute_simulation_global_to_local_jacobian(const Simulables &simulables, const PhysicsState state, SparseMat &ptm) -{ - std::vector jac_triplets; - for (unsigned int i = 0; i < simulables.particles.size(); i++) { - const unsigned int idx = simulables.particles[i].index; - for (unsigned int j = 0; j < 3; j++) - jac_triplets.emplace_back(idx + j, idx + j, 1); // Diagonal matrix - } - for (unsigned int i = 0; i < simulables.rigid_bodies.size(); i++) { - const RigidBody &rb = simulables.rigid_bodies[i]; - const Vec3 theta = rb.get_axis_angle(state.x); - const Mat3 jac = compute_global_to_local_axis_angle_jacobian(theta); - - const unsigned int idx = rb.index + 3; - - for (unsigned int j = 0; j < 3; j++) { - jac_triplets.emplace_back(rb.index + j, rb.index + j, 1); // Translation diagonal - for (unsigned int k = 0; k < 3; k++) { - jac_triplets.emplace_back(idx + j, idx + k, jac(j, k)); - } - } - } - ptm.setFromTriplets(jac_triplets.begin(), jac_triplets.end()); -} - -inline Vec approximate_Hz_finite_diff(const Simulation &simulation, - const PhysicsState &state, - const PhysicsState &state0, - const Vec &grad0, - const Vec &z, - Scalar dx) -{ - PhysicsState diff_state = state; - update_simulation_state(simulation.TimeStep, simulation.energies, dx * z, diff_state, state0); - Vec diff_grad = compute_energy_gradient(simulation.TimeStep, simulation.energies, diff_state, state0); - return (diff_grad - grad0) / dx; -} - -Vec compute_loss_function_gradient_backpropagation(const Simulation &simulation, - const std::vector &trajectory, - const LossFunctionAndDerivatives &loss, - const Mat &dx0_dp, - const Mat &dv0_dp, - const unsigned int maxIterations) -{ - const unsigned int nParameters = static_cast(loss.loss_parameter_partial_derivative.size()); - const unsigned int nDoF = static_cast(trajectory.at(0).x.size()); - const unsigned int nStates = static_cast(trajectory.size()); - const unsigned int nSteps = nStates - 1; - - // Initialize the loss function gradients dg_dp, dg_dx and dg_dv - // ------------------------------------------------------------------------- - Vec loss_gradient = loss.loss_parameter_partial_derivative; - Vec loss_position_gradient = loss.loss_position_partial_derivative[nSteps]; - Vec loss_velocity_gradient = loss.loss_velocity_partial_derivative[nSteps]; - - // Backward loop - // ------------------------------------------------------------------------- - const Scalar one_over_h = 1.0f / simulation.TimeStep; - for (int i = nSteps - 1; i >= 0; i--) { - const PhysicsState &state = trajectory[i + 1]; - const PhysicsState &state0 = trajectory[i]; - - // Compute the hessian matrix and other sparse matrices - // ------------------------------------------------------------------------- - - SparseMat hessian(nDoF, nDoF), dgradE_dx0(nDoF, nDoF), dgradE_dv0(nDoF, nDoF); - compute_gradient_partial_derivatives( - simulation.TimeStep, simulation.energies, state, state0, hessian, dgradE_dx0, dgradE_dv0); - - // NOTE: This is not always zero - Mat dgradE_dp = Mat::Zero(nDoF, nParameters); - - // Solve the linear system - // ------------------------------------------------------------------------- -#ifdef DIFF_LOCAL_RB - Eigen::ConjugateGradient solver; - // Eigen::SparseLU solver; - SparseMat jac_global_to_local(nDoF, nDoF); - compute_simulation_global_to_local_jacobian(simulation.simulables, state, jac_global_to_local); - const Vec equation_vector = - -(loss_position_gradient + one_over_h * loss_velocity_gradient).transpose() * jac_global_to_local; -#else - const Vec equation_vector = -(loss_position_gradient + one_over_h * loss_velocity_gradient); - Eigen::SparseLU solver; -#endif - solver.compute(hessian.transpose()); - - const Vec grad0 = compute_energy_gradient(simulation.TimeStep, simulation.energies, state, state0); - - Vec h_grad = equation_vector; - - const Scalar dx = 1e-6; - Vec z = solver.solve(equation_vector); - Vec dz = Vec::Zero(nDoF); - - for (unsigned int j = 0; j < maxIterations; j++) { - // Compute Hz with finite differences - - Vec Hz = approximate_Hz_finite_diff(simulation, state, state0, grad0, z, dx); - Scalar h = 0.5 * z.transpose() * Hz - equation_vector.dot(z); - // DEBUG_LOG(h); - - // Evaluate the gradient - h_grad = -Hz + equation_vector; - // DEBUG_LOG(h_grad.norm()); - // DEBUG_LOG(h); - - // Solve the system - dz = solver.solve(h_grad); - // Line search - Vec z_line_search = z + dz; - Hz = approximate_Hz_finite_diff(simulation, state, state0, grad0, z_line_search, dx); - Scalar h_new = 0.5 * z_line_search.transpose() * Hz - equation_vector.dot(z_line_search); - Scalar alpha = 1.0; - while (h_new > h) { - alpha /= 2.0; - z_line_search = z + alpha * dz; - Hz = approximate_Hz_finite_diff(simulation, state, state0, grad0, z_line_search, dx); - h_new = 0.5 * z_line_search.transpose() * Hz - equation_vector.dot(z_line_search); - } - // DEBUG_LOG(alpha); - z += alpha * dz; - // z += dz; - if (h_grad.hasNaN()) - spdlog::warn("DIFFERENTIABLE: h_grad has NaN"); - if (equation_vector.hasNaN()) - spdlog::warn("DIFFERENTIABLE: equation_vector has NaN"); - if (Hz.hasNaN()) - spdlog::warn("DIFFERENTIABLE: Hz has NaN"); - if (dz.hasNaN()) - spdlog::warn("DIFFERENTIABLE: dz has NaN"); - if (z.hasNaN()) - spdlog::info("DIFFERENTIABLE: z has NaN"); - } - // DEBUG_LOG("-------------"); - - // Update the loss function gradients - // ------------------------------------------------------------------------- - loss_position_gradient = (loss.loss_position_partial_derivative[i].transpose() - - one_over_h * loss_velocity_gradient.transpose() + z.transpose() * dgradE_dx0); - loss_velocity_gradient = (loss.loss_velocity_partial_derivative[i].transpose() + z.transpose() * dgradE_dv0); - - loss_gradient += z.transpose() * dgradE_dp; - } - - // Add the initial conditions term - // ------------------------------------------------------------------------- - loss_gradient += loss_position_gradient.transpose() * dx0_dp + loss_velocity_gradient.transpose() * dv0_dp; - - return loss_gradient; -} - -Vec compute_loss_function_gradient_backpropagation_1_step_velocity(const Simulation &simulation, - const PhysicsState state0, - const PhysicsState state1, - const Vec dg_dphi, - const Vec dg_dphi_dot) -{ - const unsigned int nDoF = static_cast(state0.x.size()); - - SparseMat hessian(nDoF, nDoF), dgradE_dx0(nDoF, nDoF), dgradE_dv0(nDoF, nDoF); - compute_gradient_partial_derivatives( - simulation.TimeStep, simulation.energies, state1, state0, hessian, dgradE_dx0, dgradE_dv0); - - const Mat dgradE_dphi = hessian.toDense(); - const Mat dgradE_dphi0 = dgradE_dx0.toDense(); - const Mat dgradE_dphi_dot0 = dgradE_dv0.toDense(); - - SparseMat ptm(nDoF, nDoF); - compute_simulation_global_to_local_jacobian(simulation.simulables, state1, ptm); - - Eigen::ConjugateGradient cg; - cg.compute(-dgradE_dphi); - const Vec eq_vec = dg_dphi.transpose() + 1.0 / simulation.TimeStep * dg_dphi_dot.transpose() * ptm; - // const Vec eq_vec = dg_dphi.transpose() + 1.0 / simulation.TimeStep * dg_dphi_dot.transpose(); - Vec adjoint = cg.solve(eq_vec); - - const Vec dgdp_adj = adjoint.transpose() * dgradE_dphi_dot0; - - return dgdp_adj; -} - -Vec compute_loss_function_gradient_backpropagation_1_step_position(const Simulation &simulation, - const PhysicsState state0, - const PhysicsState state1, - const Vec dg_dphi, - const Vec dg_dphi_dot) -{ - const unsigned int nDoF = static_cast(state0.x.size()); - - SparseMat hessian(nDoF, nDoF), dgradE_dx0(nDoF, nDoF), dgradE_dv0(nDoF, nDoF); - compute_gradient_partial_derivatives( - simulation.TimeStep, simulation.energies, state1, state0, hessian, dgradE_dx0, dgradE_dv0); - - const Mat dgradE_dphi = hessian.toDense(); - const Mat dgradE_dphi0 = dgradE_dx0.toDense(); - const Mat dgradE_dphi_dot0 = dgradE_dv0.toDense(); - - SparseMat ptm(nDoF, nDoF); - compute_simulation_global_to_local_jacobian(simulation.simulables, state1, ptm); - - Eigen::ConjugateGradient cg; - cg.compute(-dgradE_dphi); - const Vec eq_vec = dg_dphi.transpose() * ptm + 1.0 / simulation.TimeStep * dg_dphi_dot.transpose() * ptm; - // const Vec eq_vec = dg_dphi; - Vec adjoint = cg.solve(eq_vec); - - const Mat dphi_dphi0 = -dgradE_dphi.inverse() * dgradE_dphi0; - - const Vec dgdp_adj = -1.0 / simulation.TimeStep * dg_dphi_dot.transpose() + adjoint.transpose() * dgradE_dphi0; - - const Vec dgdp = dg_dphi.transpose() * dphi_dphi0; - - // return dgdp; - return dgdp_adj; -} - -} // namespace mandos::core diff --git a/src/Mandos/Core/differentiable.hpp b/src/Mandos/Core/differentiable.hpp deleted file mode 100644 index 284f47bb..00000000 --- a/src/Mandos/Core/differentiable.hpp +++ /dev/null @@ -1,37 +0,0 @@ -#ifndef MANDOS_DIFFERENTIABLE_H_ -#define MANDOS_DIFFERENTIABLE_H_ - -#include - -namespace mandos::core -{ - -// struct LossFunctionAndDerivatives { -// Scalar loss; // g -// std::vector loss_position_partial_derivative; // dg_dx -// std::vector loss_velocity_partial_derivative; // dg_dv -// Vec loss_parameter_partial_derivative; // dg_dp -// }; - -// Vec compute_loss_function_gradient_backpropagation(const Simulation& simulation, -// const std::vector& trajectory, -// const LossFunctionAndDerivatives& loss, -// const Mat& dx0_dp, -// const Mat& dv0_dp, -// const unsigned int maxIterations = 0); - -// Vec compute_loss_function_gradient_backpropagation_1_step_velocity(const Simulation& simulation, -// const PhysicsState state0, -// const PhysicsState state1, -// const Vec dg_dphi, -// const Vec dg_dphi_dot); - -// Vec compute_loss_function_gradient_backpropagation_1_step_position(const Simulation& simulation, -// const PhysicsState state0, -// const PhysicsState state1, -// const Vec dg_dphi, -// const Vec dg_dphi_dot); - -} // namespace mandos::core - -#endif // MANDOS_DIFFERENTIABLE_H_ diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index cbca41bf..fd73bd06 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -5,3 +5,4 @@ add_subdirectory(FixedProjection) add_subdirectory(StableNeoHookean) add_subdirectory(LumpedMassInertia) add_subdirectory(CosseratRods) +add_subdirectory(RotationMatrixDerivatives) diff --git a/tests/RotationMatrixDerivatives/CMakeLists.txt b/tests/RotationMatrixDerivatives/CMakeLists.txt new file mode 100644 index 00000000..4f7e7468 --- /dev/null +++ b/tests/RotationMatrixDerivatives/CMakeLists.txt @@ -0,0 +1,8 @@ +add_executable(tst_RotationMatrixDerivatives tst_RotationMatrixDerivatives.cpp) +target_link_libraries( + tst_RotationMatrixDerivatives + PUBLIC Mandos::Core + Catch2::Catch2WithMain +) + +add_test(NAME tst_RotationMatrixDerivatives COMMAND tst_RotationMatrixDerivatives) diff --git a/tests/RotationMatrixDerivatives/tst_RotationMatrixDerivatives.cpp b/tests/RotationMatrixDerivatives/tst_RotationMatrixDerivatives.cpp new file mode 100644 index 00000000..59b91742 --- /dev/null +++ b/tests/RotationMatrixDerivatives/tst_RotationMatrixDerivatives.cpp @@ -0,0 +1,221 @@ +#include + +#include + +#include +#include +#include +#include + +#include + +namespace +{ +/** + * Compute the derivative of a rotation matrix with respect to local axis angle, evaluated at theta. + */ +inline Eigen::Matrix computeVectorizedRotationMatrixDerivativeLocal( + const mandos::core::Vec3 &theta) +{ + const mandos::core::Mat3 R = mandos::core::rotationExpMap(theta); + return mandos::core::vectorizedLeviCivita() * mandos::core::blockMatrix(R); +} + +/** + * Compute the derivative of a rotation matrix with respect to the global axis angle theta. + */ +inline Eigen::Matrix computeVectorizedRotationMatrixDerivativeGlobal( + const mandos::core::Vec3 &theta) +{ + const mandos::core::Mat3 jac = mandos::core::computeLocalToGlobalAxisAngleJacobian(theta); + return jac.transpose() * computeVectorizedRotationMatrixDerivativeLocal(theta); +} + +inline Eigen::Matrix computeVectorizedRotationMatrixDerivativeLocalFinite( + const mandos::core::Vec3 &theta, + mandos::core::Scalar dx) +{ + Eigen::Matrix dRdtheta; + const mandos::core::Mat3 R = mandos::core::rotationExpMap(theta); + const mandos::core::Vec9 vecR = mandos::core::vectorizeMatrix<3>(R); + for (int i = 0; i < 3; ++i) { + mandos::core::Vec3 dvx = mandos::core::Vec3::Zero(); + dvx[i] = dx; + const mandos::core::Mat3 Rnew = mandos::core::rotationExpMap(dvx) * R; + const mandos::core::Vec9 vecRnew = mandos::core::vectorizeMatrix<3>(Rnew); + dRdtheta.row(i) = (vecRnew - vecR) / dx; + } + return dRdtheta; +} + +inline Eigen::Matrix computeVectorizedRotationMatrixDerivativeGlobalFinite( + const mandos::core::Vec3 &theta, + mandos::core::Scalar dx) +{ + Eigen::Matrix dRdtheta; + const mandos::core::Vec9 vecR = mandos::core::vectorizeMatrix<3>(mandos::core::rotationExpMap(theta)); + for (int i = 0; i < 3; ++i) { + mandos::core::Vec3 dvx = mandos::core::Vec3::Zero(); + dvx[i] = dx; + const mandos::core::Vec9 vecRnew = mandos::core::vectorizeMatrix<3>(mandos::core::rotationExpMap(theta + dvx)); + dRdtheta.row(i) = (vecRnew - vecR) / dx; + } + return dRdtheta; +} +} // namespace + +TEST_CASE("ComputeRotationMatrixDerivatives") +{ + using mandos::core::Mat3; + using mandos::core::Scalar; + using mandos::core::Vec3; + using mandos::core::Vec9; + + constexpr Scalar dx = 1e-8; + constexpr Scalar tol = 1e-6; + SECTION("GLOBAL DERIVATIVE") + { + const Vec3 theta = Vec3::Random(); + auto dRdtheta = computeVectorizedRotationMatrixDerivativeGlobal(theta); + auto dRdthetaFinite = computeVectorizedRotationMatrixDerivativeGlobalFinite(theta, dx); + std::cout << dRdtheta << "\n"; + std::cout << "\n"; + std::cout << dRdthetaFinite << "\n"; + std::cout << "\n"; + REQUIRE(dRdtheta.isApprox(dRdthetaFinite, tol)); + } + + SECTION("LOCAL DERIVATIVE") + { + const Vec3 theta = Vec3::Random(); + auto dRdtheta = computeVectorizedRotationMatrixDerivativeLocal(theta); + auto dRdthetaFinite = computeVectorizedRotationMatrixDerivativeLocalFinite(theta, dx); + std::cout << dRdtheta << "\n"; + std::cout << "\n"; + std::cout << dRdthetaFinite << "\n"; + std::cout << "\n"; + REQUIRE(dRdtheta.isApprox(dRdthetaFinite, tol)); + } + + SECTION("BLOCK AND VECTORIZED MATRIX MULTIPLICATION") + { + const Mat3 a = Mat3::Random(); + const Mat3 b = Mat3::Random(); + const Vec9 result1 = mandos::core::vectorizeMatrix<3>(a * b); + const Vec9 result2 = mandos::core::blockDiagonalMatrix<3>(a) * mandos::core::vectorizeMatrix<3>(b); + const Vec9 result3 = mandos::core::vectorizeMatrix<3>(a).transpose() * mandos::core::blockMatrix<3>(b); + std::cout << "result1\t" << result1.transpose() << "\n"; + std::cout << "result2\t" << result2.transpose() << "\n"; + std::cout << "result3\t" << result3.transpose() << "\n"; + std::cout << "\n"; + REQUIRE(result1.isApprox(result2, tol)); + REQUIRE(result1.isApprox(result3, tol)); + } + + SECTION("LAGGED DERIVATIVES") + { + const Vec3 theta = Vec3::Random(); + const Vec3 theta0 = Vec3::Random(); + const Vec3 omega0 = Vec3::Random(); + const Mat3 M = Vec3::Random().asDiagonal(); + const Scalar h = 0.1; + + using RbFuncT = Mat3(const Vec3 &, const Vec3 &, const Vec3 &); + const std::function computeRguess = [h](const Vec3 &, const Vec3 &x0, const Vec3 &v0) { + const Mat3 Rold = mandos::core::rotationExpMap(x0); + const Mat3 deltaR = mandos::core::rotationExpMap(h * v0); + const Mat3 Rguess = Rold + (Rold - deltaR.transpose() * Rold); + return Rguess; + }; + + const std::function computeRMRguess = [M, h](const Vec3 &x, const Vec3 &x0, const Vec3 &v0) { + const Mat3 R = mandos::core::rotationExpMap(x); + const Mat3 Rold = mandos::core::rotationExpMap(x0); + const Mat3 deltaR = mandos::core::rotationExpMap(h * v0); + const Mat3 Rguess = Rold + (Rold - deltaR.transpose() * Rold); + const Mat3 RMRguess = R * M * Rguess.transpose(); + return RMRguess; + }; + + const Mat3 R = mandos::core::rotationExpMap(theta); + const Mat3 Rold = mandos::core::rotationExpMap(theta0); + const Mat3 deltaR = mandos::core::rotationExpMap(h * omega0); + + const Mat3 Rguess = computeRguess(theta, theta0, omega0); + const Mat3 Rguess_x0 = computeRguess(theta, theta0 + Vec3(dx, 0, 0), omega0); + + // TEST Rguess derivative wrt theta0 + const Eigen::Matrix dRguess_dx0 = + mandos::core::blockDiagonalMatrix<3>(2 * Mat3::Identity() - deltaR.transpose()) * + computeVectorizedRotationMatrixDerivativeGlobal(theta0).transpose(); + const Vec9 dRguess_dx0_finite = + (mandos::core::vectorizeMatrix(Rguess_x0) - mandos::core::vectorizeMatrix(Rguess)) / dx; + REQUIRE(dRguess_dx0_finite.isApprox(dRguess_dx0.col(0), tol)); + std::cout << "Rguess finite\t" << dRguess_dx0_finite.transpose() << "\n"; + std::cout << "Rguess normal\t" << dRguess_dx0.col(0).transpose() << "\n"; + std::cout << "\n"; + + // TEST RMRguess derivative + const Mat3 RMRguess = computeRMRguess(theta, theta0, omega0); + const Vec9 vecRMRguess = mandos::core::vectorizeMatrix(RMRguess); + const Mat3 RMRguess_x0 = computeRMRguess(theta, theta0 + Vec3(dx, 0, 0), omega0); + const Vec9 vecRMRguess_x0 = mandos::core::vectorizeMatrix<3>(RMRguess_x0); + const Vec9 dvecRMRguess_dx0 = (vecRMRguess_x0 - vecRMRguess) / dx; + const Eigen::Matrix dRMRguess_dx0 = + mandos::core::blockDiagonalMatrix<3>(R * M) * + mandos::core::transposeVectorizedMatrixRows<9, 3>(dRguess_dx0); + REQUIRE(dvecRMRguess_dx0.isApprox(dRMRguess_dx0.col(0), tol)); + std::cout << "RMRguess finite\t" << dvecRMRguess_dx0.transpose() << "\n"; + std::cout << "RMRguess normal\t" << dRMRguess_dx0.col(0).transpose() << "\n"; + std::cout << "\n"; + + { + // TEST transpose & matrix products + const Mat3 R0old = deltaR.transpose() * Rold; + const Mat3 R0oldT = Rold.transpose() * deltaR; + const Vec9 vR0old = mandos::core::transposeVectorizedVector<9>( + mandos::core::blockDiagonalMatrix<3>(Rold.transpose()) * mandos::core::vectorizeMatrix(deltaR)); + const Vec9 vR0oldT = + mandos::core::blockDiagonalMatrix<3>(Rold.transpose()) * mandos::core::vectorizeMatrix(deltaR); + std::cout << "vR0old 1" << mandos::core::vectorizeMatrix(R0old).transpose() << "\n"; + std::cout << "vR0old 2" << vR0old.transpose() << "\n"; + std::cout << "vR0old 1T" << mandos::core::vectorizeMatrix(R0oldT).transpose() << "\n"; + std::cout << "vR0old 2T" << vR0oldT.transpose() << "\n"; + std::cout << "\n"; + REQUIRE(vR0old.reshaped().isApprox(vR0old, tol)); + REQUIRE(vR0oldT.reshaped().isApprox(vR0oldT, tol)); + } + + // TEST derivative Rguess wrt omega0 + const Mat3 Rguess_v0 = computeRguess(theta, theta0, omega0 + Vec3(dx, 0, 0)); + const Vec9 dvecRguess_dv0_finite = + (mandos::core::vectorizeMatrix(Rguess_v0) - mandos::core::vectorizeMatrix(Rguess)) / dx; + const Eigen::Matrix dvecRguess_dv0 = + -h * mandos::core::transposeVectorizedMatrixRows<9, 3>( + mandos::core::blockDiagonalMatrix<3>(Rold.transpose()) * + computeVectorizedRotationMatrixDerivativeGlobal(h * omega0).transpose()); + std::cout << "Rguess omega finite\t" << dvecRguess_dv0_finite.transpose() << "\n"; + std::cout << "Rguess omega normal\t" << dvecRguess_dv0.col(0).transpose() << "\n"; + std::cout << "\n"; + REQUIRE(dvecRguess_dv0_finite.isApprox(dvecRguess_dv0.col(0), tol)); + } +} + +TEST_CASE("TANGENTIAL TRANSFORMATION") +{ + using mandos::core::computeLocalToGlobalAxisAngleJacobian; + using mandos::core::computeTangentialTransform; + using mandos::core::Mat3; + using mandos::core::Vec3; + const Vec3 phi = Vec3::Random(); + SECTION("Local To Global equivalence") + { + const Mat3 LocalToGlobal = computeLocalToGlobalAxisAngleJacobian(phi); + const Mat3 T = computeTangentialTransform(phi); + std::cout << "Local To Global Jacobian \n" << LocalToGlobal << "\n"; + std::cout << "\n"; + std::cout << "Tangent transformation jacobian\n" << T << "\n"; + std::cout << "\n"; + REQUIRE(LocalToGlobal.isApprox(T, 1e-8)); + } +} -- GitLab From 28942dac5b26ecc18de9a1b960cd9ae17157d4ec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mag=C3=AD=20Romany=C3=A0?= Date: Thu, 5 Dec 2024 12:56:43 +0100 Subject: [PATCH 2/2] Add Differentiability Python Support & Examples MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Magí Romanyà --- .clang-format | 1 + bindings/Mandos/python/CMakeLists.txt | 1 + bindings/Mandos/python/Differentiable.cpp | 43 +++ bindings/Mandos/python/Differentiable.hpp | 13 + bindings/Mandos/python/Mandos.cpp | 6 +- bindings/Mandos/python/Model.cpp | 14 +- .../Mappings/rigid_body_point_mapping.py | 8 +- examples/python/RigidBody/local_vs_global.py | 38 ++- examples/python/RigidBody/rigidbody.py | 11 +- examples/python/diff/diff.py | 107 +++++++ examples/python/diff/diffRB.py | 153 ++++++++++ examples/python/diff/heart.obj | 144 +++++++++ .../python/diff/rigid-body-optimization.py | 133 ++++++++ examples/python/diff/rod-optimization.py | 158 ++++++++++ examples/python/diff/rotation_utilities.py | 31 ++ examples/python/diff/tkdraw.py | 285 ++++++++++++++++++ 16 files changed, 1123 insertions(+), 23 deletions(-) create mode 100644 bindings/Mandos/python/Differentiable.cpp create mode 100644 bindings/Mandos/python/Differentiable.hpp create mode 100644 examples/python/diff/diff.py create mode 100644 examples/python/diff/diffRB.py create mode 100644 examples/python/diff/heart.obj create mode 100644 examples/python/diff/rigid-body-optimization.py create mode 100644 examples/python/diff/rod-optimization.py create mode 100644 examples/python/diff/rotation_utilities.py create mode 100644 examples/python/diff/tkdraw.py diff --git a/.clang-format b/.clang-format index bd8105da..62a209ef 100644 --- a/.clang-format +++ b/.clang-format @@ -15,6 +15,7 @@ AllowShortCaseLabelsOnASingleLine: false AllowShortFunctionsOnASingleLine: false AllowShortIfStatementsOnASingleLine: false AllowShortLoopsOnASingleLine: false +AllowShortLambdasOnASingleLine: true AlwaysBreakAfterDefinitionReturnType: None AlwaysBreakAfterReturnType: None AlwaysBreakBeforeMultilineStrings: true diff --git a/bindings/Mandos/python/CMakeLists.txt b/bindings/Mandos/python/CMakeLists.txt index 4c26ced5..c169721c 100644 --- a/bindings/Mandos/python/CMakeLists.txt +++ b/bindings/Mandos/python/CMakeLists.txt @@ -5,6 +5,7 @@ set(SOURCES Deformable3D.cpp RigidBody.cpp RigidBodyPointMapping.cpp + Differentiable.cpp ) set(HEADERS Model.hpp Mesh.hpp Deformable3D.hpp RigidBody.hpp RigidBodyPointMapping.hpp) diff --git a/bindings/Mandos/python/Differentiable.cpp b/bindings/Mandos/python/Differentiable.cpp new file mode 100644 index 00000000..d96712c1 --- /dev/null +++ b/bindings/Mandos/python/Differentiable.cpp @@ -0,0 +1,43 @@ +#include +#include + +#include +#include + +namespace py = pybind11; + +namespace mandos::py +{ +void wrapDifferentiable(::py::module_ &m) +{ + ::py::class_(m, "Trajectory") + .def(::py::init()) + .def_readwrite("positions", &core::Trajectory::positions) + .def_readwrite("velocities", &core::Trajectory::velocities) + .def("get_n_dof", &core::Trajectory::getNDof) + .def("get_n_states", &core::Trajectory::getNStates) + .def("get_n_steps", &core::Trajectory::getNSteps) + .def("append_state", [](core::Trajectory &trajectory, core::Model &model) { + // Get state + core::Vec x(model.nDof()); + core::Vec v(model.nDof()); + model.state(x, v); + // Append state + trajectory.positions.push_back(std::move(x)); + trajectory.velocities.push_back(std::move(v)); + }); + + ::py::class_(m, "LossFunctionAndGradients") + .def(::py::init()) + .def_readwrite("loss", &core::LossFunctionAndGradients::loss) + .def_readwrite("loss_position_partial_derivative", + &core::LossFunctionAndGradients::lossPositionPartialDerivative) + .def_readwrite("loss_velocity_partial_derivative", + &core::LossFunctionAndGradients::lossVelocityPartialDerivative) + .def_readwrite("loss_parameter_partial_derivative", + &core::LossFunctionAndGradients::lossParameterPartialDerivative) + .def("get_n_parameters", &core::LossFunctionAndGradients::getNParameters); + + m.def("compute_loss_function_gradient_backpropagation", &core::computeLossFunctionGradientBackpropagation); +} +} // namespace mandos::py diff --git a/bindings/Mandos/python/Differentiable.hpp b/bindings/Mandos/python/Differentiable.hpp new file mode 100644 index 00000000..d7bf6c22 --- /dev/null +++ b/bindings/Mandos/python/Differentiable.hpp @@ -0,0 +1,13 @@ +#ifndef MANDOS_PY_DIFFERENTIABLE_HPP +#define MANDOS_PY_DIFFERENTIABLE_HPP + +#include + +namespace py = pybind11; + +namespace mandos::py +{ +void wrapDifferentiable(::py::module_ &m); +} // namespace mandos::py + +#endif // MANDOS_PY_DIFFERENTIABLE_HPP diff --git a/bindings/Mandos/python/Mandos.cpp b/bindings/Mandos/python/Mandos.cpp index a174773f..016af29f 100644 --- a/bindings/Mandos/python/Mandos.cpp +++ b/bindings/Mandos/python/Mandos.cpp @@ -3,6 +3,7 @@ #include #include +#include #include #include #include @@ -15,6 +16,7 @@ PYBIND11_MODULE(pymandos, m) mandos::py::wrapRigidBody(m); mandos::py::wrapRigidBodyPointMapping(m); mandos::py::wrapSurfaceMesh(m); + mandos::py::wrapDifferentiable(m); ::py::class_(m, "StepParameters") .def(::py::init()) @@ -29,7 +31,9 @@ PYBIND11_MODULE(pymandos, m) .value("Success", mandos::core::SimulationStepResult::Success) .value("LineSearchFailed", mandos::core::SimulationStepResult::LineSearchFailed) .value("NewtonFailed", mandos::core::SimulationStepResult::NewtonFailed) - .export_values(); + .export_values() + .def("__bool__", + [](mandos::core::SimulationStepResult s) { return s == mandos::core::SimulationStepResult::Success; }); m.def("step", [](mandos::py::Model &model, mandos::core::StepParameters stepParameters) { return mandos::core::step(model, stepParameters); diff --git a/bindings/Mandos/python/Model.cpp b/bindings/Mandos/python/Model.cpp index 9d6ce0d2..3373f7a3 100644 --- a/bindings/Mandos/python/Model.cpp +++ b/bindings/Mandos/python/Model.cpp @@ -1,5 +1,6 @@ #include "Mandos/Core/Model.hpp" +#include #include #include @@ -177,8 +178,17 @@ void mandos::py::wrapModel(::py::module_ &m) ///////////////// Visitors ////////////////////// - .def_property_readonly("energy", - [](const Model &model, const mandos::core::Scalar h) { return model.computeEnergy(h); }) + .def_property_readonly( + "energy", + [](const mandos::core::Model &model, const mandos::core::Scalar h) { return model.computeEnergy(h); }) + .def_property("state", + &mandos::core::Model::setState, + [](mandos::core::Model &model) { + mandos::core::Vec x(model.nDof()); + mandos::core::Vec v(model.nDof()); + model.state(x, v); + return std::pair(x, v); + }) .def("apply", [](Model &model) { auto nDof = model.nDof(); diff --git a/examples/python/Mappings/rigid_body_point_mapping.py b/examples/python/Mappings/rigid_body_point_mapping.py index 097dbddd..669bde87 100644 --- a/examples/python/Mappings/rigid_body_point_mapping.py +++ b/examples/python/Mappings/rigid_body_point_mapping.py @@ -27,8 +27,8 @@ def create_model(): rigid_body_cloud.v = velocities # Disable gravity - rigid_body_cloud.disable_gravity() - rigid_body_cloud.fix_rigid_body_center_of_mass(1) + # rigid_body_cloud.disable_gravity() + rigid_body_cloud.fix_translation(1) # Set up rigid body mass and inertia rigid_body_cloud.mass = np.ones(n_points) @@ -38,8 +38,12 @@ def create_model(): mapping = model.add_rigidbody_point_mapping(rigid_body_cloud); mapping.add_particle([2.0,0.0,0.0], 0); mapping.add_particle([0.0,2.0,0.0], 1); + # mapping.add_particle([-4.0, 0.0,0.0], 0); + mass_spring = mapping.deformable.mass_spring mass_spring.add_element((0,1), 10.0, 0.0) + mapping.deformable.particle_mass = np.zeros(mapping.deformable.size) + mapping.deformable.particle_mass = np.ones(mapping.deformable.size) model.compute_dag() diff --git a/examples/python/RigidBody/local_vs_global.py b/examples/python/RigidBody/local_vs_global.py index c183e0fd..7784c420 100644 --- a/examples/python/RigidBody/local_vs_global.py +++ b/examples/python/RigidBody/local_vs_global.py @@ -5,7 +5,8 @@ import numpy as np import meshio import polyscope as ps -import pymandos +import pymandos + def create_model(): model = pymandos.Model() @@ -38,32 +39,41 @@ def create_model(): # Render mesh mesh = meshio.read("tennis.obj") - ps.register_surface_mesh("rb_local_mesh", mesh.points, mesh.cells_dict.get("triangle")) - ps.register_surface_mesh("rb_global_mesh", mesh.points, mesh.cells_dict.get("triangle")) + ps.register_surface_mesh( + "rb_local_mesh", mesh.points, mesh.cells_dict.get("triangle") + ) + ps.register_surface_mesh( + "rb_global_mesh", mesh.points, mesh.cells_dict.get("triangle") + ) ps.get_surface_mesh("rb_local_mesh").set_transform(rigid_body_local.get_transform()) - ps.get_surface_mesh("rb_global_mesh").set_transform(rigid_body_global.get_transform()) + ps.get_surface_mesh("rb_global_mesh").set_transform( + rigid_body_global.get_transform() + ) return model, rigid_body_local, rigid_body_global + def simulate_callback(model, rigid_body_local, rigid_body_global): stepParameters = pymandos.StepParameters() - stepParameters.h = 0.01 - stepParameters.newton_iterations =5 - stepParameters.cg_iterations = 20 - stepParameters.cg_error = 1e-4 - stepParameters.grad_norm = 1e-3 - stepParameters.line_search_iterations = 5 + stepParameters.h = 0.1 + stepParameters.newton_iterations = 5 + stepParameters.cg_error = 1e-7 + stepParameters.grad_norm = 1e-1 + stepParameters.enable_line_search = False pymandos.step(model, stepParameters) ps.get_surface_mesh("rb_local_mesh").set_transform(rigid_body_local.get_transform()) - ps.get_surface_mesh("rb_global_mesh").set_transform(rigid_body_global.get_transform()) + ps.get_surface_mesh("rb_global_mesh").set_transform( + rigid_body_global.get_transform() + ) if __name__ == "__main__": ps.init() - ps.set_ground_plane_mode("none") # Disable the ground + ps.set_ground_plane_mode("none") # Disable the ground model, rigid_body_local, rigid_body_global = create_model() - ps.set_user_callback(lambda : simulate_callback(model, rigid_body_local, rigid_body_global)) + ps.set_user_callback( + lambda: simulate_callback(model, rigid_body_local, rigid_body_global) + ) ps.show() - diff --git a/examples/python/RigidBody/rigidbody.py b/examples/python/RigidBody/rigidbody.py index 3fa78d55..5180257b 100644 --- a/examples/python/RigidBody/rigidbody.py +++ b/examples/python/RigidBody/rigidbody.py @@ -5,6 +5,7 @@ import polyscope as ps import pymandos + def create_model(): model = pymandos.Model() rigid_body = model.add_rigidbody() @@ -31,21 +32,23 @@ def create_model(): return model, rigid_body + def simulate_callback(model, rigid_body): stepParameters = pymandos.StepParameters() stepParameters.h = 0.01 - stepParameters.newton_iterations =5 + stepParameters.newton_iterations = 5 stepParameters.cg_iterations = 20 stepParameters.cg_error = 1e-4 stepParameters.grad_norm = 1e-3 - stepParameters.line_search_iterations = 5 + stepParameters.enable_line_search = False pymandos.step(model, stepParameters) ps.get_surface_mesh("rb_mesh").set_transform(rigid_body.get_transform()) + if __name__ == "__main__": ps.init() - ps.set_ground_plane_mode("none") # Disable ground rendering + ps.set_ground_plane_mode("none") # Disable ground rendering model, rigid_body = create_model() - ps.set_user_callback(lambda : simulate_callback(model, rigid_body)) + ps.set_user_callback(lambda: simulate_callback(model, rigid_body)) ps.show() diff --git a/examples/python/diff/diff.py b/examples/python/diff/diff.py new file mode 100644 index 00000000..df4185cb --- /dev/null +++ b/examples/python/diff/diff.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python3 + +import numpy as np +import polyscope as ps +import matplotlib.pyplot as plt +import scipy as sp + +import pymandos + +TARGET = np.array([1.0, 0.0, 1.0]) +stepParameters = pymandos.StepParameters() +stepParameters.h = 0.01 +stepParameters.newton_iterations = 5 +stepParameters.cg_iterations = 20 +stepParameters.cg_error = 1e-8 +stepParameters.grad_norm = 1e-5 +stepParameters.line_search_iterations = 0 + + +def create_model(initial_vel: list): + model = pymandos.Model() + particles = model.add_deformable_3d() + particles.size = 1 + particles.x = np.array([[0, 0, 0]]) + particles.v = np.array([initial_vel]) + + particles.particle_mass = [1.0] + + model.commit() + + points = np.array((particles.x[0], TARGET)) + ps.register_point_cloud("particles", points, radius=0.1) + return model, particles + + +def simulate_callback(model, particles): + pymandos.step(model, stepParameters) + points = np.array((particles.x[0], TARGET)) + ps.get_point_cloud("particles").update_point_positions(points) + + +def simulate(initial_vel: list, diff_frames: int): + model, particles = create_model(initial_vel) + trajectory = pymandos.Trajectory() + trajectory.append_state(model) + for _ in range(diff_frames): + pymandos.step(model, stepParameters) + trajectory.append_state(model) + return trajectory, model + + +def compute_loss(trajectory: pymandos.Trajectory) -> pymandos.LossFunctionAndGradients: + loss = pymandos.LossFunctionAndGradients() + delta = trajectory.positions[-1] - TARGET + loss.loss = np.dot(delta, delta) + loss_position_partial_derivative = [ + np.zeros(3) for _ in range(len(trajectory.positions)) + ] + loss_position_partial_derivative[-1] = 2.0 * delta + loss.loss_position_partial_derivative = loss_position_partial_derivative + loss.loss_velocity_partial_derivative = [ + np.zeros(3) for _ in range(len(trajectory.positions)) + ] + loss.loss_parameter_partial_derivative = np.zeros( + 3 + ) # 3 parameters (one for each initial velocity vector coefficient) + return loss + + +def sim_wrapper(v0): + trajectory, model = simulate(v0, 100) + loss = compute_loss(trajectory) + dx0_dp = np.zeros((3, 3)) + dv0_dp = np.identity(3) + gradient = pymandos.compute_loss_function_gradient_backpropagation( + stepParameters.h, model, trajectory, loss, dx0_dp, dv0_dp, 0 + ) + return loss.loss, gradient + + +if __name__ == "__main__": + ps.init() + ps.set_up_dir("z_up") + + losses = [] + gradients = [] + for vel in np.linspace(1,5, 100): + trajectory, model = simulate([0,0,vel], 100) + loss = compute_loss(trajectory) + dx0_dp = np.zeros((3,3)) + dv0_dp = np.identity(3) + gradient = pymandos.compute_loss_function_gradient_backpropagation(stepParameters.h, model, trajectory, loss, dx0_dp, dv0_dp, 0) + losses.append(loss.loss) + gradients.append(gradient[-1]) + + gradients2 = np.gradient(np.array(losses), 4.0/100.0) + plt.plot(gradients2) + plt.plot(gradients) + plt.show() + + result = sp.optimize.minimize(sim_wrapper, [0, 0, 0], jac=True) + # print(result) + + model, particles = create_model(result.x) + # model, particles = create_model([0,0,0]) + ps.set_user_callback(lambda: simulate_callback(model, particles)) + ps.show() diff --git a/examples/python/diff/diffRB.py b/examples/python/diff/diffRB.py new file mode 100644 index 00000000..fac95f40 --- /dev/null +++ b/examples/python/diff/diffRB.py @@ -0,0 +1,153 @@ +#!/usr/bin/env python3 + +import numpy as np +import matplotlib.pyplot as plt +import scipy as sp +from tqdm import tqdm + +import pymandos +import rotation_utilities + +stepParameters = pymandos.StepParameters() +stepParameters.h = 0.1 +stepParameters.newton_iterations = 10 +stepParameters.cg_error = 1e-8 +stepParameters.grad_norm = 1e-6 +stepParameters.line_search_iterations = 0 + + +def compute_J_inertia_tensor(I): + Ix = I[0, 0] + Iy = I[1, 1] + Iz = I[2, 2] + return 0.5 * np.diag((-Ix + Iy + Iz, Ix - Iy + Iz, Ix + Iy - Iz)) + + +def create_model(initial_pos: list, initial_vel: list): + model = pymandos.Model() + rb = model.add_rigidbody() + rb.x = np.array(initial_pos) + rb.v = np.array(initial_vel) + + # Disable gravity + rb.disable_gravity() + + # Inertia + rb.mass = 2.0 + rb.inertiaTensor = compute_J_inertia_tensor(np.diag([0.5, 1.0, 2.0])) + + model.commit() + + return model, rb + + +def simulate_callback(model, rigid_body): + pymandos.step(model, stepParameters) + + +def simulate(initial_pos: list, initial_vel: list, diff_frames: int): + model, rb = create_model(initial_pos, initial_vel) + trajectory = pymandos.Trajectory() + trajectory.append_state(model) + for _ in range(diff_frames): + pymandos.step(model, stepParameters) + trajectory.append_state(model) + return trajectory, model + + +def compute_loss(trajectory: pymandos.Trajectory) -> pymandos.LossFunctionAndGradients: + loss = pymandos.LossFunctionAndGradients() + + lin_pos = trajectory.positions[-1][:3] + target_pos = np.array([1.0, 1.0, 1.0]) + delta = lin_pos - target_pos + loss_position = np.dot(delta, delta) + loss_position_gradient = 2.0 * delta + + axis_angle = trajectory.positions[-1][3:] + R = rotation_utilities.rotation_exp_map(axis_angle) + R_target = rotation_utilities.rotation_exp_map(np.array([1.0, 0.0, 1.0])) + R_diff = R.T @ R_target + loss_rotation = 0.5 * np.trace(np.identity(3) - R_diff) + + loss_rotation_gradient = -0.5 * rotation_utilities.unskew(R_diff - R_diff.T).T @ R.T + + dof = trajectory.get_n_dof() + loss.loss = loss_position + loss_rotation + loss_position_partial_derivative = [np.zeros(dof) for _ in trajectory.positions] + loss_position_partial_derivative[-1][:3] = loss_position_gradient + loss_position_partial_derivative[-1][3:] = loss_rotation_gradient + loss.loss_position_partial_derivative = loss_position_partial_derivative + loss.loss_velocity_partial_derivative = [ + np.zeros(dof) for _ in trajectory.positions + ] + loss.loss_parameter_partial_derivative = np.zeros(2 * dof) + return loss + + +def sim_wrapper(x0: list, v0: list): + trajectory, model = simulate(x0, v0, 100) + loss = compute_loss(trajectory) + dof = trajectory.get_n_dof() + dx0_dp = np.block( + [ + [np.eye(dof), np.zeros((dof, dof))], + ] + ) + dv0_dp = np.block( + [ + [np.zeros((dof, dof)), np.eye(dof)], + ] + ) + # dx0_dp = np.zeros((dof, dof)) + # dv0_dp = np.identity(dof) + gradient = pymandos.compute_loss_function_gradient_backpropagation( + stepParameters.h, model, trajectory, loss, dx0_dp, dv0_dp, 10 + ) + return loss.loss, gradient + + +if __name__ == "__main__": + positions = np.linspace(0.0, 2.0, 500 + 1) + velocities = np.linspace(0.0, 2.0, 500 + 1) + dx = positions[1] - positions[0] + dv = velocities[1] - velocities[0] + losses = [] + gradients = [] + for x in tqdm(positions): + pos = [0.0, 0.0, 0.0, 1.0, x, 0] + vel = [0.0, 0.0, 0.0, 1.0, 0, 0] + loss, gradient = sim_wrapper(pos, vel) + losses.append(loss) + gradients.append(gradient[3 + 1]) + + diff_gradient = np.gradient(losses, dx) + + """ Plot results """ + plt.grid() + plt.title("Gradient with respect to initial orientation") + plt.plot(positions, losses, label="Loss function") + plt.plot(positions, gradients, label="BP") + plt.plot(positions, diff_gradient, label="FD") + plt.legend() + plt.show() + + losses = [] + gradients = [] + for v in tqdm(velocities): + pos = [0.0, 0.0, 0.0, 1.0, 0, 0] + vel = [0.0, 0.0, 0.0, 1.0, v, 0] + loss, gradient = sim_wrapper(pos, vel) + losses.append(loss) + gradients.append(gradient[6 + 3 + 1]) + + diff_gradient = np.gradient(losses, dv) + + """ Plot results """ + plt.grid() + plt.title("Gradient with respect to initial angular velocity") + plt.plot(velocities, losses, label="Loss function") + plt.plot(velocities, gradients, label="BP") + plt.plot(velocities, diff_gradient, label="FD") + plt.legend() + plt.show() diff --git a/examples/python/diff/heart.obj b/examples/python/diff/heart.obj new file mode 100644 index 00000000..ffef9ad5 --- /dev/null +++ b/examples/python/diff/heart.obj @@ -0,0 +1,144 @@ +# Blender 4.2.1 LTS +# www.blender.org +o heart +v -0.006683 -0.231520 0.000000 +v -0.282312 -0.017188 0.000000 +v -0.555923 0.199717 -0.000000 +v -0.826888 0.419915 -0.000000 +v -1.094211 0.644523 -0.000000 +v -1.357785 0.873518 -0.000000 +v -1.617431 1.106961 -0.000000 +v -1.872526 1.345366 -0.000000 +v -2.122321 1.589314 -0.000000 +v -2.366040 1.839340 -0.000000 +v -2.602702 2.096050 -0.000000 +v -2.831054 2.360183 -0.000000 +v -3.049483 2.632579 -0.000000 +v -3.255916 2.914176 -0.000000 +v -3.447637 3.205986 -0.000000 +v -3.621047 3.509036 -0.000000 +v -3.771323 3.824199 -0.000000 +v -3.892076 4.151812 -0.000000 +v -3.975240 4.490924 -0.000000 +v -4.011520 4.838195 -0.000000 +v -3.991634 5.186799 -0.000000 +v -3.910133 5.526312 -0.000000 +v -3.771561 5.846779 -0.000000 +v -3.589630 6.144796 -0.000000 +v -3.375668 6.420726 -0.000000 +v -3.130190 6.669042 -0.000001 +v -2.847747 6.874338 -0.000001 +v -2.528708 7.016223 -0.000001 +v -2.185264 7.079073 -0.000001 +v -1.836581 7.060945 -0.000001 +v -1.498443 6.973983 -0.000001 +v -1.177488 6.836415 -0.000001 +v -0.875480 6.661122 -0.000001 +v -0.588794 6.462618 -0.000000 +v -0.325224 6.228934 -0.000000 +v 0.000000 6.113363 -0.000000 +v 0.325224 6.228934 -0.000000 +v 0.588794 6.462618 -0.000000 +v 0.875480 6.661122 -0.000001 +v 1.177488 6.836415 -0.000001 +v 1.498443 6.973983 -0.000001 +v 1.836581 7.060945 -0.000001 +v 2.185264 7.079073 -0.000001 +v 2.528708 7.016223 -0.000001 +v 2.847747 6.874338 -0.000001 +v 3.130190 6.669042 -0.000001 +v 3.375668 6.420726 -0.000000 +v 3.589630 6.144796 -0.000000 +v 3.771561 5.846779 -0.000000 +v 3.910133 5.526312 -0.000000 +v 3.991634 5.186799 -0.000000 +v 4.011520 4.838195 -0.000000 +v 3.975240 4.490924 -0.000000 +v 3.892076 4.151812 -0.000000 +v 3.771323 3.824199 -0.000000 +v 3.621047 3.509036 -0.000000 +v 3.447637 3.205986 -0.000000 +v 3.255916 2.914176 -0.000000 +v 3.049483 2.632579 -0.000000 +v 2.831054 2.360183 -0.000000 +v 2.602702 2.096050 -0.000000 +v 2.366040 1.839340 -0.000000 +v 2.122321 1.589314 -0.000000 +v 1.872526 1.345366 -0.000000 +v 1.617431 1.106961 -0.000000 +v 1.357785 0.873518 -0.000000 +v 1.094211 0.644523 -0.000000 +v 0.826888 0.419915 -0.000000 +v 0.555923 0.199717 -0.000000 +v 0.282312 -0.017188 0.000000 +v 0.006683 -0.231520 0.000000 +l 1 2 +l 2 3 +l 3 4 +l 4 5 +l 5 6 +l 6 7 +l 7 8 +l 8 9 +l 9 10 +l 10 11 +l 11 12 +l 12 13 +l 13 14 +l 14 15 +l 15 16 +l 16 17 +l 17 18 +l 18 19 +l 19 20 +l 20 21 +l 21 22 +l 22 23 +l 23 24 +l 24 25 +l 25 26 +l 26 27 +l 27 28 +l 28 29 +l 29 30 +l 30 31 +l 31 32 +l 32 33 +l 33 34 +l 34 35 +l 35 36 +l 36 37 +l 37 38 +l 38 39 +l 39 40 +l 40 41 +l 41 42 +l 42 43 +l 43 44 +l 44 45 +l 45 46 +l 46 47 +l 47 48 +l 48 49 +l 49 50 +l 50 51 +l 51 52 +l 52 53 +l 53 54 +l 54 55 +l 55 56 +l 56 57 +l 57 58 +l 58 59 +l 59 60 +l 60 61 +l 61 62 +l 62 63 +l 63 64 +l 64 65 +l 65 66 +l 66 67 +l 67 68 +l 68 69 +l 69 70 +l 70 71 diff --git a/examples/python/diff/rigid-body-optimization.py b/examples/python/diff/rigid-body-optimization.py new file mode 100644 index 00000000..a86cbeec --- /dev/null +++ b/examples/python/diff/rigid-body-optimization.py @@ -0,0 +1,133 @@ +#!/usr/bin/env python3 + +import numpy as np +from scipy.optimize import minimize +import polyscope as ps +import meshio + +import pymandos +import rotation_utilities + +stepParameters = pymandos.StepParameters() +stepParameters.h = 0.1 +stepParameters.newton_iterations = 10 +stepParameters.cg_error = 1e-8 +stepParameters.grad_norm = 1e-6 +stepParameters.line_search_iterations = 10 + +target_pos = 3.0 * (2.0 * np.random.rand(3) - 1.0) +R_target = rotation_utilities.rotation_exp_map(2.0 * np.random.rand(3) - 1.0) + + +def compute_J_inertia_tensor(I): + Ix = I[0, 0] + Iy = I[1, 1] + Iz = I[2, 2] + return 0.5 * np.diag((-Ix + Iy + Iz, Ix - Iy + Iz, Ix + Iy - Iz)) + + +def create_model(initial_vel: list): + model = pymandos.Model() + rb = model.add_rigidbody() + rb.x = np.array([0.0, 0.0, 0.0, 1.0, 0.0, 0.0]) + rb.v = np.array(initial_vel) + + # Inertia + rb.mass = 2.0 + rb.inertiaTensor = compute_J_inertia_tensor(np.diag([0.5, 1.0, 2.0])) + rb.disable_gravity() + + model.compute_dag() + + return model, rb + + +def simulate(initial_vel: list, diff_frames: int): + model, rb = create_model(initial_vel) + trajectory = pymandos.Trajectory() + trajectory.append_state(model) + for _ in range(diff_frames): + pymandos.step(model, stepParameters) + trajectory.append_state(model) + return trajectory, model + + +def compute_loss(trajectory: pymandos.Trajectory) -> pymandos.LossFunctionAndGradients: + loss = pymandos.LossFunctionAndGradients() + + lin_pos = trajectory.positions[-1][:3] + delta = lin_pos - target_pos + loss_position = np.dot(delta, delta) + loss_position_gradient = 2.0 * delta + + axis_angle = trajectory.positions[-1][3:] + R = rotation_utilities.rotation_exp_map(axis_angle) + R_diff = R.T @ R_target + loss_rotation = 0.5 * np.trace(np.identity(3) - R_diff) + + loss_rotation_gradient = -0.5 * rotation_utilities.unskew(R_diff - R_diff.T).T @ R.T + + dof = trajectory.get_n_dof() + loss.loss = loss_position + loss_rotation + loss_position_partial_derivative = [np.zeros(dof) for _ in trajectory.positions] + loss_position_partial_derivative[-1][:3] = loss_position_gradient + loss_position_partial_derivative[-1][3:] = loss_rotation_gradient + loss.loss_position_partial_derivative = loss_position_partial_derivative + loss.loss_velocity_partial_derivative = [ + np.zeros(dof) for _ in trajectory.positions + ] + loss.loss_parameter_partial_derivative = np.zeros(dof) + return loss + + +def sim_wrapper(v0): + trajectory, model = simulate(v0, 100) + loss = compute_loss(trajectory) + dof = trajectory.get_n_dof() + dx0_dp = np.zeros((dof, dof)) + dv0_dp = np.identity(dof) + gradient = pymandos.compute_loss_function_gradient_backpropagation( + stepParameters.h, model, trajectory, loss, dx0_dp, dv0_dp, 0 + ) + return loss.loss, gradient + + +if __name__ == "__main__": + v0 = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]) + res = minimize( + sim_wrapper, + v0, + jac=True, + method="L-BFGS-B", + options={"disp": True, "maxiter": 100}, + ) + + trajectory, model = simulate(res.x, 100) + R = rotation_utilities.rotation_exp_map(trajectory.positions[-1][3:].copy()) + print(f"R target =\n{R_target}, ") + print(f"R optimized =\n{R}") + print(f"Initial velocity vector: {res.x}") + + ps.init() + ps.set_ground_plane_mode("none") # Disable ground rendering + ps.set_up_dir("z_up") + ps.look_at((0.0, 20.0, 4.0), (0.0, 0.0, 0.0)) + mesh = meshio.read("../RigidBody/tennis.obj") + ps.register_surface_mesh("rb_mesh", mesh.points, mesh.cells_dict.get("triangle")) + ps.register_surface_mesh( + "rb_mesh_target", mesh.points, mesh.cells_dict.get("triangle") + ) + transform = np.eye(4) + transform[:3, :3] = R_target + transform[:3, 3] = target_pos + print(transform) + ps.get_surface_mesh("rb_mesh_target").set_transform(transform) + + model, rb = create_model(res.x) + + def simulate_callback(model, rigid_body): + pymandos.step(model, stepParameters) + ps.get_surface_mesh("rb_mesh").set_transform(rigid_body.get_transform()) + + ps.set_user_callback(lambda: simulate_callback(model, rb)) + ps.show() diff --git a/examples/python/diff/rod-optimization.py b/examples/python/diff/rod-optimization.py new file mode 100644 index 00000000..6dceaae5 --- /dev/null +++ b/examples/python/diff/rod-optimization.py @@ -0,0 +1,158 @@ +#!/usr/bin/env python3 + +import numpy as np +import meshio +import polyscope as ps +from scipy.optimize import minimize +import os, sys + +sys.path.append(os.getcwd() + "/../Rod") +from parallel_transport import compute_rotations_parallel_transport +import pymandos + +curve = meshio.read("heart.obj") +points = curve.points +n_points = curve.points.shape[0] +indices = np.array([(i, i + 1) for i in range(n_points - 1)]) +delta_x = np.linalg.norm(points[1] - points[0]) +origin = np.array(np.mean(points, axis=0)) +print(f"Origin {origin}") +print(f"Curve with {n_points} points, and with a distance of {delta_x}") + +linear_dof_matrix = [] +for i in range(n_points): + index = i * 6 + for j in range(3): + row = np.zeros(6 * n_points) + row[index + j] = 1.0 + linear_dof_matrix.append(row) +linear_dof_matrix = np.array(linear_dof_matrix) + +stepParameters = pymandos.StepParameters() +stepParameters.h = 0.01 +stepParameters.newton_iterations = 10 +stepParameters.cg_error = 1e-8 +stepParameters.grad_norm = 1e-4 +stepParameters.line_search_iterations = 10 + +backward_iterations = 5 + + +def create_model(initial_velocities): + model = pymandos.Model() + rigid_body_cloud = model.add_rigidbody_cloud() + + # Set initial conditions + positions = np.zeros((n_points, 6)) + velocities = np.zeros((n_points, 6)) + for i, position in enumerate(positions): + position[0] += delta_x * i - 0.5 * n_points * delta_x + position[0:3] += origin + for i, vel in enumerate(velocities): + vel[0:3] = initial_velocities[3 * i : 3 * i + 3] + + lin_pos = np.array([pos[:3] for pos in positions]) + rotvecs = compute_rotations_parallel_transport(lin_pos) + for rotvec, pos in zip(rotvecs, positions): + pos[3:] = rotvec + rigid_body_cloud.size = n_points + rigid_body_cloud.x = positions + rigid_body_cloud.v = velocities + rigid_body_cloud.disable_gravity() + + # Set up rigid body mass and inertia + rigid_body_cloud.mass = 1.0 * np.ones(n_points) + rigid_body_cloud.inertiaTensor = [np.diag([1.0, 1.0, 1.0]) for _ in range(n_points)] + + # Rod energy + Ks = [2000.0 for _ in range(n_points-1)] + stiffness_tensor = [100.0 * np.ones(3) for _ in range(n_points-2)] + cosserat_stiffness = [1000.0 for _ in range(n_points-1)] + cosserat_rod = rigid_body_cloud.cosserat_rod + cosserat_rod.set_parameters(Ks, cosserat_stiffness, stiffness_tensor) + + model.commit() + + # Render + positions = np.array([x[:3] for x in rigid_body_cloud.x]) + indices = np.array([(i, i + 1) for i in range(n_points - 1)]) + ps.register_curve_network("rod_curve", positions, indices, radius=0.01) + ps.register_curve_network("heart", points, indices, radius=0.01) + + return model, rigid_body_cloud + + +def compute_loss(trajectory: pymandos.Trajectory) -> pymandos.LossFunctionAndGradients: + loss = pymandos.LossFunctionAndGradients() + dof = trajectory.get_n_dof() + x_final = trajectory.positions[-1] + delta = linear_dof_matrix @ x_final.copy() - points.flatten() + loss.loss = np.dot(delta, delta) + loss_position_partial_derivative = [np.zeros(dof) for _ in trajectory.positions] + loss_position_partial_derivative[-1] = 2.0 * linear_dof_matrix.T @ delta + loss.loss_position_partial_derivative = loss_position_partial_derivative + loss.loss_velocity_partial_derivative = [ + np.zeros(dof) for _ in trajectory.positions + ] + loss.loss_parameter_partial_derivative = np.zeros(int(dof / 6 * 3)) + return loss + + +def simulate(initial_vel: list, diff_frames: int): + model, rb_cloud = create_model(initial_vel) + trajectory = pymandos.Trajectory() + trajectory.append_state(model) + for _ in range(diff_frames): + converged = pymandos.step(model, stepParameters) + if not converged: + print("Not converged!") + trajectory.append_state(model) + # Render + positions = np.array([x[:3] for x in rb_cloud.x]) + ps.get_curve_network("rod_curve").update_node_positions(positions) + ps.frame_tick() + return trajectory, model + + +def sim_wrapper(v0): + trajectory, model = simulate(v0, 100) + loss = compute_loss(trajectory) + dof = trajectory.get_n_dof() + dx0_dp = np.zeros((dof, int(dof / 6 * 3))) + dv0_dp = linear_dof_matrix.copy().T + # dv0_dp = np.diag( + # np.array( + # [[1.0, 1.0, 1.0, 0.0, 0.0, 0.0] for _ in range(int(dof / 6))] + # ).flatten() + # ) + + gradient = pymandos.compute_loss_function_gradient_backpropagation( + stepParameters.h, model, trajectory, loss, dx0_dp, dv0_dp, backward_iterations + ) + return loss.loss, gradient + + +def simulate_callback(model, rigid_body_cloud): + pymandos.step(model, stepParameters) + + positions = np.array([x[:3] for x in rigid_body_cloud.x]) + ps.get_curve_network("rod_curve").update_node_positions(positions) + + +if __name__ == "__main__": + ps.init() + ps.set_up_dir("z_up") + ps.look_at((0, -15.0, 8.0), (0, 0.0, 0.0)) + ps.set_ground_plane_mode("none") # Disable ground rendering + initial_velocities = np.zeros(n_points * 3) + + minimize( + sim_wrapper, + initial_velocities, + method="L-BFGS-B", + jac=True, + options={"disp": True}, + ) + model, rigid_body_cloud = create_model(initial_velocities) + ps.set_user_callback(lambda: simulate_callback(model, rigid_body_cloud)) + ps.show() diff --git a/examples/python/diff/rotation_utilities.py b/examples/python/diff/rotation_utilities.py new file mode 100644 index 00000000..31dc9d44 --- /dev/null +++ b/examples/python/diff/rotation_utilities.py @@ -0,0 +1,31 @@ +#!/usr/bin/env python3 + +from scipy.spatial.transform import Rotation +import numpy as np + + +def rotation_exp_map(axis_angle: np.array) -> np.array: + return Rotation.from_rotvec(axis_angle.copy()).as_matrix() + + +def skew(v: np.array) -> np.array: + return np.array( + [ + [0.0, -v[2], v[1]], + [v[2], 0.0, -v[0]], + [-v[1], v[0], 0.0], + ] + ) + + +def unskew(m: np.array) -> np.array: + return np.array([m[2][1], m[0][2], m[1][0]]) + + +def expMapJacobian(axis_angle: np.array) -> np.array: + angle = np.linalg.norm(axis_angle) + return ( + np.identity(3) + + (1 - np.cos(angle)) / angle**2 * skew(axis_angle) + - (angle - np.sin(angle)) / angle**3 * skew(axis_angle) @ skew(axis_angle) + ) diff --git a/examples/python/diff/tkdraw.py b/examples/python/diff/tkdraw.py new file mode 100644 index 00000000..83294728 --- /dev/null +++ b/examples/python/diff/tkdraw.py @@ -0,0 +1,285 @@ +#!/usr/bin/env python3 + +import tkinter as tk +import numpy as np + +import polyscope as ps +from scipy.optimize import minimize +import os, sys + +sys.path.append(os.getcwd() + "/../Rod") +from parallel_transport import compute_rotations_parallel_transport +import pymandos + + +def compute_length(a, b): + return ((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2) ** 0.5 + + +class DrawingApp: + def __init__(self, root): + self.root = root + self.root.title("Draw") + + # Create a canvas widget + self.canvas = tk.Canvas(self.root, bg="white", width=800, height=600) + self.canvas.pack() + + # Create Widget + self.save_button = tk.Button( + self.root, text="Simulate", command=self.compute_points + ) + self.save_button.pack() + + # Variables to track the current position + self.last_x, self.last_y = None, None + + # Bind mouse events to canvas + self.canvas.bind("", self.draw) + self.canvas.bind("", self.reset) + self.canvas.bind("", self.clear_canvas) + + self.positions = [] + + def draw(self, event): + if self.last_x and self.last_y: + # Create a line from the last position to the current position + self.canvas.create_line( + (self.last_x, self.last_y, event.x, event.y), fill="black", width=2 + ) + self.positions.append((event.x, event.y)) + + # Update the last position + self.last_x, self.last_y = event.x, event.y + + def reset(self, event): + # Reset the last position when the mouse button is released + self.last_x, self.last_y = None, None + + def clear_canvas(self, event): + self.canvas.delete("all") + self.last_x, self.last_y = None, None + self.positions = [] + + def compute_points(self): + def compute_lengths(positions): + lengths = [] + for i in range(len(self.positions) - 1): + p1 = self.positions[i] + p2 = self.positions[i + 1] + lengths.append(compute_length(p1, p2)) + return lengths + + lengths = compute_lengths(self.positions) + + lengths = np.array(lengths) + comulated_lengths = [] + total_len = 0 + for length in lengths: + comulated_lengths.append(total_len) + total_len += length + if total_len < 300: + print("Too short") + return + + def sample_curve(s): + if s > total_len: + raise ValueError( + "The desired distance is bigger than the total curve distance!" + ) + for i, length in enumerate(comulated_lengths): + if s <= length: + p1 = self.positions[i - 1] + p2 = self.positions[i] + norm = compute_length(p1, p2) + if norm > 1e-8: + t = (length - s) / norm + else: + t = 0.0 + return t * np.array(p1) + (1 - t) * np.array(p2) + raise ValueError("Unreachable") + + target_points = 71 + delta = float(total_len) / float(target_points) + positions = [] + for i in range(target_points): + p = sample_curve(i * delta) + positions.append(p) + radius = 5 + self.canvas.create_oval( + int(p[0]) - radius, + int(p[1]) - radius, + int(p[0]) + radius, + int(p[1]) + radius, + fill="blue", + ) + desired_delta = 0.4 + scaling = desired_delta / delta + positions = scaling * np.array(positions) + positions = np.hstack((positions, np.zeros((target_points, 1)))) + self.root.destroy() + self.positions = positions + + +# Create the Tkinter window +root = tk.Tk() + +# Create the drawing app +app = DrawingApp(root) + +# Run the application +root.mainloop() +print(app.positions) + +direction = -(app.positions[0] - app.positions[-1]) / np.linalg.norm( + app.positions[0] - app.positions[-1] +) +direction = np.array((direction[0], direction[1], 0)) + +points = app.positions +n_points = points.shape[0] +indices = np.array([(i, i + 1) for i in range(n_points - 1)]) +delta_x = np.linalg.norm(points[1] - points[0]) +origin = np.array(np.mean(points, axis=0)) +print(f"Origin {origin}") +print(f"Curve with {n_points} points, and with a distance of {delta_x}") + +linear_dof_matrix = [] +for i in range(n_points): + index = i * 6 + for j in range(3): + row = np.zeros(6 * n_points) + row[index + j] = 1.0 + linear_dof_matrix.append(row) +linear_dof_matrix = np.array(linear_dof_matrix) + +stepParameters = pymandos.StepParameters() +stepParameters.h = 0.01 +stepParameters.newton_iterations = 10 +stepParameters.cg_error = 1e-8 +stepParameters.grad_norm = 1e-7 +stepParameters.line_search_iterations = 0 + +backward_iterations = 0 + + +def create_model(initial_velocities): + model = pymandos.Model() + rigid_body_cloud = model.add_rigidbody_cloud() + + # Set initial conditions + positions = np.zeros((n_points, 6)) + velocities = np.zeros((n_points, 6)) + for i, position in enumerate(positions): + position[0:3] += (delta_x * i - 0.5 * n_points * delta_x) * direction + position[0:3] += origin + for i, vel in enumerate(velocities): + vel[0:3] = initial_velocities[3 * i : 3 * i + 3] + + lin_pos = np.array([pos[:3] for pos in positions]) + rotvecs = compute_rotations_parallel_transport(lin_pos) + for rotvec, pos in zip(rotvecs, positions): + pos[3:] = rotvec + rigid_body_cloud.size = n_points + rigid_body_cloud.x = positions + rigid_body_cloud.v = velocities + rigid_body_cloud.disable_gravity() + + # Set up rigid body mass and inertia + rigid_body_cloud.mass = 1.0 * np.ones(n_points) + rigid_body_cloud.inertiaTensor = [np.diag([1.0, 1.0, 1.0]) for _ in range(n_points)] + + # Rod energy + Ks = [3000.0 for _ in range(n_points-1)] + stiffness_tensor = [100.0 * np.ones(3) for _ in range(n_points-2)] + cosserat_stiffness = [1000.0 for _ in range(n_points-1)] + cosserat_rod = rigid_body_cloud.cosserat_rod + cosserat_rod.set_parameters(Ks, cosserat_stiffness, stiffness_tensor) + + model.compute_dag() + + # Render + positions = np.array([x[:3] for x in rigid_body_cloud.x]) + indices = np.array([(i, i + 1) for i in range(n_points - 1)]) + ps.register_curve_network("rod_curve", positions, indices, radius=0.01) + ps.register_curve_network("heart", points, indices, radius=0.01) + + return model, rigid_body_cloud + + +def compute_loss(trajectory: pymandos.Trajectory) -> pymandos.LossFunctionAndGradients: + loss = pymandos.LossFunctionAndGradients() + dof = trajectory.get_n_dof() + x_final = trajectory.positions[-1] + delta = linear_dof_matrix @ x_final.copy() - points.flatten() + loss.loss = np.dot(delta, delta) + loss_position_partial_derivative = [np.zeros(dof) for _ in trajectory.positions] + loss_position_partial_derivative[-1] = 2.0 * linear_dof_matrix.T @ delta + loss.loss_position_partial_derivative = loss_position_partial_derivative + loss.loss_velocity_partial_derivative = [ + np.zeros(dof) for _ in trajectory.positions + ] + loss.loss_parameter_partial_derivative = np.zeros(int(dof / 6 * 3)) + return loss + + +def simulate(initial_vel: list, diff_frames: int): + model, rb_cloud = create_model(initial_vel) + trajectory = pymandos.Trajectory() + trajectory.append_state(model) + for _ in range(diff_frames): + converged = pymandos.step(model, stepParameters) + if not converged: + print("Not converged:", converged) + trajectory.append_state(model) + # Render + positions = np.array([x[:3] for x in rb_cloud.x]) + ps.get_curve_network("rod_curve").update_node_positions(positions) + if ps.window_requests_close(): + exit() + ps.frame_tick() + return trajectory, model + + +def sim_wrapper(v0): + trajectory, model = simulate(v0, 100) + loss = compute_loss(trajectory) + dof = trajectory.get_n_dof() + dx0_dp = np.zeros((dof, int(dof / 6 * 3))) + dv0_dp = linear_dof_matrix.copy().T + # dv0_dp = np.diag( + # np.array( + # [[1.0, 1.0, 1.0, 0.0, 0.0, 0.0] for _ in range(int(dof / 6))] + # ).flatten() + # ) + + gradient = pymandos.compute_loss_function_gradient_backpropagation( + stepParameters.h, model, trajectory, loss, dx0_dp, dv0_dp, backward_iterations + ) + return loss.loss, gradient + + +def simulate_callback(model, rigid_body_cloud): + pymandos.step(model, stepParameters) + + positions = np.array([x[:3] for x in rigid_body_cloud.x]) + ps.get_curve_network("rod_curve").update_node_positions(positions) + + +if __name__ == "__main__": + ps.init() + ps.set_up_dir("z_up") + ps.look_at(np.array((0, 1.0, 18.0)) + origin, origin) + ps.set_ground_plane_mode("none") # Disable ground rendering + initial_velocities = np.zeros(n_points * 3) + + minimize( + sim_wrapper, + initial_velocities, + method="L-BFGS-B", + jac=True, + options={"disp": True}, + ) + model, rigid_body_cloud = create_model(initial_velocities) + ps.set_user_callback(lambda: simulate_callback(model, rigid_body_cloud)) + ps.show() -- GitLab