feat: final parity changes with MATLAB

news/warnings-fix.rst

@@ -0,0 +1,24 @@
+**Added:**
+
+* Use faster iterations when possible
+* Allow storing of objective log
+
+**Changed:**
+
+* <news item>
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* Remove warning by explicitly interpreting cubic root as complex
+
+**Security:**
+
+* <news item>

src/diffpy/stretched_nmf/snmf_class.py

@@ -1,3 +1,4 @@
+import csv
 import time
 
 import cvxpy as cp
@@ -77,14 +78,16 @@ class SNMFOptimizer:
     def __init__(
         self,
         n_components=None,
-        max_iter=500,
+        max_iter=200,
         min_iter=20,
         tol=5e-7,
         rho=0,
         eta=0,
         random_state=None,
         show_plots=False,
         verbose=False,
+        stretch_max_iter=8,
+        stretch_slow_iter=200,
     ):
         """Initialize an instance of sNMF with estimator
         hyperparameters.
@@ -116,9 +119,22 @@ def __init__(
             created by the decomposition. Optional.
         show_plots : bool
             Enables plotting at each step of the decomposition. Optional.
+        stretch_max_iter : int
+            Maximum number of projected-gradient stretch steps per outer-loop
+            iteration. Optional.
+        stretch_slow_iter : int
+            Number of initial outer-loop stretch updates to solve with the
+            slower constrained optimizer before switching to projected-gradient
+            stretch updates. Optional.
         """
         if n_components is not None and n_components < 1:
             raise ValueError("n_components must be a positive integer.")
+        if not isinstance(stretch_max_iter, int) or stretch_max_iter < 1:
+            raise ValueError("stretch_max_iter must be a positive integer.")
+        if not isinstance(stretch_slow_iter, int) or stretch_slow_iter < 0:
+            raise ValueError(
+                "stretch_slow_iter must be a non-negative integer."
+            )
 
         self.n_components = n_components
         self.max_iter = max_iter
@@ -129,10 +145,13 @@ def __init__(
         self.random_state = random_state
         self.show_plots = show_plots
         self.verbose = verbose
+        self.stretch_max_iter = stretch_max_iter
+        self.stretch_slow_iter = stretch_slow_iter
 
         self._rng = np.random.default_rng(self.random_state)
         self._plotter = SNMFPlotter() if self.show_plots else None
         self._fill_tail_zero = False
+        self._stretch_step_size = None
 
     def _initialize_factors(
         self,
@@ -213,6 +232,7 @@ def _initialize_factors(
         self._init_weights = self.weights_.copy()
         self._init_stretch = self.stretch_.copy()
         self._fill_tail_zero = False
+        self._stretch_step_size = None
 
         # Second-order spline: Tridiagonal (-2 on diags, 1 on sub/superdiags)
         self._spline_smooth_operator = 0.25 * diags(
@@ -389,6 +409,43 @@ def fit(
 
         return self
 
+    def save_objective_log(self, filename):
+        """Save the objective log to a tab-delimited text file.
+
+        Parameters
+        ----------
+        filename : str or path-like
+            Output file path.
+        """
+        if not hasattr(self, "objective_log"):
+            raise ValueError("Cannot save objective_log before calling fit.")
+
+        fieldnames = ("step", "iteration", "objective", "dt_ms")
+        with open(filename, "w", newline="") as fp:
+            writer = csv.DictWriter(
+                fp,
+                fieldnames=fieldnames,
+                delimiter="\t",
+            )
+            writer.writeheader()
+            previous_timestamp = None
+            for row in self.objective_log:
+                timestamp = row["timestamp"]
+                dt_ms = (
+                    0.0
+                    if previous_timestamp is None
+                    else 1000 * (timestamp - previous_timestamp)
+                )
+                writer.writerow(
+                    {
+                        "step": row["step"],
+                        "iteration": row["iteration"],
+                        "objective": f"{row['objective']:.6E}",
+                        "dt_ms": f"+{dt_ms:.3g}",
+                    }
+                )
+                previous_timestamp = timestamp
+
     def _normalize_results(self):
         if self.verbose:
             print("\nNormalizing results after convergence...")
@@ -410,6 +467,7 @@ 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()
@@ -418,7 +476,9 @@ def _normalize_results(self):
             self._objective_history = [self.objective_function_]
             self._outer_iter = 0
             self._inner_iter = 0
-            for outiter in range(self.max_iter):
+
+            normalization_max_iter = max(self.max_iter, 100)
+            for outiter in range(normalization_max_iter):
                 self._outer_iter = outiter
                 if outiter == 1:
                     self._inner_iter = (
@@ -451,7 +511,7 @@ def _normalize_results(self):
                     print(
                         f"\n--- Iteration {outiter} after normalization---"
                         f"\nTotal Objective   : {self.objective_function_:.5e}"
-                        "\nConvergence Check : Δ "
+                        "\nConvergence Check : Delta "
                         f"({self.objective_difference_:.2e})"
                         f" < Threshold ({convergence_threshold:.2e})\n"
                     )
@@ -1122,7 +1182,22 @@ def _regularize_function_hessian(self, stretch):
 
         return 0.5 * (hessian + hessian.T)
 
-    def _update_stretch(self):
+    @staticmethod
+    def _project_stretch(stretch, lower_bound=0.1):
+        return np.maximum(stretch, lower_bound)
+
+    def _initial_stretch_step_size(self, stretch, gradient):
+        if self._stretch_step_size is not None:
+            return self._stretch_step_size
+
+        gradient_norm = np.linalg.norm(gradient, "fro")
+        if gradient_norm == 0 or not np.isfinite(gradient_norm):
+            return 1.0
+
+        stretch_norm = max(np.linalg.norm(stretch, "fro"), 1.0)
+        return 0.05 * stretch_norm / gradient_norm
+
+    def _update_stretch_trust_constr(self):
         """Updates stretching matrix using constrained optimization
         (equivalent to fmincon in MATLAB)."""
         if self.verbose:
@@ -1182,6 +1257,67 @@ def hessian(stretch_vec):
 
         # Update stretch with the optimized values
         self.stretch_ = result.x.reshape(self.stretch_.shape)
+        self._stretch_step_size = None
+
+    def _update_stretch_projected_gradient(self):
+        if self.verbose:
+            print("Updating stretch factors...")
+
+        for _ in range(self.stretch_max_iter):
+            stretch = self.stretch_
+            current_objective, gradient = self._regularize_function(stretch)
+            step_size = self._initial_stretch_step_size(stretch, gradient)
+            best_stretch = stretch
+            best_objective = current_objective
+
+            for _ in range(20):
+                candidate_stretch = self._project_stretch(
+                    stretch - step_size * gradient
+                )
+                step = candidate_stretch - stretch
+                step_norm_sq = np.linalg.norm(step, "fro") ** 2
+                if step_norm_sq == 0:
+                    step_size *= 0.5
+                    continue
+
+                candidate_residuals = self._get_residual_matrix(
+                    stretch=candidate_stretch
+                )
+                candidate_objective = self._get_objective_function(
+                    residuals=candidate_residuals,
+                    stretch=candidate_stretch,
+                )
+                if candidate_objective < best_objective:
+                    best_stretch = candidate_stretch
+                    best_objective = candidate_objective
+
+                sufficient_decrease = (
+                    current_objective - 1e-4 * step_norm_sq / step_size
+                )
+                if candidate_objective <= sufficient_decrease:
+                    break
+
+                step_size *= 0.5
+
+            self._stretch_step_size = step_size
+            if best_objective >= current_objective:
+                break
+
+            self.stretch_ = best_stretch
+
+    def _update_stretch(self):
+        """Update stretching factors with a hybrid strategy.
+
+        The first ``stretch_slow_iter`` outer-loop updates use the original
+        constrained nonlinear optimizer to find a good non-convex basin. Later
+        updates switch to the lightweight Algorithm 2 style path: compute the
+        stretch gradient, take linearized proximal steps, then project back
+        onto the feasible stretch range.
+        """
+        if getattr(self, "_outer_iter", 0) < self.stretch_slow_iter:
+            self._update_stretch_trust_constr()
+        else:
+            self._update_stretch_projected_gradient()
 
     @staticmethod
     def _compute_objective_function(
@@ -1255,9 +1391,9 @@ def _cubic_largest_real_root(p, q):
     # Compute discriminant
     delta = (q / 2) ** 2 + (p / 3) ** 3
 
-    # Compute square root of delta safely
-    d = np.where(delta >= 0, np.sqrt(delta), np.sqrt(np.abs(delta)) * 1j)
-    # TODO: this line causes a warning but results seem correct
+    # Match the MATLAB helper's real-root branch without evaluating invalid
+    # real square roots for entries that are handled in complex arithmetic.
+    d = np.sqrt(delta.astype(complex))
 
     # Compute cube roots safely
     a1 = (-q / 2 + d) ** (1 / 3)
@@ -1273,9 +1409,7 @@ 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.max(r_roots, axis=0) * (
-        delta < 0
-    )  # Keep only real roots when delta < 0
+    y = np.where(delta < 0, np.max(r_roots, axis=0), 0.0)
 
     return y