diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md index 6edfbe1..8e9efb1 100644 --- a/RELEASE_NOTES.md +++ b/RELEASE_NOTES.md @@ -1,3 +1,14 @@ +## Unreleased + +### Changed + +- `ellphi.grad.pdist_tangency_grad(...)` now dispatches the batched + distance/gradient computation to the C++ backend when it is available + (same auto-selection as `pdist_tangency`), giving orders-of-magnitude + speedups for gradient-based workflows. The pure-Python implementation + remains as the fallback and the public API is unchanged. Thanks to + koki3070 and collaborators (TDA-ML) for the prototype and benchmarks. + ## 0.1.2 - 2026-03-25 ### Added diff --git a/src/ellphi/_tangency_cpp.py b/src/ellphi/_tangency_cpp.py index 9049573..f315c95 100644 --- a/src/ellphi/_tangency_cpp.py +++ b/src/ellphi/_tangency_cpp.py @@ -152,6 +152,10 @@ def _raise_backend_error(message: str) -> None: raise ValueError(message) if "Degenerate conic" in message: raise ZeroDivisionError(message) + if "Derivative with respect to mu is numerically zero" in message: + raise ZeroDivisionError(message) + if "Tangency distance is zero" in message: + raise ZeroDivisionError(message) raise RuntimeError(message or "Unknown C++ backend error") @@ -301,3 +305,88 @@ def pdist_tangency(coef: numpy.ndarray) -> numpy.ndarray: _raise_backend_error(message) return output + + +def has_pdist_tangency_grad() -> bool: + """Checks if the C++ backend exports the batched gradient kernel. + + Returns: + `True` if the C++ backend is loaded and exports the + ``pdist_tangency_grad`` symbol, `False` otherwise (e.g. for a stale + library built from an older source tree). + """ + return _LIB is not None and hasattr(_LIB, "pdist_tangency_grad") + + +def pdist_tangency_grad( + coef: numpy.ndarray, +) -> Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]: + """Computes pairwise tangency distances and per-pair gradient blocks. + + Args: + coef: A NumPy array of shape `(m, n)` containing the coefficient + vectors for `m` ellipses. + + Returns: + A tuple `(dists, dt_dp, dt_dq)` where `dists` is the condensed + distance matrix of tangency distances and `dt_dp` / `dt_dq` hold, for + each pair, the gradient of the pair distance with respect to the + first / second ellipse coefficients (shape `(n_pairs, n)`). + + Raises: + RuntimeError: If the C++ backend is unavailable, or if the loaded + library is stale and lacks the ``pdist_tangency_grad`` export. + """ + lib = _ensure_available() + if not has_pdist_tangency_grad(): + raise RuntimeError( + "C++ backend is missing the 'pdist_tangency_grad' export; the " + "loaded library is stale. " + "Rebuild the extension yourself: in a development checkout " + "run 'uv sync --reinstall-package ellphi' or " + "'python scripts/build_tangency_cpp.py'; with pip, run " + "'pip install --force-reinstall .' from the source tree. " + "It is not rebuilt automatically during import." + ) + + func = lib.pdist_tangency_grad + func.restype = ctypes.c_int + func.argtypes = [ + ctypes.POINTER(ctypes.c_double), + ctypes.c_size_t, + ctypes.c_size_t, + ctypes.POINTER(ctypes.c_double), + ctypes.POINTER(ctypes.c_double), + ctypes.POINTER(ctypes.c_double), + ctypes.c_void_p, + ctypes.c_size_t, + ] + + coef_arr = numpy.ascontiguousarray(coef, dtype=float) + if coef_arr.ndim != 2: + raise ValueError("Coefficient array must have shape (m, n)") + coef_length = coef_arr.shape[1] + infer_dim_from_coef_length(coef_length) + m = coef_arr.shape[0] + n = m * (m - 1) // 2 + dists = numpy.empty(n, dtype=float) + dt_dp = numpy.empty((n, coef_length), dtype=float) + dt_dq = numpy.empty((n, coef_length), dtype=float) + error_buffer = ctypes.create_string_buffer(_ERROR_BUFFER) + + status = func( + coef_arr.ctypes.data_as(ctypes.POINTER(ctypes.c_double)), + ctypes.c_size_t(m), + ctypes.c_size_t(coef_length), + dists.ctypes.data_as(ctypes.POINTER(ctypes.c_double)), + dt_dp.ctypes.data_as(ctypes.POINTER(ctypes.c_double)), + dt_dq.ctypes.data_as(ctypes.POINTER(ctypes.c_double)), + ctypes.cast(error_buffer, ctypes.c_void_p), + ctypes.c_size_t(_ERROR_BUFFER), + ) + + if status != 0: + message = error_buffer.value.decode("utf-8", errors="ignore") + _raise_backend_error(message) + + return dists, dt_dp, dt_dq diff --git a/src/ellphi/_tangency_cpp.pyi b/src/ellphi/_tangency_cpp.pyi index aca13b9..8897148 100644 --- a/src/ellphi/_tangency_cpp.pyi +++ b/src/ellphi/_tangency_cpp.pyi @@ -17,3 +17,7 @@ def tangency( failsafe: bool, ) -> TangencyResult: ... def pdist_tangency(coef: FloatArray) -> FloatArray: ... +def has_pdist_tangency_grad() -> bool: ... +def pdist_tangency_grad( + coef: FloatArray, +) -> Tuple[FloatArray, FloatArray, FloatArray]: ... diff --git a/src/ellphi/_tangency_cpp_impl.cpp b/src/ellphi/_tangency_cpp_impl.cpp index afbfc30..3382e37 100644 --- a/src/ellphi/_tangency_cpp_impl.cpp +++ b/src/ellphi/_tangency_cpp_impl.cpp @@ -1100,6 +1100,186 @@ double solve_mu( raise("Unknown method"); } +struct GradWorkspace { + std::vector pencil_coef; + DecodedConic conic; + std::vector chol; + std::vector rhs; + std::vector center; + std::vector diff_coef; + DecodedConic diff_dec; + std::vector residual; + std::vector solved; + std::vector base; + std::vector chain; + std::vector basis_rhs; +}; + +// Tangency distance and its gradients dt/dp, dt/dq for one pair. +// +// Mirrors the Python reference (ellphi.grad.tangency_grad backed by +// ellphi.differentiable_solver.solve_mu_gradients): the envelope theorem +// removes the chain through the tangent point, and d_mu/dp, d_mu/dq come +// from implicit differentiation of the optimality condition F(mu, p, q) = 0. +void tangency_grad_pair( + const std::vector& p, + const std::vector& q, + int dim, + GradWorkspace& ws, + double* out_dist, + double* out_dt_dp, + double* out_dt_dq +) { + auto hybrid_iters = default_hybrid_iterations(dim); + double mu = solve_mu( + p, + q, + "brentq+newton", + {0.0, 1.0}, + false, + 0.0, + hybrid_iters.first, + hybrid_iters.second, + true // failsafe enabled by default for pdist + ); + + pencil_into(p, q, mu, ws.pencil_coef); + decode_conic_into(ws.pencil_coef, ws.conic); + + // The gradient requires the Cholesky factor of the pencil quadratic + // form; failure signals a degenerate configuration and surfaces as + // ZeroDivisionError, matching the Python reference. + cholesky_factor_into(ws.conic.quad, dim, ws.chol); + + if (ws.rhs.size() != static_cast(dim)) { + ws.rhs.resize(dim); + } + for (int i = 0; i < dim; ++i) { + ws.rhs[static_cast(i)] = + -ws.conic.linear[static_cast(i)]; + } + solve_with_cholesky_into(ws.chol, ws.rhs, dim, ws.center); + + double value = quad_eval(ws.conic, ws.center); + if (value < 0.0) { + value = 0.0; + } + double t = std::sqrt(value); + if (t == 0.0) { + raise("Tangency distance is zero (gradient undefined)"); + } + + const std::size_t coef_length = p.size(); + if (ws.diff_coef.size() != coef_length) { + ws.diff_coef.resize(coef_length); + } + for (std::size_t i = 0; i < coef_length; ++i) { + ws.diff_coef[i] = p[i] - q[i]; + } + decode_conic_into(ws.diff_coef, ws.diff_dec); + + // residual = -(diff_quad @ center + diff_linear) + matvec_into(ws.diff_dec.quad, ws.center, dim, ws.residual); + for (int i = 0; i < dim; ++i) { + ws.residual[static_cast(i)] = -( + ws.residual[static_cast(i)] + + ws.diff_dec.linear[static_cast(i)] + ); + } + + // dF/dmu = 2 residual . K^{-1} residual (implicit differentiation) + solve_with_cholesky_into(ws.chol, ws.residual, dim, ws.solved); + double dF_dmu = 0.0; + for (int i = 0; i < dim; ++i) { + dF_dmu += ws.residual[static_cast(i)] * + ws.solved[static_cast(i)]; + } + dF_dmu *= 2.0; + // Same threshold as numpy.isclose(dF_dmu, 0.0) in the Python reference. + if (std::abs(dF_dmu) <= 1e-8) { + raise("Derivative with respect to mu is numerically zero"); + } + + // Monomial basis at the center: [x_i*x_j (i<=j), 2*x_k, 1] + if (ws.base.size() != coef_length) { + ws.base.resize(coef_length); + } + std::size_t idx = 0; + for (int i = 0; i < dim; ++i) { + for (int j = i; j < dim; ++j) { + ws.base[idx++] = (i == j) + ? ws.center[static_cast(i)] * + ws.center[static_cast(i)] + : 2.0 * ws.center[static_cast(i)] * + ws.center[static_cast(j)]; + } + } + for (int i = 0; i < dim; ++i) { + ws.base[idx++] = 2.0 * ws.center[static_cast(i)]; + } + ws.base[idx] = 1.0; + + // chain[k] = (dx*/dr_k) . phi_x with phi_x = -2 residual and Jacobian + // row -K^{-1} rhs_k, hence chain[k] = 2 (K^{-1} rhs_k) . residual. + if (ws.chain.size() != coef_length) { + ws.chain.resize(coef_length); + } + if (ws.basis_rhs.size() != static_cast(dim)) { + ws.basis_rhs.resize(dim); + } + idx = 0; + for (int i = 0; i < dim; ++i) { + for (int j = i; j < dim; ++j) { + std::fill(ws.basis_rhs.begin(), ws.basis_rhs.end(), 0.0); + if (i == j) { + ws.basis_rhs[static_cast(i)] = + ws.center[static_cast(i)]; + } else { + ws.basis_rhs[static_cast(i)] = + ws.center[static_cast(j)]; + ws.basis_rhs[static_cast(j)] = + ws.center[static_cast(i)]; + } + solve_with_cholesky_into(ws.chol, ws.basis_rhs, dim, ws.solved); + double dot = 0.0; + for (int k = 0; k < dim; ++k) { + dot += ws.solved[static_cast(k)] * + ws.residual[static_cast(k)]; + } + ws.chain[idx++] = 2.0 * dot; + } + } + for (int axis = 0; axis < dim; ++axis) { + std::fill(ws.basis_rhs.begin(), ws.basis_rhs.end(), 0.0); + ws.basis_rhs[static_cast(axis)] = 1.0; + solve_with_cholesky_into(ws.chol, ws.basis_rhs, dim, ws.solved); + double dot = 0.0; + for (int k = 0; k < dim; ++k) { + dot += ws.solved[static_cast(k)] * + ws.residual[static_cast(k)]; + } + ws.chain[idx++] = 2.0 * dot; + } + // The constant term does not move the center. + ws.chain[idx] = 0.0; + + double scalar = 0.0; + for (std::size_t k = 0; k < coef_length; ++k) { + scalar += ws.base[k] * (q[k] - p[k]); + } + const double inv2t = 0.5 / t; + const double one_minus_mu = 1.0 - mu; + for (std::size_t k = 0; k < coef_length; ++k) { + double dF_dp = ws.base[k] + one_minus_mu * ws.chain[k]; + double dF_dq = -ws.base[k] + mu * ws.chain[k]; + double d_mu_dp = -dF_dp / dF_dmu; + double d_mu_dq = -dF_dq / dF_dmu; + out_dt_dp[k] = inv2t * (one_minus_mu * ws.base[k] + scalar * d_mu_dp); + out_dt_dq[k] = inv2t * (mu * ws.base[k] + scalar * d_mu_dq); + } + out_dist[0] = t; +} + void copy_error(char* buffer, std::size_t size, const std::string& message) { if (buffer == nullptr || size == 0) { return; @@ -1223,3 +1403,52 @@ ELLPHI_EXPORT extern "C" int pdist_tangency( return 1; } } + +ELLPHI_EXPORT extern "C" int pdist_tangency_grad( + const double* coef, + std::size_t m, + std::size_t coef_length, + double* out_dists, + double* out_dt_dp, + double* out_dt_dq, + char* err_buffer, + std::size_t err_buffer_len +) { + try { + int dim = infer_dim_from_coef_length(coef_length); + std::vector> conics( + m, + std::vector(coef_length, 0.0) + ); + for (std::size_t i = 0; i < m; ++i) { + const double* start = coef + i * coef_length; + std::copy(start, start + coef_length, conics[i].begin()); + } + + GradWorkspace ws{}; + std::size_t idx = 0; + for (std::size_t i = 0; i < m; ++i) { + const std::vector& p = conics[i]; + for (std::size_t j = i + 1; j < m; ++j) { + const std::vector& q = conics[j]; + tangency_grad_pair( + p, + q, + dim, + ws, + out_dists + idx, + out_dt_dp + idx * coef_length, + out_dt_dq + idx * coef_length + ); + ++idx; + } + } + return 0; + } catch (const std::exception& ex) { + copy_error(err_buffer, err_buffer_len, ex.what()); + return 1; + } catch (...) { + copy_error(err_buffer, err_buffer_len, "Unknown error"); + return 1; + } +} diff --git a/src/ellphi/grad.py b/src/ellphi/grad.py index 92cc0e5..1619e8c 100644 --- a/src/ellphi/grad.py +++ b/src/ellphi/grad.py @@ -14,7 +14,9 @@ from typing import Callable import numpy as np +import numpy.typing as npt +from . import _tangency_cpp as _cpp from .differentiable_solver import solve_mu_gradients from .solver import tangency @@ -90,9 +92,32 @@ def tangency_grad(p: np.ndarray, q: np.ndarray, **solver_kwargs) -> TangencyGrad ) +def _pdist_tangency_grad_python( + coefs: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Pure-Python reference: one ``tangency_grad`` call per pair. + + Returns the condensed distances plus the per-pair gradient blocks + ``dt_dp`` / ``dt_dq`` of shape ``(n_pairs, m)``. + """ + N, m = coefs.shape + n_pairs = N * (N - 1) // 2 + dists = np.empty(n_pairs) + dt_dp = np.empty((n_pairs, m)) + dt_dq = np.empty((n_pairs, m)) + + for k, (i, j) in enumerate(combinations(range(N), 2)): + g = tangency_grad(coefs[i], coefs[j]) + dists[k] = g.t + dt_dp[k] = g.dt_dp + dt_dq[k] = g.dt_dq + + return dists, dt_dp, dt_dq + + def pdist_tangency_grad( coefs: np.ndarray, -) -> tuple[np.ndarray, Callable[[np.ndarray], np.ndarray]]: +) -> tuple[np.ndarray, Callable[[npt.ArrayLike], np.ndarray]]: """Pairwise tangency distances and a VJP (pullback) for all pairs. Computes the same condensed distance array as @@ -100,6 +125,11 @@ def pdist_tangency_grad( and additionally returns a VJP function that maps upstream gradient vectors back to per-ellipsoid coefficient gradients. + The batched distance/gradient computation runs on the C++ backend when it + is available (mirroring ``pdist_tangency``) and falls back to the + pure-Python reference implementation otherwise. The VJP accumulation + itself always runs in NumPy. + Args: coefs: Array of shape ``(N, m)`` containing ``N`` ellipsoid coefficient vectors. @@ -111,6 +141,9 @@ def pdist_tangency_grad( tangency distances in the same order as ``scipy.spatial.distance.pdist``. - ``vjp`` is a callable ``(grad_dists,) -> grad_coefs`` that accumulates upstream gradients into an array of shape ``(N, m)``. + ``grad_dists`` must be broadcastable to shape ``(N*(N-1)//2,)`` + (one cotangent per pair; a scalar is broadcast to all pairs), + otherwise a ``ValueError`` is raised. Examples: >>> import numpy as np @@ -130,20 +163,28 @@ def pdist_tangency_grad( if coefs.ndim != 2: raise ValueError("Expected coefficient array with shape (N, m)") N = len(coefs) - pairs = list(combinations(range(N), 2)) - dists = np.empty(len(pairs)) - store: list[tuple[int, int, np.ndarray, np.ndarray]] = [] - - for k, (i, j) in enumerate(pairs): - g = tangency_grad(coefs[i], coefs[j]) - dists[k] = g.t - store.append((i, j, g.dt_dp, g.dt_dq)) - def vjp(grad_dists: np.ndarray) -> np.ndarray: + if _cpp.has_pdist_tangency_grad(): + dists, dt_dp, dt_dq = _cpp.pdist_tangency_grad(coefs) + else: + dists, dt_dp, dt_dq = _pdist_tangency_grad_python(coefs) + + n_pairs = N * (N - 1) // 2 + idx_i, idx_j = np.triu_indices(N, k=1) + + def vjp(grad_dists: npt.ArrayLike) -> np.ndarray: + gd = np.asarray(grad_dists, dtype=float) + try: + gd = np.broadcast_to(gd, (n_pairs,)) + except ValueError: + raise ValueError( + "grad_dists must be broadcastable to shape " + f"({n_pairs},) (one cotangent per ellipsoid pair); " + f"got shape {gd.shape}" + ) from None g_coefs = np.zeros_like(coefs) - for k, (i, j, dt_dp, dt_dq) in enumerate(store): - g_coefs[i] += grad_dists[k] * dt_dp - g_coefs[j] += grad_dists[k] * dt_dq + np.add.at(g_coefs, idx_i, gd[:, None] * dt_dp) + np.add.at(g_coefs, idx_j, gd[:, None] * dt_dq) return g_coefs return dists, vjp diff --git a/tests/test_grad_cpp.py b/tests/test_grad_cpp.py new file mode 100644 index 0000000..c58a9e3 --- /dev/null +++ b/tests/test_grad_cpp.py @@ -0,0 +1,191 @@ +"""Tests for the C++ backend of pdist_tangency_grad.""" + +from __future__ import annotations + +import types + +import numpy as np +import pytest + +import ellphi._tangency_cpp as _cpp +from ellphi import pdist_tangency, pdist_tangency_grad +from ellphi.geometry import coef_from_cov +from ellphi.grad import _pdist_tangency_grad_python + +from .factories import random_covariance, rotation_matrix + + +requires_cpp_grad = pytest.mark.skipif( + not _cpp.has_pdist_tangency_grad(), + reason="C++ pdist_tangency_grad kernel not available", +) + + +def _random_coefs(rng, n, dim=2): + means = rng.uniform(-20.0, 20.0, size=(n, dim)) + covs = np.stack([random_covariance(rng, dim=dim) for _ in range(n)]) + return coef_from_cov(means, covs) + + +def _eccentric_coefs(rng, n): + """Nearly degenerate 2-D ellipses: high aspect ratio, close centres.""" + means = rng.uniform(-1.0, 1.0, size=(n, 2)) + covs = [] + for _ in range(n): + axes = np.array([rng.uniform(1e-3, 1e-2), rng.uniform(1.0, 2.0)]) + rot = rotation_matrix(rng.uniform(0.0, np.pi)) + covs.append(rot @ np.diag(axes) @ rot.T) + return coef_from_cov(means, np.stack(covs)) + + +@requires_cpp_grad +@pytest.mark.parametrize("n,dim", [(4, 2), (10, 2), (25, 2), (4, 3), (10, 3)]) +def test_cpp_grad_matches_python_reference(rng, n, dim): + """C++ kernel agrees with the pure-Python reference to <= 1e-10.""" + coefs = _random_coefs(rng, n, dim=dim) + + dists_cpp, dt_dp_cpp, dt_dq_cpp = _cpp.pdist_tangency_grad(coefs) + dists_py, dt_dp_py, dt_dq_py = _pdist_tangency_grad_python(coefs) + + np.testing.assert_allclose(dists_cpp, dists_py, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dt_dp_cpp, dt_dp_py, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dt_dq_cpp, dt_dq_py, rtol=1e-10, atol=1e-10) + + +@requires_cpp_grad +def test_cpp_grad_matches_python_reference_eccentric(rng): + """Agreement holds for nearly degenerate (highly eccentric) ellipses.""" + coefs = _eccentric_coefs(rng, 6) + + dists_cpp, dt_dp_cpp, dt_dq_cpp = _cpp.pdist_tangency_grad(coefs) + dists_py, dt_dp_py, dt_dq_py = _pdist_tangency_grad_python(coefs) + + np.testing.assert_allclose(dists_cpp, dists_py, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dt_dp_cpp, dt_dp_py, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(dt_dq_cpp, dt_dq_py, rtol=1e-10, atol=1e-10) + + +@requires_cpp_grad +def test_cpp_grad_dists_match_pdist_tangency(rng): + """Distances from the grad kernel equal the forward pdist_tangency.""" + coefs = _random_coefs(rng, 8) + dists, _ = pdist_tangency_grad(coefs) + ref = pdist_tangency(coefs) + np.testing.assert_allclose(dists, ref, rtol=0, atol=1e-12) + + +@requires_cpp_grad +def test_vjp_matches_python_path(rng, monkeypatch): + """The public VJP gives the same pullback on both backends.""" + coefs = _random_coefs(rng, 7) + cotangent = rng.standard_normal(7 * 6 // 2) + + dists_cpp, vjp_cpp = pdist_tangency_grad(coefs) + grad_cpp = vjp_cpp(cotangent) + + monkeypatch.setattr(_cpp, "has_pdist_tangency_grad", lambda: False) + dists_py, vjp_py = pdist_tangency_grad(coefs) + grad_py = vjp_py(cotangent) + + np.testing.assert_allclose(dists_cpp, dists_py, rtol=1e-10, atol=1e-10) + np.testing.assert_allclose(grad_cpp, grad_py, rtol=1e-10, atol=1e-10) + + +def test_python_fallback_when_backend_missing(rng, monkeypatch): + """Without the C++ library the Python path is selected and works.""" + coefs = _random_coefs(rng, 5) + expected_dists, expected_dt_dp, expected_dt_dq = _pdist_tangency_grad_python(coefs) + + monkeypatch.setattr(_cpp, "_LIB", None) + assert _cpp.is_available() is False + assert _cpp.has_pdist_tangency_grad() is False + + dists, vjp = pdist_tangency_grad(coefs) + np.testing.assert_allclose(dists, expected_dists, rtol=1e-10, atol=1e-10) + + cotangent = rng.standard_normal(len(dists)) + g_coefs = vjp(cotangent) + expected = np.zeros_like(coefs) + for k, (i, j) in enumerate((i, j) for i in range(5) for j in range(i + 1, 5)): + expected[i] += cotangent[k] * expected_dt_dp[k] + expected[j] += cotangent[k] * expected_dt_dq[k] + np.testing.assert_allclose(g_coefs, expected, rtol=1e-10, atol=1e-12) + + +def test_stale_library_without_grad_symbol(monkeypatch): + """A library lacking the new export reports the grad kernel as missing.""" + monkeypatch.setattr(_cpp, "_LIB", types.SimpleNamespace()) + assert _cpp.has_pdist_tangency_grad() is False + + +def test_stale_library_direct_call_raises(monkeypatch): + """Calling the internal binding on a stale library raises RuntimeError.""" + monkeypatch.setattr(_cpp, "_LIB", types.SimpleNamespace()) + with pytest.raises(RuntimeError, match="pdist_tangency_grad' export"): + _cpp.pdist_tangency_grad(np.zeros((2, 6))) + + +def test_vjp_accepts_valid_1d_cotangent(rng): + """A 1-D cotangent of length n_pairs is accepted.""" + coefs = _random_coefs(rng, 4) + dists, vjp = pdist_tangency_grad(coefs) + g_coefs = vjp(np.ones(len(dists))) + assert g_coefs.shape == coefs.shape + + +def test_vjp_accepts_scalar_cotangent(rng): + """A scalar cotangent broadcasts to all pairs.""" + coefs = _random_coefs(rng, 4) + dists, vjp = pdist_tangency_grad(coefs) + expected = vjp(2.0 * np.ones(len(dists))) + np.testing.assert_allclose(vjp(2.0), expected) + np.testing.assert_allclose(vjp(np.asarray(2.0)), expected) + + +def test_vjp_rejects_wrong_length_cotangent(rng): + """Cotangents of the wrong length raise a clear ValueError.""" + coefs = _random_coefs(rng, 4) + dists, vjp = pdist_tangency_grad(coefs) + with pytest.raises(ValueError, match="broadcastable to shape"): + vjp(np.ones(len(dists) - 1)) + with pytest.raises(ValueError, match="broadcastable to shape"): + vjp(np.ones(len(dists) + 1)) + + +def test_vjp_rejects_2d_cotangent(rng): + """2-D cotangents raise a clear ValueError.""" + coefs = _random_coefs(rng, 4) + dists, vjp = pdist_tangency_grad(coefs) + with pytest.raises(ValueError, match="broadcastable to shape"): + vjp(np.ones((len(dists), coefs.shape[1]))) + with pytest.raises(ValueError, match="broadcastable to shape"): + vjp(np.ones((len(dists), 1))) + + +@requires_cpp_grad +def test_cpp_grad_identical_ellipsoids_raise(): + """Identical ellipsoids are degenerate and raise ZeroDivisionError.""" + means = np.array([[0.3, -0.2], [0.3, -0.2]]) + covs = np.stack([np.eye(2), np.eye(2)]) + coefs = coef_from_cov(means, covs) + with pytest.raises(ZeroDivisionError): + _cpp.pdist_tangency_grad(coefs) + + +@requires_cpp_grad +def test_cpp_grad_concentric_ellipsoids_raise(): + """Concentric nested ellipsoids raise ZeroDivisionError as in Python.""" + means = np.array([[0.5, 0.5], [0.5, 0.5]]) + covs = np.stack([np.eye(2), 4.0 * np.eye(2)]) + coefs = coef_from_cov(means, covs) + with pytest.raises(ZeroDivisionError): + _cpp.pdist_tangency_grad(coefs) + + +def test_single_ellipsoid_empty_result(rng): + """N=1 produces empty distances and a zero VJP on any backend.""" + coefs = _random_coefs(rng, 1) + dists, vjp = pdist_tangency_grad(coefs) + assert dists.shape == (0,) + g_coefs = vjp(np.zeros(0)) + np.testing.assert_array_equal(g_coefs, np.zeros_like(coefs))