TriOptimizerMesh.h

#pragma once
 
#include <wmtk/OptimizerParameters.h>
#include <wmtk/RationalPositions.h>
#include <wmtk/SurfaceTagAttributes.h>
#include <wmtk/TriMesh.h>
#include <wmtk/ExactPosition.hpp>
#include <wmtk/Types.hpp>
#include <wmtk/envelope/Envelope.hpp>
#include <wmtk/history/History.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 <optional>
#include <set>
#include <vector>
 
namespace wmtk {
 
class TriOptimizerMesh : public wmtk::TriMesh, public wmtk::RationalPositions
{
public:
    struct VertexAttributes
    {
        Vector2d m_posf; // position as double
        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;
 
        Vector2r pos() const { return m_exact.value(m_posf); }
        bool pos_coord_equals(int k, const Rational& r) const
        {
            return m_exact.coord_equals(k, m_posf, r);
        }
        void set_pos(const Vector2r& p) { m_exact.store(p); }
        void set_pos_to_posf() { m_exact.clear(); }
        ExactPosition<Vector2r, Vector2d> m_exact;
        VertexAttributes() {}
        VertexAttributes(const Vector2d& p)
            : m_posf(p)
            , m_is_rounded(true)
        {}
        VertexAttributes(const Vector2r& p)
            : m_posf(to_double(p))
        {
            set_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;
 
    bool m_tracks_orientation = false;
 
    std::vector<std::array<size_t, 2>> oriented_tracked_edges() const;
 
    std::vector<std::pair<size_t, int>> tracked_curve_boundary() const;
    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<Envelope> 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;
 
    static constexpr size_t NO_FEATURE_POINT = std::numeric_limits<size_t>::max();
    struct FeaturePointAttributes
    {
        size_t m_feature_point_id = NO_FEATURE_POINT;
    };
    using FeaturePointAttCol = AttributeCollection<FeaturePointAttributes>;
    FeaturePointAttCol m_feature_point_attribute;
    std::vector<Vector2d> m_feature_points;
    double m_feature_eps = -1;
    bool m_track_feature_points = false;
    bool m_preserve_feature_points = false;
    std::atomic<size_t> m_feature_rejects = 0;
 
    bool m_in_operation_pass = false;
    void throw_if_in_operation_pass(const char* call) const
    {
        if (m_in_operation_pass) {
            log_and_throw_error(
                "{}() called while an operation pass runs; call it before or after the pass.",
                call);
        }
    }
    void enable_feature_point_tracking(const bool preserve)
    {
        throw_if_in_operation_pass("enable_feature_point_tracking");
        if (!m_track_feature_points) {
            m_vertex_attr_group.add(&m_feature_point_attribute);
            m_feature_point_attribute.resize(vert_storage_capacity());
        }
        m_track_feature_points = true;
        m_preserve_feature_points = preserve;
    }
 
    bool collapse_breaks_feature_point(const size_t v1_id, const size_t v2_id) const;
 
    bool feature_point_ball_contains(const size_t vid, const Vector2d& p) const;
 
    std::pair<size_t, size_t> feature_retention(double* worst_ratio = nullptr) const;
    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 storage_grows_on_demand() const override { return true; }
 
protected:
    void on_slot_storage_grown() override
    {
        if (!m_high_valence_claim) return;
        const size_t n = std::max(vert_capacity(), m_vertex_attribute.size());
        if (n <= m_high_valence_claim_size) return;
        auto grown = std::make_unique<std::atomic<int>[]>(n);
        for (size_t i = 0; i < m_high_valence_claim_size; ++i) {
            grown[i].store(m_high_valence_claim[i].load(std::memory_order_relaxed));
        }
        m_high_valence_claim = std::move(grown);
        m_high_valence_claim_size = n;
    }
 
public:
    bool split_edge_is_due(const Tuple& e) const;
    size_t m_split_order_waits = 0;
    size_t m_split_order_wait_defects = 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 use_colored_smoothing() const;
    void smooth_vertices_colored(const std::vector<std::pair<std::string, Tuple>>& ops);
    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(const size_t vid, const Vector2d& p) const
    {
        return feature_point_ball_contains(vid, p);
    }
 
    virtual std::shared_ptr<Envelope> 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<Envelope> smoothing_containment_envelope(
        const size_t /*vid*/,
        const std::array<size_t, 2>& edge,
        const EdgeAttributes& attr) const
    {
        return surface_envelope_for_edge(edge, attr);
    }
 
    virtual std::shared_ptr<polysolve::nonlinear::Problem> smoothing_extra_energy(
        const size_t vid) const
    {
        return nullptr;
    }
 
    virtual double smoothing_cell_energy(const size_t /*fid*/, const double quality) const
    {
        return quality;
    }
 
    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 swap_quality_allowed(double after, double before, bool /*is_surface_flip*/) const
    {
        return after < before;
    }
 
    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_tracked_edge(const Tuple& e) const override { return is_edge_on_surface(e); }
    std::vector<size_t> tracked_feature_points() const override;
 
    using CurveSources = std::map<std::array<size_t, 2>, std::vector<uint64_t>>;
    void record_curve_history(const std::string& metadata, const CurveSources& sources = {});
    std::optional<history::CurveHistory> finish_curve_history(
        const history::CurveHistory* input = nullptr);
 
private:
    std::shared_ptr<history::TrackedCurveRecorder> m_recorded_curves;
    std::vector<std::vector<uint64_t>> m_recorded_curve_sources;
 
public:
    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<Envelope> surface_envelope_for_edge(
        const std::array<size_t, 2>& vids,
        const EdgeAttributes& attr) const
    {
        return m_envelope;
    }
 
    bool surface_segment_is_outside(const size_t a, const size_t b, const EdgeAttributes& attr)
        const
    {
        const std::shared_ptr<Envelope> env = surface_envelope_for_edge({{a, b}}, attr);
        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::pair<std::array<size_t, 2>, EdgeAttributes>> surface_edges;
        std::vector<std::pair<std::array<size_t, 2>, int>> renamed_orientations;
        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