ManExtractMesh.h

#pragma once
 
#include <igl/Timer.h>
#include <wmtk/TetMesh.h>
#include <wmtk/utils/Morton.h>
#include <wmtk/utils/PartitionMesh.h>
#include <bitset>
#include <wmtk/envelope/Envelope.hpp>
#include <wmtk/simplex/Simplex.hpp>
#include "Parameters.h"
 
// clang-format off
#include <wmtk/utils/DisableWarnings.hpp>
#include <tbb/concurrent_map.h>
#include <tbb/parallel_sort.h>
#include <wmtk/utils/EnableWarnings.hpp>
// clang-format on
 
// #include <igl/remove_unreferenced.h>
// #include <memory>
 
 
namespace wmtk::components::manifold_extraction {
 
 
// for all attributes:
// label: 0=default, 1=input, 2=offset
class VertexAttributes
{
public:
    Vector3d m_posf;
    int label = 0;
    bool on_surf = false;
 
    VertexAttributes() {};
    VertexAttributes(const Vector3d& p);
};
 
 
class EdgeAttributes
{
public:
    int label = 0;
    bool on_surf = false;
};
 
 
class FaceAttributes
{
public:
    int label = 0;
    bool on_surf = false;
};
 
 
class TetAttributes
{
public:
    bool in_out = false; // in or out of mesh body
    int label = 0;
    std::set<int64_t> tag;
};
 
 
class ManExtractMesh : public wmtk::TetMesh
{
public:
    int m_vtu_counter = 0;
    int m_surfvtu_counter = 0;
    std::array<size_t, 4> init_counts = {{0, 0, 0, 0}};
    bool m_boundary_input_vert_found = false;
 
    Parameters& m_params;
    size_t m_tags_count;
 
    // tag mapping stuff
    std::map<std::string, int64_t> m_tag_name_to_id;
    std::map<int64_t, std::string> m_tag_id_to_name;
    std::set<int64_t> m_in_tag_ids;
    std::set<int64_t> m_replace_tag_ids;
 
    // dont actually use this
    bool m_has_envelope = false;
    MatrixXd m_V_envelope;
    MatrixXi m_F_envelope;
 
    using VertAttCol = wmtk::AttributeCollection<VertexAttributes>;
    using EdgeAttCol = wmtk::AttributeCollection<EdgeAttributes>;
    using FaceAttCol = wmtk::AttributeCollection<FaceAttributes>;
    using TetAttCol = wmtk::AttributeCollection<TetAttributes>;
    VertAttCol m_vertex_attribute;
    EdgeAttCol m_edge_attribute;
    FaceAttCol m_face_attribute;
    TetAttCol m_tet_attribute;
 
    ManExtractMesh(Parameters& _m_params, int _num_threads = 0)
        : m_params(_m_params)
    {
        NUM_THREADS = _num_threads;
        p_vertex_attrs = &m_vertex_attribute;
        p_edge_attrs = &m_edge_attribute;
        p_face_attrs = &m_face_attribute;
        p_tet_attrs = &m_tet_attribute;
    }
 
    ~ManExtractMesh() {}
 
    // splitting/invariants
    bool split_edge_before(const Tuple& t) override;
    bool split_edge_after(const Tuple& t) override;
    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>& t) override; // this is now automatically checked
 
private:
    struct EdgeSplitCache
    {
        size_t v1_id;
        size_t v2_id;
        VertexAttributes new_v;
        EdgeAttributes split_e; // split edge
        std::map<size_t, EdgeAttributes> internal_e; // internal edges
        std::map<simplex::Edge, EdgeAttributes> external_e; // edge is boundary edge (not link)
        std::map<simplex::Edge, EdgeAttributes> link_e;
        std::map<size_t, FaceAttributes> split_f; // faces incident to split edge
        std::map<simplex::Edge, FaceAttributes> internal_f; // new internal faces
        std::map<std::pair<simplex::Edge, size_t>, FaceAttributes> external_f; // retained faces
        std::map<simplex::Edge, TetAttributes> tets; // tets incident to split edge
    };
    tbb::enumerable_thread_specific<EdgeSplitCache> edge_split_cache;
 
    struct FaceSplitCache
    {
        size_t v1_id;
        size_t v2_id;
        size_t v3_id;
        std::map<simplex::Edge, EdgeAttributes> existing_e; // retained edges
        std::map<simplex::Face, FaceAttributes> existing_f; // retained faces
        int splitf_label; // split face
        std::map<size_t, TetAttributes> tets; // incident tets
    };
    tbb::enumerable_thread_specific<FaceSplitCache> face_split_cache;
 
    struct TetSplitCache
    {
        std::array<size_t, 4> v_ids;
        std::map<simplex::Edge, EdgeAttributes> existing_e; // retained edges
        std::map<simplex::Face, FaceAttributes> existing_f; // retained faces
        TetAttributes tet; // split tet
    };
    tbb::enumerable_thread_specific<TetSplitCache> tet_split_cache;
 
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 label_surface_simplices(bool need_to_clear);
 
    std::pair<size_t, size_t> label_non_manifold(); // return # of non manifold edges, verts
 
    bool edge_is_manifold(const Tuple& t) const;
 
    void edge_dfs_helper(std::set<size_t>& visited_tids, const Tuple& t) const;
 
    bool vertex_is_manifold(const Tuple& t) const;
 
    void vertex_dfs_helper(
        std::set<size_t>& visited_tids,
        const Tuple& t,
        const bool include,
        const std::vector<simplex::Face>& b_out_faces) const;
 
    bool is_boundary_vertex(size_t vid) const;
 
    std::vector<simplex::Face> get_boundary_faces_for_out_tets(size_t vid) const;
 
    bool is_simplicially_embedded() const;
 
    bool tet_is_simp_emb(const Tuple& t) const;
 
    void simplicial_embedding();
 
    void perform_offset();
 
    void set_offset_tags();
 
    void extract_surface_mesh(MatrixXd& V, MatrixXi& F);
 
    void write_input_complex(const std::string& path);
 
    void write_surface(const std::string& path);
 
    void write_vtu(const std::string& path); // debugging, write .vtu of tet mesh
 
    void write_msh_groups(const std::string& path);
 
private: // helpers
    // /**
    //  * @brief return true if all tags in tag1 are present in tag2. If tag2 is empty, only return
    //  * true if tag1 empty (ambient).
    //  */
    // bool all_tags_present(std::set<int64_t> tag1, std::set<int64_t> tag2)
    // {
    //     if (tag2.empty()) {
    //         return tag1.empty();
    //     }
    //     if (tag1.empty()) { // tag1 is ambient and tag2 is not.
    //         return false;
    //     }
 
    //     for (const int64_t& i : tag1) {
    //         if (tag2.find(i) == tag2.end()) {
    //             return false;
    //         }
    //     }
    //     return true;
    // }
 
    bool any_tag_present(const std::set<int64_t>& tag1, const std::set<int64_t>& 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;
    }
};
 
} // namespace wmtk::components::manifold_extraction