TriOptimizerMesh.h

#pragma once
 
#include <wmtk/OptimizerParameters.h>
#include <wmtk/RationalPositions.h>
#include <wmtk/SurfaceTagAttributes.h>
#include <wmtk/TriMesh.h>
#include <wmtk/Types.hpp>
#include <wmtk/envelope/Envelope.hpp>
#include <wmtk/optimization/SmoothVertex.hpp>
#include <wmtk/optimization/solver.hpp>
#include <wmtk/threading/enumerable_thread_specific.hpp>
 
#include <spdlog/fmt/bundled/format.h>
#include <polysolve/nonlinear/Problem.hpp>
 
#include <atomic>
#include <limits>
#include <map>
#include <memory>
#include <set>
#include <vector>
 
namespace wmtk {
 
class TriOptimizerMesh : public wmtk::TriMesh, public wmtk::RationalPositions
{
public:
    struct VertexAttributes
    {
        Vector2d m_posf; // position as double
        Vector2r m_pos; // exact position in rational
        bool m_is_rounded = false;
 
        bool m_is_on_surface = false;
        std::vector<int> on_bbox_faces;
 
        double m_sizing_scalar = 1;
 
        size_t partition_id = 0;
 
        VertexAttributes() {}
        VertexAttributes(const Vector2d& p)
            : m_posf(p)
            , m_pos(to_rational(p))
            , m_is_rounded(true)
        {}
        VertexAttributes(const Vector2r& p)
            : m_posf(to_double(p))
            , m_pos(p)
        {}
    };
 
    using EdgeAttributes = wmtk::SurfaceTagAttributes;
 
    struct FaceAttributes
    {
        double m_quality;
        double m_winding_number = 0;
        std::set<int64_t> tags;
        int part_id = -1;
    };
 
    using VertAttCol = AttributeCollection<VertexAttributes>;
    using EdgeAttCol = AttributeCollection<EdgeAttributes>;
    using FaceAttCol = AttributeCollection<FaceAttributes>;
 
    VertAttCol m_vertex_attribute;
    EdgeAttCol m_edge_attribute;
    FaceAttCol m_face_attribute;
 
    AttributeContainerGroup m_vertex_attr_group;
 
    AttributeContainerGroup m_edge_attr_group;
 
    AttributeContainerGroup m_face_attr_group;
 
    static constexpr double MAX_ENERGY = 1e50;
 
    OptimizerParameters& m_params;
 
    std::shared_ptr<SampleEnvelope> m_envelope;
    double m_envelope_eps = -1;
 
    double m_s_amips = -1;
    double m_s_envelope = -1;
 
    wmtk::threading::enumerable_thread_specific<std::unique_ptr<polysolve::nonlinear::Solver>>
        m_solver;
 
    optimization::SmoothRejectCounters m_smooth_rejects;
 
    bool m_collapse_limit_length = true;
    int m_debug_print_counter = 0;
    std::string m_debug_pass_name;
 
    size_t m_tags_count = 0;
    std::map<int64_t, std::string> m_tag_id_to_name;
    std::map<std::string, int64_t> m_tag_name_to_id;
 
    explicit TriOptimizerMesh(OptimizerParameters& params)
        : m_params(params)
    {
        m_vertex_attr_group.add(&m_vertex_attribute);
        p_vertex_attrs = &m_vertex_attr_group;
        m_edge_attr_group.add(&m_edge_attribute);
        p_edge_attrs = &m_edge_attr_group;
        m_face_attr_group.add(&m_face_attribute);
        p_face_attrs = &m_face_attr_group;
 
        m_s_amips = 1.;
        m_s_envelope = 1. / (m_params.eps * m_params.eps);
    }
    ~TriOptimizerMesh() override = default;
 
    size_t get_partition_id(const Tuple& loc) const
    {
        return m_vertex_attribute[loc.vid(*this)].partition_id;
    }
    void partition_mesh();
    void partition_mesh_morton();
 
    double get_length2(const Tuple& l) const;
 
    bool is_inverted(const std::array<size_t, 3>& vs) const;
    bool is_inverted(const Tuple& loc) const;
    bool is_inverted(const size_t fid) const;
    bool is_inverted_f(const Tuple& loc) const;
    bool is_inverted_f(const size_t fid) const;
 
    double get_quality(const std::array<size_t, 3>& vs) const;
    double get_quality(const Tuple& loc) const;
    double get_quality(const size_t fid) const;
 
    std::tuple<double, double> get_max_avg_energy();
 
    virtual void update_attributes() {}
 
    int m_iterations_used = 0;
    void mesh_improvement(int max_its = 80);
    std::tuple<double, double> local_operations(
        const std::array<int, 4>& ops,
        bool collapse_limit_length = true);
 
    std::set<simplex::Edge> m_force_split_edges;
    size_t m_force_split_count = 0;
    std::unique_ptr<std::atomic<int>[]> m_high_valence_claim;
    size_t m_high_valence_claim_size = 0;
    std::atomic<size_t> m_high_valence_rejects = 0;
 
    bool is_force_split_edge(const size_t v1, const size_t v2) const
    {
        return m_force_split_edges.find(simplex::Edge(v1, v2)) != m_force_split_edges.end();
    }
 
    void split_all_edges();
    bool split_edge_before(const Tuple& t) override;
    bool split_edge_after(const Tuple& loc) override;
 
    void collapse_all_edges(bool is_limit_length = true);
    bool collapse_edge_before(const Tuple& t) override;
    bool collapse_edge_after(const Tuple& t) override;
 
    size_t coarsen_mesh();
 
    bool coarsen_collapse_edge(const Tuple& e, std::vector<Tuple>& new_tris);
 
    struct CoarsenStats
    {
        size_t accepted = 0;
        size_t cells_before = 0;
        size_t cells_after = 0;
        double max_energy_before = 0.;
        double max_energy_after = 0.;
    };
    CoarsenStats m_coarsen_stats;
 
    size_t swap_all_edges();
    double swap_weight(const Tuple& t) const;
    bool swap_edge_before(const Tuple& t) override;
    bool swap_edge_after(const Tuple& t) override;
 
    void smooth_all_vertices(size_t n_iters = 1);
    bool smooth_before(const Tuple& t) override;
    bool smooth_after(const Tuple& t) override;
 
    Vector2d smoothing_position(size_t vid) const;
    void set_smoothing_position(size_t vid, const Vector2d& p);
    virtual bool smoothing_position_is_allowed(size_t vid, const Vector2d& p) const = 0;
 
    virtual std::shared_ptr<SampleEnvelope> smoothing_energy_envelope(const size_t vid) const
    {
        return m_vertex_attribute.at(vid).m_is_on_surface ? m_envelope : nullptr;
    }
    virtual std::shared_ptr<SampleEnvelope> smoothing_containment_envelope(const size_t vid) const
    {
        return m_vertex_attribute.at(vid).m_is_on_surface ? m_envelope : nullptr;
    }
 
    virtual std::shared_ptr<polysolve::nonlinear::Problem> smoothing_extra_energy(
        const size_t vid) const
    {
        return nullptr;
    }
 
    double active_quality_threshold() const
    {
        return m_params.skip_good_regions_margin * m_params.stop_energy;
    }
    virtual std::vector<size_t> active_vertices() const;
 
    virtual double quality_rel(const size_t fid) const
    {
        return m_face_attribute.at(fid).m_quality / m_params.stop_energy;
    }
 
    bool round(const Tuple& v);
 
protected:
    virtual std::tuple<double, double> optimization_quality_stats();
    virtual double optimization_stop_metric() const { return m_params.stop_energy; }
    virtual size_t refine_sizing_around_worst(double max_metric) = 0;
 
    virtual bool optimization_bare_coarsen_passes() const { return true; }
    uint32_t m_op_epoch = 0;
    virtual bool optimization_stop_at_float() const { return false; }
 
    virtual bool optimization_stalled(double prev, double cur)
    {
        return (prev - cur) <= m_params.stuck_refine_stall_eps * (cur - optimization_stop_metric());
    }
 
    virtual void optimization_debug_checkpoint()
    {
        if (m_params.debug_output) {
            write_smoothing_debug_output(fmt::format("debug_{}", m_debug_print_counter++));
        }
    }
 
    virtual void collapse_pass_begin() {}
    virtual void collapse_pass_end(size_t) {}
    virtual bool collapse_before_vertex(size_t, size_t) { return true; }
    virtual bool collapse_quality_allowed(size_t v1, size_t, double q, double ring_max) const
    {
        return !m_vertex_attribute.at(v1).m_is_rounded || q <= m_params.stop_energy ||
               q <= ring_max;
    }
    virtual void collapse_after_vertex(size_t, size_t) {}
 
    virtual bool split_adjust_position(size_t, const std::vector<Tuple>&) { return true; }
    virtual void split_after_vertex(size_t) {}
 
    virtual void write_smoothing_debug_output(const std::string& path) const = 0;
 
    // RationalPositions supplies round_all_vertices() and round_and_check_all_rounded() on
    // top of these three.
    std::vector<size_t> all_vertex_ids() const override;
    bool vertex_is_rounded(const size_t vid) const override
    {
        return m_vertex_attribute.at(vid).m_is_rounded;
    }
    // Equivalent to round() on the tuple get_vertices() would have yielded: round() reads only
    // the vid and the one-ring, and the one-ring query is answered from
    // m_vertex_connectivity[vid] regardless of which cell the tuple names.
    bool round_vertex(const size_t vid) override { return round(tuple_from_vertex(vid)); }
 
public:
    bool is_edge_on_surface(const Tuple& loc) const;
    bool is_edge_on_surface(const std::array<size_t, 2>& vids) const;
    bool is_edge_on_bbox(const Tuple& loc) const;
    bool is_edge_on_bbox(const std::array<size_t, 2>& vids) const;
 
    bool vertex_is_on_surface(const size_t vid) const override
    {
        return m_vertex_attribute.at(vid).m_is_on_surface ||
               !m_vertex_attribute.at(vid).on_bbox_faces.empty();
    }
    bool edge_is_on_surface(const std::array<size_t, 2>& vids) const override
    {
        if (!vertex_is_on_surface(vids[0]) || !vertex_is_on_surface(vids[1])) {
            return false;
        }
 
        const auto [_, eid] = tuple_from_edge(vids);
        bool on_surface = m_edge_attribute.at(eid).m_is_surface_fs;
        bool on_bbox = m_edge_attribute.at(eid).m_is_bbox_fs >= 0;
        return on_surface || on_bbox;
    }
 
    virtual std::shared_ptr<SampleEnvelope> surface_envelope_for_edge(
        const std::array<size_t, 2>& vids) const
    {
        return m_envelope;
    }
 
    bool surface_segment_is_outside(const size_t a, const size_t b) const
    {
        const std::shared_ptr<SampleEnvelope> env = surface_envelope_for_edge({{a, b}});
        if (!env) return false;
        const auto& VA = m_vertex_attribute;
        return env->is_outside(std::array<Vector2d, 2>{{VA[a].m_posf, VA[b].m_posf}});
    }
 
    std::vector<std::array<size_t, 2>> get_edges_by_condition(
        std::function<bool(const EdgeAttributes&)> cond) const;
 
    void gradation_smooth_sizing(double grade, const std::vector<size_t>& seeds);
 
protected:
    struct SplitInfoCache
    {
        size_t v1_id = 0;
        size_t v2_id = 0;
        EdgeAttributes old_e_attrs;
        std::map<simplex::Edge, EdgeAttributes> changed_edges;
        std::map<size_t, FaceAttributes> faces;
    };
    wmtk::threading::enumerable_thread_specific<SplitInfoCache> split_cache;
 
    struct CollapseInfoCache
    {
        size_t v1_id = 0;
        size_t v2_id = 0;
        double max_energy = 0.;
        double edge_length = 0.;
        std::vector<std::pair<EdgeAttributes, std::array<size_t, 2>>> changed_edges;
        std::vector<std::array<size_t, 2>> surface_edges;
        std::vector<size_t> changed_fids;
        std::vector<double> changed_energies;
        double region_max_rel_before = 0.;
    };
    wmtk::threading::enumerable_thread_specific<CollapseInfoCache> collapse_cache;
 
    bool m_coarsen_mode = false;
 
private:
    struct SwapInfoCache
    {
        double max_energy;
        std::map<simplex::Edge, EdgeAttributes> changed_edges;
        std::set<int64_t> face_tags;
    };
    wmtk::threading::enumerable_thread_specific<SwapInfoCache> swap_cache;
 
    struct CoarsenScratch
    {
        std::vector<size_t> ring; // vertices to re-smooth, BFS order
        std::vector<size_t> frontier;
        std::vector<size_t> next;
        std::vector<size_t> one_ring;
        std::vector<uint32_t> stamp;
        uint32_t epoch = 0;
        VertexAttributes saved_vertex;
        std::vector<std::pair<size_t, double>> saved_qualities;
    };
    wmtk::threading::enumerable_thread_specific<CoarsenScratch> coarsen_scratch;
 
    const std::vector<size_t>&
    collect_vertex_ball(const size_t* seeds, size_t n_seeds, int n, CoarsenScratch& scr) const;
 
    double region_max_quality_rel(const std::vector<size_t>& vids) const;
 
    bool smooth_vertex_reversible(size_t vid, CoarsenScratch& scr);
 
    size_t collapse_all_edges_impl(bool is_limit_length, int lock_ring, size_t max_passes = 0);
};
 
} // namespace wmtk