SmoothVertex.hpp

#pragma once
 
#include <wmtk/optimization/AMIPSEnergy.hpp>
#include <wmtk/optimization/EnergySum.hpp>
#include <wmtk/optimization/EnvelopeEnergy.hpp>
#include <wmtk/optimization/solver.hpp>
 
#include <wmtk/Types.hpp>
#include <wmtk/envelope/Envelope.hpp>
#include <wmtk/simplex/SimplexCollection.hpp>
#include <wmtk/utils/Logger.hpp>
#include <wmtk/utils/TetraQualityUtils.hpp>
 
#include <array>
#include <atomic>
#include <limits>
#include <memory>
#include <vector>
 
namespace wmtk::optimization {
 
struct SmoothRejectCounters
{
    std::atomic<size_t> already_inverted{0}; 
    std::atomic<size_t> envelope{0}; 
    std::atomic<size_t> inverted{0}; 
    std::atomic<size_t> quality{0}; 
    std::atomic<size_t> accepted{0};
 
    void reset()
    {
        already_inverted = 0;
        envelope = 0;
        inverted = 0;
        quality = 0;
        accepted = 0;
    }
 
    std::string to_string() const
    {
        return fmt::format(
            "accepted {} | rejected: pre-inverted {}, envelope {}, inverted {}, quality {}",
            accepted.load(),
            already_inverted.load(),
            envelope.load(),
            inverted.load(),
            quality.load());
    }
};
 
struct NewtonCounters
{
    static constexpr size_t kIterBuckets = 12;
    static constexpr size_t kStatuses = 14;
    static_assert(
        int(polysolve::nonlinear::Status::UpdateDirectionFailed) == int(kStatuses) - 2,
        "polysolve's Status enum changed; update kStatuses and status_name()");
 
    std::array<std::atomic<size_t>, kIterBuckets> iterations{};
    std::array<std::atomic<size_t>, kStatuses> status{};
    std::atomic<size_t> iterations_total{0};
    std::atomic<size_t> threw{0};
 
    void record(const polysolve::nonlinear::Solver& s, const bool did_throw)
    {
        const size_t it = s.current_criteria().iterations;
        iterations[std::min(it, kIterBuckets - 1)].fetch_add(1, std::memory_order_relaxed);
        iterations_total.fetch_add(it, std::memory_order_relaxed);
        const int st = int(s.status()) + 1;
        if (st >= 0 && size_t(st) < kStatuses) {
            status[size_t(st)].fetch_add(1, std::memory_order_relaxed);
        }
        if (did_throw) threw.fetch_add(1, std::memory_order_relaxed);
    }
 
    size_t solves() const
    {
        size_t n = 0;
        for (const auto& c : iterations) n += c.load();
        return n;
    }
 
    void reset()
    {
        for (auto& c : iterations) c = 0;
        for (auto& c : status) c = 0;
        iterations_total = 0;
        threw = 0;
    }
 
    static const char* status_name(const size_t i)
    {
        static constexpr std::array<const char*, kStatuses> names = {
            {"NotStarted",
             "Continue",
             "IterationLimit",
             "XDeltaTolerance",
             "RelXDeltaTolerance",
             "FDeltaTolerance",
             "GradNormTolerance",
             "RelGradNormTolerance",
             "NewtonDecrementTolerance",
             "ObjectiveCustomStop",
             "NanEncountered",
             "NotDescentDirection",
             "LineSearchFailed",
             "UpdateDirectionFailed"}};
        return i < kStatuses ? names[i] : "?";
    }
 
    std::string to_string() const
    {
        const size_t n = solves();
        std::string out = fmt::format(
            "solves={} mean_iterations={:.3g} | iterations",
            n,
            n > 0 ? double(iterations_total.load()) / double(n) : 0.);
        for (size_t i = 0; i < kIterBuckets; ++i) {
            out += i + 1 < kIterBuckets ? fmt::format(" {}={}", i, iterations[i].load())
                                        : fmt::format(" >{}={}", i - 1, iterations[i].load());
        }
        out += " | stop";
        for (size_t i = 0; i < kStatuses; ++i) {
            if (status[i].load() != 0)
                out += fmt::format(" {}={}", status_name(i), status[i].load());
        }
        out += fmt::format(" | threw={}", threw.load());
        return out;
    }
};
 
struct SmoothVertexOptions
{
    double w_amips = 1e-4;
    double w_envelope = 1.0 - 1e-4;
    double s_amips = 1.0;
    double s_envelope = 1.0;
 
    bool two_stage = true;
 
    bool quality_veto_on_surface = true;
 
    enum class SmoothingMode { Projected, Exact };
    SmoothingMode smoothing_mode = SmoothingMode::Projected;
 
    int project_line_search_steps = 12;
 
    bool quality_veto = true;
 
 
    int project_line_search_nested_steps = 0;
};
 
template <class Mesh>
bool smooth_vertex_3d(
    Mesh& m,
    const typename Mesh::Tuple& t,
    const SmoothVertexOptions& opts,
    std::unique_ptr<polysolve::nonlinear::Solver>& solver,
    SmoothRejectCounters* counters = nullptr,
    NewtonCounters* newton = nullptr)
{
    using Tuple = typename Mesh::Tuple;
 
    const size_t vid = t.vid(m);
    auto& VA = m.m_vertex_attribute;
    const auto locs = m.get_one_ring_tets_for_vertex(t);
    assert(!locs.empty());
 
    double max_quality = 0.;
    for (const Tuple& tet : locs) {
        max_quality = std::max(max_quality, m.cell_quality(tet.tid(m)));
        if (m.is_inverted_f(tet)) {
            // A neighbour that is not rounded can leave a tet inverted in floats even
            // though it is fine in exact arithmetic; there is nothing to optimize from.
            if (counters) ++counters->already_inverted;
            return false;
        }
    }
 
    // AMIPS wants each tet as 12 doubles with the moving vertex first, so that the solver
    // can overwrite the leading xyz.
    std::vector<std::array<double, 12>> assembles(locs.size());
    for (size_t i = 0; i < locs.size(); ++i) {
        std::array<size_t, 4> local_verts = m.oriented_tet_vids(locs[i].tid(m));
        local_verts = wmtk::orient_preserve_tet_reorder(local_verts, vid);
        for (int k = 0; k < 4; k++) {
            for (int j = 0; j < 3; j++) {
                assembles[i][k * 3 + j] = VA[local_verts[k]].m_posf[j];
            }
        }
    }
 
    if (!solver) {
        solver = create_basic_solver();
    }
 
    const double amips_w = opts.w_amips > 0 ? opts.s_amips * opts.w_amips : 1.0;
    auto amips_energy = std::make_shared<AMIPSEnergy3D>(assembles, amips_w);
 
    // An application-supplied term for this vertex. It carries its own weight, exactly as
    // ExactDistanceEnergy3D below does, so it is summed here at 1. Null for TetWild and SimWild,
    // which leaves `base_energy` as `amips_energy` and every expression below exactly what it
    // was.
    const std::shared_ptr<polysolve::nonlinear::Problem> extra_energy =
        m.smoothing_extra_energy(vid);
    std::shared_ptr<polysolve::nonlinear::Problem> base_energy = amips_energy;
    if (extra_energy) {
        auto sum = std::make_shared<EnergySum>();
        if (opts.w_amips > 0) sum->add_energy(amips_energy);
        sum->add_energy(extra_energy);
        base_energy = sum;
    }
    std::shared_ptr<polysolve::nonlinear::Problem> total_energy = base_energy;
 
    auto solve = [&]() {
        VectorXd x = VA[vid].m_posf;
        bool threw = false;
        try {
            solver->minimize(*total_energy, x);
        } catch (const std::exception&) {
            // polysolve reports a failed line search by throwing; the position it reached
            // is still the best it found, and the checks below decide whether to keep it.
            threw = true;
        }
        if (newton) newton->record(*solver, threw);
        VA[vid].m_posf = x;
    };
 
    const std::shared_ptr<SampleEnvelope> pull_env =
        VA[vid].m_is_on_surface ? m.smoothing_energy_envelope(vid) : nullptr;
 
    if (pull_env && opts.smoothing_mode == SmoothVertexOptions::SmoothingMode::Projected) {
        // Smooth as if the vertex were interior, then walk back onto the input.
        const Vector3d x_orig = VA[vid].m_posf;
        total_energy = base_energy;
        solve();
        const Vector3d x_new = VA[vid].m_posf;
 
        // Place a candidate and report the worst incident quality, or infinity if it
        // inverts. is_inverted is exact, so the rational position tracks every candidate.
        const auto worst_at = [&](const Vector3d& p) {
            VA[vid].m_posf = p;
            VA[vid].m_pos = to_rational(p);
            double mq = 0.;
            for (const Tuple& loc : locs) {
                if (m.is_inverted(loc)) {
                    return std::numeric_limits<double>::infinity();
                }
                mq = std::max(mq, m.get_quality(loc));
            }
            return mq;
        };
 
        bool accepted = false;
        std::vector<Vector3d> interp, proj; // kept for the nested pass below
        interp.reserve(opts.project_line_search_steps);
        proj.reserve(opts.project_line_search_steps);
        for (int k = 0; k < opts.project_line_search_steps; ++k) {
            const Vector3d p = x_orig + std::pow(0.5, k) * (x_new - x_orig);
            Vector3d q;
            pull_env->nearest_point(p, q);
            interp.push_back(p);
            proj.push_back(q);
            if (worst_at(q) < max_quality) {
                accepted = true;
                break;
            }
        }
 
        // Nothing ON the input was acceptable anywhere, so settle for getting as close to
        // it as the one-ring allows. For each candidate, bisect the segment from the
        // interpolated point (s = 0) to its projection (s = 1) for the LARGEST acceptable s.
        // s = 1 is already known to fail -- that is what the first pass just established --
        // so this brackets the boundary and converges up to it from below. Halving s down
        // from 1 instead would cap the result at the midpoint and leave the vertex needlessly
        // far from the input.
        if (!accepted && opts.project_line_search_nested_steps > 0) {
            for (size_t k = 0; k < proj.size() && !accepted; ++k) {
                double lo = 0.0, hi = 1.0; // lo: best acceptable so far, hi: known bad
                Vector3d best;
                bool found = false;
                for (int j = 0; j < opts.project_line_search_nested_steps; ++j) {
                    const double mid = 0.5 * (lo + hi);
                    const Vector3d cand = interp[k] + mid * (proj[k] - interp[k]);
                    if (worst_at(cand) < max_quality) {
                        lo = mid;
                        best = cand;
                        found = true;
                    } else {
                        hi = mid;
                    }
                }
                if (found) {
                    // worst_at left the vertex at the last candidate tried, which is not
                    // necessarily the best one.
                    worst_at(best);
                    accepted = true;
                }
            }
        }
        if (!accepted) {
            VA[vid].m_posf = x_orig;
            VA[vid].m_pos = to_rational(x_orig);
            if (counters) ++counters->quality;
            return false;
        }
    } else if (pull_env) {
        auto envelope_energy =
            std::make_shared<ExactDistanceEnergy3D>(pull_env, opts.s_envelope * opts.w_envelope);
 
        if (opts.two_stage) {
            auto warmup = std::make_shared<EnergySum>();
            if (opts.w_amips > 0) warmup->add_energy(amips_energy, 1. / opts.w_amips);
            if (opts.w_envelope > 0) warmup->add_energy(envelope_energy, 1. / opts.w_envelope);
            if (extra_energy) warmup->add_energy(extra_energy);
            total_energy = warmup;
            solve();
        }
 
        auto weighted = std::make_shared<EnergySum>();
        if (opts.w_amips > 0) weighted->add_energy(amips_energy);
        if (opts.w_envelope > 0) weighted->add_energy(envelope_energy);
        if (extra_energy) weighted->add_energy(extra_energy);
        total_energy = weighted;
        solve();
    } else {
        solve();
    }
 
    // Containment: every surface triangle at this vertex must still be inside. Checked
    // against the containment envelope, which is not necessarily the one it was pulled to.
    const std::shared_ptr<SampleEnvelope> check_env =
        VA[vid].m_is_on_surface ? m.smoothing_containment_envelope(vid) : nullptr;
    if (check_env) {
        const simplex::SimplexCollection surf = m.get_surface_faces_for_vertex(vid);
        for (const simplex::Face& f : surf.faces()) {
            const std::array<Eigen::Vector3d, 3> face = {
                {VA[f.vertices()[0]].m_posf,
                 VA[f.vertices()[1]].m_posf,
                 VA[f.vertices()[2]].m_posf}};
            if (check_env->is_outside(face)) {
                if (counters) ++counters->envelope;
                return false;
            }
        }
    }
 
    // The rational position must be current before the exact inversion test.
    VA[vid].m_pos = to_rational(VA[vid].m_posf);
 
    double max_after_quality = 0.;
    for (const Tuple& loc : locs) {
        if (m.is_inverted(loc)) {
            if (counters) ++counters->inverted;
            return false;
        }
        const size_t tid = loc.tid(m);
        const double quality = m.get_quality(loc);
        m.set_cell_quality(tid, quality);
        max_after_quality = std::max(max_after_quality, quality);
    }
 
    if (opts.quality_veto && (!VA[vid].m_is_on_surface || opts.quality_veto_on_surface)) {
        if (max_after_quality > max_quality) {
            if (counters) ++counters->quality;
            return false;
        }
    }
 
    if (counters) ++counters->accepted;
    return true;
}
 
 
template <class Mesh>
bool smooth_vertex_2d(
    Mesh& m,
    const typename Mesh::Tuple& t,
    const SmoothVertexOptions& opts,
    std::unique_ptr<polysolve::nonlinear::Solver>& solver,
    SmoothRejectCounters* counters = nullptr)
{
    const size_t vid = t.vid(m);
    auto& VA = m.m_vertex_attribute;
    auto& FA = m.m_face_attribute;
 
    const std::vector<size_t>& locs = m.get_one_ring_fids_for_vertex(t);
    assert(!locs.empty());
 
    double max_quality = 0.;
    for (const size_t fid : locs) {
        max_quality = std::max(max_quality, FA[fid].m_quality);
        if (m.is_inverted_f(fid)) {
            // Nothing to optimize from: a neighbour that is not rounded can leave a face
            // inverted in floats even when it is fine exactly.
            if (counters) ++counters->already_inverted;
            return false;
        }
    }
 
    // AMIPS wants each face as 6 doubles with the moving vertex first, keeping the winding.
    std::vector<std::array<double, 6>> assembles;
    assembles.reserve(locs.size());
    for (const size_t fid : locs) {
        std::array<size_t, 3> vs = m.oriented_tri_vids(fid);
        size_t v_loc = 0;
        for (size_t i = 0; i < 3; ++i) {
            if (vs[i] == vid) {
                v_loc = i;
                break;
            }
        }
        const std::array<size_t, 3> buf = vs;
        vs[0] = buf[v_loc];
        vs[1] = buf[(v_loc + 1) % 3];
        vs[2] = buf[(v_loc + 2) % 3];
 
        std::array<double, 6> T;
        for (int i = 0; i < 3; i++) {
            const Vector2d p = m.smoothing_position(vs[i]);
            T[i * 2] = p[0];
            T[i * 2 + 1] = p[1];
        }
        assembles.push_back(T);
    }
 
    if (!solver) {
        solver = create_basic_solver();
    }
 
    const double amips_w = opts.w_amips > 0 ? opts.s_amips * opts.w_amips : 1.0;
    auto amips_energy = std::make_shared<AMIPSEnergy2D>(assembles, amips_w);
 
    // An application-supplied term for this vertex. It carries its own weight, exactly as
    // ExactDistanceEnergy2D below does, so it is summed here at 1. Null for TriWild and
    // SimWild, which leaves `base_energy` as `amips_energy` and every expression below exactly
    // what it was.
    const std::shared_ptr<polysolve::nonlinear::Problem> extra_energy =
        m.smoothing_extra_energy(vid);
    std::shared_ptr<polysolve::nonlinear::Problem> base_energy = amips_energy;
    if (extra_energy) {
        auto sum = std::make_shared<EnergySum>();
        if (opts.w_amips > 0) sum->add_energy(amips_energy);
        sum->add_energy(extra_energy);
        base_energy = sum;
    }
    std::shared_ptr<polysolve::nonlinear::Problem> total_energy = base_energy;
 
    auto solve = [&]() {
        VectorXd x = m.smoothing_position(vid);
        try {
            solver->minimize(*total_energy, x);
        } catch (const std::exception&) {
            // A failed line search is reported by throwing; the position reached is still
            // the best found, and the checks below decide whether to keep it.
        }
        m.set_smoothing_position(vid, Vector2d(x));
    };
 
    // Neighbours along the incident surface edges, captured before the solve. Only `vid`
    // moves, so their positions are the same either way, but taking them first matches what
    // both applications did and keeps the assert below meaningful.
    std::vector<Vector2d> surf_neighbors;
    if (VA[vid].m_is_on_surface) {
        const simplex::SimplexCollection es = m.get_surface_edges_for_vertex(vid);
        surf_neighbors.reserve(es.edges().size());
        for (const simplex::Edge& e : es.edges()) {
            const auto& evs = e.vertices();
            surf_neighbors.push_back(m.smoothing_position(evs[0] != vid ? evs[0] : evs[1]));
        }
        assert(!surf_neighbors.empty());
    }
 
    // Per-vertex rather than the mesh's single m_envelope, and split into the PULL and the
    // CONTAINER exactly as the 3D path above: see TriOptimizerMesh::smoothing_energy_envelope.
    // The defaults return the old expression, so TriWild and SimWild are unaffected.
    const std::shared_ptr<SampleEnvelope> envelope = m.smoothing_energy_envelope(vid);
 
    if (envelope && opts.smoothing_mode == SmoothVertexOptions::SmoothingMode::Projected) {
        const Vector2d x_orig = m.smoothing_position(vid);
        total_energy = base_energy;
        solve();
        const Vector2d x_new = m.smoothing_position(vid);
 
        // set_smoothing_position keeps the rational position in step, which the exact
        // is_inverted below depends on.
        const auto worst_at = [&](const Vector2d& p) {
            m.set_smoothing_position(vid, p);
            double mq = 0.;
            for (const size_t fid : locs) {
                if (m.is_inverted(fid)) {
                    return std::numeric_limits<double>::infinity();
                }
                mq = std::max(mq, m.get_quality(fid));
            }
            return mq;
        };
 
        bool accepted = false;
        std::vector<Vector2d> interp, proj; // kept for the nested pass below
        interp.reserve(opts.project_line_search_steps);
        proj.reserve(opts.project_line_search_steps);
        for (int k = 0; k < opts.project_line_search_steps; ++k) {
            const Vector2d p = x_orig + std::pow(0.5, k) * (x_new - x_orig);
            Vector2d q;
            envelope->nearest_point(p, q);
            interp.push_back(p);
            proj.push_back(q);
            if (worst_at(q) < max_quality) {
                accepted = true;
                break;
            }
        }
 
        // Nothing ON the input was acceptable anywhere, so settle for getting as close to
        // it as the one-ring allows. For each candidate, bisect the segment from the
        // interpolated point (s = 0) to its projection (s = 1) for the LARGEST acceptable s.
        // s = 1 is already known to fail -- that is what the first pass just established --
        // so this brackets the boundary and converges up to it from below. Halving s down
        // from 1 instead would cap the result at the midpoint and leave the vertex needlessly
        // far from the input.
        if (!accepted && opts.project_line_search_nested_steps > 0) {
            for (size_t k = 0; k < proj.size() && !accepted; ++k) {
                double lo = 0.0, hi = 1.0; // lo: best acceptable so far, hi: known bad
                Vector2d best;
                bool found = false;
                for (int j = 0; j < opts.project_line_search_nested_steps; ++j) {
                    const double mid = 0.5 * (lo + hi);
                    const Vector2d cand = interp[k] + mid * (proj[k] - interp[k]);
                    if (worst_at(cand) < max_quality) {
                        lo = mid;
                        best = cand;
                        found = true;
                    } else {
                        hi = mid;
                    }
                }
                if (found) {
                    // worst_at left the vertex at the last candidate tried, which is not
                    // necessarily the best one.
                    worst_at(best);
                    accepted = true;
                }
            }
        }
        if (!accepted) {
            m.set_smoothing_position(vid, x_orig);
            if (counters) ++counters->quality;
            return false;
        }
    } else if (envelope) {
        auto envelope_energy =
            std::make_shared<ExactDistanceEnergy2D>(envelope, opts.s_envelope * opts.w_envelope);
 
        if (opts.two_stage) {
            auto warmup = std::make_shared<EnergySum>();
            if (opts.w_amips > 0) warmup->add_energy(amips_energy, 1. / opts.w_amips);
            if (opts.w_envelope > 0) warmup->add_energy(envelope_energy, 1. / opts.w_envelope);
            if (extra_energy) warmup->add_energy(extra_energy);
            total_energy = warmup;
            solve();
        }
 
        auto weighted = std::make_shared<EnergySum>();
        if (opts.w_amips > 0) weighted->add_energy(amips_energy);
        if (opts.w_envelope > 0) weighted->add_energy(envelope_energy);
        if (extra_energy) weighted->add_energy(extra_energy);
        total_energy = weighted;
        solve();
    } else {
        solve();
    }
 
    // Per-vertex positional constraint, on top of the envelope. A mesh uses this to pin a
    // vertex to a 0-dimensional feature it stands for -- within a ball, so the vertex is
    // still free to move and improve quality, it just cannot walk away from the feature.
    // Meshes with no such features answer true unconditionally.
    if (!m.smoothing_position_is_allowed(vid, m.smoothing_position(vid))) {
        if (counters) ++counters->envelope;
        return false;
    }
 
    // Containment, edge by edge rather than face by face as in 3D -- and against the
    // CONTAINMENT envelope, which is not necessarily the one the vertex was pulled to.
    if (const std::shared_ptr<SampleEnvelope> hold_env = m.smoothing_containment_envelope(vid)) {
        const Vector2d p = m.smoothing_position(vid);
        for (const Vector2d& q : surf_neighbors) {
            const std::array<Eigen::Vector2d, 2> edge = {{p, q}};
            if (hold_env->is_outside(edge)) {
                if (counters) ++counters->envelope;
                return false;
            }
        }
    }
 
    double max_after_quality = 0.;
    for (const size_t fid : locs) {
        if (m.is_inverted(fid)) {
            if (counters) ++counters->inverted;
            return false;
        }
        const double q = m.get_quality(fid);
        FA[fid].m_quality = q;
        max_after_quality = std::max(max_after_quality, q);
    }
 
    if (opts.quality_veto && (!VA[vid].m_is_on_surface || opts.quality_veto_on_surface)) {
        if (max_after_quality > max_quality) {
            if (counters) ++counters->quality;
            return false;
        }
    }
 
    if (counters) ++counters->accepted;
    return true;
}
 
} // namespace wmtk::optimization