ImageSimulationMesh.h

#pragma once
 
#include <igl/Timer.h>
#include <wmtk/TetMesh.h>
#include <wmtk/utils/Morton.h>
#include <wmtk/utils/PartitionMesh.h>
#include <polysolve/nonlinear/Problem.hpp>
#include <wmtk/envelope/Envelope.hpp>
#include <wmtk/optimization/solver.hpp>
#include <wmtk/simplex/Simplex.hpp>
 
#include "ConnectedComponent.hpp"
#include "Parameters.h"
 
// clang-format off
#include <wmtk/utils/DisableWarnings.hpp>
#include <fastenvelope/FastEnvelope.h>
#include <tbb/concurrent_queue.h>
#include <tbb/concurrent_priority_queue.h>
#include <tbb/concurrent_vector.h>
#include <tbb/concurrent_map.h>
#include <tbb/concurrent_unordered_map.h>
#include <tbb/enumerable_thread_specific.h>
#include <tbb/parallel_for.h>
#include <tbb/task_arena.h>
#include <tbb/parallel_sort.h>
#include <wmtk/utils/EnableWarnings.hpp>
#include <VolumeRemesher/embed.h>
// clang-format on
 
#include <igl/remove_unreferenced.h>
#include <memory>
 
#include "expression_parser/Expression.hpp"
 
namespace wmtk::components::image_simulation {
 
class VertexAttributes
{
public:
    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; // same as is_bbox_fs?
 
    double m_sizing_scalar = 1;
 
    size_t partition_id = 0;
 
    VertexAttributes() {};
    VertexAttributes(const Vector3r& p);
};
 
 
// class EdgeAttributes
// {
// public:
//     bool m_is_on_open_boundary = false;
// };
 
class FaceAttributes
{
public:
    bool m_is_surface_fs = false;
 
    int m_is_bbox_fs = -1; //-1; 0~5
 
    void reset()
    {
        m_is_surface_fs = false;
        m_is_bbox_fs = -1;
    }
 
    void merge(const FaceAttributes& attr)
    {
        m_is_surface_fs = m_is_surface_fs || attr.m_is_surface_fs;
        if (attr.m_is_bbox_fs >= 0) m_is_bbox_fs = attr.m_is_bbox_fs;
    }
};
 
class TetAttributes
{
public:
    double m_quality;
    CellTag tags;
};
 
class ImageSimulationMesh : public wmtk::TetMesh
{
public:
    int m_debug_print_counter = 0;
    size_t m_tags_count = 0;
    std::map<int64_t, std::string> m_tag_id_to_name;
    std::map<std::string, int64_t> m_tag_name_to_id;
 
    double time_env = 0.0;
    igl::Timer isout_timer;
    const double MAX_ENERGY = std::numeric_limits<double>::max();
 
    Parameters& m_params;
    std::vector<Vector3d> m_V_envelope;
    std::vector<Vector3i> m_F_envelope;
    std::shared_ptr<SampleEnvelope> m_envelope;
    std::shared_ptr<SampleEnvelope> m_envelope_orig;
    double m_envelope_eps = -1;
 
    using ExprPtr = expression_parser::ExpressionPtr;
    std::vector<std::tuple<ExprPtr, double>> m_length_regions;
 
    bool m_collapse_check_quality = true;
 
    // for open boundary
    std::shared_ptr<SampleEnvelope> m_order_2_edge_envelope; // todo: add sample envelope option
 
    tbb::enumerable_thread_specific<std::unique_ptr<polysolve::nonlinear::Solver>> m_solver;
 
    // scaling factors
    double m_s_amips = -1;
    double m_s_envelope = -1;
 
    // When set, split_edge_after binary-searches vmid onto the zero-crossing of this function.
    // Negative = stays on v1 side, positive = stays on v2 side.
    std::function<double(const Vector3d&)> m_voronoi_split_fn = nullptr;
 
    ImageSimulationMesh(Parameters& _m_params, double envelope_eps, int _num_threads = 0)
        : m_params(_m_params)
        , m_envelope_eps(envelope_eps)
    {
        NUM_THREADS = _num_threads;
        p_vertex_attrs = &m_vertex_attribute;
        p_face_attrs = &m_face_attribute;
        p_tet_attrs = &m_tet_attribute;
        m_collapse_check_link_condition = false;
        m_collapse_check_manifold = false;
 
        // solver is lazily created on first use
 
        optimization::deactivate_opt_logger();
 
        m_s_amips = 1.;
        m_s_envelope = 1. / (m_params.diag_l * m_params.eps * m_params.eps);
 
        double& wa = m_params.w_amips;
        double& we = m_params.w_envelope;
        we = 1 - wa;
        logger().info("w_envelope = {}", we);
    }
 
    ~ImageSimulationMesh() {}
    using VertAttCol = wmtk::AttributeCollection<VertexAttributes>;
    using FaceAttCol = wmtk::AttributeCollection<FaceAttributes>;
    using TetAttCol = wmtk::AttributeCollection<TetAttributes>;
    // using EdgeAttCol = wmtk::AttributeCollection<EdgeAttributes>;
    VertAttCol m_vertex_attribute;
    FaceAttCol m_face_attribute;
    TetAttCol m_tet_attribute;
    // EdgeAttCol m_edge_attribute;
 
    // only used with unit tests
    void create_mesh_attributes(
        const std::vector<VertexAttributes>& _vertex_attribute,
        const std::vector<TetAttributes>& _tet_attribute)
    {
        auto n_tet = _tet_attribute.size();
        m_vertex_attribute.resize(_vertex_attribute.size());
        m_face_attribute.resize(4 * n_tet);
        m_tet_attribute.resize(n_tet);
 
        for (auto i = 0; i < _vertex_attribute.size(); i++) {
            m_vertex_attribute[i] = _vertex_attribute[i];
        }
        m_tet_attribute.m_attributes = tbb::concurrent_vector<TetAttributes>(_tet_attribute.size());
        for (auto i = 0; i < _tet_attribute.size(); i++) {
            m_tet_attribute[i] = _tet_attribute[i];
        }
        for (auto i = 0; i < _tet_attribute.size(); i++) {
            m_tet_attribute[i].m_quality = get_quality(tuple_from_tet(i));
        }
    }
 
    // TODO This should not be here but inside wmtk
    void compute_vertex_partition()
    {
        auto partition_id = partition_TetMesh(*this, NUM_THREADS);
        for (auto i = 0; i < vert_capacity(); i++)
            m_vertex_attribute[i].partition_id = partition_id[i];
    }
 
    // TODO This should not be here but inside wmtk
    void compute_vertex_partition_morton()
    {
        if (NUM_THREADS == 0) return;
 
        wmtk::logger().info("Number of parts: {} by morton", NUM_THREADS);
 
        tbb::task_arena arena(NUM_THREADS);
 
        arena.execute([&] {
            std::vector<Eigen::Vector3d> V_v(vert_capacity());
 
            tbb::parallel_for(
                tbb::blocked_range<int>(0, V_v.size()),
                [&](tbb::blocked_range<int> r) {
                    for (int i = r.begin(); i < r.end(); i++) {
                        V_v[i] = m_vertex_attribute[i].m_posf;
                    }
                });
 
 
            struct sortstruct
            {
                int order;
                Resorting::MortonCode64 morton;
            };
 
            std::vector<sortstruct> list_v;
            list_v.resize(V_v.size());
            // since the morton code requires a correct scale of input vertices,
            //  we need to scale the vertices if their coordinates are out of range
            std::vector<Eigen::Vector3d> V = V_v; // this is for rescaling vertices
            Eigen::Vector3d vmin, vmax;
            vmin = V.front();
            vmax = V.front();
 
            for (size_t j = 0; j < V.size(); j++) {
                for (int i = 0; i < 3; i++) {
                    vmin(i) = std::min(vmin(i), V[j](i));
                    vmax(i) = std::max(vmax(i), V[j](i));
                }
            }
 
            // get_bb_corners(V, vmin, vmax);
            Eigen::Vector3d center = (vmin + vmax) / 2;
 
            tbb::parallel_for(tbb::blocked_range<int>(0, V.size()), [&](tbb::blocked_range<int> r) {
                for (int i = r.begin(); i < r.end(); i++) {
                    V[i] = V[i] - center;
                }
            });
 
            Eigen::Vector3d scale_point =
                vmax - center; // after placing box at origin, vmax and vmin are symetric.
 
            double xscale, yscale, zscale;
            xscale = fabs(scale_point[0]);
            yscale = fabs(scale_point[1]);
            zscale = fabs(scale_point[2]);
            double scale = std::max(std::max(xscale, yscale), zscale);
            if (scale > 300) {
                tbb::parallel_for(
                    tbb::blocked_range<int>(0, V.size()),
                    [&](tbb::blocked_range<int> r) {
                        for (int i = r.begin(); i < r.end(); i++) {
                            V[i] = V[i] / scale;
                        }
                    });
            }
 
            constexpr int multi = 1000;
            tbb::parallel_for(tbb::blocked_range<int>(0, V.size()), [&](tbb::blocked_range<int> r) {
                for (int i = r.begin(); i < r.end(); i++) {
                    list_v[i].morton = Resorting::MortonCode64(
                        int(V[i][0] * multi),
                        int(V[i][1] * multi),
                        int(V[i][2] * multi));
                    list_v[i].order = i;
                }
            });
 
            const auto morton_compare = [](const sortstruct& a, const sortstruct& b) {
                return (a.morton < b.morton);
            };
 
            tbb::parallel_sort(list_v.begin(), list_v.end(), morton_compare);
 
            int interval = list_v.size() / NUM_THREADS + 1;
 
            tbb::parallel_for(
                tbb::blocked_range<int>(0, list_v.size()),
                [&](tbb::blocked_range<int> r) {
                    for (int i = r.begin(); i < r.end(); i++) {
                        m_vertex_attribute[list_v[i].order].partition_id = i / interval;
                    }
                });
        });
    }
 
    size_t get_partition_id(const Tuple& loc) const
    {
        return m_vertex_attribute[loc.vid(*this)].partition_id;
    }
 
    void init_envelope(const MatrixXd& V, const MatrixXi& F, const bool use_exact);
 
    CellTag string_set_to_cell_tag(const std::set<std::string>& str_set);
 
    void set_length_regions(const nlohmann::json& length_region_json);
 
    double get_length2(const Tuple& l) const;
 
 
    void write_msh(std::string file, const bool write_envelope = true);
    void output_faces(std::string file, std::function<bool(const FaceAttributes&)> cond);
 
public:
    void split_all_edges();
    bool split_edge_before(const Tuple& t) override;
    bool split_edge_after(const Tuple& loc) override;
 
    void smooth_all_vertices(const size_t n_iters);
    bool smooth_before(const Tuple& t) override;
    bool smooth_after(const Tuple& t) override;
 
    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;
 
    void simplify();
 
    size_t swap_all_edges_44();
    bool swap_edge_44_before(const Tuple& t) override;
    double swap_edge_44_energy(const std::vector<std::array<size_t, 4>>& tets, const int op_case)
        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;
    double swap_edge_56_energy(const std::vector<std::array<size_t, 4>>& tets, const int op_case)
        override;
    bool swap_edge_56_after(const Tuple& t) override;
 
    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_faces();
    bool swap_face_before(const Tuple& t) override;
    bool swap_face_after(const Tuple& t) override;
 
    size_t swap_all_edges_all();
 
    bool is_inverted_f(const Tuple& loc) const;
    bool is_inverted(const std::array<size_t, 4>& vs) const;
    bool is_inverted(const Tuple& loc) const;
    double get_quality(const std::array<size_t, 4>& vs) const;
    double get_quality(const Tuple& loc) const;
 
    std::vector<std::array<double, 12>> get_amips_assembles(const Tuple& t) const;
 
    std::shared_ptr<polysolve::nonlinear::Problem> get_amips_energy(const Tuple& t) const;
    std::shared_ptr<polysolve::nonlinear::Problem> get_envelope_energy(const Tuple& t) const;
 
    bool round(const Tuple& loc);
 
    bool all_rounded() const;
 
    //
    bool is_edge_on_surface(const Tuple& loc);
    bool is_edge_on_bbox(const Tuple& loc);
    //
    void mesh_improvement(int max_its = 80);
    std::tuple<double, double> local_operations(
        const std::array<int, 4>& ops,
        bool collapse_limit_length = true);
    std::tuple<double, double> get_max_avg_energy();
 
    bool check_attributes();
 
    std::vector<std::array<size_t, 3>> get_faces_by_condition(
        std::function<bool(const FaceAttributes&)> cond) const;
 
    bool invariants(const std::vector<Tuple>& t) override; // this is now automatically checked
 
    // debug use
    std::atomic<int> cnt_split = 0, cnt_collapse = 0, cnt_swap = 0;
 
private:
    // tags: correspondence map from new tet-face node indices to in-triangle ids.
    // built up while triangles are inserted.
    tbb::concurrent_map<std::array<size_t, 3>, std::vector<int>> tet_face_tags;
 
 
    struct SplitInfoCache
    {
        //        VertexAttributes vertex_info;
        size_t v_new;
        size_t v1_id;
        size_t v2_id;
        bool is_edge_on_surface = false;
        bool is_edge_open_boundary = false;
        std::vector<size_t> v1_param_type;
        std::vector<size_t> v2_param_type;
 
        std::vector<std::pair<FaceAttributes, std::array<size_t, 3>>> changed_faces;
 
        std::map<simplex::Edge, TetAttributes> tets;
    };
    tbb::enumerable_thread_specific<SplitInfoCache> split_cache;
 
    struct CollapseInfoCache
    {
        size_t v1_id;
        size_t v2_id;
        double max_energy;
        double edge_length;
        bool is_limit_length;
 
        std::vector<std::pair<FaceAttributes, std::array<size_t, 3>>> changed_faces;
        // all faces incident to the delete vertex (v1) that are on the tracked surface
        std::vector<std::array<size_t, 3>> surface_faces;
        // all edges incident to the deleted vertex(v1) that are on the open boundary
        std::vector<std::array<size_t, 2>> boundary_edges;
        std::vector<size_t> changed_tids;
        std::vector<double> changed_energies;
    };
    tbb::enumerable_thread_specific<CollapseInfoCache> collapse_cache;
 
 
    struct SwapInfoCache
    {
        double max_energy;
        std::map<std::array<size_t, 3>, FaceAttributes> changed_faces;
        CellTag tet_tags;
    };
    tbb::enumerable_thread_specific<SwapInfoCache> swap_cache;
 
public:
    void init_from_image(
        const MatrixXr& V,
        const MatrixXi& T,
        const MatrixSi& T_tags,
        const std::vector<std::string>& tag_names);
    void init_from_image(
        const MatrixXd& V,
        const MatrixXi& T,
        const MatrixSi& T_tags,
        const std::vector<std::string>& tag_names);
 
    void init_surfaces_and_boundaries();
 
    std::vector<std::array<size_t, 3>> triangulate_polygon_face(std::vector<Vector3r> points);
 
    bool adjust_sizing_field_serial(double max_energy);
 
    void find_order_2_edges();
    bool is_order_2_edge(const Tuple& e) const;
    bool is_order_2_edge(const std::array<size_t, 2>& e) const;
 
    void write_vtu(const std::string& path);
 
    void write_surface(const std::string& path) const;
 
public:
    // substructure functions
 
    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();
 
public:
    // Annotations
 
    double tet_volume(const size_t tid) const;
 
    std::vector<ConnectedComponent> compute_connected_components(const CellTag& tag_in) const;
 
    std::vector<ConnectedComponent> find_holes(const std::vector<CellTag>& tag_in) const;
 
    void compute_tag_boundary(const CellTag& tag, MatrixXd& V, MatrixXi& F) const;
 
    void keep_largest_connected_component(
        const std::vector<CellTag>& lcc_tags,
        const size_t n_lcc = 1);
 
    void fill_holes_topo(
        const std::vector<CellTag>& fill_holes_tags,
        double threshold = std::numeric_limits<double>::infinity());
 
    void seal_connected_components(
        const std::vector<CellTag>& tag_sets,
        const std::vector<ConnectedComponent>& components);
 
    void tight_seal_topo(
        const std::vector<std::vector<CellTag>>& tight_seal_tag_sets,
        double threshold = std::numeric_limits<double>::infinity());
 
    void resolve_intersections(const std::vector<CellTag>& intersecting_tags);
 
    void replace_tags(const std::vector<CellTag>& tags_in, const std::vector<CellTag>& tags_out);
 
    void tag_priority(const std::vector<int64_t>& tags);
};
 
} // namespace wmtk::components::image_simulation