diff --git a/src/pyrecest/smoothers/so3_chordal_mean_smoother.py b/src/pyrecest/smoothers/so3_chordal_mean_smoother.py index 25f9d6d09a..f72c4bf685 100644 --- a/src/pyrecest/smoothers/so3_chordal_mean_smoother.py +++ b/src/pyrecest/smoothers/so3_chordal_mean_smoother.py @@ -5,9 +5,10 @@ from operator import index as operator_index from typing import Sequence +import numpy as np + # pylint: disable=no-member -from pyrecest import backend -from pyrecest.backend import asarray, diag, isfinite, linalg, ndim, zeros +from pyrecest.backend import asarray, diag, linalg, ndim, to_numpy, zeros from .abstract_smoother import AbstractSmoother @@ -101,20 +102,32 @@ def _normalize_weight_vector( raise ValueError(f"{name} must be one-dimensional.") if weights_array.shape[0] != length: raise ValueError(f"{name} must have length {length}.") - if not bool(backend.all(isfinite(weights_array))): - raise ValueError(f"{name} must contain only finite values.") - for idx in range(length): - if weights_array[idx] < 0.0: - raise ValueError(f"{name} must be nonnegative.") + # Validate and scale using the stored host values. Some XLA operations + # flush subnormal operands, and division by a near-maximum finite value + # can be lowered through a reciprocal that underflows to zero. Moving + # only the resulting O(1) ratios back to the active backend avoids both + # failure modes while preserving the intended relative weights. + try: + host_weights = np.asarray(to_numpy(weights_array)) + finite_weights = np.isfinite(host_weights) + except (TypeError, ValueError, OverflowError, RuntimeError) as exc: + raise ValueError(f"{name} must contain real numeric values.") from exc + if np.iscomplexobj(host_weights): + raise ValueError(f"{name} must contain real numeric values.") + if not np.all(finite_weights): + raise ValueError(f"{name} must contain only finite values.") + if np.any(host_weights < 0.0): + raise ValueError(f"{name} must be nonnegative.") - weight_scale = backend.max(weights_array) + weight_scale = float(np.max(host_weights)) if weight_scale <= 0.0: raise ValueError(f"{name} must contain at least one positive entry.") - scaled_weights = weights_array / weight_scale + scaled_host_weights = host_weights / weight_scale + scaled_weights = asarray(scaled_host_weights) if normalize: - return scaled_weights / backend.sum(scaled_weights) + return scaled_weights / float(np.sum(scaled_host_weights)) return scaled_weights @staticmethod diff --git a/tests/smoothers/test_so3_chordal_mean_backend_weight_stability.py b/tests/smoothers/test_so3_chordal_mean_backend_weight_stability.py new file mode 100644 index 0000000000..03d00e7709 --- /dev/null +++ b/tests/smoothers/test_so3_chordal_mean_backend_weight_stability.py @@ -0,0 +1,81 @@ +"""Cross-backend regressions for extreme SO(3) chordal smoother weights.""" + +import numpy as np +import numpy.testing as npt +import pytest +from pyrecest.backend import array, cos, eye, sin, to_numpy +from pyrecest.smoothers import SO3ChordalMeanSmoother + +# pylint: disable=protected-access + + +def _active_dtype(): + return to_numpy(array([0.0], dtype=float)).dtype + + +def _z_rotation(angle): + return array( + [ + [cos(angle), -sin(angle), 0.0], + [sin(angle), cos(angle), 0.0], + [0.0, 0.0, 1.0], + ] + ) + + +def test_chordal_mean_preserves_largest_finite_weight_ratio_across_backends(): + dtype = _active_dtype() + largest = np.finfo(dtype).max + weights = array(np.asarray([largest, largest / 2.0], dtype=dtype), dtype=float) + + mean_rotation = SO3ChordalMeanSmoother.chordal_mean( + [eye(3), _z_rotation(0.5 * np.pi)], + weights=weights, + ) + + npt.assert_allclose( + to_numpy(mean_rotation), + to_numpy(_z_rotation(np.arctan(0.5))), + atol=1.0e-6, + ) + + +def test_weight_scaling_preserves_positive_subnormal_ratio_across_backends(): + dtype = _active_dtype() + smallest = np.finfo(dtype).smallest_subnormal + weights = array( + np.asarray([2.0 * smallest, smallest], dtype=dtype), + dtype=float, + ) + + scaled_weights = SO3ChordalMeanSmoother._normalize_weight_vector( + weights, + 2, + "weights", + normalize=False, + ) + normalized_weights = SO3ChordalMeanSmoother._normalize_weight_vector( + weights, + 2, + "weights", + normalize=True, + ) + + npt.assert_allclose(to_numpy(scaled_weights), [1.0, 0.5]) + npt.assert_allclose(to_numpy(normalized_weights), [2.0 / 3.0, 1.0 / 3.0]) + + +def test_rejects_negative_subnormal_weight_across_backends(): + dtype = _active_dtype() + smallest = np.finfo(dtype).smallest_subnormal + weights = array( + np.asarray([1.0, -smallest], dtype=dtype), + dtype=float, + ) + + with pytest.raises(ValueError, match="nonnegative"): + SO3ChordalMeanSmoother._normalize_weight_vector( + weights, + 2, + "weights", + )