EdgeBasedGeodesicSolver.cpp

// BSD 3-Clause License
//
// Copyright (c) 2019, Jiong Tao, Bailin Deng, Yue Peng
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are met:
//
// * Redistributions of source code must retain the above copyright notice, this
//   list of conditions and the following disclaimer.
//
// * Redistributions in binary form must reproduce the above copyright notice,
//   this list of conditions and the following disclaimer in the documentation
//   and/or other materials provided with the distribution.
//
// * Neither the name of the copyright holder nor the names of its
//   contributors may be used to endorse or promote products derived from
//   this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

#include "EdgeBasedGeodesicSolver.h"
#include "surface_mesh/IO.h"
#include "OMPHelper.h"
#include <iostream>
#include <utility>
#include <limits>

EdgeBasedGeodesicSolver::EdgeBasedGeodesicSolver()
    : model_scaling_factor(1.0),
      bfs_laplacian_coef(NULL),
      current_SX(NULL),
      prev_SX(NULL),
      need_compute_residual_norms(false),
      n_vertices(0),
      n_faces(0),
      n_edges(0),
      n_halfedges(0),
      n_interior_edges(0),
      iter_num(0),
      primal_residual_sqr_norm(0),
      dual_residual_sqr_norm(0),
      primal_residual_sqr_norm_threshold(0),
      dual_residual_sqr_norm_threshold(0),
      output_progress(false),
      optimization_converge(false),
      optimization_end(false) {
}

const Eigen::VectorXd& EdgeBasedGeodesicSolver::get_distance_values() {
  return geod_dist_values;
}

bool EdgeBasedGeodesicSolver::solve(const char *mesh_file,
                                    const Parameters& para) {
  param = para;

  std::cout << "Reading triangle mesh......" << std::endl;

  if (!load_input(mesh_file)) {
    return false;
  }

  normalize_mesh();

  std::cout << "Initialize BFS path......" << std::endl;

  Timer timer;
  Timer::EventID start = timer.get_time();

  // Precompute breadth-first propagation order
  init_bfs_paths();
  std::cout << "Gauss-Seidel initilization of gradients......" << std::endl;

  Timer::EventID before_GS = timer.get_time();

  gauss_seidel_init_gradients();

  Timer::EventID before_ADMM = timer.get_time();
  std::cout << "ADMM solver for integrable gradients......" << std::endl;

  prepare_integrate_geodesic_distance();

  compute_integrable_gradients();

  Timer::EventID after_ADMM = timer.get_time();
  std::cout << "Recovery of geodesic distance......" << std::endl;

  integrate_geodesic_distance();

  Timer::EventID end = timer.get_time();

  std::cout << std::endl;
  std::cout << "====== Timing ======" << std::endl;
  std::cout << "Pre-computation of BFS paths: "
            << timer.elapsed_time(start, before_GS) << " seconds" << std::endl;
  std::cout << "Gauss-Seidel initialization of gradients: "
            << timer.elapsed_time(before_GS, before_ADMM) << " seconds"
            << std::endl;
  std::cout << "ADMM solver for integrable gradients: "
            << timer.elapsed_time(before_ADMM, after_ADMM) << " seconds"
            << std::endl;
  std::cout << "Integration of gradients: "
            << timer.elapsed_time(after_ADMM, end) << " seconds" << std::endl;
  std::cout << "Total time: " << timer.elapsed_time(start, end) << " seconds"
            << std::endl;

  return true;
}

void EdgeBasedGeodesicSolver::init_bfs_paths() {
  bfs_vertex_list.resize(n_vertices);
  bfs_vertex_list.fill(-1);
  transition_halfedge_idx.setConstant(n_vertices, -1);

  bfs_laplacian_coef_addr.resize(n_vertices + 1);
  bfs_laplacian_coef_addr(0) = 0;

  // Store source vertices as the first layer
  std::vector<bool> visited(n_vertices, false);
  std::vector<int> front1 = param.source_vertices, front2;
  std::vector<int> *current_front = &front1, *next_front = &front2;

  int n_sources = param.source_vertices.size();

  std::vector<int> bfs_segment_addr_vec;
  bfs_segment_addr_vec.push_back(0);
  bfs_segment_addr_vec.push_back(n_sources);

  int id = 0;
  for (; id < n_sources; ++id) {
    int current_source_vtx = param.source_vertices[id];
    visited[current_source_vtx] = true;
    bfs_vertex_list(id) = current_source_vtx;
    bfs_laplacian_coef_addr(id + 1) = bfs_laplacian_coef_addr(id)
        + mesh.valence(MeshType::Vertex(current_source_vtx)) + 1;
  }

  while (!current_front->empty()) {

    next_front->clear();

    for (int k = 0; k < static_cast<int>(current_front->size()); ++k) {
      MeshType::Vertex vh(current_front->at(k));
      MeshType::Halfedge_around_vertex_circulator vhc, vhc_end;
      vhc = vhc_end = mesh.halfedges(vh);

      do {
        MeshType::Halfedge heh = *vhc;
        MeshType::Vertex next_vh = mesh.to_vertex(heh);
        int next_v = next_vh.idx();

        if (!visited[next_v]) {
          next_front->push_back(next_v);
          bfs_vertex_list(id) = next_v;
          // Each segment stores the weights for neighbors and the current vertex for GS update
          bfs_laplacian_coef_addr(id + 1) = bfs_laplacian_coef_addr(id)
              + mesh.valence(next_vh) + 1;
          transition_halfedge_idx(id) = heh.idx();
          id++;
        }

        visited[next_v] = true;

      } while (++vhc != vhc_end);
    }

    bfs_segment_addr_vec.push_back(
        bfs_segment_addr_vec.back() + next_front->size());
    std::swap(current_front, next_front);
  }

  bfs_segment_addr = Eigen::Map < Eigen::VectorXi
      > (bfs_segment_addr_vec.data(), bfs_segment_addr_vec.size());
}

bool EdgeBasedGeodesicSolver::load_input(const char* mesh_file) {
  if (!surface_mesh::read_mesh(mesh, mesh_file)) {
    std::cerr << "Error: unable to read input mesh from the file " << mesh_file
              << std::endl;
    return false;
  }

  mesh.free_memory();  // Free unused memory

  n_vertices = mesh.n_vertices();
  n_faces = mesh.n_faces();
  n_edges = mesh.n_edges();
  n_halfedges = mesh.n_halfedges();

  if (n_vertices == 0 || n_faces == 0 || n_edges == 0) {
    std::cerr << "Error: zero mesh element count " << std::endl;
    return false;
  }

  for (int i = 0; i < static_cast<int>(param.source_vertices.size()); ++i) {
    if (param.source_vertices[i] < 0
        || param.source_vertices[i] >= n_vertices) {
      std::cerr << "Error: invalid source vertex index "
                << param.source_vertices[i] << std::endl;
      return false;
    }
  }

  return true;
}

void EdgeBasedGeodesicSolver::normalize_mesh() {
  std::vector<surface_mesh::Point> &pos = mesh.points();

  surface_mesh::Point min_coord = pos.front(), max_coord = pos.front();
  std::vector<surface_mesh::Point>::iterator iter = pos.begin(), iter_end = pos
      .end();

  for (++iter; iter != iter_end; ++iter) {
    surface_mesh::Point &coord = *iter;
    min_coord.minimize(coord);
    max_coord.maximize(coord);
  }

  model_scaling_factor = surface_mesh::norm(max_coord - min_coord);
  surface_mesh::Point center_pos = (min_coord + max_coord) * 0.5;

  for (iter = pos.begin(); iter != iter_end; ++iter) {
    surface_mesh::Point &coord = *iter;
    coord -= center_pos;
    coord /= model_scaling_factor;
  }
}

void EdgeBasedGeodesicSolver::gauss_seidel_init_gradients() {
  edge_vector.resize(3, n_edges);
  DenseVector edge_sqr_length;
  edge_sqr_length.setZero(n_edges);

  face_area.resize(n_faces);

  int n_halfedges = mesh.n_halfedges();
  DenseVector halfedge_halfcot;  // Collect the half cotan value for edge halfedge, to be used for computing cotan Laplacian weights
  halfedge_halfcot.setZero(n_halfedges);

  double step_length = 0;
  HeatScalar init_source_val = 1;
  VectorHS current_d;
  VectorHS temp_d;
  VectorHS vertex_area;
  int gs_iter = 0;
  int segment_count = 0;
  int n_segments = 0;
  int segment_begin_addr = 0, segment_end_addr = 0;
  bool end_gs_loop = false;
  bool reset_iter = true;
  bool need_check_residual = false;
  HeatScalar eps = 0;
  VectorHS heatflow_residuals;

  OMP_PARALLEL
  {
    // Compute Laplacian weights
    OMP_FOR
    for (int i = 0; i < n_edges; ++i) {
      // Precompute edge vectors and squared edge length,
      // to be used later for computing cotan weights and areas
      MeshType::Halfedge heh = mesh.halfedge(MeshType::Edge(i), 0);
      Eigen::Vector3d edge_vec = to_eigen_vec3d(
          mesh.position(mesh.to_vertex(heh))
              - mesh.position(mesh.from_vertex(heh)));
      double l2 = edge_vec.squaredNorm();
      edge_vector.col(i) = edge_vec;
      edge_sqr_length(i) = l2;
    }

    OMP_SINGLE
    {
      // Compute heat flow step size
      double h = edge_sqr_length.array().sqrt().mean();
      step_length = h * h;
    }

    OMP_FOR
    for (int i = 0; i < n_faces; ++i) {
      // Compute face areas and half-cotan weights for halfedges
      Eigen::Vector3i fh_idx, fe_idx;
      Eigen::Vector3d edge_l2;
      int k = 0;

      MeshType::Halfedge_around_face_circulator fhc, fhc_end;
      fhc = fhc_end = mesh.halfedges(MeshType::Face(i));
      do {
        MeshType::Halfedge heh = *fhc;
        fh_idx(k) = heh.idx();
        fe_idx(k) = mesh.edge(heh).idx();
        edge_l2(k) = edge_sqr_length(fe_idx(k));
        k++;
      } while (++fhc != fhc_end);

      double area = edge_vector.col(fe_idx(0)).cross(edge_vector.col(fe_idx(1)))
          .norm() * 0.5;
      for (int j = 0; j < 3; ++j) {
        halfedge_halfcot(fh_idx(j)) = 0.125
            * (edge_l2((j + 1) % 3) + edge_l2((j + 2) % 3) - edge_l2(j)) / area;
      }

      face_area(i) = area;
    }

    OMP_SINGLE
    {
      // Allocate arrays for storing relevant vertices and weights for Laplacian operator at each vertex
      edge_sqr_length.resize(0);
      int n_laplacian_vertices = n_edges * 2 + n_vertices;
      bfs_laplacian_coef = new std::pair<int, double>[n_laplacian_vertices];
      vertex_area.setZero(n_vertices);
    }

    OMP_FOR
    for (int i = 0; i < n_vertices; ++i) {
      // Compute and store vertex indices and weights for Laplacian operators
      int start_addr = bfs_laplacian_coef_addr(i), end_addr =
          bfs_laplacian_coef_addr(i + 1);

      DenseVector weights;
      IndexVector vtx_idx;
      int n = end_addr - start_addr;
      weights.setZero(n);
      vtx_idx.setZero(n);

      int v_idx = bfs_vertex_list(i);
      MeshType::Vertex vh(v_idx);
      int k = 0;
      MeshType::Halfedge_around_vertex_circulator vhc, vhc_end;
      vhc = vhc_end = mesh.halfedges(vh);

      do {
        MeshType::Halfedge heh = *vhc;
        double w = halfedge_halfcot(heh.idx())
            + halfedge_halfcot(mesh.opposite_halfedge(heh).idx());
        vtx_idx(k) = mesh.to_vertex(heh).idx();
        weights(k) = w;
        k++;
      } while (++vhc != vhc_end);

      vtx_idx(k) = v_idx;
      weights(k) = weights.head(k).sum();  // Store the sum of neighbor weights, to be used for Gauss-Seidel update
      weights *= step_length;

      double A = 0;
      MeshType::Face_around_vertex_circulator vfc, vfc_end;
      vfc = vfc_end = mesh.faces(vh);
      do {
        A += face_area((*vfc).idx());
      } while (++vfc != vfc_end);

      double vertex_A = A / 3.0;
      vertex_area(v_idx) = vertex_A;
      weights(k) += vertex_A;

      for (int j = start_addr; j < end_addr; ++j) {
        bfs_laplacian_coef[j] = std::pair<int, double>(vtx_idx(j - start_addr),
                                                       weights(j - start_addr));
      }
    }

    OMP_SINGLE
    {
      // Set up heat value arrays
      halfedge_halfcot.resize(0);

      int n_sources = param.source_vertices.size();
      HeatScalar total_source_area = 0;
      for (int i = 0; i < n_sources; ++i) {
        total_source_area += vertex_area(param.source_vertices[i]);
      }
      init_source_val = std::sqrt(
          std::min(HeatScalar(n_vertices) / HeatScalar(n_sources),
                   vertex_area.sum() / total_source_area));
      vertex_area.resize(0);

      current_d.setZero(n_vertices);
      for (int i = 0; i < n_sources; ++i) {
        current_d(param.source_vertices[i]) = init_source_val;
      }

      n_segments = bfs_segment_addr.size() - 1;
      int buffer_size = (Eigen::Map < IndexVector
          > (&(bfs_segment_addr[1]), n_segments) - Eigen::Map < IndexVector
          > (&(bfs_segment_addr[0]), n_segments)).maxCoeff();
      temp_d.setZero(buffer_size);

      heatflow_residuals.setZero(n_vertices);
    }

    compute_heatflow_residual(current_d, init_source_val, heatflow_residuals);

    OMP_SINGLE
    {
      // Rescale heat source values to make the initial residual norm close to 1
      HeatScalar init_residual_norm = heatflow_residuals.norm();
      eps = std::max(HeatScalar(1e-16),
                     init_residual_norm * HeatScalar(param.heat_solver_eps));
      std::cout << "Initial residual: " << init_residual_norm << ", threshold: "
                << eps << std::endl;
    }

  }

  while (!end_gs_loop) {
    OMP_PARALLEL
    {
      // Gauss-Seidel update of heat values in breadth-first order
      OMP_SINGLE
      {
        segment_begin_addr = bfs_segment_addr(segment_count);
        segment_end_addr = bfs_segment_addr(segment_count + 1);
      }

      OMP_FOR
      for (int i = segment_begin_addr; i < segment_end_addr; ++i) {
        int lap_coef_begin_addr = bfs_laplacian_coef_addr(i);
        int lap_coef_end_addr = bfs_laplacian_coef_addr(i + 1);

        HeatScalar new_heat_value = 0;
        if (segment_count == 0) {  // Check whether the current vertex is a source
          new_heat_value += init_source_val;
        }

        for (int j = lap_coef_begin_addr; j < lap_coef_end_addr - 1; ++j) {
          std::pair<int, double> &coef = bfs_laplacian_coef[j];
          new_heat_value += current_d(coef.first) * coef.second;
        }

        temp_d(i - segment_begin_addr) = new_heat_value
            / bfs_laplacian_coef[lap_coef_end_addr - 1].second;
      }

      OMP_FOR
      for (int i = segment_begin_addr; i < segment_end_addr; ++i) {
        current_d(bfs_vertex_list(i)) = temp_d(i - segment_begin_addr);
      }

      OMP_SINGLE
      {
        segment_count++;
        reset_iter = (segment_count == n_segments);
        if (reset_iter) {
          gs_iter++;
          segment_count = 0;
        }

        end_gs_loop = gs_iter >= param.heat_solver_max_iter;
        need_check_residual =
            end_gs_loop
                || (reset_iter
                    && (gs_iter % param.heat_solver_convergence_check_frequency
                        == 0));
      }

      if (need_check_residual) {
        compute_heatflow_residual(current_d, init_source_val,
                                  heatflow_residuals);

        OMP_SINGLE
        {
          HeatScalar residual_norm = heatflow_residuals.norm();
          std::cout << "Gauss-Seidel iteration " << gs_iter
                    << ", current residual: " << residual_norm
                    << ", threshold: " << eps << std::endl;

          if (residual_norm <= eps) {
            end_gs_loop = true;
          }
        }
      }
    }
  }

  OMP_PARALLEL
  {
    OMP_SINGLE
    {
      temp_d.resize(0);
      heatflow_residuals.resize(0);
      vertex_area.resize(0);
      delete[] bfs_laplacian_coef;
      bfs_laplacian_coef_addr.resize(0);
      init_grad.resize(3, n_faces);
    }

    // Compute initial gradient and get target edge difference.
    OMP_FOR
    for (int i = 0; i < n_faces; ++i) {
      Matrix3HS edge_vecs;
      Vector3HS heat_vals;
      int k = 0;

      MeshType::Halfedge_around_face_circulator fhc, fhc_end;
      fhc = fhc_end = mesh.halfedges(MeshType::Face(i));

      do {
        MeshType::Halfedge heh = *fhc;
        MeshType::Edge eh = mesh.edge(heh);
        Eigen::Vector3d current_edge = edge_vector.col(eh.idx());
        if (mesh.halfedge(eh, 0) != heh) {
          current_edge *= -1;
        }

        edge_vecs(0, k) = HeatScalar(current_edge[0]);
        edge_vecs(1, k) = HeatScalar(current_edge[1]);
        edge_vecs(2, k) = HeatScalar(current_edge[2]);
        heat_vals(k) = current_d(mesh.to_vertex(heh).idx());
        ++k;

      } while (++fhc != fhc_end);

      heat_vals.normalize();
      edge_vecs.normalize();

      Vector3HS N = edge_vecs.col(0).cross(edge_vecs.col(1)).normalized();
      Vector3HS V = edge_vecs.col(0) * heat_vals(1)
          + edge_vecs.col(1) * heat_vals(2) + edge_vecs.col(2) * heat_vals(0);
      Vector3HS grad_vec = V.cross(N).normalized();
      init_grad(0, i) = grad_vec(0);
      init_grad(1, i) = grad_vec(1);
      init_grad(2, i) = grad_vec(2);
    }
  }
}

void EdgeBasedGeodesicSolver::compute_heatflow_residual(
    const VectorHS &heat_values, HeatScalar init_source_val,
    VectorHS &residuals) {
  OMP_FOR
  for (int i = 0; i < n_vertices; ++i) {
    int lap_coef_begin_addr = bfs_laplacian_coef_addr(i);
    int lap_coef_end_addr = bfs_laplacian_coef_addr(i + 1);

    HeatScalar res = 0;
    if (i < static_cast<int>(param.source_vertices.size())) {  // Check whether the current vertex is a source
      res += init_source_val;
    }

    for (int j = lap_coef_begin_addr; j < lap_coef_end_addr; ++j) {
      std::pair<int, double> &coef = bfs_laplacian_coef[j];
      res += heat_values(coef.first) * coef.second
          * ((j == (lap_coef_end_addr - 1)) ? (-1) : 1);
    }

    residuals(i) = res;
  }
}

void EdgeBasedGeodesicSolver::prepare_integrate_geodesic_distance() {
  transition_from_vtx.setConstant(n_vertices, -1);
  transition_edge_idx.setConstant(n_vertices, -1);

  IndexVector num_rows;
  num_rows.setZero(n_edges);

  Z.resize(3 * n_faces);
  S.resize(3, n_faces);
  Q.resize(3, n_faces);
  edges_Y_index.setConstant(2, n_edges, -1);  // the set of rows in Y associated with each edge.

  OMP_PARALLEL
  {
    // Set up incident relation between edges and faces.
    OMP_SINGLE
    {
      for (int i = 0; i < n_faces; ++i) {
        MeshType::Face f = MeshType::Face(i);
        MeshType::Halfedge_around_face_circulator fhc, fhc_end;
        fhc = fhc_end = mesh.halfedges(f);
        int k = 0;

        do {
          MeshType::Edge e = mesh.edge(*fhc);
          int edge_index = mesh.edge(*fhc).idx();
          Eigen::Vector3d e_vector = to_eigen_vec3d(
              mesh.position(mesh.from_vertex(*fhc))
                  - mesh.position(mesh.to_vertex(*fhc)));
          if (*fhc == mesh.halfedge(e, 0))  // the halfedge with index 0 as orientation halfedge
                                    {
            Q(k, i) = 1;
            Z(3 * i + k) = init_grad.col(i).dot(e_vector);

          } else {
            Q(k, i) = -1;
            Z(3 * i + k) = init_grad.col(i).dot(-e_vector);
          }

          S(k, i) = edge_index;
          edges_Y_index(num_rows(edge_index)++, edge_index) = 3 * i + k;
          k++;

        } while (++fhc != fhc_end);
      }
    }

    // Set up transition vector needed in recovering distance step.
    OMP_FOR
    for (int i = 0; i < n_vertices; ++i) {
      MeshType::Halfedge heh(transition_halfedge_idx(i));
      if (heh.is_valid()) {
        MeshType::Vertex from_vh = mesh.from_vertex(heh);
        MeshType::Edge e = mesh.edge(heh);
        transition_from_vtx(i) = from_vh.idx();

        if (heh == mesh.halfedge(e, 0)) {
          transition_edge_idx(i) = -e.idx() - 1;
        } else {
          transition_edge_idx(i) = e.idx();
        }
      }
    }

    OMP_SINGLE
    {
      mesh.clear();
      transition_halfedge_idx.resize(0);
      init_grad.resize(3, 0);

      D.setZero(3 * n_faces);
      X.resize(n_edges);
      Y.setZero(3 * n_faces);

      SX1.setZero(3 * n_faces);
      SX2.setZero(3 * n_faces);
      current_SX = &SX1;
      prev_SX = &SX2;

      primal_residual_sqr_norm_threshold = param.grad_solver_eps
          * param.grad_solver_eps;
      dual_residual_sqr_norm_threshold = param.grad_solver_eps
          * param.grad_solver_eps;
      (*prev_SX) = (*current_SX);
    }

    // Initialize X.
    OMP_FOR
    for (int i = 0; i < n_edges; ++i) {
      int n_var = 0;
      double r = 0;
      for (int j = 0; j < 2; ++j) {
        int index = edges_Y_index(j, i);
        if (index >= 0) {
          r += Z(index);
          n_var++;
        }
      }

      X(i) = r / n_var;
    }

    // Initialize SX.
    OMP_FOR
    for (int i = 0; i < n_faces; ++i) {
      (*prev_SX)(3 * i) = X(S(0, i));
      (*prev_SX)(3 * i + 1) = X(S(1, i));
      (*prev_SX)(3 * i + 2) = X(S(2, i));
    }
  }
}

void EdgeBasedGeodesicSolver::compute_integrable_gradients() {
  optimization_end = false;
  iter_num = 0;
  OMP_PARALLEL
  {
    while (!optimization_end) {
      update_Y();
      update_X();
      update_dual_variables();
    }
  }
}

void EdgeBasedGeodesicSolver::integrate_geodesic_distance() {
  geod_dist_values.setZero(n_vertices);
  int n_segments = bfs_segment_addr.size() - 1;
  int segment_count = 1;  // We update the distance values strting from the second layer of BFS vertex list
  int segment_begin_addr, segment_end_addr;
  bool end_propagation = false;

  OMP_PARALLEL
  {
    while (!end_propagation) {
      OMP_SINGLE
      {
        segment_begin_addr = bfs_segment_addr(segment_count);
        segment_end_addr = bfs_segment_addr(segment_count + 1);
      }

      OMP_FOR
      for (int i = segment_begin_addr; i < segment_end_addr; ++i) {
        double from_d = geod_dist_values(transition_from_vtx(i));
        int edge_index = transition_edge_idx(i);
        if (edge_index >= 0) {
          geod_dist_values(bfs_vertex_list(i)) = from_d + X(edge_index);
        } else {
          geod_dist_values(bfs_vertex_list(i)) = from_d - X(-(edge_index + 1));
        }
      }

      OMP_SINGLE
      {
        segment_count++;
        end_propagation = (segment_count >= n_segments);
      }
    }
  }

  // Recover geodesic distance in the original scale
  geod_dist_values *= model_scaling_factor;
}

void EdgeBasedGeodesicSolver::update_Y() {
  OMP_FOR
  for (int i = 0; i < n_faces; ++i) {
    Eigen::Vector3d y = prev_SX->segment(3 * i, 3) - D.segment(3 * i, 3);
    Eigen::Vector3d q = Q.col(i).cast<double>();
    Y.segment(3 * i, 3) = y - 1.0 / 3 * q.dot(y) * q;
  }
}

void EdgeBasedGeodesicSolver::update_X() {
  OMP_FOR
  for (int i = 0; i < n_edges; ++i) {
    int n_aux_var = 0;
    double r = 0;
    for (int j = 0; j < 2; ++j) {
      int index = edges_Y_index(j, i);
      if (index >= 0) {
        r += param.penalty * (Y(index) + D(index)) + Z(index);
        n_aux_var++;
      }
    }

    X(i) = r / ((param.penalty + 1) * n_aux_var);
  }
}

void EdgeBasedGeodesicSolver::update_dual_variables() {
  OMP_FOR
  for (int i = 0; i < n_faces; ++i) {
    (*current_SX)(3 * i) = X(S(0, i));
    (*current_SX)(3 * i + 1) = X(S(1, i));
    (*current_SX)(3 * i + 2) = X(S(2, i));
  }

  OMP_SINGLE
  {
    need_compute_residual_norms = ((iter_num + 1)
        % param.grad_solver_convergence_check_frequency == 0);
  }

  OMP_SECTIONS
  {
    OMP_SECTION
    {
      if (need_compute_residual_norms) {
        primal_residual_sqr_norm = (Y - (*current_SX)).squaredNorm();
      }
    }

    OMP_SECTION
    {
      if (need_compute_residual_norms) {
        dual_residual_sqr_norm = ((*current_SX) - (*prev_SX)).squaredNorm()
            * param.penalty * param.penalty;
      }
    }

    OMP_SECTION
    {
      D += Y - (*current_SX);
    }
  }

  OMP_SINGLE
  {
    iter_num++;
    optimization_converge = need_compute_residual_norms
        && (primal_residual_sqr_norm <= primal_residual_sqr_norm_threshold
            && dual_residual_sqr_norm <= dual_residual_sqr_norm_threshold);
    optimization_end = optimization_converge
        || iter_num >= param.grad_solver_max_iter;
    output_progress = need_compute_residual_norms
        && (iter_num % param.grad_solver_output_frequency == 0);

    if (optimization_converge) {
      std::cout << "Solver converged." << std::endl;
    } else if (optimization_end) {
      std::cout << "Maximum number of iterations reached." << std::endl;
    }

    if (output_progress || optimization_end) {
      std::cout << "Iteration " << iter_num << ":" << std::endl;
      std::cout << "Primal residual squared norm: " << primal_residual_sqr_norm
                << ",  threshold:" << primal_residual_sqr_norm_threshold
                << std::endl;
      std::cout << "Dual residual squared norm: " << dual_residual_sqr_norm
                << ",  threshold:" << dual_residual_sqr_norm_threshold
                << std::endl;
    }

    std::swap(current_SX, prev_SX);
  }
}