diff --git a/tda_ml/topology.py b/tda_ml/topology.py index f2f529c..9389a4d 100644 --- a/tda_ml/topology.py +++ b/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 diff --git a/tests/test_topology.py b/tests/test_topology.py index 8031254..bb98b1e 100644 --- a/tests/test_topology.py +++ b/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() diff --git a/tests/test_topology_gradients.py b/tests/test_topology_gradients.py index a8715fa..a06670b 100644 --- a/tests/test_topology_gradients.py +++ b/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()