fix: issue #61 — 対角 self-distance の sqrt NaN 勾配を構造的に修正

tda_ml/topology.py

@@ -124,6 +124,21 @@ def compute_anisotropic_distance_matrix(
             f"symmetrize must be 'max' or 'min', got {symmetrize!r}"
         )
 
-    # sqrt'(0) is undefined; diagonal entries are exactly zero, so add NUMERICAL_EPS
-    # inside sqrt for a bounded derivative (declared L2 floor, not a modeling prior).
-    return torch.sqrt(dist_sq + NUMERICAL_EPS)
+    return _sqrt_off_diagonal_only(dist_sq)
+
+
+def _sqrt_off_diagonal_only(dist_sq: torch.Tensor) -> torch.Tensor:
+    """
+    Mahalanobis distances with zero diagonal without sqrt backward on self-distances.
+
+    Diagonal squared distances are identically zero (self-distance). Differentiating
+    ``sqrt(x)`` at ``x=0`` yields NaN gradients. VR / persistence uses off-diagonal
+    entries only, so diagonals stay zero and ``sqrt`` is applied only where
+    ``dist_sq > 0`` (never to zero squared distances, including self-distance and
+    coincident off-diagonal pairs).
+    """
+    dist_sq = torch.clamp(dist_sq, min=0.0)
+    dist = torch.zeros_like(dist_sq)
+    positive = dist_sq > 0
+    dist[positive] = torch.sqrt(dist_sq[positive])
+    return dist

tests/test_topology.py

@@ -1,9 +1,7 @@
-import math
 import unittest
 
 import torch
 
-from tda_ml.numerical_eps import NUMERICAL_EPS
 from tda_ml.topology import compute_anisotropic_distance_matrix
 
 
@@ -19,13 +17,12 @@ def _make_batch(self):
         )
         return points, params
 
-    def test_diagonal_uses_numerical_sqrt_floor(self):
+    def test_diagonal_is_exactly_zero(self):
         points, params = self._make_batch()
         dist = compute_anisotropic_distance_matrix(points, params, symmetrize="max")
-        expected = math.sqrt(NUMERICAL_EPS)
         n = points.shape[1]
         for i in range(n):
-            self.assertAlmostEqual(dist[0, i, i].item(), expected, places=8)
+            self.assertEqual(dist[0, i, i].item(), 0.0)
 
     def test_symmetry_max_and_min(self):
         points, params = self._make_batch()
@@ -53,8 +50,7 @@ def test_probs_changes_off_diagonal(self):
                 atol=1e-8,
             )
         )
-        expected_diag = math.sqrt(NUMERICAL_EPS)
-        self.assertAlmostEqual(dist_probs[0, 0, 0].item(), expected_diag, places=8)
+        self.assertEqual(dist_probs[0, 0, 0].item(), 0.0)
 
     def test_invalid_symmetrize_raises(self):
         points, params = self._make_batch()

tests/test_topology_gradients.py

@@ -1,32 +1,115 @@
+"""Gradient health for anisotropic distance matrices (issue #61)."""
+
 import torch
 
 from tda_ml.topology import compute_anisotropic_distance_matrix
 
 
+def test_anisotropic_distance_diagonal_is_exactly_zero():
+    batch_size, num_points = 1, 5
+    points = torch.randn(batch_size, num_points, 2, dtype=torch.float32)
+    params = torch.tensor(
+        [
+            [
+                [1.0, 0.5, 0.0],
+                [1.0, 0.5, 0.0],
+                [1.0, 0.5, 0.0],
+                [1.0, 0.5, 0.0],
+                [1.0, 0.5, 0.0],
+            ]
+        ],
+        dtype=torch.float32,
+    )
+    probs = torch.full((batch_size, num_points), 0.2, dtype=torch.float32)
+
+    dist = compute_anisotropic_distance_matrix(
+        points, params, probs=probs, symmetrize="max"
+    )
+    diag = dist.diagonal(dim1=-2, dim2=-1)
+    assert torch.all(diag == 0.0)
+    assert torch.all(dist[..., ~torch.eye(num_points, dtype=torch.bool)] > 0)
+
+
 def test_anisotropic_distance_backward_finite_at_zero_diagonal():
-  """Diagonal distances are zero; backward must not emit inf/nan gradients."""
-  batch_size, num_points = 1, 5
-  points = torch.randn(batch_size, num_points, 2, dtype=torch.float32)
-  params = torch.tensor(
-      [
-          [
-              [1.0, 0.5, 0.0],
-              [1.0, 0.5, 0.0],
-              [1.0, 0.5, 0.0],
-              [1.0, 0.5, 0.0],
-              [1.0, 0.5, 0.0],
-          ]
-      ],
-      dtype=torch.float32,
-      requires_grad=True,
-  )
-  probs = torch.full((batch_size, num_points), 0.2, dtype=torch.float32)
-
-  dist = compute_anisotropic_distance_matrix(
-      points, params, probs=probs, symmetrize="max"
-  )
-  assert torch.all(dist.diagonal(dim1=-2, dim2=-1) > 0)
-
-  dist.sum().backward()
-  assert params.grad is not None
-  assert torch.isfinite(params.grad).all()
+    """Self-distance diagonal is zero; backward must not emit inf/nan gradients."""
+    batch_size, num_points = 1, 5
+    points = torch.randn(batch_size, num_points, 2, dtype=torch.float32)
+    params = torch.tensor(
+        [
+            [
+                [1.0, 0.5, 0.0],
+                [1.0, 0.5, 0.0],
+                [1.0, 0.5, 0.0],
+                [1.0, 0.5, 0.0],
+                [1.0, 0.5, 0.0],
+            ]
+        ],
+        dtype=torch.float32,
+        requires_grad=True,
+    )
+    probs = torch.full((batch_size, num_points), 0.2, dtype=torch.float32)
+
+    dist = compute_anisotropic_distance_matrix(
+        points, params, probs=probs, symmetrize="max"
+    )
+    assert torch.all(dist.diagonal(dim1=-2, dim2=-1) == 0.0)
+
+    dist.sum().backward()
+    assert params.grad is not None
+    assert torch.isfinite(params.grad).all()
+
+
+def test_anisotropic_distance_backward_no_sqrt_anomaly_on_diagonal():
+    with torch.autograd.detect_anomaly():
+        points = torch.randn(1, 8, 2, dtype=torch.float32, requires_grad=True)
+        params = torch.rand(1, 8, 3, dtype=torch.float32, requires_grad=True)
+        probs = torch.full((1, 8), 0.3, dtype=torch.float32)
+        dist = compute_anisotropic_distance_matrix(
+            points, params, probs=probs, symmetrize="max"
+        )
+        dist.sum().backward()
+    assert torch.isfinite(points.grad).all()
+    assert torch.isfinite(params.grad).all()
+
+
+def test_anisotropic_distance_backward_finite_with_coincident_points():
+    """Off-diagonal dist_sq is zero when two centers coincide; backward must stay finite."""
+    points = torch.tensor(
+        [[[0.0, 0.0], [1.0, 0.0], [0.0, 0.0]]],
+        dtype=torch.float32,
+        requires_grad=True,
+    )
+    params = torch.tensor(
+        [[[1.0, 0.5, 0.0], [1.0, 0.5, 0.0], [1.0, 0.5, 0.0]]],
+        dtype=torch.float32,
+        requires_grad=True,
+    )
+
+    dist = compute_anisotropic_distance_matrix(points, params, symmetrize="max")
+    assert dist[0, 0, 2].item() == 0.0
+    assert dist[0, 2, 0].item() == 0.0
+
+    dist.sum().backward()
+    assert torch.isfinite(points.grad).all()
+    assert torch.isfinite(params.grad).all()
+
+
+def test_anisotropic_distance_backward_finite_with_close_points():
+    """Near-coincident pairs have small but positive off-diagonal distances."""
+    base = torch.randn(1, 6, 2, dtype=torch.float32)
+    points = base.clone().requires_grad_(True)
+    with torch.no_grad():
+        points[0, 1] = points[0, 0] + 1e-4
+
+    params = torch.rand(1, 6, 3, dtype=torch.float32, requires_grad=True)
+    probs = torch.full((1, 6), 0.25, dtype=torch.float32)
+
+    dist = compute_anisotropic_distance_matrix(
+        points, params, probs=probs, symmetrize="max"
+    )
+    offdiag = ~torch.eye(6, dtype=torch.bool)
+    assert torch.all(dist[0, offdiag] > 0)
+
+    dist.sum().backward()
+    assert torch.isfinite(points.grad).all()
+    assert torch.isfinite(params.grad).all()