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 <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;
 
    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;
    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);
 
    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; }
 
    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