fix: maintain valid state through the algorithm

news/logic-fix.rst

@@ -0,0 +1,23 @@
+**Added:**
+
+* <news item>
+
+**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):