OffsetPotential.hpp

#pragma once
 
#include <functional>
 
#include <wmtk/Types.hpp>
#include <wmtk/envelope/Envelope.hpp>
 
#include <wmtk/threading/enumerable_thread_specific.hpp>
#include "ExactBandIntegral.hpp"
#include "SimplicialComplexBVH.hpp"
 
#include <polysolve/nonlinear/Problem.hpp>
 
#include <array>
#include <atomic>
#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 void value_gradient(const VecD& p, double& v, VecD& g) const
    {
        v = value(p);
        g = gradient(p);
    }
 
    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;
    void value_gradient(const VecD& p, double& v, VecD& g) 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;
 
    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<Envelope>& 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<Envelope> 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,
        bool one_sided = 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 m_one_sided;
    void residual(const VecD& p, double& r, VecD& dr) const;
};
 
using OffsetEnergy2D = OffsetEnergy<2>;
using OffsetEnergy3D = OffsetEnergy<3>;
 
class StencilEnergy2D : 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;
    };
    struct Edge
    {
        Eigen::Vector2d q1;
        std::vector<Sample> samples;
    };
 
    StencilEnergy2D(
        const std::shared_ptr<const OffsetPotential2D>& potential,
        std::vector<Edge> edges,
        double weight,
        bool gauss_newton = 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:
    struct Reading
    {
        double r = 0.;
        Eigen::Vector2d dr = Eigen::Vector2d::Zero();
        bool r_ok = false;
        bool dr_ok = false;
    };
 
    const std::vector<Reading>& readings_at(const Eigen::Vector2d& x, bool need_dr) const;
 
    std::shared_ptr<const OffsetPotential2D> m_potential;
    std::vector<Edge> m_edges;
    double m_weight;
    bool m_gauss_newton;
    double m_c = 1.; 
 
    mutable std::vector<Reading> m_readings;
    mutable Eigen::Vector2d m_readings_x;
    mutable bool m_readings_valid = false;
    mutable bool m_readings_have_dr = false;
};
 
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 CubedAMIPSEnergy3D : public polysolve::nonlinear::Problem
{
public:
    using typename polysolve::nonlinear::Problem::Scalar;
    using typename polysolve::nonlinear::Problem::THessian;
    using typename polysolve::nonlinear::Problem::TVector;
    CubedAMIPSEnergy3D(std::vector<std::array<double, 12>> 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:
    void at_point(const TVector& x);
 
    std::vector<std::array<double, 12>> m_cells;
    double m_weight;
    bool m_x_set = false;
    Eigen::Vector3d m_x = Eigen::Vector3d::Zero();
    std::vector<double> m_a; 
    bool m_value_set = false;
    double m_value = 0.;
    bool m_grad_set = false;
    TVector m_grad;
};
 
class VolAMIPSEnergy3D : 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
    {
        std::array<Eigen::Vector3d, 3> q;
        Eigen::Matrix3d rest_inv;
        Eigen::Vector3d s; 
        double c; 
    };
    static bool cell(const std::array<Eigen::Vector3d, 3>& q, const Eigen::Matrix3d& R, Cell& out);
    static Eigen::Matrix3d regular_rest();
    static double value_of(const std::array<Eigen::Vector3d, 4>& p, const Eigen::Matrix3d& R);
 
    VolAMIPSEnergy3D(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:
    double eval(const Eigen::Vector3d& x, int need, Eigen::Vector3d& g, Eigen::Matrix3d& H) const;
    std::vector<Cell> m_cells;
    double m_weight;
};
 
double input_pad(const Eigen::MatrixXd& V);
 
int64_t nearest_primitive(const Envelope& env, double pad, const Eigen::Vector3d& p);
 
double min_mean_distance(
    const Envelope& env,
    double pad,
    const Eigen::Vector3d* p,
    int n,
    int64_t& best,
    int64_t hint = -1);
 
struct BallCandidate
{
    int64_t tri = -1;
    double d1 = 0., d2 = 0., d3 = 0.;
};
void ball_candidates(
    const Envelope& env,
    double pad,
    const Eigen::Vector3d& c,
    double r,
    const std::array<Eigen::Vector3d, 3>& q,
    std::vector<BallCandidate>& out);
 
double min_triangle_distance(
    const Envelope& env,
    double pad,
    const Eigen::Vector3d& a,
    const Eigen::Vector3d& b,
    const Eigen::Vector3d& c);
 
double min_max_corner_distance(const Envelope& env, double pad, const Eigen::Vector3d* p, int n);
 
double face_mean_sq_cos(
    const Envelope& env,
    double pad,
    const Eigen::Vector3d& a,
    const Eigen::Vector3d& b,
    const Eigen::Vector3d& c,
    const Eigen::Vector3d& n,
    double delta);
 
class BandVolumeEnergy3D : 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;
    };
    BandVolumeEnergy3D(
        std::shared_ptr<const Envelope> env,
        double pad,
        std::vector<Cell> cells,
        const Eigen::Vector3d& x0,
        double delta);
 
    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 {}
    void set_exact_integral(std::shared_ptr<const ExactBandIntegral> eb, double delta);
    struct Face
    {
        Eigen::Vector3d a, b;
    };
    void set_front_flux(
        std::shared_ptr<const ExactBandIntegral> eb,
        double delta,
        std::vector<Face> faces);
    static std::string& front_flux_error();
    static double swept_change(
        const ExactBandIntegral& eb,
        double delta,
        const Eigen::Vector3d& x0,
        const Eigen::Vector3d& x1,
        const std::vector<std::array<Eigen::Vector3d, 2>>& faces,
        ExactBandIntegral::Stats* st = nullptr);
    struct SweptDerivCounts
    {
        std::atomic<size_t> evals{0}, ring_cells{0};
    };
    void set_swept(
        std::shared_ptr<const ExactBandIntegral> eb,
        double delta,
        std::vector<std::array<Eigen::Vector3d, 2>> faces,
        SweptDerivCounts* counts);
    std::pair<Eigen::Vector3d, double> certified_ball() const { return {m_ball_c, m_ball_r}; }
 
private:
    double eval(const Eigen::Vector3d& x, int need, Eigen::Vector3d& g, Eigen::Matrix3d& H) const;
    void exact_minima(const Eigen::Vector3d& x) const;
    bool certify_ball(const Eigen::Vector3d& x) const;
    std::shared_ptr<const Envelope> m_env; 
    double m_pad = 0.; 
    double m_delta = 1.;
    mutable std::vector<int64_t> m_hint; 
    mutable bool m_min_set = false;
    mutable Eigen::Vector3d m_min_x = Eigen::Vector3d::Zero();
    mutable std::vector<double> m_min_md;
    mutable std::vector<int64_t> m_min_tri;
    mutable Eigen::Vector3d m_ball_c = Eigen::Vector3d::Zero();
    mutable double m_ball_r = -1.;
    mutable std::vector<std::vector<BallCandidate>> m_ball_cand;
    std::shared_ptr<const ExactBandIntegral> m_eb;
    mutable std::vector<ExactBandIntegral::Partition> m_eb_parts;
    mutable bool m_eb_val_set = false, m_eb_der_set = false;
    mutable Eigen::Vector3d m_eb_val_x = Eigen::Vector3d::Zero();
    mutable Eigen::Vector3d m_eb_der_x = Eigen::Vector3d::Zero();
    mutable double m_eb_E = 0.;
    mutable Eigen::Vector3d m_eb_g = Eigen::Vector3d::Zero();
    mutable Eigen::Matrix3d m_eb_H = Eigen::Matrix3d::Zero();
    bool m_ff = false;
    std::vector<Face> m_ff_faces;
    mutable std::vector<ExactBandIntegral::FaceParts> m_ff_fparts;
    mutable std::vector<double> m_ff_U;
    mutable std::vector<Eigen::Vector3d> m_ff_Ux;
    mutable std::vector<char> m_ff_Ux_set;
    double ff_eval(const Eigen::Vector3d& x, int need, Eigen::Vector3d& g, Eigen::Matrix3d& H)
        const;
    bool m_sw = false;
    std::vector<std::array<Eigen::Vector3d, 2>> m_sw_faces;
    mutable std::vector<ExactBandIntegral::FaceParts> m_sw_fparts;
    SweptDerivCounts* m_sw_counts = nullptr;
    double sw_eval(const Eigen::Vector3d& x, int need, Eigen::Vector3d& g, Eigen::Matrix3d& H)
        const;
    Eigen::Vector3d m_x0 = Eigen::Vector3d::Zero();
 
    std::vector<Cell> m_cells;
};
 
const std::vector<std::array<double, 3>>& subdivision_face_stencil(int order);
 
const std::vector<std::array<double, 3>>& interior_face_stencil(int order);
 
const std::vector<std::array<double, 3>>& fem_face_stencil(int order);
 
const std::vector<double>& fem_face_stencil_weights(int order);
 
class SurfaceOffsetError3D : 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;
    };
    SurfaceOffsetError3D(
        std::shared_ptr<const OffsetPotential3D> pot,
        std::vector<Face> faces,
        int stencil_order,
        double weight);
    SurfaceOffsetError3D(
        std::shared_ptr<const OffsetPotential3D> pot,
        std::vector<Face> faces,
        const std::vector<std::array<double, 3>>& stencil,
        const std::vector<double>* point_weights,
        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 {}
 
    static double face_value(
        const OffsetPotential3D& pot,
        const std::vector<std::array<double, 3>>& stencil,
        const std::vector<double>* point_weights,
        double weight,
        const Eigen::Vector3d& a,
        const Eigen::Vector3d& b,
        const Eigen::Vector3d& c);
 
private:
    double eval(const Eigen::Vector3d& x, int need, Eigen::Vector3d& g, Eigen::Matrix3d& H) const;
    std::shared_ptr<const OffsetPotential3D> m_pot;
    std::vector<Face> m_faces;
    const std::vector<std::array<double, 3>>& m_stencil;
    const std::vector<double>* m_point_weights = nullptr;
    double m_weight;
};
 
} // namespace wmtk::components::topological_offset