TopoOffsetTetMesh.h

#pragma once
 
#include <wmtk/TetMesh.h>
#include <wmtk/TetOptimizerMesh.h>
#include <wmtk/envelope/Envelope.hpp>
#include <wmtk/optimization/solver.hpp>
#include <wmtk/simplex/Simplex.hpp>
#include <wmtk/threading/enumerable_thread_specific.hpp>
#include "Parameters.h"
#include "SimplicialComplexBVH.hpp"
 
// clang-format off
#include <wmtk/utils/DisableWarnings.hpp>
#include <wmtk/utils/EnableWarnings.hpp>
// clang-format on
 
using CellTag = std::set<int64_t>;
 
 
namespace wmtk::components::topological_offset {
 
 
const int64_t TEMP_OFFSET_TET_TAG = -1;
const CellTag TEMP_OFFSET_TET_TAG_SET{TEMP_OFFSET_TET_TAG};
 
 
// for all attributes:
// label: 0=default, 1=input, 2=offset
 
class VertexExtra
{
public:
    int label = 0;
    size_t component_id = 0;
    bool m_is_on_input = false; // is this vertex on the input complex
    bool m_is_on_offset = false; // is this vertex on the offset surface
};
 
 
class EdgeAttributes
{
public:
    int label = 0; // label: 0=default, 1=input, 2=offset
};
 
 
class FaceExtra
{
public:
    int label = 0; // label: 0=default, 1=input, 2=offset
};
 
 
class TetAttributes
{
public:
    int label = 0; // label: 0=default, 1=input, 2=offset
    CellTag tag;
    double m_quality = 0; // AMIPS energy, kept up to date by smoothing
};
 
 
struct OffsetSurfaceSample
{
    Vector3d point; // the sample point, on (or near a corner of) the face
    Vector3d nearest; // nearest point on the input complex to `point`
    Vector3d normal; // (point - nearest).normalized()
    double weight; // 1 for the face centroid, 0.1 for the 3 near-corner samples
};
 
 
class TopoOffsetTetMesh : public wmtk::TetOptimizerMesh
{
public:
    static constexpr int OFFSET_SURFACE_CLASS = 1;
 
public: // mode for splitting in marching tets
    enum class EdgeSplitMode {
        Midpoint = 0, // used for simplicial embedding steps
        Initial = 1, // this is used to initialize the complex. Its a little hacky
 
        // only one of these is used, hard coded in execute_offset()
        BinarySearch = 2, // bisection root finding algo
        LogRootFind = 3, // 'custom' root finding, using the fact that d(x) - d* < 0 at first vertex
        SphereTracing = 4, // use sphere tracing to compute the zero of the distance field
 
        Optimization = 5 // this is used in the optimization phase of the algorithm
    };
 
public:
    int m_vtu_counter = 0;
    std::array<size_t, 4> m_init_counts = {{0, 0, 0, 0}};
    size_t m_tags_count;
    SimplicialComplexBVH m_input_complex_bvh;
    EdgeSplitMode m_edge_split_mode = EdgeSplitMode::Midpoint;
 
    // tag map stuff
    std::map<std::string, int64_t> m_tag_name_to_id;
    std::map<int64_t, std::string> m_tag_id_to_name;
    CellTag m_offset_output_tag_ids;
 
    // if in 'singlebody' mode
    bool m_singlebody = false;
    int64_t m_single_tag;
 
    // dont actually use, just for retaining in output
    bool m_has_envelope;
    std::vector<Vector3d> m_V_envelope;
    std::vector<Vector3i> m_F_envelope;
    // m_envelope itself lives on the base, which is what checks tracked-surface triangles
    // against it.
    double m_envelope_eps = -1;
 
    Parameters& m_offset_params;
 
    using VertexExtraCol = wmtk::AttributeCollection<VertexExtra>;
    using EdgeAttCol = wmtk::AttributeCollection<EdgeAttributes>;
    using FaceExtraCol = wmtk::AttributeCollection<FaceExtra>;
    using TetAttCol = wmtk::AttributeCollection<TetAttributes>;
    // m_vertex_attribute and m_face_attribute are the base's; these two are registered
    // alongside them in its attribute groups.
    VertexExtraCol m_vertex_extra;
    FaceExtraCol m_face_extra;
    EdgeAttCol m_edge_attribute;
    TetAttCol m_tet_attribute;
 
    struct FaceSnapshot
    {
        FaceAttributes tags;
        FaceExtra extra;
    };
    FaceSnapshot face_snapshot(const size_t fid) const
    {
        return FaceSnapshot{m_face_attribute[fid], m_face_extra[fid]};
    }
    void restore_face(const size_t fid, const FaceSnapshot& s)
    {
        m_face_attribute[fid] = s.tags;
        m_face_extra[fid] = s.extra;
    }
 
    void set_vertex_position(const size_t vid, const Vector3d& p)
    {
        m_vertex_attribute[vid].m_posf = p;
        m_vertex_attribute[vid].m_pos = to_rational(p);
        m_vertex_attribute[vid].m_is_rounded = true;
    }
 
    bool cell_in_region(const size_t tid) const
    {
        const CellTag& tags = m_tet_attribute[tid].tag;
        for (const int64_t t : m_offset_output_tag_ids) {
            if (tags.count(t) != 0) return true; // in the offset band
        }
        const ExpressionPtr& expr = m_offset_params.offset_selection;
        return expr && expr->eval(tags); // in the input complex it wraps
    }
 
    bool face_is_offset(const size_t fid) const
    {
        return m_face_attribute[fid].m_is_surface_fs &&
               m_face_attribute[fid].m_surface_class == OFFSET_SURFACE_CLASS;
    }
    bool face_is_input(const size_t fid) const
    {
        return m_face_attribute[fid].m_is_surface_fs &&
               m_face_attribute[fid].m_surface_class != OFFSET_SURFACE_CLASS;
    }
 
    // debug use
    std::atomic<int> cnt_split = 0, cnt_collapse = 0;
 
    TopoOffsetTetMesh(Parameters& _m_offset_params, int _num_threads = 0)
        : wmtk::TetOptimizerMesh(_m_offset_params, nullptr)
        , m_offset_params(_m_offset_params)
    {
        NUM_THREADS = _num_threads;
        // The base owns the vertex and face slots; register the offset's own data with its
        // groups so it is resized, protected and rolled back with them.
        m_vertex_attr_group.add(&m_vertex_extra);
        m_face_attr_group.add(&m_face_extra);
        p_edge_attrs = &m_edge_attribute;
        p_tet_attrs = &m_tet_attribute;
 
        m_collapse_check_link_condition = false;
        m_collapse_check_manifold = false;
 
        optimization::deactivate_opt_logger();
    }
 
    ~TopoOffsetTetMesh() override = default;
 
 
    double cell_quality(const size_t tid) const override { return m_tet_attribute[tid].m_quality; }
    void set_cell_quality(const size_t tid, const double q) override
    {
        m_tet_attribute[tid].m_quality = q;
    }
 
    std::shared_ptr<SampleEnvelope> surface_envelope_for_face(
        const std::array<size_t, 3>& vids) const override
    {
        for (const size_t v : vids) {
            if (!m_vertex_extra[v].m_is_on_input) return nullptr;
        }
        return m_envelope;
    }
 
    bool allow_surface_swap() const override { return true; }
    bool check_surface_topology() const override { return m_offset_params.perform_sanity_checks; }
 
    std::shared_ptr<SampleEnvelope> smoothing_energy_envelope(const size_t) const override
    {
        return nullptr;
    }
 
    void write_optimization_debug_output(const std::string& path) override { write_vtu(path); }
 
    size_t refine_sizing_around_worst(double) override { return update_sizing_field(); }
 
    bool swap_before_interior(const std::vector<size_t>& tids) override;
    bool swap_before_surface(
        const std::vector<size_t>& tids,
        size_t a,
        size_t b,
        size_t c,
        size_t d) override;
    bool swap_after_cells(const std::vector<size_t>& tids, bool is_surface_flip) override;
 
    bool collapse_edge_before(const Tuple& t) override;
    bool collapse_before_vertex(size_t v1, size_t v2, double edge_length) override;
    bool collapse_after_connectivity(
        size_t v1,
        size_t v2,
        const std::vector<std::array<size_t, 2>>& boundary_edges) override;
    bool collapse_is_order_2_edge(const std::array<size_t, 2>& e) override
    {
        return is_order_2_edge(e);
    }
    void collapse_after_vertex(size_t v1, size_t v2) override;
 
    bool split_before_cells(const Tuple& edge, const std::vector<Tuple>& parents) override;
    bool split_after_cells(size_t v1, size_t v2, size_t v_new, const std::vector<Tuple>& children)
        override;
    void split_after_vertex(size_t v_new, bool is_edge_open_boundary) override;
 
    bool is_open_boundary_edge(const Tuple& e) override { return is_order_2_edge(e); }
 
private:
    struct OptSplitCache
    {
        bool is_edge_on_input = false;
        bool is_edge_on_offset = false;
        std::map<simplex::Edge, TetAttributes> tets;
    };
    wmtk::threading::enumerable_thread_specific<OptSplitCache> m_opt_split_cache;
 
public:
private:
    wmtk::threading::enumerable_thread_specific<double> m_collapse_nd_before;
 
public:
private:
    bool swap_capture_tag(const std::vector<size_t>& tids);
    wmtk::threading::enumerable_thread_specific<CellTag> m_swap_tag;
 
public:
    void init_from_image(
        const MatrixXd& V,
        const MatrixXi& T,
        const MatrixSi& T_tags,
        const MatrixXd& V_env,
        const MatrixXi F_env,
        const std::vector<std::string>& tag_names);
 
    void init_surfaces_and_boundaries();
 
    bool is_edge_on_input(const Tuple& loc);
    bool is_edge_on_offset(const Tuple& loc);
 
    bool ambient_assert();
 
    void label_input_complex();
 
    bool empty_input_complex();
 
    void init_input_complex_bvh();
 
    void edge_split_binary_search(const size_t v1, const size_t v2, Vector3d& p_new) const;
    void edge_split_binary_search(const Vector3d& v1_pos, const Vector3d& v2_pos, Vector3d& p_new)
        const;
 
    void edge_split_log_root_find(const size_t v1, const size_t v2, Vector3d& p_new) const;
 
    void edge_split_sphere_tracing(const size_t v1, const size_t v2, Vector3d& p_new) const;
 
    size_t flood_fill();
 
    std::vector<std::array<size_t, 3>> get_faces_by_condition(
        std::function<bool(const FaceAttributes&)> cond) const;
 
    bool split_edge_before(const Tuple& t) override;
    bool split_edge_after(const Tuple& t) override;
    bool marching_split_edge_before(const Tuple& t);
    bool marching_split_edge_after(const Tuple& t);
    bool split_face_before(const Tuple& t) override;
    bool split_face_after(const Tuple& t) override;
    bool split_tet_before(const Tuple& t) override;
    bool split_tet_after(const Tuple& t) override;
    bool invariants(const std::vector<Tuple>& tets) override;
    bool smooth_before(const Tuple& t) override;
    bool smooth_after(const Tuple& t) override;
 
    void execute_offset(const std::filesystem::path& output_file);
 
    void optimize_offset(const std::filesystem::path& output_file);
 
 
    bool is_offset_face(const Tuple& f) const;
    bool is_offset_face(const size_t fid) const;
 
    std::vector<Tuple> get_offset_surface_faces_for_vertex(const Tuple& t) const;
 
    std::array<OffsetSurfaceSample, 4> offset_surface_samples(const Tuple& f) const;
 
    bool smooth_after_offset_surface(const Tuple& t);
 
 
    // max angle (degrees, 0-90, orientation independent) allowed between an offset-surface
    // face's own normal and the input-complex normal it is supposed to approximate, checked by
    // collapse_edge_before/after and offset_swap_normal_deviation_ok(). Exposed to the user as
    // /max_normal_deviation_deg (see Parameters::max_normal_deviation_deg). See
    // face_normal_deviation() and
    // https://github.com/wildmeshing/topological-offsets/blob/main/components/topological_offsets/wmtk/components/topological_offsets/internal/invariants/NormalDeviationAfterInvariant.cpp
    // and .../invariants/OffsetCollapseBeforeInvariant.cpp
    double m_max_normal_deviation_swap_max_deg = 75.0;
 
    double face_normal_deviation(const Tuple& f) const;
 
    double max_offset_surface_normal_deviation_at_vertex(size_t vid) const;
 
    double collapse_normal_deviation(size_t v_from, size_t v_to, size_t remove_vid) const;
 
 
    // Bounds and thresholds for the sizing field are user-configurable; see
    // Parameters::min_sizing_scalar / max_sizing_scalar / sizing_mrm_threshold /
    // sizing_gradation. Refinement stops at min_sizing_scalar; coarsening never exceeds
    // max_sizing_scalar (default: the base target itself, since a sizing_scalar > 1 would mean
    // "coarser than what the user asked for", which update_sizing_field() never has a reason
    // to produce).
 
    static double mean_ratio_metric(const Vector3d& p0, const Vector3d& p1, const Vector3d& p2);
 
    size_t update_sizing_field();
 
 
 
    bool offset_swap_normal_deviation_ok(
        const Tuple& face_abc,
        const Tuple& face_abd,
        size_t a,
        size_t b,
        size_t c,
        size_t d) const;
 
    void marching_tets();
 
 
 
    bool is_simplicially_embedded() const;
 
    bool tet_is_simp_emb(const Tuple& t) const;
 
    void simplicial_embedding();
 
 
    bool tag_tet_consistent_topology(size_t t_id, int64_t tag) const;
 
    bool offset_tet_consistent_topology(const size_t t_id) const;
 
    bool tet_is_in_offset_conservative(const size_t t_id, const double threshold_r) const;
 
    bool tet_is_in_offset_aggressive(const size_t t_id) const;
 
    void grow_offset_conservative();
 
    void grow_offset_aggressive();
 
    void set_offset_tet_tags();
 
    bool offset_is_manifold();
 
    void write_input_complex(const std::string& path); // write components labeled to be offset
    void write_vtu(const std::string& path);
    // void write_msh(const std::string& file);
    void write_msh_groups(const std::string& file);
 
private:
    struct EdgeSplitCache
    {
        size_t v1_id;
        size_t v2_id;
        // The marching-tets placement and the offset's own data for the new vertex. The
        // shared attributes (position, rounding, bbox, order) are written in `after`.
        Vector3d new_v_pos;
        VertexExtra new_v_extra;
 
        bool is_edge_on_input = false;
        bool is_edge_on_offset = false;
        bool is_edge_open_boundary = false;
 
        std::vector<std::pair<FaceAttributes, std::array<size_t, 3>>> changed_faces;
 
        // cache edge attributes
        EdgeAttributes split_e;
        std::map<size_t, EdgeAttributes> internal_e;
        std::map<simplex::Edge, EdgeAttributes> external_e; // edge is boundary edge (not link)
        std::map<simplex::Edge, EdgeAttributes> link_e; // link edge around splitted edge
 
        // cache face attributes
        std::map<size_t, FaceSnapshot> split_f; // splitted faces
        std::map<simplex::Edge, FaceSnapshot> internal_f; // new faces created by split
        std::map<std::pair<simplex::Edge, size_t>, FaceSnapshot>
            external_f; // closed star boundary faces of splitted edge
 
        // cache tet attributes
        std::map<simplex::Edge, TetAttributes> tets;
    };
    wmtk::threading::enumerable_thread_specific<EdgeSplitCache> edge_split_cache;
 
    struct FaceSplitCache
    {
        size_t v1_id;
        size_t v2_id;
        size_t v3_id;
 
        // cache edge attributes
        std::map<simplex::Edge, EdgeAttributes> existing_e;
 
        // cache face attributes
        std::map<simplex::Face, FaceSnapshot> existing_f;
        int splitf_label;
 
        // cache tet attributes
        std::map<size_t, TetAttributes> tets;
    };
    wmtk::threading::enumerable_thread_specific<FaceSplitCache> face_split_cache;
 
    struct TetSplitCache
    {
        std::array<size_t, 4> v_ids;
 
        // cache retained edge attributes
        std::map<simplex::Edge, EdgeAttributes> existing_e;
 
        // cache retained face attributes
        std::map<simplex::Face, FaceSnapshot> existing_f;
 
        // cache tet attribute
        TetAttributes tet;
    };
    wmtk::threading::enumerable_thread_specific<TetSplitCache> tet_split_cache;
 
 
public:
    // substructure functions
 
    bool is_order_2_edge(const Tuple& e) const;
    bool is_order_2_edge(const std::array<size_t, 2>& e) const;
 
    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;
    void init_vertex_order();
 
private: // helpers
    bool any_tag_present(const CellTag& tag1, const CellTag& tag2)
    {
        if (tag2.empty()) {
            return tag1.empty();
        }
        if (tag1.empty()) { // tag1 is ambient, tag2 is not
            return false;
        }
 
        for (const int64_t& i : tag1) {
            if (tag2.find(i) != tag2.end()) {
                return true;
            }
        }
        return false;
    }
 
    void sort_edges_by_length(std::vector<simplex::Edge>& edges)
    {
        std::sort(
            edges.begin(),
            edges.end(),
            [this](const simplex::Edge& e1, const simplex::Edge& e2) {
                double len1 = (m_vertex_attribute[e1.vertices()[0]].m_posf -
                               m_vertex_attribute[e1.vertices()[1]].m_posf)
                                  .norm();
                double len2 = (m_vertex_attribute[e2.vertices()[0]].m_posf -
                               m_vertex_attribute[e2.vertices()[1]].m_posf)
                                  .norm();
                return len1 > len2;
            });
    }
 
public: // helpers
    void compute_vertex_partition();
 
    size_t get_partition_id(const Tuple& loc) const
    {
        return m_vertex_attribute[loc.vid(*this)].partition_id;
    }
 
    std::vector<Tuple> get_face_adjacent_tets(const Tuple& t) const
    {
        std::vector<Tuple> adj_tets;
        auto tet_1 = t.switch_tetrahedron(*this);
        if (tet_1) {
            adj_tets.push_back(tet_1.value());
        }
        auto tet_2 = t.switch_face(*this).switch_tetrahedron(*this);
        if (tet_2) {
            adj_tets.push_back(tet_2.value());
        }
        auto tet_3 = t.switch_edge(*this).switch_face(*this).switch_tetrahedron(*this);
        if (tet_3) {
            adj_tets.push_back(tet_3.value());
        }
        auto tet_4 =
            t.switch_vertex(*this).switch_edge(*this).switch_face(*this).switch_tetrahedron(*this);
        if (tet_4) {
            adj_tets.push_back(tet_4.value());
        }
        return adj_tets;
    }
 
    std::vector<size_t> connected_components_helper(const size_t& v_id)
    {
        auto onering_v_ids = get_one_ring_vids_for_vertex(v_id);
        std::vector<size_t> ret_v_ids;
        for (const size_t& other_v_id : onering_v_ids) {
            size_t e_id = tuple_from_edge({{v_id, other_v_id}}).eid(*this);
            if (m_edge_attribute[e_id].label != 0) { // edge labelled 1 or 2
                ret_v_ids.push_back(other_v_id);
            }
        }
        return ret_v_ids;
    }
 
    void reset_connected_components()
    {
        auto verts = get_vertices();
        for (const Tuple& v : verts) {
            size_t v_id = v.vid(*this);
            m_vertex_extra[v_id].component_id = 0;
        }
    }
};
 
 
} // namespace wmtk::components::topological_offset