OffsetPotential.hpp

#pragma once
 
#include <wmtk/Types.hpp>
#include <wmtk/envelope/Envelope.hpp>
 
#include <wmtk/threading/enumerable_thread_specific.hpp>
#include "SimplicialComplexBVH.hpp"
 
#include <polysolve/nonlinear/Problem.hpp>
 
#include <cmath>
#include <memory>
#include <string>
#include <vector>
 
namespace wmtk::components::topological_offset {
 
template <int DIM>
class OffsetPotential
{
    static_assert(DIM == 2 || DIM == 3, "the offset potential exists in 2D and 3D only");
 
public:
    using VecD = Eigen::Matrix<double, DIM, 1>;
    using MatD = Eigen::Matrix<double, DIM, DIM>;
 
    virtual ~OffsetPotential();
 
    double target_level() const { return m_c; }
 
    double delta() const { return m_delta; }
 
    double dhat() const { return m_dhat; }
 
    double level_set_slope() const { return m_grad_ref; }
 
    virtual double value(const VecD& p) const = 0;
    virtual bool is_euclidean() const { return false; }
    virtual VecD gradient(const VecD& p) const = 0;
    virtual MatD hessian(const VecD& p) const = 0;
 
    virtual double residual_length(const VecD& p) const = 0;
 
    virtual double relative_residual(const VecD& p) const = 0;
 
    virtual bool within_support(const VecD& p) const = 0;
 
    virtual bool is_inside_offset(const VecD& p) const = 0;
 
    virtual std::string describe_active(const VecD& p) const = 0;
 
protected:
    OffsetPotential(const double delta, const double dhat)
        : m_delta(delta)
        , m_dhat(dhat)
    {}
 
    double m_delta = 0.;
    double m_dhat = 0.;
    double m_c = 0.;
    double m_grad_ref = 1.;
};
 
using OffsetPotential2D = OffsetPotential<2>;
using OffsetPotential3D = OffsetPotential<3>;
 
 
template <int DIM>
class SmoothOffsetPotential : public OffsetPotential<DIM>
{
    static_assert(DIM == 2 || DIM == 3, "the offset potential exists in 2D and 3D only");
 
public:
    using VecD = typename OffsetPotential<DIM>::VecD;
    using MatD = typename OffsetPotential<DIM>::MatD;
 
    SmoothOffsetPotential(
        const MatrixXd& V,
        const MatrixXi& E,
        const MatrixXi& F,
        const std::vector<int>& P,
        double delta,
        double dhat_factor);
 
    ~SmoothOffsetPotential() override;
 
    double value(const VecD& p) const override;
    VecD gradient(const VecD& p) const override;
    MatD hessian(const VecD& p) const override;
 
    double residual_length(const VecD& p) const override;
 
    double relative_residual(const VecD& p) const override;
 
    bool within_support(const VecD& p) const override { return value(p) > 0.; }
 
    bool is_inside_offset(const VecD& p) const override { return value(p) >= m_c; }
 
    std::string describe_active(const VecD& p) const override;
 
private:
    // Dependent base members: name them unqualified in the definitions below.
    using OffsetPotential<DIM>::m_delta;
    using OffsetPotential<DIM>::m_dhat;
    using OffsetPotential<DIM>::m_c;
    using OffsetPotential<DIM>::m_grad_ref;
 
    SmoothOffsetPotential(double delta, double dhat_factor, int /*calibration tag*/);
 
    void build(const MatrixXd& V, const MatrixXi& E, const MatrixXi& F, const std::vector<int>& P);
 
    bool level_set_distance(const VecD& p, double& t) const;
 
    void value_gradient(const VecD& p, double& v, VecD& g) const;
 
    struct Impl;
    std::unique_ptr<Impl> m_impl;
};
 
using SmoothOffsetPotential2D = SmoothOffsetPotential<2>;
using SmoothOffsetPotential3D = SmoothOffsetPotential<3>;
 
 
template <int DIM>
class EuclideanOffsetPotential : public OffsetPotential<DIM>
{
    static_assert(DIM == 2 || DIM == 3, "the offset potential exists in 2D and 3D only");
 
public:
    bool is_euclidean() const override { return true; }
    using VecD = typename OffsetPotential<DIM>::VecD;
    using MatD = typename OffsetPotential<DIM>::MatD;
 
    EuclideanOffsetPotential(const std::shared_ptr<SampleEnvelope>& envelope, double delta);
 
    EuclideanOffsetPotential(const std::shared_ptr<SimplicialComplexBVH>& bvh, double delta);
 
    double value(const VecD& p) const override;
    VecD gradient(const VecD& p) const override;
    MatD hessian(const VecD& p) const override;
 
    double residual_length(const VecD& p) const override
    {
        return std::abs(value(p) - m_c) / m_grad_ref;
    }
 
    double relative_residual(const VecD& p) const override { return (value(p) - m_c) / m_c; }
 
    bool within_support(const VecD& p) const override { return true; }
 
    bool is_inside_offset(const VecD& p) const override { return value(p) <= m_c; }
 
    std::string describe_active(const VecD& p) const override;
 
private:
    void nearest_feature(const VecD& p, VecD& foot, int& dim, VecD& dir) const;
 
    using OffsetPotential<DIM>::m_delta;
    using OffsetPotential<DIM>::m_grad_ref;
    using OffsetPotential<DIM>::m_dhat;
    using OffsetPotential<DIM>::m_c;
 
    std::shared_ptr<SampleEnvelope> m_envelope;
    std::shared_ptr<SimplicialComplexBVH> m_bvh;
};
 
using EuclideanOffsetPotential2D = EuclideanOffsetPotential<2>;
using EuclideanOffsetPotential3D = EuclideanOffsetPotential<3>;
 
 
template <int DIM>
class OffsetEnergy : public polysolve::nonlinear::Problem
{
public:
    using typename polysolve::nonlinear::Problem::Scalar;
    using typename polysolve::nonlinear::Problem::THessian;
    using typename polysolve::nonlinear::Problem::TVector;
 
    using VecD = Eigen::Matrix<double, DIM, 1>;
    using MatD = Eigen::Matrix<double, DIM, DIM>;
 
    OffsetEnergy(
        const std::shared_ptr<const OffsetPotential<DIM>>& potential,
        double weight = 1.,
        bool gauss_newton = true,
        bool distance_residual = false);
 
    double value(const TVector& x) override;
    void gradient(const TVector& x, TVector& gradv) override;
    void hessian(const TVector& x, THessian& hessian) override
    {
        log_and_throw_error("Sparse functions do not exist, use dense solver");
    }
    void hessian(const TVector& x, MatrixXd& hessian) override;
 
    void solution_changed(const TVector& new_x) override {}
 
private:
    std::shared_ptr<const OffsetPotential<DIM>> m_potential;
    double m_weight;
    bool m_gauss_newton;
    bool m_distance_residual;
    bool root_distance(const VecD& p, double& s, VecD& n) const;
    void residual(const VecD& p, double& r, VecD& dr) const;
    mutable double m_last_root = 0.; 
    mutable bool m_cache_valid = false;
    mutable VecD m_cache_p;
    mutable double m_cache_r = 0.;
    mutable VecD m_cache_dr;
};
 
using OffsetEnergy2D = OffsetEnergy<2>;
 
class AlignEnergy2D : public polysolve::nonlinear::Problem
{
public:
    using typename polysolve::nonlinear::Problem::Scalar;
    using typename polysolve::nonlinear::Problem::THessian;
    using typename polysolve::nonlinear::Problem::TVector;
    struct Edge
    {
        Eigen::Vector2d q; 
        double sigma; 
        double agree = 1.;
    };
    AlignEnergy2D(
        const std::shared_ptr<const OffsetPotential2D>& potential,
        std::vector<Edge> edges,
        double outward_sign,
        double weight);
    void residual(const Eigen::Vector2d& x, const Edge& e, double& r, Eigen::Vector2d& J) const;
    double value(const TVector& x) override;
    void gradient(const TVector& x, TVector& gradv) override;
    void hessian(const TVector& x, THessian& hessian) override
    {
        log_and_throw_error("Sparse functions do not exist, use dense solver");
    }
    void hessian(const TVector& x, MatrixXd& hessian) override;
    void solution_changed(const TVector& new_x) override {}
 
private:
    std::shared_ptr<const OffsetPotential2D> m_potential;
    std::vector<Edge> m_edges;
    double m_sign, m_weight;
};
using OffsetEnergy3D = OffsetEnergy<3>;
 
class AlignEnergy3D : public polysolve::nonlinear::Problem
{
public:
    using typename polysolve::nonlinear::Problem::Scalar;
    using typename polysolve::nonlinear::Problem::THessian;
    using typename polysolve::nonlinear::Problem::TVector;
    struct Face
    {
        Eigen::Vector3d q1, q2; 
        double sigma; 
        double agree = 1.;
    };
    AlignEnergy3D(
        const std::shared_ptr<const OffsetPotential3D>& potential,
        std::vector<Face> faces,
        double outward_sign,
        double weight);
    void residual(const Eigen::Vector3d& x, const Face& f, double& r, Eigen::Vector3d& J) const;
    double value(const TVector& x) override;
    void gradient(const TVector& x, TVector& gradv) override;
    void hessian(const TVector& x, THessian& hessian) override
    {
        log_and_throw_error("Sparse functions do not exist, use dense solver");
    }
    void hessian(const TVector& x, MatrixXd& hessian) override;
    void solution_changed(const TVector& new_x) override {}
 
private:
    std::shared_ptr<const OffsetPotential3D> m_potential;
    std::vector<Face> m_faces;
    double m_sign, m_weight;
};
 
class StencilEnergy3D : public polysolve::nonlinear::Problem
{
public:
    using typename polysolve::nonlinear::Problem::Scalar;
    using typename polysolve::nonlinear::Problem::THessian;
    using typename polysolve::nonlinear::Problem::TVector;
 
    struct Sample
    {
        double a, b, c;
    };
    struct Face
    {
        Eigen::Vector3d q1, q2;
        std::vector<Sample> samples;
        double weight = 1.;
    };
 
    StencilEnergy3D(
        const std::shared_ptr<const OffsetPotential3D>& potential,
        std::vector<Face> faces,
        double weight);
 
    double value(const TVector& x) override;
    void gradient(const TVector& x, TVector& gradv) override;
    void hessian(const TVector& x, THessian& hessian) override
    {
        log_and_throw_error("Sparse functions do not exist, use dense solver");
    }
    void hessian(const TVector& x, MatrixXd& hessian) override;
    void solution_changed(const TVector& new_x) override {}
 
private:
    bool residual_at(const Eigen::Vector3d& p, double& r, Eigen::Vector3d* dr) const;
 
    std::shared_ptr<const OffsetPotential3D> m_potential;
    std::vector<Face> m_faces;
    double m_weight;
    double m_c = 1.; 
};
 
class RestAMIPSEnergy2D : public polysolve::nonlinear::Problem
{
public:
    using typename polysolve::nonlinear::Problem::Scalar;
    using typename polysolve::nonlinear::Problem::THessian;
    using typename polysolve::nonlinear::Problem::TVector;
    struct Cell
    {
        Eigen::Vector2d q1, q2; 
        Eigen::Matrix2d rest_inv; 
    };
    RestAMIPSEnergy2D(std::vector<Cell> cells, double weight);
 
    double value(const TVector& x) override;
    void gradient(const TVector& x, TVector& gradv) override;
    void hessian(const TVector& x, THessian& hessian) override
    {
        log_and_throw_error("Sparse functions do not exist, use dense solver");
    }
    void hessian(const TVector& x, MatrixXd& hessian) override;
    void solution_changed(const TVector& new_x) override {}
    bool is_step_valid(const TVector& x0, const TVector& x1) override;
 
private:
    bool cell_F(const Eigen::Vector2d& x, const Cell& c, Eigen::Matrix2d& F, double& d) const;
    std::vector<Cell> m_cells;
    double m_weight;
};
 
class RestAMIPSEnergy3D : public polysolve::nonlinear::Problem
{
public:
    using typename polysolve::nonlinear::Problem::Scalar;
    using typename polysolve::nonlinear::Problem::THessian;
    using typename polysolve::nonlinear::Problem::TVector;
    struct Cell
    {
        Eigen::Vector3d q1, q2, q3; 
        Eigen::Matrix3d rest_inv; 
    };
    RestAMIPSEnergy3D(std::vector<Cell> cells, double weight);
 
    double value(const TVector& x) override;
    void gradient(const TVector& x, TVector& gradv) override;
    void hessian(const TVector& x, THessian& hessian) override
    {
        log_and_throw_error("Sparse functions do not exist, use dense solver");
    }
    void hessian(const TVector& x, MatrixXd& hessian) override;
    void solution_changed(const TVector& new_x) override {}
    bool is_step_valid(const TVector& x0, const TVector& x1) override;
 
private:
    bool cell_F(const Eigen::Vector3d& x, const Cell& c, Eigen::Matrix3d& F, double& d) const;
    std::vector<Cell> m_cells;
    double m_weight;
};
 
} // namespace wmtk::components::topological_offset