Add warnings in dtm.py for DistanceToMeasure and DTMDensity

Add DTMDensity warning test

2 files changed, 20 insertions, 1 deletions

diff --git a/src/python/gudhi/point_cloud/dtm.py b/src/python/gudhi/point_cloud/dtm.py
index 55ac58e6..96a9e7bf 100644
--- a/src/python/gudhi/point_cloud/dtm.py
+++ b/src/python/gudhi/point_cloud/dtm.py
@@ -9,6 +9,7 @@
 
 from .knn import KNearestNeighbors
 import numpy as np
+import warnings
 
 __author__ = "Marc Glisse"
 __copyright__ = "Copyright (C) 2020 Inria"
@@ -66,6 +67,11 @@ class DistanceToMeasure:
         distances = distances ** self.q
         dtm = distances.sum(-1) / self.k
         dtm = dtm ** (1.0 / self.q)
+        with warnings.catch_warnings():
+            import torch
+            if isinstance(dtm, torch.Tensor):
+                if not(torch.isfinite(dtm).all()):
+                    warnings.warn("Overflow/infinite value encountered while computing 'dtm'", RuntimeWarning)
         # We compute too many powers, 1/p in knn then q in dtm, 1/q in dtm then q or some log in the caller.
         # Add option to skip the final root?
         return dtm
@@ -163,6 +169,11 @@ class DTMDensity:
             distances = self.knn.transform(X)
         distances = distances ** q
         dtm = (distances * weights).sum(-1)
+        with warnings.catch_warnings():
+            import torch
+            if isinstance(dtm, torch.Tensor):
+                if not(torch.isfinite(dtm).all()):
+                    warnings.warn("Overflow/infinite value encountered while computing 'dtm' for density", RuntimeWarning)
         if self.normalize:
             dtm /= (np.arange(1, k + 1) ** (q / dim) * weights).sum()
         density = dtm ** (-dim / q)
diff --git a/src/python/test/test_dtm.py b/src/python/test/test_dtm.py
index c29471cf..52468d0f 100755
--- a/src/python/test/test_dtm.py
+++ b/src/python/test/test_dtm.py
@@ -97,7 +97,15 @@ def test_dtm_overflow_warnings():
         for impl in impl_warn:
             dtm = DistanceToMeasure(2, q=10000, implementation=impl)
             r = dtm.fit_transform(pts)
-        assert len(w) == 2
+        assert len(w) == 3
         for i in range(len(w)):
             assert issubclass(w[i].category, RuntimeWarning)
             assert "Overflow" in str(w[i].message)
+
+def test_density_overflow_warning():
+    distances = numpy.array([[10., 100.], [10000000000000., 10.]])
+    with warnings.catch_warnings(record=True) as w:
+        density = DTMDensity(k=2, q=100000, implementation="keops", dim=1).fit_transform(distances)
+        assert len(w) == 1
+        assert issubclass(w[0].category, RuntimeWarning)
+        assert "Overflow" in str(w[0].message)