Skip to content

Commit 3b05f43

Browse files
eurunuelaclaude
andcommitted
Make _has_converged numerically safe for sparse coefficients
Dividing the iterate change by |s_old| produced NaN/inf wherever a coefficient was zero (the common case under sparsity), so jnp.all(... <= tol) was never true and FISTA always ran to max_iter. Guard the denominator with maximum(|s_old|, 1e-10) so a 0 -> 0 entry reads as converged and 0 -> nonzero as a real change, letting early stopping actually fire. Outputs stay within atol=1e-6 of the references at the pinned jax (full suite passes), so no reference regeneration is needed. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
1 parent a33d8ce commit 3b05f43

1 file changed

Lines changed: 5 additions & 3 deletions

File tree

pySPFM/_solvers/fista.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -148,11 +148,13 @@ def _has_converged(s, s_old, tol=1e-6):
148148
bool
149149
True if FISTA has converged, False otherwise
150150
"""
151-
# Calculate normalized error between current and previous estimate
152-
estimate_error = jnp.abs(s - s_old) / jnp.abs(s_old)
151+
# Normalized change between iterates. Guard the denominator against the
152+
# zeros that sparsity produces: 0 -> 0 reads as converged, 0 -> nonzero as a
153+
# real change (a bare |s_old| denominator yields NaN/inf and blocks stopping).
154+
estimate_error = jnp.abs(s - s_old) / jnp.maximum(jnp.abs(s_old), 1e-10)
153155

154156
# Check if the error is smaller than the tolerance for all voxels
155-
return jnp.all(jnp.abs(estimate_error) <= tol).astype(jnp.bool_)
157+
return jnp.all(estimate_error <= tol).astype(jnp.bool_)
156158

157159

158160
def fista(

0 commit comments

Comments
 (0)