TetWildMesh.h

#pragma once
 
#include <wmtk/MeshData.hpp>
 
#include <igl/Timer.h>
#include <wmtk/SurfaceTagAttributes.h>
#include <wmtk/TetMesh.h>
#include <wmtk/TetOptimizerMesh.h>
#include <wmtk/utils/PartitionMesh.h>
#include <algorithm>
#include <wmtk/envelope/CompositeEnvelopes.hpp>
#include <wmtk/envelope/Envelope.hpp>
#include <wmtk/optimization/SmoothVertex.hpp>
#include <wmtk/threading/concurrent_map.hpp>
#include <wmtk/threading/enumerable_thread_specific.hpp>
#include <wmtk/threading/parallel_for.hpp>
 
#include "Parameters.h"
 
// clang-format off
#include <wmtk/utils/DisableWarnings.hpp>
#include <fastenvelope/FastEnvelope.h>
#include <wmtk/utils/EnableWarnings.hpp>
// clang-format on
 
#include <igl/remove_unreferenced.h>
#include <memory>
#include <set>
#include <unordered_set>
#include <utility>
#include <wmtk/utils/EmbedTriangles.hpp>
#include <wmtk/utils/SurfaceTopology.hpp>
#include <wmtk/utils/partition_utils.hpp>
 
namespace wmtk::components::tetwild {
 
using VertexAttributes = wmtk::TetOptimizerMesh::VertexAttributes;
using FaceAttributes = wmtk::TetOptimizerMesh::FaceAttributes;
 
class TetAttributes
{
public:
    double m_quality;
};
 
class TetFinalizeAttributes
{
public:
    double m_winding_number_input = 0; // winding number w.r.t. the input
    double m_winding_number_tracked = 0; // winding number w.r.t. the tracked surface
    std::vector<double> m_winding_number_per_input;
    int part_id = -1; // flood fill ID
};
 
class TetWildMesh : public wmtk::TetOptimizerMesh
{
public:
    struct VertexExtras
    {
        bool m_is_on_open_boundary = false;
        bool m_is_on_feature_curve = false;
    };
    wmtk::AttributeCollection<VertexExtras> m_vertex_extra;
 
    // The 0-dimensional features -- the input feature points, and the feature-curve
    // endpoints (junctions too, when allow_junction_cleanup is off) -- are anchored by the
    // base (TetOptimizerMesh's feature-point anchors), turned on in the constructor. The
    // driver sets m_feature_eps alongside m_feature_envelope; init_from_Volumeremesher fills
    // the anchors.
 
    Parameters& m_tet_params;
 
    std::shared_ptr<Envelope> m_order2_envelope;
 
    wmtk::InputEnvelopes m_input_envelopes;
 
    std::shared_ptr<Envelope> surface_envelope_for_face(
        const std::array<size_t, 3>& vids,
        const FaceAttributes& attr) const override
    {
        if (m_input_envelopes.empty()) return m_envelope;
        return m_input_envelopes.containment(attr.input_set());
    }
 
    std::vector<std::string> m_input_names;
 
    using TetAttCol = wmtk::AttributeCollection<TetAttributes>;
    TetAttCol m_tet_attribute;
 
    void enable_finalize_attributes()
    {
        if (m_tet_finalize_enabled) return;
        m_tet_finalize.resize(m_tet_attribute.size());
        m_tet_attr_group.add(&m_tet_finalize);
        m_tet_finalize_enabled = true;
    }
    TetFinalizeAttributes& tet_finalize(const size_t tid)
    {
        enable_finalize_attributes();
        return m_tet_finalize[tid];
    }
    const TetFinalizeAttributes& tet_finalize(const size_t tid) const
    {
        static const TetFinalizeAttributes defaults;
        return m_tet_finalize_enabled ? m_tet_finalize[tid] : defaults;
    }
private:
    wmtk::AttributeCollection<TetFinalizeAttributes> m_tet_finalize;
    wmtk::AttributeContainerGroup m_tet_attr_group;
    bool m_tet_finalize_enabled = false;
 
public:
    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;
    }
    bool allow_surface_swap() const override { return m_tet_params.allow_surface_swap; }
    bool check_surface_topology() const override { return m_tet_params.check_surface_topology; }
    void split_after_vertex(const size_t vid, const bool is_open_boundary) override
    {
        m_vertex_extra[vid].m_is_on_open_boundary = is_open_boundary;
    }
    bool boundary_edges_at_vertex_inside(size_t vid) const override
    {
        if (!m_vertex_extra[vid].m_is_on_open_boundary || !m_order2_envelope) {
            return true;
        }
        for (const size_t u : get_one_ring_vids_for_vertex(vid)) {
            if (!m_vertex_extra[u].m_is_on_open_boundary) {
                continue;
            }
            // Exact membership: an open boundary edge has exactly one incident tracked
            // face. Flags alone would also catch interior chords between two boundary
            // vertices (e.g. across a corner), which must not be tube-tested.
            if (get_num_surface_faces_for_edge({{vid, u}}) != 1) {
                continue;
            }
            if (m_order2_envelope->is_outside(
                    std::array<Eigen::Vector3d, 2>{
                        {m_vertex_attribute[vid].m_posf, m_vertex_attribute[u].m_posf}})) {
                return false;
            }
        }
        return true;
    }
    bool collapse_before_vertex(size_t v1, size_t v2, double edge_length) override
    {
        if (edge_length <= 0 || !m_vertex_extra[v1].m_is_on_open_boundary) return true;
        return m_vertex_extra[v2].m_is_on_open_boundary ||
               !m_order2_envelope->is_outside(m_vertex_attribute[v2].m_posf);
    }
 
    bool collapse_is_order_2_edge(const std::array<size_t, 2>& e) override
    {
        return is_open_boundary_edge(e);
    }
    bool collapse_after_connectivity(
        size_t v1,
        size_t v2,
        const std::vector<std::array<size_t, 2>>&) override
    {
        m_vertex_extra[v2].m_is_on_open_boundary =
            m_vertex_extra[v1].m_is_on_open_boundary || m_vertex_extra[v2].m_is_on_open_boundary;
        m_vertex_extra[v2].m_is_on_feature_curve =
            m_vertex_extra[v1].m_is_on_feature_curve || m_vertex_extra[v2].m_is_on_feature_curve;
        return true;
    }
    void collapse_after_vertex(size_t, size_t v2) override
    {
        if (m_vertex_extra[v2].m_is_on_open_boundary && !is_vertex_on_boundary(v2)) {
            m_vertex_extra[v2].m_is_on_open_boundary = false;
        }
    }
    void collapse_after_features(size_t, size_t v2) override
    {
        // Same cleanup for the feature flag: v1's flag was OR-ed in (see
        // collapse_after_connectivity), but if the survivor has no incident tagged edge the
        // curve rerouted around it and the flag is stale.
        if (m_vertex_extra[v2].m_is_on_feature_curve && !vertex_has_feature_edge(v2)) {
            m_vertex_extra[v2].m_is_on_feature_curve = false;
        }
    }
 
    void split_after_vertex_feature(size_t vid, bool on_feature) override
    {
        m_vertex_extra[vid].m_is_on_feature_curve = on_feature;
    }
 
 
    std::shared_ptr<Envelope> smoothing_energy_envelope(const size_t vid) const override;
    TetWildMesh(Parameters& _m_params, std::shared_ptr<Envelope> _m_envelope, int _num_threads = 1)
        : wmtk::TetOptimizerMesh(_m_params, std::move(_m_envelope))
        , m_tet_params(_m_params)
    {
        m_vertex_attr_group.add(&m_vertex_extra);
        // Registered unconditionally, as the ids always were: the driver exports and audits
        // the anchors whether or not preserve_feature_points turns the policy on.
        enable_feature_point_tracking(m_tet_params.preserve_feature_points);
        NUM_THREADS = _num_threads;
        m_tet_attr_group.add(&m_tet_attribute);
        p_tet_attrs = &m_tet_attr_group;
        m_collapse_check_link_condition = false;
        m_collapse_check_manifold = false;
 
        optimization::deactivate_opt_logger();
        m_s_envelope = 1. / (m_params.eps * m_params.eps);
        m_params.w_envelope = 1. - m_params.w_amips;
    }
 
    ~TetWildMesh() {}
 
    // only used with unit tests
    void create_mesh_attributes(
        const std::vector<VertexAttributes>& _vertex_attribute,
        const std::vector<TetAttributes>& _tet_attribute)
    {
        const size_t n_tet = _tet_attribute.size();
        m_vertex_attribute.resize(_vertex_attribute.size());
        // The extras carry no data a caller supplies -- every field defaults -- so they only
        // have to be sized alongside the shared collection.
        m_vertex_extra.resize(_vertex_attribute.size());
        m_feature_point_attribute.resize(_vertex_attribute.size());
        m_face_attribute.resize(4 * n_tet);
        m_tet_attribute.resize(n_tet);
 
        // new for edge
        // m_edge_attribute.resize(6 * n_tet);
 
        for (size_t i = 0; i < _vertex_attribute.size(); i++)
            m_vertex_attribute[i] = _vertex_attribute[i];
 
        // Keep whatever init() reserved. AttributeCollection::resize only ever grows, so the
        // calls above cannot shrink a collection -- but assigning m_attributes directly does,
        // and it used to drop the tet attributes to exactly n_tet. The attributes have to stay
        // at least as large as the connectivity: an operation that creates a new element
        // indexes them by its id, and a 5->6 swap creates one. That left m_tet_attribute one
        // short of the tet the swap adds, and AttributeCollection::operator[] read past the end
        // of it -- ASan reports a heap-buffer-overflow reached from swap_edge_56_after, and
        // libstdc++ then aborts on the corrupted heap a few allocations later, while libc++
        // carries on and the tests pass.
        const size_t tcap = std::max(n_tet, m_tet_attribute.size());
        m_tet_attribute.m_attributes = std::vector<TetAttributes>(tcap);
        for (size_t i = 0; i < n_tet; i++) m_tet_attribute[i] = _tet_attribute[i];
        for (size_t i = 0; i < n_tet; i++)
            m_tet_attribute[i].m_quality = get_quality(tuple_from_tet(i));
    }
 
 
    void output_mesh(std::string file);
    MeshData tet_mesh_data();
    void init_from_delaunay_box_mesh(const std::vector<Eigen::Vector3d>& vertices);
 
public:
    //
    bool is_vertex_on_boundary(const size_t vid);
    //
    void mesh_improvement_legacy(int max_its = 80);
 
    Eigen::MatrixXd tet_barycenters(const std::vector<Tuple>& tets) const;
 
    void compute_winding_number(
        const std::vector<Tuple>& tets,
        const Eigen::MatrixXd& barycenters,
        const std::vector<Vector3d>& vertices = {},
        const std::vector<std::array<size_t, 3>>& faces = {});
 
    void compute_input_winding_numbers(
        const std::vector<Tuple>& tets,
        const Eigen::MatrixXd& barycenters,
        const std::vector<Vector3d>& vertices,
        const std::vector<std::array<size_t, 3>>& faces,
        const std::vector<int>& face_input,
        int n_inputs);
 
    bool m_has_input_winding_number = false;
    bool m_has_tracked_winding_number = false;
    bool m_has_parts = false;
 
    void filter_with_input_surface_winding_number();
    void filter_with_tracked_surface_winding_number();
    void filter_with_flood_fill();
 
 
    // debug use
    std::atomic<int> cnt_split = 0, cnt_collapse = 0;
 
private:
    // tags: correspondence map from new tet-face node indices to in-triangle ids.
    // built up while triangles are inserted.
    wmtk::threading::concurrent_map<std::array<size_t, 3>, std::vector<int>> tet_face_tags;
 
    struct TriangleInsertionLocalInfoCache
    {
        // local info: for each face insertion
        int face_id;
        std::vector<std::array<size_t, 3>> old_face_vids;
    };
    wmtk::threading::enumerable_thread_specific<TriangleInsertionLocalInfoCache>
        triangle_insertion_local_cache;
 
    // for incremental tetwild
public:
    void insertion_by_volumeremesher(
        const std::vector<Vector3d>& vertices,
        const std::vector<std::array<size_t, 3>>& faces,
        std::vector<Vector3r>& v_rational,
        std::vector<std::array<size_t, 3>>& facets_after,
        std::vector<bool>& is_v_on_input,
        std::vector<std::array<size_t, 4>>& tets_after,
        std::vector<bool>& tet_face_on_input_surface,
        std::vector<int8_t>* tet_face_orientation = nullptr,
        const std::vector<Vector3d>& feature_edge_vertices = {},
        const std::vector<std::array<size_t, 2>>& feature_edges = {},
        const std::vector<Vector3d>& feature_points = {},
        utils::EmbedFeaturesResult* features_out = nullptr,
        const std::vector<int>* face_input = nullptr,
        std::vector<InputSet>* facet_input_set = nullptr,
        std::vector<std::vector<uint32_t>>* facet_sources = nullptr);
 
    void init_from_Volumeremesher(
        const std::vector<Vector3r>& v_rational,
        const std::vector<std::array<size_t, 3>>& facets,
        const std::vector<bool>& is_v_on_input,
        const std::vector<std::array<size_t, 4>>& tets,
        const std::vector<bool>& tet_face_on_input_surface,
        const std::vector<int8_t>* tet_face_orientation = nullptr,
        const utils::EmbedFeaturesResult* features = nullptr,
        const std::vector<Vector3d>* feature_anchors = nullptr,
        const std::vector<InputSet>* facet_input_set = nullptr);
 
    void init_from_file(std::string input_dir);
 
 
    size_t refine_sizing_around_worst(double max_energy) override;
 
    // for open boundary
    void find_open_boundary();
    bool is_open_boundary_edge(const Tuple& e) override;
    bool is_open_boundary_edge(const std::array<size_t, 2>& e);
 
public:
    // substructure functions
 
    void init_vertex_order();
 
public:
    // for boolean operations
    int flood_fill();
 
    void save_paraview(const std::string& path, const bool use_hdf5, MeshMap* output = nullptr);
    void write_optimization_debug_output(const std::string& path) override
    {
        save_paraview(path, false);
    }
    void optimization_sanity_checks_extra() override;
 
public:
    struct ExportStruct
    {
        // tet mesh
        MatrixXd V;
        MatrixXi T;
        // tracked surface
        MatrixXi F;
        // attributes; the ones the finalization did not compute (see m_has_parts and the
        // m_has_*_winding_number flags) hold their defaults: 0, and -1 for the part
        VectorXd t_amips;
        VectorXd t_winding_number_input;
        VectorXd t_winding_number_tracked;
        MatrixXd t_winding_number_per_input;
        VectorXi t_part;
    };
    // export functionality
    ExportStruct export_mesh_data() const;
};
 
 
} // namespace wmtk::components::tetwild