Skip to content

Fix coef_from_cov_grad: validate inputs and handle singular covariances #102

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
t-uda merged 6 commits into from
Mar 23, 2026
Merged

Fix coef_from_cov_grad: validate inputs and handle singular covariances #102

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
ffdd80e
Fix coef_from_cov_grad: validate inputs and handle singular covariances
t-uda Mar 23, 2026
679d154
Fix black formatting in grad.py
t-uda Mar 23, 2026
38f914b
Add tests for coef_from_cov_grad input validation and singular fallback
t-uda Mar 23, 2026
f567009
Move pytest import to top of file (flake8 E402)
t-uda Mar 23, 2026
ff2616d
Fix black formatting in test_coef_from_cov_grad.py
t-uda Mar 23, 2026
78d2300
Return NaN gradients (not zeros) for singular covariances
t-uda Mar 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
28 changes: 26 additions & 2 deletions src/ellphi/grad.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -183,11 +183,35 @@ def coef_from_cov_grad(
if cov_arr.ndim == 2:
cov_arr = cov_arr[np.newaxis, :, :]

if centers.shape[0] != cov_arr.shape[0]:
raise ValueError("Mismatch between number of centres and covariance matrices")
if cov_arr.shape[-1] != cov_arr.shape[-2]:
Comment on lines +186 to +188

@chatgpt-codex-connector chatgpt-codex-connector Bot Mar 23, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Validate input rank before indexing shape tuple

The new validation path indexes centers.shape[0] and cov_arr.shape[-2] without first checking rank, so malformed inputs like scalar X or 1-D cov now raise IndexError instead of the intended ValueError. This makes error handling inconsistent for callers that rely on catching ValueError for bad shapes and contradicts the commit’s goal of robust input validation.

Useful? React with 👍 / 👎.

@t-uda t-uda Mar 23, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is already handled. The rank promotion at lines 181–184 ensures centers is at least 2-D and cov_arr is at least 3-D before the validation runs:

if centers.ndim == 1:
    centers = centers[np.newaxis, :]
if cov_arr.ndim == 2:
    cov_arr = cov_arr[np.newaxis, :, :]

Inputs with ranks outside the documented contract ((n, d) / (d,) for X, (n, d, d) / (d, d) for cov) are not supported — this matches coef_from_cov() in geometry.py:318-324 which uses identical promotion logic and has no additional rank checks.

raise ValueError("Covariance matrices must be square")

n, d = centers.shape

if cov_arr.shape[-1] != d:
raise ValueError("Centre dimensionality and covariance size must agree")
s2 = scale**2

# Forward pass
inv_cov = np.linalg.inv(cov_arr) # (n, d, d)
# Forward pass – handle singular covariances the same way coef_from_cov
# does: return NaN coefficients so callers can detect degeneracy.
try:
inv_cov = np.linalg.inv(cov_arr) # (n, d, d)
except np.linalg.LinAlgError:
m = (d + 1) * (d + 2) // 2
coefs = np.full((n, m), np.nan, dtype=float)

def vjp_nan(
grad_coefs: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
return (
np.full_like(centers, np.nan),
np.full_like(cov_arr, np.nan),
)

return coefs, vjp_nan

A = inv_cov / s2 # (n, d, d)
b = -np.einsum("nij,nj->ni", A, centers) # (n, d)
c = np.einsum("ni,nij,nj->n", centers, A, centers) # (n,)
Expand Down
38 changes: 38 additions & 0 deletions tests/test_coef_from_cov_grad.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from __future__ import annotations

import numpy as np
import pytest

from ellphi import coef_from_cov_grad
from ellphi.geometry import coef_from_cov
Expand Down Expand Up @@ -185,3 +186,40 @@ def test_grad_shapes(rng):
grad_X, grad_cov = vjp(g)
assert grad_X.shape == (n, d)
assert grad_cov.shape == (n, d, d)


# ---------------------------------------------------------------------------
# Input validation
# ---------------------------------------------------------------------------


def test_batch_size_mismatch():
"""Mismatched centre/covariance batch sizes raise ValueError."""
with pytest.raises(ValueError, match="Mismatch"):
coef_from_cov_grad(np.zeros((3, 2)), np.eye(2)[np.newaxis])


def test_non_square_cov():
"""Non-square covariance raises ValueError."""
with pytest.raises(ValueError, match="square"):
coef_from_cov_grad(np.zeros((2, 2)), np.zeros((2, 3, 2)))


def test_dimension_mismatch():
"""Centre/covariance dimension mismatch raises ValueError."""
with pytest.raises(ValueError, match="dimensionality"):
coef_from_cov_grad(np.zeros((2, 2)), np.eye(3)[np.newaxis].repeat(2, axis=0))


def test_singular_cov_returns_nan():
"""Singular covariance returns NaN coefficients, matching coef_from_cov."""
centers = np.zeros((1, 2))
cov_sing = np.zeros((1, 2, 2))

coefs, vjp = coef_from_cov_grad(centers, cov_sing)
ref = coef_from_cov(centers, cov_sing)
np.testing.assert_array_equal(np.isnan(coefs), np.isnan(ref))

grad_X, grad_cov = vjp(np.ones_like(coefs))
assert np.all(np.isnan(grad_X))
assert np.all(np.isnan(grad_cov))
Loading