fix: maintain valid state through the algorithm

news/logic-fix.rst

@@ -0,0 +1,24 @@
+**Added:**
+
+* Add test for recovery from failed component update
+* Add a regression test for the cubic root solver
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* Maintain valid state through the algorithm
+
+**Security:**
+
+* <news item>

src/diffpy/stretched_nmf/snmf_class.py

@@ -467,7 +467,6 @@ def _normalize_results(self):
         self._prev_grad_components = np.zeros_like(
             self.components_
         )  # Previous gradient of X (zeros for now)
-
         self._fill_tail_zero = True
         try:
             self.residuals_ = self._get_residual_matrix()
@@ -476,7 +475,6 @@ def _normalize_results(self):
             self._objective_history = [self.objective_function_]
             self._outer_iter = 0
             self._inner_iter = 0
-
             normalization_max_iter = max(self.max_iter, 100)
             for outiter in range(normalization_max_iter):
                 self._outer_iter = outiter
@@ -686,7 +684,9 @@ def _reconstruct_from_stretched_components(
             order="F",
         )
 
-    def _get_objective_function(self, residuals=None, stretch=None):
+    def _get_objective_function(
+        self, residuals=None, stretch=None, components=None
+    ):
         """Return the objective value, passing stored attributes or
         overrides to _compute_objective_function().
 
@@ -696,14 +696,16 @@ def _get_objective_function(self, residuals=None, stretch=None):
             Residual matrix to use instead of self.residuals_.
         stretch : ndarray, optional
             Stretch matrix to use instead of self.stretch_.
+        components : ndarray, optional
+            Component matrix to use instead of self.components_.
 
         Returns
         -------
         float
             Current objective function value.
         """
         return SNMFOptimizer._compute_objective_function(
-            components=self.components_,
+            components=self.components_ if components is None else components,
             residuals=self.residuals_ if residuals is None else residuals,
             stretch=self.stretch_ if stretch is None else stretch,
             rho=self.rho,
@@ -1032,33 +1034,40 @@ def _update_components(self):
                 self._prev_components - self._grad_components / step_size
             )
             # Solve x^3 + p*x + q = 0 for the largest real root
-            self.components_ = np.square(
+            candidate_components = np.square(
                 _cubic_largest_real_root(
                     -components_step, self.eta / (2 * step_size)
                 )
             )
             # Mask values that should be set to zero
             mask = (
-                self.components_**2 * step_size / 2
-                - step_size * self.components_ * components_step
-                + self.eta * np.sqrt(self.components_)
+                candidate_components**2 * step_size / 2
+                - step_size * candidate_components * components_step
+                + self.eta * np.sqrt(candidate_components)
                 < 0
             )
-            self.components_ = mask * self.components_
+            candidate_components = mask * candidate_components
 
             objective_improvement = (
                 self.objective_function_
                 - self._get_objective_function(
-                    residuals=self._get_residual_matrix()
+                    components=candidate_components,
+                    residuals=self._get_residual_matrix(
+                        components=candidate_components
+                    ),
                 )
             )
 
-            # Check if objective function improves
-            if objective_improvement > 0:
+            if (
+                np.isfinite(objective_improvement)
+                and objective_improvement > 0
+            ):
+                self.components_ = candidate_components
                 break
             # If not, increase step_size (step size)
             step_size *= 2
             if np.isinf(step_size):
+                self.components_ = self._prev_components
                 break
 
     def _update_weights(self):
@@ -1383,10 +1392,12 @@ def _compute_objective_function(
 def _cubic_largest_real_root(p, q):
     """Solves x^3 + p*x + q = 0 element-wise for matrices, returning the
     largest real root."""
-    # Handle special case where q == 0
-    y = np.where(
-        q == 0, np.maximum(0, -p) ** 0.5, np.zeros_like(p)
-    )  # q=0 case
+    # For q == 0, the non-negative solution is available directly.  Keep
+    # this branch separate: the general complex-root calculation below is
+    # numerically unstable at this degenerate cubic and previously overwrote
+    # the exact result.
+    q_is_zero = q == 0
+    zero_q_root = np.maximum(0, -p) ** 0.5
 
     # Compute discriminant
     delta = (q / 2) ** 2 + (p / 3) ** 3
@@ -1409,9 +1420,9 @@ def _cubic_largest_real_root(p, q):
 
     # Take the largest real root element-wise when delta < 0
     r_roots = np.stack([np.real(y1), np.real(y2), np.real(y3)], axis=0)
-    y = np.where(delta < 0, np.max(r_roots, axis=0), 0.0)
+    general_root = np.where(delta < 0, np.max(r_roots, axis=0), 0.0)
 
-    return y
+    return np.where(q_is_zero, zero_q_root, general_root)
 
 
 def _reconstruct_matrix(components, weights, stretch):

tests/test_snmf_optimizer.py

@@ -2,7 +2,10 @@
 import pytest
 from scipy.sparse import csr_matrix
 
-from diffpy.stretched_nmf.snmf_class import SNMFOptimizer
+from diffpy.stretched_nmf.snmf_class import (
+    SNMFOptimizer,
+    _cubic_largest_real_root,
+)
 
 
 def test_fit_recovers_rank_one_factors():
@@ -40,6 +43,40 @@ def test_fit_recovers_rank_one_factors():
     assert np.allclose(model.weights_, expected_weights, rtol=0.2, atol=0.1)
 
 
+def test_cubic_largest_real_root_preserves_tiny_zero_q_root():
+    root = _cubic_largest_real_root(np.array([[-1e-300]]), np.zeros((1, 1)))
+
+    np.testing.assert_allclose(root, [[1e-150]], rtol=1e-12, atol=0)
+
+
+def test_failed_component_update_restores_previous_components():
+    model = SNMFOptimizer(n_components=1, eta=0.0)
+    model.signal_length_ = model.n_signals_ = model.n_components_ = 1
+    model.components_ = np.array([[1.0]])
+    max_float = np.finfo(float).max
+    model.weights_ = np.array([[np.sqrt(max_float)]])
+    model.stretch_ = np.ones((1, 1))
+    model._source_matrix = np.zeros((1, 1))
+    model._fill_tail_zero = True
+    model._outer_iter = model._inner_iter = 0
+    model.objective_function_ = 0.0
+    model._compute_stretched_components = lambda: (
+        np.zeros((1, 1)),
+        None,
+        None,
+    )
+    model._compute_component_gradient_zero_tail = lambda residuals: np.array(
+        [[np.finfo(float).max]]
+    )
+    model._get_residual_matrix = lambda **kwargs: np.zeros((1, 1))
+    model._get_objective_function = lambda **kwargs: 1.0
+
+    with np.errstate(over="ignore"):
+        model._update_components()
+
+    np.testing.assert_array_equal(model.components_, [[1.0]])
+
+
 @pytest.mark.parametrize(
     "inputs, expected",
     # inputs tuple: