1 files changed, 20 insertions, 0 deletions
diff --git a/geom_matching/wasserstein/include/diagram_reader.h b/geom_matching/wasserstein/include/diagram_reader.h
index 84bf49c..55228b4 100644
--- a/geom_matching/wasserstein/include/diagram_reader.h
+++ b/geom_matching/wasserstein/include/diagram_reader.h
@@ -324,6 +324,26 @@ void remove_duplicates(ContType& dgm_A, ContType& dgm_B)
     }
 }
 
+template<class RealType = double>
+int finitize(RealType finitization, std::vector<std::pair<RealType, RealType> >& diagram)
+{
+  if (finitization == std::numeric_limits<double>::infinity())
+    return 0;
+  
+  int altered = 0;
+  
+  for (auto it = diagram.begin(); it != diagram.end(); ++it)
+  {
+    if (it->second > finitization)
+    {
+      it->second = finitization;
+      ++altered;
+    }
+  }
+  
+  return altered;
+}
+  
 
 #ifdef WASSERSTEIN_PURE_GEOM