TetOptimizerMesh.h

#pragma once
 
#include <wmtk/OptimizerParameters.h>
#include <wmtk/RationalPositions.h>
#include <wmtk/SurfaceTagAttributes.h>
#include <wmtk/TetMesh.h>
#include <wmtk/Types.hpp>
#include <wmtk/envelope/Envelope.hpp>
#include <wmtk/optimization/SmoothVertex.hpp>
#include <wmtk/optimization/solver.hpp>
#include <wmtk/simplex/Simplex.hpp>
#include <wmtk/threading/enumerable_thread_specific.hpp>
#include <wmtk/utils/SurfaceTopology.hpp>
 
#include <igl/Timer.h>
 
#include <array>
#include <atomic>
#include <cmath>
#include <functional>
#include <map>
#include <memory>
#include <set>
#include <vector>
 
namespace wmtk {
 
class TetOptimizerMesh : public wmtk::TetMesh, public wmtk::RationalPositions
{
public:
    struct VertexAttributes
    {
        Vector3r m_pos; // exact position in rational
        Vector3d m_posf; // position as double
        bool m_is_rounded = false;
 
        bool m_is_on_surface = false;
        size_t m_order = 0;
        std::vector<int> on_bbox_faces;
 
        double m_sizing_scalar = 1;
 
        size_t partition_id = 0;
 
        VertexAttributes() {}
        VertexAttributes(const Vector3r& p);
    };
 
    using FaceAttributes = wmtk::SurfaceTagAttributes;
 
    using VertAttCol = AttributeCollection<VertexAttributes>;
    using FaceAttCol = AttributeCollection<FaceAttributes>;
 
    VertAttCol m_vertex_attribute;
    FaceAttCol m_face_attribute;
 
    AttributeContainerGroup m_vertex_attr_group;
 
    AttributeContainerGroup m_face_attr_group;
 
    static constexpr double MAX_ENERGY = 1e50;
 
    OptimizerParameters& m_params;
 
    std::shared_ptr<SampleEnvelope> m_envelope;
 
    double m_s_amips = 1.;
    double m_s_envelope = -1.;
 
    double time_env = 0.0;
    igl::Timer isout_timer;
 
    wmtk::threading::enumerable_thread_specific<std::unique_ptr<polysolve::nonlinear::Solver>>
        m_solver;
 
    optimization::SmoothRejectCounters m_smooth_rejects;
    optimization::NewtonCounters m_newton;
 
    bool m_collapse_limit_length = true;
 
    std::set<simplex::Edge> m_force_split_edges;
    size_t m_force_split_count = 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();
    }
 
    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;
 
    explicit TetOptimizerMesh(OptimizerParameters& params, std::shared_ptr<SampleEnvelope> env)
        : m_params(params)
        , m_envelope(std::move(env))
    {
        m_vertex_attr_group.add(&m_vertex_attribute);
        p_vertex_attrs = &m_vertex_attr_group;
        m_face_attr_group.add(&m_face_attribute);
        p_face_attrs = &m_face_attr_group;
    }
    ~TetOptimizerMesh() override = default;
 
    virtual double cell_quality(const size_t tid) const = 0;
    virtual void set_cell_quality(const size_t tid, const double q) = 0;
 
    virtual double quality_rel(const size_t tid) const
    {
        return std::cbrt(cell_quality(tid)) / m_params.stop_energy;
    }
 
    // TODO This should not be here but inside wmtk
    void compute_vertex_partition();
    void compute_vertex_partition_morton();
 
    size_t get_partition_id(const Tuple& loc) const
    {
        return m_vertex_attribute[loc.vid(*this)].partition_id;
    }
 
    double get_length2(const Tuple& l) const;
 
    bool is_inverted(const std::array<size_t, 4>& vs) const;
    bool is_inverted(const Tuple& loc) const;
    bool is_inverted_f(const Tuple& loc) const;
 
    double get_quality(const std::array<size_t, 4>& vs) const;
    double get_quality(const Tuple& loc) const;
    std::tuple<double, double> get_max_avg_energy();
 
    virtual void update_attributes() {}
 
    int m_iterations_used = 0;
    int m_debug_print_counter = 0;
    std::string m_debug_pass_name;
    void mesh_improvement(int max_its = 80);
    std::tuple<double, double> local_operations(
        const std::array<int, 4>& ops,
        bool collapse_limit_length = true);
 
    bool round(const Tuple& v);
 
protected:
    // 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);
    bool is_edge_on_bbox(const Tuple& loc);
    virtual bool is_open_boundary_edge(const Tuple& e) { return false; }
    int edge_incident_surface_face_count(const Tuple& e);
 
    virtual std::shared_ptr<SampleEnvelope> surface_envelope_for_face(
        const std::array<size_t, 3>& vids) const
    {
        return m_envelope;
    }
 
    bool surface_triangle_is_outside(const size_t a, const size_t b, const size_t c) const
    {
        const std::shared_ptr<SampleEnvelope> env = surface_envelope_for_face({{a, b, c}});
        if (!env) return false;
        const auto& VA = m_vertex_attribute;
        return env->is_outside({{VA[a].m_posf, VA[b].m_posf, VA[c].m_posf}});
    }
 
    bool vertex_is_on_surface(const size_t vid) const override;
    bool face_is_on_surface(const size_t fid) const override;
    size_t get_order_of_vertex(const size_t vid) const override;
 
    std::vector<std::array<size_t, 3>> get_faces_by_condition(
        std::function<bool(const FaceAttributes&)> cond) const;
 
    void output_faces(std::string file, std::function<bool(const FaceAttributes&)> cond);
 
    void gradation_smooth_sizing(double grade, const std::vector<size_t>& seeds);
 
    double active_quality_threshold() const
    {
        const double e = m_params.skip_good_regions_margin * m_params.stop_energy;
        return e * e * e;
    }
 
    virtual std::vector<size_t> active_vertices() const;
 
    // TetWild's split engine. SimWild inherits it and customizes only application data,
    // annotation-only placement, and order-2 metadata through the hooks below.
    void split_all_edges();
    bool split_edge_before(const Tuple& t) override;
    bool split_edge_after(const Tuple& loc) override;
 
    // TetWild's collapse engine. SimWild customizes only application policy/metadata.
    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_tets);
 
    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;
 
    // TetWild's complete swap engine. SimWild inherits these operations and customizes only
    // the tag bookkeeping through the hooks below.
    size_t swap_all_edges_32();
    bool swap_edge_before(const Tuple& t) override;
    bool swap_edge_after(const Tuple& t) override;
 
    size_t swap_all_edges_44();
    bool swap_edge_44_before(const Tuple& t) override;
    bool swap_edge_44_accept_case(const std::array<size_t, 2>& new_edge) override;
    bool swap_edge_44_after(const Tuple& t) override;
 
    size_t swap_all_edges_56();
    bool swap_edge_56_before(const Tuple& t) override;
    bool swap_edge_56_accept_case(const std::array<size_t, 3>& new_face) override;
    bool swap_edge_56_after(const Tuple& t) override;
 
    size_t swap_all_faces();
    bool swap_face_before(const Tuple& t) override;
    bool swap_face_after(const Tuple& t) override;
    size_t swap_all_edges_all();
 
    bool prepare_surface_flip(const Tuple& t, const std::vector<size_t>& incident_tets);
 
    using SurfaceTopoSignature = wmtk::utils::SurfaceTopoSignature;
    SurfaceTopoSignature surface_topology_signature() const
    {
        return wmtk::utils::surface_topology_signature(*this, [this](size_t fid) {
            return m_face_attribute[fid].m_is_surface_fs;
        });
    }
    void warn_if_surface_topology_changed(const SurfaceTopoSignature& before, const char* where)
        const
    {
        wmtk::utils::warn_if_surface_topology_changed(before, surface_topology_signature(), where);
    }
 
    // Operation diagnostics shared by TetWild and SimWild.
    std::atomic<int> cnt_swap = 0;
    std::atomic<int> cnt_surface_swap = 0;
    std::atomic<int> cnt_surface_swap_32 = 0, cnt_surface_swap_44 = 0, cnt_surface_swap_56 = 0;
 
    double swap_edge_44_energy(const std::vector<std::array<size_t, 4>>& tets, const int op_case)
        override;
    double swap_edge_56_energy(const std::vector<std::array<size_t, 4>>& tets, const int op_case)
        override;
 
    virtual std::shared_ptr<SampleEnvelope> smoothing_containment_envelope(const size_t vid) const;
    virtual std::shared_ptr<SampleEnvelope> smoothing_energy_envelope(const size_t vid) const = 0;
 
    virtual std::shared_ptr<polysolve::nonlinear::Problem> smoothing_extra_energy(
        const size_t /*vid*/) const
    {
        return nullptr;
    }
 
    bool smooth_before(const Tuple& t) override;
    bool smooth_after(const Tuple& t) override;
    void smooth_all_vertices(const size_t n_iters = 1);
    virtual void log_smoothing_pass_accounting() {}
 
    bool invariants(const std::vector<Tuple>& t) override;
 
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; }
    virtual void write_optimization_debug_output(const std::string& path) = 0;
    virtual void optimization_sanity_checks_extra() {}
 
    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 bool collapse_before_vertex(size_t, size_t, double) { return true; }
    virtual bool collapse_quality_allowed(size_t v1, double quality, double ring_max) const
    {
        return !m_vertex_attribute.at(v1).m_is_rounded || quality <= ring_max;
    }
    virtual bool collapse_is_order_2_edge(const std::array<size_t, 2>&) { return false; }
    virtual bool
    collapse_after_connectivity(size_t, size_t, const std::vector<std::array<size_t, 2>>&)
    {
        return true;
    }
    virtual void collapse_after_vertex(size_t, size_t) {}
 
    virtual bool split_before_cells(const Tuple&, const std::vector<Tuple>&) { return true; }
    virtual bool split_after_cells(size_t, size_t, size_t, const std::vector<Tuple>&)
    {
        return true;
    }
    virtual bool split_adjust_position(size_t, const std::vector<Tuple>&) { return true; }
    virtual void split_after_vertex(size_t, bool) {}
 
    virtual bool allow_surface_swap() const = 0;
    virtual bool check_surface_topology() const = 0;
 
    virtual bool swap_before_interior(const std::vector<size_t>&) { return true; }
    virtual bool swap_before_surface(const std::vector<size_t>&, size_t, size_t, size_t, size_t)
    {
        return true;
    }
    virtual bool swap_after_cells(const std::vector<size_t>&, bool) { return true; }
 
    // ---- SWAP REJECTION INSTRUMENTATION (debug; see swap_reject_report()) ----
    // Every place a proposed swap can be turned down, counted. Added to find out why an offset
    // surface flip is never accepted; the counters cost one relaxed atomic increment on a path
    // that is already refusing the operation, and nothing reads them unless a caller asks.
    //
    // The refusals TetMesh makes itself, below every hook (the valence re-test, an open ring, a
    // face or edge that already exists, no allowed 4-4 / 5-6 case beating the current cells, out
    // of slots, a duplicate tet), were once visible here only as before_pass - after_enter. They
    // are now reported one by one through TetMesh::op_event(); see op_accounting_report(). The
    // per-turn line of this block still prints that difference as `connectivity`.
    //
    // Face swaps count their own reasons (face_*, and base_before / interior_hook / base_after /
    // after_cells on their path) in the per-kind table alone -- swap_reject_kind_only() -- so
    // swap_reject_report(), which has only ever described edge swaps, keeps meaning what it meant.
    enum class SwapReject : int {
        base_before, // TetMesh::swap_*_before said no
        base_after, // TetMesh::swap_*_after said no
        valence, // wrong incident-tet count for this swap kind
        bbox, // edge on the bounding box
        surface_not_allowed, // has surface faces and allow_surface_swap() is false
        prepare_flip, // prepare_surface_flip() refused (see the flip_* reasons)
        interior_hook, // swap_before_interior() refused
        flip_open_boundary, // flipping would change the surface's boundary loops
        flip_not_two_surf, // not exactly two incident surface faces
        flip_nonmanifold_edge, // more than two incident surface faces
        flip_cd_nonmanifold, // a surface face already sits on the new edge (c,d)
        flip_new_face_surface, // (a,c,d) or (b,c,d) is already tagged surface
        flip_app_refused, // swap_before_surface() refused (see the app_* reasons)
        flip_wrong_case, // 4-4 / 5-6: the retetrahedralization is not the one that makes
        // the surface diagonal (c,d), so this case is not the flip
        app_capture_label, // application: the swap ring spans two construction labels
        app_capture_tag, // application: the swap ring spans two cell tags
        app_fid_missing, // application: a face of the pair could not be found
        app_class_mismatch, // application: the two faces are different tracked surfaces
        app_mask_mismatch, // application: the two faces carry different boundary masks
        app_sag_raised, // application: the flip would raise the local sag
        after_inverted, // a created cell is inverted
        after_quality, // swap_quality_allowed() said no
        after_cells, // swap_after_cells() refused (see the app_after_* reasons)
        app_after_side_conflict, // application: two ring vertices of a new cell sit on two sides
        app_after_no_side, // application: a new cell holds no ring vertex to take a side from
        after_envelope, // a new surface triangle left the envelope
        face_tracked_surface, // face swap: the face is a tracked surface face
        face_tracked_bbox, // face swap: the face is on the bounding box
        face_inverted, // face swap: one of the three new cells would be inverted
        face_not_better, // face swap: a new cell is not strictly better than the two old ones
        COUNT
    };
    enum class SwapStage : int {
        attempt,
        surface_attempt,
        before_pass,
        after_enter,
        accepted,
        COUNT
    };
 
    static const char* swap_reject_name(SwapReject r);
    static const char* swap_stage_name(SwapStage s);
    bool swap_reject(SwapReject r) const
    {
        m_swap_reject[size_t(r)].fetch_add(1, std::memory_order_relaxed);
        return swap_reject_kind_only(r);
    }
    bool swap_reject_kind_only(SwapReject r) const
    {
        m_swap_reject_by_kind[size_t(current_op_kind())][size_t(r)].fetch_add(
            1,
            std::memory_order_relaxed);
        return false;
    }
    void swap_stage(SwapStage s) const
    {
        m_swap_stage[size_t(s)].fetch_add(1, std::memory_order_relaxed);
    }
    void swap_counters_reset();
    std::string swap_reject_report() const;
 
    mutable std::array<std::atomic<long>, size_t(SwapReject::COUNT)> m_swap_reject{};
    mutable std::array<std::atomic<long>, size_t(SwapStage::COUNT)> m_swap_stage{};
 
    // ---- COLLAPSE REJECTION INSTRUMENTATION (see op_accounting_report()) ----
    // Every place a collapse is turned down in a hook, counted once, at the place that decided.
    // Before this, 85% of the collapse attempts on the deliverable cube (target_distance_rel 1e-2
    // / front_conv_rel 1e-4, turn 5: 213k of 232k) were refused with nothing recorded but the
    // ops guard's 367. What TetMesh refuses below the hooks (the link condition, a duplicate tet)
    // arrives through op_event(). Before-hook reasons come first, after-hook reasons from
    // after_base on.
    enum class CollapseReject : int {
        coarsen_length, // coarsening pass only: the edge is at or above the collapse target
        bbox, // v1 sits on a bounding-box face that v2 does not
        surface_leaves_envelope, // v1 on the surface, v2 not, and v2 not inside the envelope
        inverted, // a cell the collapse reshapes would invert
        quality, // collapse_quality_allowed() said no
        substructure_link, // preserve_topology, both ends on the surface: substructure link
        app_both_surfaces, // application: an end on the input complex and an end on the offset
        app_front_unlimited, // application: an unlimited-length pass may not remove a front vertex
        app_leaves_input, // application: v1 on the input complex, v2 not
        app_leaves_offset, // application: v1 on the offset surface, v2 not
        app_leaves_region, // application: v1 on a region boundary, v2 not
        app_order2, // application: an order-2 vertex would merge into a lower-order one
        app_substructure_link, // application: the substructure link condition, unconditional
        app_ops_guard, // application: the ops divergence guard
        after_base, // TetMesh::collapse_edge_after said no
        after_connectivity, // collapse_after_connectivity() said no
        after_envelope, // a surface face the survivor keeps left the envelope
        after_face_attribute, // a face attribute found no face left to land on
        after_coarsen_region, // coarsening only: the region's worst quality rose
        COUNT
    };
    static const char* collapse_reject_name(CollapseReject r);
    bool collapse_reject(CollapseReject r) const
    {
        m_collapse_reject[size_t(r)].fetch_add(1, std::memory_order_relaxed);
        return false;
    }
    mutable std::array<std::atomic<long>, size_t(CollapseReject::COUNT)> m_collapse_reject{};
 
    // ---- OPERATION ACCOUNTING (counting only) ----
    // TetMesh reports every operation's attempt and its one outcome through op_event(); this
    // counts them per kind. op_accounting_report() then checks, per kind, that the outcomes add
    // up to the attempts, and that the before- and after-hook refusals add up to the reasons the
    // hooks counted (SwapReject / CollapseReject). A nonzero remainder is a refusal nobody named.
    void op_event(OpKind k, OpEvent e) const override
    {
        if (e == OpEvent::attempt) current_op_kind() = k;
        m_op_events[size_t(k)][size_t(e)].fetch_add(1, std::memory_order_relaxed);
    }
    static OpKind& current_op_kind()
    {
        static thread_local OpKind k = OpKind::COUNT;
        return k;
    }
    std::string op_accounting_report(OpKind k) const;
    bool op_accounting_unexplained(OpKind k) const;
    void op_accounting_reset();
    static const char* op_kind_name(OpKind k);
    static const char* op_event_name(OpEvent e);
 
    enum class SwapRejectRole { before_leaf, after_leaf, wrapper, per_case };
    static SwapRejectRole swap_reject_role(SwapReject r);
    struct OpHookReasons
    {
        std::string before, after;
        long before_sum = 0, after_sum = 0;
        bool counted = true; // false: this kind's hooks count no reasons (split)
    };
    OpHookReasons op_hook_reasons(OpKind k) const;
    long op_event_count(const OpKind k, const OpEvent e) const
    {
        return m_op_events[size_t(k)][size_t(e)].load(std::memory_order_relaxed);
    }
    long op_unaccounted(OpKind k) const;
 
    mutable std::array<std::array<std::atomic<long>, size_t(OpEvent::COUNT)>, size_t(OpKind::COUNT)>
        m_op_events{};
    mutable std::
        array<std::array<std::atomic<long>, size_t(SwapReject::COUNT)>, size_t(OpKind::COUNT) + 1>
            m_swap_reject_by_kind{};
 
    virtual bool swap_quality_allowed(double after, double before, bool /*is_surface_flip*/) const
    {
        return after < before;
    }
 
    struct SwapInfoCache
    {
        double max_energy;
        std::map<std::array<size_t, 3>, FaceAttributes> changed_faces;
 
        bool is_surface_flip = false;
        size_t sf_a = 0, sf_b = 0, sf_c = 0, sf_d = 0;
        FaceAttributes sf_face_attr;
    };
    wmtk::threading::enumerable_thread_specific<SwapInfoCache> swap_cache;
 
    struct SplitInfoCache
    {
        size_t v1_id = 0;
        size_t v2_id = 0;
        bool is_edge_on_surface = false;
        bool is_edge_open_boundary = false;
        size_t edge_order = 0;
        std::vector<std::pair<FaceAttributes, std::array<size_t, 3>>> changed_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<FaceAttributes, std::array<size_t, 3>>> changed_faces;
        std::vector<std::array<size_t, 3>> surface_faces;
        std::vector<std::array<size_t, 2>> boundary_edges;
        std::vector<size_t> changed_tids;
        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 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,
        bool exact_ball_lock = false);
};
 
} // namespace wmtk