diff --git a/src/pyrecest/distributions/hypertorus/_tensor_train.py b/src/pyrecest/distributions/hypertorus/_tensor_train.py index ca71eb5ccd..59afff0ea1 100644 --- a/src/pyrecest/distributions/hypertorus/_tensor_train.py +++ b/src/pyrecest/distributions/hypertorus/_tensor_train.py @@ -65,21 +65,40 @@ def _check_dense_validation_size(size, max_entries): ) +def _stable_vector_norm(values): + """Return a scale-safe Euclidean norm for finite array values.""" + magnitudes = np.abs(np.asarray(values)) + if magnitudes.size == 0: + return 0.0 + scale = float(np.max(magnitudes)) + if scale == 0.0 or not np.isfinite(scale): + return scale + return scale * sqrt(float(np.sum((magnitudes / scale) ** 2))) + + def _choose_rank(singular_values, max_rank, local_tolerance): max_rank = _normalize_max_rank(max_rank) full_rank = singular_values.size if local_tolerance <= 0: rank = full_rank else: - squared_tail = np.cumsum(singular_values[::-1] ** 2)[::-1] - rank = full_rank - for candidate in range(1, full_rank + 1): - tail = ( - 0.0 if candidate == full_rank else sqrt(float(squared_tail[candidate])) - ) - if tail <= local_tolerance: - rank = candidate - break + scale = float(np.max(np.abs(singular_values))) + if scale == 0.0: + rank = 1 + else: + scaled_values = singular_values / scale + squared_tail = np.cumsum(scaled_values[::-1] ** 2)[::-1] + scaled_tolerance = local_tolerance / scale + rank = full_rank + for candidate in range(1, full_rank + 1): + tail = ( + 0.0 + if candidate == full_rank + else sqrt(float(squared_tail[candidate])) + ) + if tail <= scaled_tolerance: + rank = candidate + break if max_rank is not None: rank = min(rank, max_rank) return max(1, rank) @@ -153,7 +172,7 @@ def from_dense(cls, tensor, *, max_rank=None, rtol=0.0, atol=0.0): if array.ndim == 1: return cls((array.reshape(1, array.shape[0], 1),)) - norm = float(np.linalg.norm(array.ravel())) + norm = _stable_vector_norm(array.ravel()) global_tolerance = max(atol, rtol * norm) local_tolerance = ( global_tolerance / sqrt(array.ndim - 1) if global_tolerance > 0 else 0.0 diff --git a/tests/distributions/test_hypertoroidal_tensor_train.py b/tests/distributions/test_hypertoroidal_tensor_train.py index 0a82bbe6e0..0631df4ce4 100644 --- a/tests/distributions/test_hypertoroidal_tensor_train.py +++ b/tests/distributions/test_hypertoroidal_tensor_train.py @@ -13,6 +13,14 @@ def test_dense_roundtrip_and_entry_access(self): self.assertEqual(tt.shape, (3, 3, 3)) npt.assert_allclose(tt.entry((1, 2, 0)), tensor[1, 2, 0], atol=1e-12) + def test_relative_truncation_handles_large_finite_values(self): + tensor = np.diag([1e200, 1e200]) + with np.errstate(over="raise", invalid="raise"): + tt = TensorTrain.from_dense(tensor, rtol=1e-12) + + self.assertEqual(tt.ranks, (1, 2, 1)) + npt.assert_allclose(tt.to_dense() / 1e200, np.eye(2), atol=1e-12) + def test_frobenius_norm(self): tensor = np.arange(9, dtype=float).reshape(3, 3) - 2.0 tt = TensorTrain.from_dense(tensor)