fix: correct FMM source propagation

src/surface/fast_marching_method.cpp

@@ -1,5 +1,8 @@
 #include "geometrycentral/surface/fast_marching_method.h"
 
+#include <algorithm>
+#include <cmath>
+#include <limits>
 #include <queue>
 #include <tuple>
 
@@ -53,7 +56,7 @@ VertexData<double> FMMDistance(IntrinsicGeometryInterface& geometry,
                                const std::vector<std::vector<std::pair<SurfacePoint, double>>>& initialDistances,
                                bool sign) {
 
-  typedef std::pair<double, Vertex> Entry;
+  typedef std::tuple<double, Vertex, int> Entry;
 
   SurfaceMesh& mesh = geometry.mesh;
   geometry.requireEdgeLengths();
@@ -69,16 +72,42 @@ VertexData<double> FMMDistance(IntrinsicGeometryInterface& geometry,
   VertexData<int> signs(mesh, 1);
   VertexData<char> finalized(mesh, false);
   VertexData<char> isSource(mesh, false);
+  VertexData<int> sourceLabels(mesh, -1);
+  EdgeData<char> isSourceBaseEdge(mesh, false);
+  EdgeData<int> sourceBaseLabels(mesh, -1);
+  int nextSourceLabel = 0;
 
   auto cmp = [&signs](Entry left, Entry right) {
-    if (signs[left.second] != signs[right.second]) {
+    Vertex leftV = std::get<1>(left);
+    Vertex rightV = std::get<1>(right);
+    if (signs[leftV] != signs[rightV]) {
       // We're looking at an edge that intersects the source, which may not have initial distance 0.
       // I *think* this can only happen if the positive vertex lies on the source, in which case it's the "closer" one.
-      return signs[left.second] != 1;
+      return signs[leftV] != 1;
     }
-    return signs[left.second] * left.first > signs[right.second] * right.first;
+    return signs[leftV] * std::get<0>(left) > signs[rightV] * std::get<0>(right);
   };
   std::priority_queue<Entry, std::vector<Entry>, decltype(cmp)> frontierPQ(cmp);
+
+  auto candidateImprovesDistance = [&](Vertex v, double candidateDist) {
+    if (!std::isfinite(candidateDist)) return false;
+    if (!std::isfinite(distances[v])) return true;
+    if (signs[v] == 0) return std::abs(candidateDist) < std::abs(distances[v]);
+    return signs[v] * candidateDist < signs[v] * distances[v];
+  };
+
+  auto addCandidateDistance = [&](Vertex v, double candidateDist, int sourceLabel) {
+    if (!candidateImprovesDistance(v, candidateDist)) return;
+    distances[v] = candidateDist;
+    sourceLabels[v] = sourceLabel;
+    frontierPQ.push(std::make_tuple(candidateDist, v, sourceLabel));
+  };
+
+  auto markSourceBaseEdge = [&](Edge e, int sourceLabel) {
+    isSourceBaseEdge[e] = true;
+    if (sourceBaseLabels[e] == -1) sourceBaseLabels[e] = sourceLabel;
+  };
+
   // Initialize signs
   if (sign) {
     for (Vertex v : mesh.vertices()) signs[v] = 0;
@@ -108,6 +137,13 @@ VertexData<double> FMMDistance(IntrinsicGeometryInterface& geometry,
         }
       }
     }
+    for (auto& curve : initialDistances) {
+      size_t nNodes = curve.size();
+      for (size_t i = 0; i < nNodes - 1; i++) {
+        Edge commonEdge = sharedEdge(curve[i].first, curve[i + 1].first);
+        if (commonEdge != Edge()) markSourceBaseEdge(commonEdge, nextSourceLabel++);
+      }
+    }
     // Vertices on the curve are always assumed to have positive sign.
     for (auto& curve : initialDistances) {
       size_t nNodes = curve.size();
@@ -140,20 +176,22 @@ VertexData<double> FMMDistance(IntrinsicGeometryInterface& geometry,
   for (auto& curve : initialDistances) {
     for (auto& x : curve) {
       const SurfacePoint& p = x.first;
+      int sourceLabel = nextSourceLabel++;
       switch (p.type) {
       case (SurfacePointType::Vertex): {
-        frontierPQ.push(std::make_pair(x.second, p.vertex));
+        addCandidateDistance(p.vertex, x.second, sourceLabel);
         isSource[p.vertex] = true;
         break;
       }
       case (SurfacePointType::Edge): {
         const Vertex& vA = p.edge.firstVertex();
         const Vertex& vB = p.edge.secondVertex();
         double l = geometry.edgeLengths[p.edge];
-        frontierPQ.push(std::make_pair(x.second + signs[vA] * p.tEdge * l, vA));
-        frontierPQ.push(std::make_pair(x.second + signs[vB] * (1. - p.tEdge) * l, vB));
+        addCandidateDistance(vA, x.second + signs[vA] * p.tEdge * l, sourceLabel);
+        addCandidateDistance(vB, x.second + signs[vB] * (1. - p.tEdge) * l, sourceLabel);
         isSource[vA] = true;
         isSource[vB] = true;
+        markSourceBaseEdge(p.edge, sourceLabel);
         break;
       }
       case (SurfacePointType::Face): {
@@ -170,20 +208,24 @@ VertexData<double> FMMDistance(IntrinsicGeometryInterface& geometry,
         double u = p.faceCoords[0];
         double v = p.faceCoords[1];
         double w = p.faceCoords[2];
-        double dist2_A = lAB2 * (v * (1. - u)) + lCA2 * (w * (1. - u)) - lAB2 * v * w; // squared distance from p to vA
+        double dist2_A = lAB2 * (v * (1. - u)) + lCA2 * (w * (1. - u)) - lBC2 * v * w; // squared distance from p to vA
         double dist2_B = lAB2 * (u * (1. - v)) + lBC2 * (w * (1. - v)) - lCA2 * u * w; // squared distance from p to vB
-        double dist2_C = lCA2 * (u * (1. - w)) + lBC2 * (v * (1. - w)) - lBC2 * u * v; // squared distance from p to vC
-        frontierPQ.push(std::make_pair(x.second + signs[vA] * std::sqrt(dist2_A), vA));
-        frontierPQ.push(std::make_pair(x.second + signs[vB] * std::sqrt(dist2_B), vB));
-        frontierPQ.push(std::make_pair(x.second + signs[vC] * std::sqrt(dist2_C), vC));
+        double dist2_C = lCA2 * (u * (1. - w)) + lBC2 * (v * (1. - w)) - lAB2 * u * v; // squared distance from p to vC
+        addCandidateDistance(vA, x.second + signs[vA] * std::sqrt(std::max(0., dist2_A)), sourceLabel);
+        addCandidateDistance(vB, x.second + signs[vB] * std::sqrt(std::max(0., dist2_B)), sourceLabel);
+        addCandidateDistance(vC, x.second + signs[vC] * std::sqrt(std::max(0., dist2_C)), sourceLabel);
         isSource[vA] = true;
         isSource[vB] = true;
         isSource[vC] = true;
+        markSourceBaseEdge(he.edge(), sourceLabel);
+        markSourceBaseEdge(he.next().edge(), sourceLabel);
+        markSourceBaseEdge(he.next().next().edge(), sourceLabel);
         break;
       }
       }
     }
   }
+
   size_t nFound = 0;
   size_t nVert = mesh.nVertices();
 
@@ -193,12 +235,13 @@ VertexData<double> FMMDistance(IntrinsicGeometryInterface& geometry,
     // Pop the nearest element
     Entry currPair = frontierPQ.top();
     frontierPQ.pop();
-    Vertex currV = currPair.second;
-    double currDist = currPair.first;
+    double currDist = std::get<0>(currPair);
+    Vertex currV = std::get<1>(currPair);
+    int currSourceLabel = std::get<2>(currPair);
 
 
     // Accept it if not stale
-    if (finalized[currV]) {
+    if (finalized[currV] || currDist != distances[currV] || currSourceLabel != sourceLabels[currV]) {
       continue;
     }
     distances[currV] = currDist;
@@ -213,13 +256,15 @@ VertexData<double> FMMDistance(IntrinsicGeometryInterface& geometry,
       if (!finalized[neighVert] && !isSource[neighVert]) {
         if (signs[neighVert] == 0) signs[neighVert] = signs[currV];
         double newDist = currDist + signs[neighVert] * geometry.edgeLengths[he.edge()];
-        if (signs[neighVert] * newDist < signs[neighVert] * distances[neighVert] || std::isinf(distances[neighVert])) {
-          frontierPQ.push(std::make_pair(newDist, neighVert));
-          distances[neighVert] = newDist;
-        }
+        addCandidateDistance(neighVert, newDist, currSourceLabel);
         continue;
       }
 
+      bool useSourceBaseEdge = sourceLabels[neighVert] != currSourceLabel && isSource[currV] && isSource[neighVert] &&
+                               isSourceBaseEdge[he.edge()];
+      if (sourceLabels[neighVert] != currSourceLabel && !useSourceBaseEdge) continue;
+      int updateSourceLabel = useSourceBaseEdge ? sourceBaseLabels[he.edge()] : currSourceLabel;
+
       // Check the third point of the "left" triangle straddling this edge
       if (he.isInterior()) {
         Vertex newVert = he.next().next().vertex();
@@ -233,10 +278,7 @@ VertexData<double> FMMDistance(IntrinsicGeometryInterface& geometry,
           double theta = geometry.cornerAngles[he.next().next().corner()];
           if (signs[newVert] == 0) signs[newVert] = (signs[currV] != 0) ? signs[currV] : signs[he.next().vertex()];
           double newDist = eikonalDistanceSubroutine(lenA, lenB, theta, distA, distB, signs[newVert]);
-          if (signs[newVert] * newDist < signs[newVert] * distances[newVert] || std::isinf(distances[newVert])) {
-            frontierPQ.push(std::make_pair(newDist, newVert));
-            distances[newVert] = newDist;
-          }
+          addCandidateDistance(newVert, newDist, updateSourceLabel);
         }
       }
 
@@ -254,10 +296,7 @@ VertexData<double> FMMDistance(IntrinsicGeometryInterface& geometry,
           double theta = geometry.cornerAngles[heT.next().next().corner()];
           if (signs[newVert] == 0) signs[newVert] = (signs[currV] != 0) ? signs[currV] : signs[he.next().vertex()];
           double newDist = eikonalDistanceSubroutine(lenA, lenB, theta, distA, distB, signs[newVert]);
-          if (signs[newVert] * newDist < signs[newVert] * distances[newVert] || std::isinf(distances[newVert])) {
-            frontierPQ.push(std::make_pair(newDist, newVert));
-            distances[newVert] = newDist;
-          }
+          addCandidateDistance(newVert, newDist, updateSourceLabel);
         }
       }
     }