snmf_class.py

import csv
import time

import cvxpy as cp
import numpy as np
from scipy.optimize import minimize
from scipy.sparse import coo_matrix, diags

from diffpy.stretched_nmf.plotter import SNMFPlotter


class SNMFOptimizer:
    """An implementation of stretched NMF, including sparse stretched
    NMF.

    Instantiate the estimator with hyperparameters, then call ``fit`` to
    optimize model factors. Trailing underscores indicate that an attribute
    was determined during the fit process.

    For more information on sNMF, please reference:
    Gu, R., Rakita, Y., Lan, L. et al.
    Stretched non-negative matrix factorization.
    npj Comput Mater 10, 193 (2024) https://doi.org/10.1038/s41524-024-01377-5

    Attributes
    ----------
    stretch_ : numpy.ndarray
        The best guess (or while running, the current guess) for the stretching
        factor matrix.
    components_ : numpy.ndarray
        The best guess (or while running, the current guess) for the matrix of
        component intensities.
    weights_ : numpy.ndarray
        The best guess (or while running, the current guess) for the matrix of
        component weights.
    rho : float
        The stretching factor that influences the decomposition. Zero
        corresponds to no stretching present. Relatively insensitive and
        typically adjusted in powers of 10.
    eta : float
        The sparsity factor that influences the decomposition. Should be set
        to zero for non-sparse data such as PDF. Can be used to improve
        results for sparse data such as XRD, but due to instability, should
        be used only after first selecting the best value for rho. Suggested
        adjustment is by powers of 2.
    max_iter : int
        The maximum number of times to update each of stretch, components,
        and weights before stopping the optimization.
    min_iter : int
        The minimum number of times to update each of stretch, components,
        and weights before terminating the optimization due to low/no
        improvement.
    tol : float
        The convergence threshold. This is the minimum fractional improvement
        in the objective function to allow without terminating the
        optimization.
    n_components : int
        The referred number of components when ``init_weights`` is not
        provided to ``fit``.
    random_state : int
        The seed for the initial guesses at the matrices (stretch, components,
        and weights) created by the decomposition.
    n_components_ : int
        The learned number of components from initialization.
    signal_length_ : int
        The number of rows in the fitted source matrix.
    n_signals_ : int
        The number of columns in the fitted source matrix.
    objective_function_ : float
        Current objective value from the most recent update.
    objective_difference_ : float
        The change in the objective function value since the last update. A
        positive value means that the result improved.
    n_iter_ : int
        The number of outer iterations completed in ``fit``.
    """

    def __init__(
        self,
        n_components=None,
        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.

        Parameters
        ----------
        n_components : int, optional
            The number of components to extract when ``init_weights`` is not
            provided to ``fit``.
        max_iter : int
            The maximum number of times to update each of A, X, and Y before
            stopping the optimization. Optional.
        min_iter : int
            The minimum number of outer-loop iterations before convergence
            checks can stop optimization. Optional.
        tol : float
            The convergence threshold. This is the minimum fractional
            improvement in the objective function to allow without terminating
            the optimization. Note that a minimum of 20 updates are run before
            this parameter is checked. Optional.
        rho : float
            The stretching regularization hyperparameter. Zero corresponds to
            no stretching.
        eta : float
            The sparsity regularization hyperparameter. Turn off for non-sparse
            data such as PDF.
        random_state : int
            The seed for the initial guesses at the matrices (A, X, and Y)
            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
        self.min_iter = min_iter
        self.tol = tol
        self.rho = rho
        self.eta = eta
        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,
        source_matrix,
        init_weights=None,
        init_components=None,
        init_stretch=None,
    ):
        self._rng = np.random.default_rng(self.random_state)
        self.signal_length_, self.n_signals_ = source_matrix.shape

        if init_weights is None and self.n_components is None:
            raise ValueError(
                "n_components must be provided when init_weights is not set."
            )

        if init_weights is None:
            n_components = self.n_components
            weights = self._rng.beta(
                a=2.0,
                b=2.0,
                size=(n_components, self.n_signals_),
            )
        else:
            weights = np.asarray(init_weights, dtype=float)
            n_components = weights.shape[0]
            if (
                self.n_components is not None
                and self.n_components != n_components
            ):
                raise ValueError(
                    "init_weights has a different number of components than "
                    "n_components."
                )

        if init_stretch is None:
            stretch = np.ones(
                (n_components, self.n_signals_)
            ) + self._rng.normal(
                0,
                1e-3,
                size=(n_components, self.n_signals_),
            )
        else:
            stretch = np.asarray(init_stretch, dtype=float)

        if init_components is None:
            components = self._rng.random((self.signal_length_, n_components))
        else:
            components = np.asarray(init_components, dtype=float)

        expected_weights_shape = (n_components, self.n_signals_)
        expected_stretch_shape = (n_components, self.n_signals_)
        expected_components_shape = (self.signal_length_, n_components)

        if weights.shape != expected_weights_shape:
            raise ValueError(
                "init_weights must have shape "
                f"{expected_weights_shape}, got {weights.shape}."
            )
        if stretch.shape != expected_stretch_shape:
            raise ValueError(
                "init_stretch must have shape "
                f"{expected_stretch_shape}, got {stretch.shape}."
            )
        if components.shape != expected_components_shape:
            raise ValueError(
                "init_components must have shape "
                f"{expected_components_shape}, got {components.shape}."
            )

        self.n_components_ = n_components
        self.weights_ = np.maximum(0, weights)
        self.stretch_ = stretch
        self.components_ = np.maximum(0, components)

        self._init_components = self.components_.copy()
        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(
            [1, -2, 1],
            offsets=[0, 1, 2],
            shape=(self.n_signals_ - 2, self.n_signals_),
            dtype=float,
        )

    def fit(
        self,
        source_matrix,
        init_weights=None,
        init_components=None,
        init_stretch=None,
        reset=True,
    ):
        """Run the sNMF optimization on ``source_matrix``.

        Parameters
        ----------
        source_matrix : ndarray of shape (signal_length, n_signals)
            The source data matrix to decompose.
        init_weights : ndarray, optional
            The initial weights matrix of shape
            ``(n_components, n_signals)``.
        init_components : ndarray, optional
            Optional initial components matrix of shape
            ``(signal_length, n_components)``.
        init_stretch : ndarray, optional
            The initial stretch matrix of shape
            ``(n_components, n_signals)``.
        reset : bool
            Whether to reinitialize model factors before fitting. If ``False``,
            the previous factor matrices are reused.
        """
        source_matrix = np.asarray(source_matrix, dtype=float)
        if source_matrix.ndim != 2:
            raise ValueError("source_matrix must be a 2D array.")
        self.converged_ = False

        self._source_matrix = source_matrix

        if reset:
            self._initialize_factors(
                source_matrix=source_matrix,
                init_weights=init_weights,
                init_components=init_components,
                init_stretch=init_stretch,
            )
        else:
            if any(
                v is not None
                for v in (init_weights, init_components, init_stretch)
            ):
                raise ValueError(
                    "init_weights, init_components, and init_stretch can only "
                    "be provided when reset=True."
                )
            if not all(
                hasattr(self, name)
                for name in (
                    "components_",
                    "weights_",
                    "stretch_",
                    "n_components_",
                    "signal_length_",
                    "n_signals_",
                    "_spline_smooth_operator",
                )
            ):
                raise ValueError(
                    "Cannot warm-start before initialization. Call fit with "
                    "reset=True first."
                )
            expected_shape = (self.signal_length_, self.n_signals_)
            if source_matrix.shape != expected_shape:
                raise ValueError(
                    "Warm-start requires source_matrix to keep the same shape "
                    f"{expected_shape}, got {source_matrix.shape}."
                )

        # Set stretch matrix to 1 if no stretching present
        if self.rho == 0:
            self.stretch_ = np.ones_like(self.stretch_)

        # Set up residual matrix, objective function, and history
        self.residuals_ = self._get_residual_matrix()
        self.objective_function_ = self._get_objective_function()
        self.best_objective_ = self.objective_function_
        self.best_matrices_ = [
            self.components_.copy(),
            self.weights_.copy(),
            self.stretch_.copy(),
        ]
        self.objective_difference_ = None
        self.objective_log = [
            {
                "step": "start",
                "iteration": 0,
                "objective": self.objective_function_,
                "timestamp": time.time(),
            }
        ]

        # Set up tracking variables for _update_components()
        self._prev_components = None
        self._grad_components = np.zeros_like(self.components_)
        self._prev_grad_components = np.zeros_like(self.components_)
        self.n_iter_ = 0

        regularization_term = (
            0.5
            * self.rho
            * np.linalg.norm(
                self._spline_smooth_operator @ self.stretch_.T, "fro"
            )
            ** 2
        )
        sparsity_term = self.eta * np.sum(
            np.sqrt(self.components_)
        )  # Square root penalty
        base_obj = (
            self.objective_function_ - regularization_term - sparsity_term
        )
        if self.verbose:
            print(
                f"\n--- Start ---"
                f"\nTotal Objective   : {self.objective_function_:.5e}"
                f"\nBase Obj (No Reg) : {base_obj:.5e}"
            )

        # Main optimization loop
        for outiter in range(self.max_iter):
            self._outer_iter = outiter
            self._outer_loop()
            self.n_iter_ = outiter + 1
            # Print diagnostics
            regularization_term = (
                0.5
                * self.rho
                * np.linalg.norm(
                    self._spline_smooth_operator @ self.stretch_.T, "fro"
                )
                ** 2
            )
            sparsity_term = self.eta * np.sum(
                np.sqrt(self.components_)
            )  # Square root penalty
            base_obj = (
                self.objective_function_ - regularization_term - sparsity_term
            )
            convergence_threshold = self.objective_function_ * self.tol
            # Convergence check: Stop if diffun is small
            # and at least min_iter iterations have passed
            if self.verbose:
                print(
                    f"\n--- Iteration {self._outer_iter} ---"
                    f"\nTotal Objective   : {self.objective_function_:.5e}"
                    f"\nBase Obj (No Reg) : {base_obj:.5e}"
                    "\nConvergence Check : Δ "
                    f"({self.objective_difference_:.2e})"
                    f" < Threshold ({convergence_threshold:.2e})\n"
                )
            if (
                self.objective_difference_ < convergence_threshold
                and outiter >= self.min_iter
            ):
                self.converged_ = True
                break

        self._normalize_results()
        self.reconstruction_err_ = np.linalg.norm(self.residuals_, "fro")

        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...")
        # Select our best results for normalization
        self.components_ = self.best_matrices_[0]
        self.weights_ = self.best_matrices_[1]
        self.stretch_ = self.best_matrices_[2]

        # Normalize weights/stretch first
        weights_row_max = np.max(self.weights_, axis=1, keepdims=True)
        self.weights_ = self.weights_ / weights_row_max
        stretch_row_max = np.max(self.stretch_, axis=1, keepdims=True)
        self.stretch_ = self.stretch_ / stretch_row_max

        # re-running with component updates only vs normalized weights/stretch
        self._grad_components = np.zeros_like(
            self.components_
        )  # Gradient of X (zeros for now)
        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()
            self.objective_function_ = self._get_objective_function()
            self.objective_difference_ = None
            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
                if outiter == 1:
                    self._inner_iter = (
                        1  # So step size can adapt without an inner loop
                    )
                self._update_components()
                self.residuals_ = self._get_residual_matrix()
                self.objective_function_ = self._get_objective_function()
                self.objective_log.append(
                    {
                        "step": "c_norm",
                        "iteration": outiter,
                        "objective": self.objective_function_,
                        "timestamp": time.time(),
                    }
                )
                self.objective_difference_ = (
                    self.objective_log[-2]["objective"]
                    - self.objective_log[-1]["objective"]
                )
                if self._plotter is not None:
                    self._plotter.update(
                        components=self.components_,
                        weights=self.weights_,
                        stretch=self.stretch_,
                        update_tag="normalize components",
                    )
                convergence_threshold = self.objective_function_ * self.tol
                if self.verbose:
                    print(
                        f"\n--- Iteration {outiter} after normalization---"
                        f"\nTotal Objective   : {self.objective_function_:.5e}"
                        "\nConvergence Check : Delta "
                        f"({self.objective_difference_:.2e})"
                        f" < Threshold ({convergence_threshold:.2e})\n"
                    )
                if (
                    self.objective_difference_ < convergence_threshold
                    and outiter >= 7
                ):
                    break
        finally:
            self._fill_tail_zero = False

    def _outer_loop(self):
        if self.verbose:
            print("Updating components and weights...")
        for inner_iter in range(4):
            self._inner_iter = inner_iter
            self._prev_grad_components = self._grad_components.copy()
            self._update_components()
            self.residuals_ = self._get_residual_matrix()
            self.objective_function_ = self._get_objective_function()
            self.objective_log.append(
                {
                    "step": "c",
                    "iteration": self._outer_iter,
                    "objective": self.objective_function_,
                    "timestamp": time.time(),
                }
            )
            self.objective_difference_ = (
                self.objective_log[-2]["objective"]
                - self.objective_log[-1]["objective"]
            )
            if self.objective_function_ < self.best_objective_:
                self.best_objective_ = self.objective_function_
                self.best_matrices_ = [
                    self.components_.copy(),
                    self.weights_.copy(),
                    self.stretch_.copy(),
                ]
            if self._plotter is not None:
                self._plotter.update(
                    components=self.components_,
                    weights=self.weights_,
                    stretch=self.stretch_,
                    update_tag="components",
                )

            self._update_weights()
            self.residuals_ = self._get_residual_matrix()
            self.objective_function_ = self._get_objective_function()
            self.objective_log.append(
                {
                    "step": "w",
                    "iteration": self._outer_iter,
                    "objective": self.objective_function_,
                    "timestamp": time.time(),
                }
            )

            self.objective_difference_ = (
                self.objective_log[-2]["objective"]
                - self.objective_log[-1]["objective"]
            )
            if self.objective_function_ < self.best_objective_:
                self.best_objective_ = self.objective_function_
                self.best_matrices_ = [
                    self.components_.copy(),
                    self.weights_.copy(),
                    self.stretch_.copy(),
                ]
            if self._plotter is not None:
                self._plotter.update(
                    components=self.components_,
                    weights=self.weights_,
                    stretch=self.stretch_,
                    update_tag="weights",
                )

            self.objective_difference_ = (
                self.objective_log[-2]["objective"]
                - self.objective_log[-1]["objective"]
            )
            if (
                self.objective_log[-3]["objective"] - self.objective_function_
                < self.objective_difference_ * 1e-3
            ):
                break

        # Skip updating stretch if no stretching factor
        if not self.rho == 0:
            self._update_stretch()
            self.residuals_ = self._get_residual_matrix()
            self.objective_function_ = self._get_objective_function()
            self.objective_log.append(
                {
                    "step": "s",
                    "iteration": self._outer_iter,
                    "objective": self.objective_function_,
                    "timestamp": time.time(),
                }
            )
            self.objective_difference_ = (
                self.objective_log[-2]["objective"]
                - self.objective_log[-1]["objective"]
            )
            if self.objective_function_ < self.best_objective_:
                self.best_objective_ = self.objective_function_
                self.best_matrices_ = [
                    self.components_.copy(),
                    self.weights_.copy(),
                    self.stretch_.copy(),
                ]
            if self._plotter is not None:
                self._plotter.update(
                    components=self.components_,
                    weights=self.weights_,
                    stretch=self.stretch_,
                    update_tag="stretch",
                )

    def _get_residual_matrix(
        self, components=None, weights=None, stretch=None
    ):
        """Return the residuals (difference) between the source matrix
        and its reconstruction.

        Parameters
        ----------
        components : (signal_len, n_components) array, optional
        weights    : (n_components, n_signals) array, optional
        stretch    : (n_components, n_signals) array, optional

        Returns
        -------
        residuals : (signal_len, n_signals) array
        """
        if components is None:
            components = self.components_
        if weights is None:
            weights = self.weights_
        if stretch is None:
            stretch = self.stretch_

        if self._fill_tail_zero:
            reconstructed_matrix = self._reconstruct_from_stretched_components(
                components=components,
                weights=weights,
                stretch=stretch,
            )
        else:
            reconstructed_matrix = _reconstruct_matrix(
                components, weights, stretch
            )
        residuals = reconstructed_matrix - self._source_matrix

        return residuals

    def _reconstruct_from_stretched_components(
        self, components=None, weights=None, stretch=None
    ):
        stretched_components, _, _ = self._compute_stretched_components(
            components=components,
            weights=weights,
            stretch=stretch,
        )
        intermediate = stretched_components.flatten(order="F").reshape(
            (self.signal_length_ * self.n_signals_, self.n_components_),
            order="F",
        )
        return intermediate.sum(axis=1).reshape(
            (self.signal_length_, self.n_signals_),
            order="F",
        )

    def _get_objective_function(
        self, residuals=None, stretch=None, components=None
    ):
        """Return the objective value, passing stored attributes or
        overrides to _compute_objective_function().

        Parameters
        ----------
        residuals : ndarray, optional
            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_ 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,
            eta=self.eta,
            spline_smooth_operator=self._spline_smooth_operator,
        )

    def _compute_stretched_components(
        self, components=None, weights=None, stretch=None
    ):
        """Interpolates each component along its sample axis according
        to per-(component, signal) stretch factors, then applies
        per-(component, signal) weights. Also computes the first and
        second derivatives with respect to stretch. Left and right,
        respectively, refer to the sample prior to and subsequent to the
        interpolated sample's position.

        Inputs
        ------
        components : array, shape (signal_len, n_components)
            Each column is a component with signal_len samples.
        weights : array, shape (n_components, n_signals)
            Per-(component, signal) weights.
        stretch : array, shape (n_components, n_signals)
            Per-(component, signal) stretch factors.

        Outputs
        -------
        stretched_components : array, shape (signal_len, n_comps * n_sigs)
            Interpolated and weighted components.
        d_stretched_components : array, shape (signal_len, n_comps * n_sigs)
            First derivatives with respect to stretch.
        dd_stretched_components : array, shape (signal_len, n_comps * n_sigs)
            Second derivatives with respect to stretch.
        """
        # --- Defaults ---
        if components is None:
            components = self.components_
        if weights is None:
            weights = self.weights_
        if stretch is None:
            stretch = self.stretch_

        # Dimensions
        signal_len = components.shape[0]  # number of samples
        n_components = components.shape[1]  # number of components
        n_signals = weights.shape[1]  # number of signals

        # Guard stretches
        eps = 1e-8
        stretch = np.clip(stretch, eps, None)
        stretch_inv = 1.0 / stretch

        # Apply stretching to the original sample indices,
        # represented as a "time-stretch"
        t = (
            np.arange(signal_len, dtype=float)[:, None, None]
            * stretch_inv[None, :, :]
        )
        # has shape (signal_len, n_components, n_signals)

        # For each stretched coordinate, find its prior integer (original)
        # index and their difference
        i0_raw = np.floor(t).astype(np.int64)  # prior original index
        alpha = t - i0_raw.astype(float)
        i1_raw = i0_raw + 1

        # Clip indices to range (0, signal_len - 1) to gather safely.
        max_idx = signal_len - 1
        i0 = np.clip(i0_raw, 0, max_idx)
        i1 = np.clip(i1_raw, 0, max_idx)

        # Gather sample values
        comps_3d = components[
            :, :, None
        ]  # expand components by a dim for broadcasting across n_signals
        c0 = np.take_along_axis(comps_3d, i0, axis=0)  # left sample values
        c1 = np.take_along_axis(comps_3d, i1, axis=0)  # right sample values
        if self._fill_tail_zero:
            c0 = np.where(i0_raw < signal_len, c0, 0.0)
            c1 = np.where(i1_raw < signal_len, c1, 0.0)

        # Linear interpolation to determine stretched sample values
        interp = c0 * (1.0 - alpha) + c1 * alpha
        interp_weighted = interp * weights[None, :, :]

        # Derivatives
        di = -t * stretch_inv[None, :, :]  # first-derivative coefficient
        ddi = (
            -di * stretch_inv[None, :, :] * 2.0
        )  # second-derivative coefficient

        d_unweighted = c0 * (-di) + c1 * di
        dd_unweighted = c0 * (-ddi) + c1 * ddi

        d_weighted = d_unweighted * weights[None, :, :]
        dd_weighted = dd_unweighted * weights[None, :, :]

        # Flatten back to expected shape (signal_len, n_components * n_signals)
        return (
            interp_weighted.reshape(signal_len, n_components * n_signals),
            d_weighted.reshape(signal_len, n_components * n_signals),
            dd_weighted.reshape(signal_len, n_components * n_signals),
        )

    def _apply_transformation_matrix(
        self, stretch=None, weights=None, residuals=None
    ):
        """Computes the transformation matrix `stretch_transformed` for
        residuals, using scaling matrix `stretch` and weight
        coefficients `weights`."""
        if stretch is None:
            stretch = self.stretch_
        if weights is None:
            weights = self.weights_
        if residuals is None:
            residuals = self.residuals_

        # Compute scaling matrix
        stretch_tiled = np.tile(
            stretch.reshape(1, self.n_signals_ * self.n_components_, order="F")
            ** -1,
            (self.signal_length_, 1),
        )

        # Compute indices
        indices = np.arange(self.signal_length_)[:, None] * stretch_tiled

        # Weighting coefficients
        weights_tiled = np.tile(
            weights.reshape(
                1, self.n_signals_ * self.n_components_, order="F"
            ),
            (self.signal_length_, 1),
        )

        # Compute floor indices
        floor_indices = np.floor(indices).astype(int)
        floor_indices_1 = np.minimum(floor_indices + 1, self.signal_length_)
        floor_indices_2 = np.minimum(floor_indices_1 + 1, self.signal_length_)

        # Compute fractional part
        fractional_indices = indices - floor_indices

        # Expand row indices
        repm = np.tile(
            np.arange(self.n_components_),
            (self.signal_length_, self.n_signals_),
        )

        # Compute transformations
        kron = np.kron(residuals, np.ones((1, self.n_components_)))
        fractional_kron = kron * fractional_indices
        fractional_weights = (fractional_indices - 1) * weights_tiled

        # Construct sparse matrices
        x2 = coo_matrix(
            (
                (-kron * fractional_weights).flatten(),
                (floor_indices_1.flatten() - 1, repm.flatten()),
            ),
            shape=(self.signal_length_ + 1, self.n_components_),
        ).tocsc()
        x3 = coo_matrix(
            (
                (fractional_kron * weights_tiled).flatten(),
                (floor_indices_2.flatten() - 1, repm.flatten()),
            ),
            shape=(self.signal_length_ + 1, self.n_components_),
        ).tocsc()

        # Combine the last row into previous, then remove the last row
        x2[self.signal_length_ - 1, :] += x2[self.signal_length_, :]
        x3[self.signal_length_ - 1, :] += x3[self.signal_length_, :]
        x2 = x2[:-1, :]
        x3 = x3[:-1, :]

        stretch_transformed = x2 + x3

        return stretch_transformed

    def _compute_component_gradient_zero_tail(
        self, stretch=None, weights=None, residuals=None
    ):
        if stretch is None:
            stretch = self.stretch_
        if weights is None:
            weights = self.weights_
        if residuals is None:
            residuals = self.residuals_

        gradient = np.zeros_like(self.components_)
        sample_indices = np.arange(self.signal_length_, dtype=float)
        for signal in range(self.n_signals_):
            for comp in range(self.n_components_):
                positions = sample_indices / stretch[comp, signal]
                left_indices = np.floor(positions).astype(int)
                alpha = positions - left_indices
                right_indices = left_indices + 1
                scaled_residual = residuals[:, signal] * weights[comp, signal]

                left_mask = left_indices < self.signal_length_
                np.add.at(
                    gradient[:, comp],
                    left_indices[left_mask],
                    scaled_residual[left_mask] * (1.0 - alpha[left_mask]),
                )

                right_mask = right_indices < self.signal_length_
                np.add.at(
                    gradient[:, comp],
                    right_indices[right_mask],
                    scaled_residual[right_mask] * alpha[right_mask],
                )

        return gradient

    def _solve_quadratic_program(self, t, m):
        """
        Solves the quadratic program for updating y in stretched NMF:

            min J(y) = 0.5 * y^T q y + d^T y
            subject to: 0 ≤ y ≤ 1

        Parameters:
        - t: (N, k) ndarray
        - source_matrix_col: (N,) column of source_matrix for the
        corresponding m

        Returns:
        - y: (k,) optimal solution
        """

        source_matrix_col = self._source_matrix[:, m]

        # Compute q and d
        q = t.T @ t  # Gram matrix (k x k)
        d = -t.T @ source_matrix_col  # Linear term (k,)

        k = q.shape[0]  # Number of variables

        # Regularize q to ensure positive semi-definiteness
        reg_factor = 1e-8 * np.linalg.norm(
            q, ord="fro"
        )  # Adaptive regularization, original was fixed
        q += np.eye(k) * reg_factor

        # Define optimization variable
        y = cp.Variable(k)

        # Define quadratic objective
        objective = cp.Minimize(0.5 * cp.quad_form(y, q) + d.T @ y)

        # Define constraints (0 ≤ y ≤ 1)
        constraints = [y >= 0, y <= 1]

        # Solve using a QP solver
        prob = cp.Problem(objective, constraints)
        prob.solve(
            solver=cp.OSQP,
            verbose=False,
            polish=False,  # TODO keep? removes polish message
            # solver_verbose=False
        )

        # Get the solution
        return np.maximum(
            y.value, 0
        )  # Ensure non-negative values in case of solver tolerance issues

    def _update_components(self):
        """Updates `components` using gradient-based optimization with
        adaptive step size."""
        # Compute stretched components using the interpolation function
        stretched_components, _, _ = (
            self._compute_stretched_components()
        )  # Discard the derivatives
        # Compute reshaped_stretched_components and component_residuals
        intermediate_reshaped = stretched_components.flatten(
            order="F"
        ).reshape(
            (self.signal_length_ * self.n_signals_, self.n_components_),
            order="F",
        )
        reshaped_stretched_components = intermediate_reshaped.sum(
            axis=1
        ).reshape((self.signal_length_, self.n_signals_), order="F")
        component_residuals = (
            reshaped_stretched_components - self._source_matrix
        )
        # Compute gradient
        if self._fill_tail_zero:
            self._grad_components = self._compute_component_gradient_zero_tail(
                residuals=component_residuals
            )
        else:
            self._grad_components = self._apply_transformation_matrix(
                residuals=component_residuals
            ).toarray()  # toarray equivalent of full, make non-sparse

        # Compute initial step size `initial_step_size`
        initial_step_size = np.linalg.eigvalsh(
            self.weights_.T @ self.weights_
        ).max() * np.max([self.stretch_.max(), 1 / self.stretch_.min()])
        # Compute adaptive step size `step_size`
        if self._outer_iter == 0 and self._inner_iter == 0:
            step_size = initial_step_size
        else:
            num = np.sum(
                (self._grad_components - self._prev_grad_components)
                * (self.components_ - self._prev_components)
            )  # Element-wise multiplication
            denom = (
                np.linalg.norm(self.components_ - self._prev_components, "fro")
                ** 2
            )  # Frobenius norm squared
            step_size = num / denom if denom > 0 else initial_step_size
            if step_size <= 0:
                step_size = initial_step_size

        # Store our old X before updating because it is used in step selection
        self._prev_components = self.components_.copy()

        while True:  # iterate updating components
            components_step = (
                self._prev_components - self._grad_components / step_size
            )
            # Solve x^3 + p*x + q = 0 for the largest real root
            candidate_components = np.square(
                _cubic_largest_real_root(
                    -components_step, self.eta / (2 * step_size)
                )
            )
            # Mask values that should be set to zero
            mask = (
                candidate_components**2 * step_size / 2
                - step_size * candidate_components * components_step
                + self.eta * np.sqrt(candidate_components)
                < 0
            )
            candidate_components = mask * candidate_components

            objective_improvement = (
                self.objective_function_
                - self._get_objective_function(
                    components=candidate_components,
                    residuals=self._get_residual_matrix(
                        components=candidate_components
                    ),
                )
            )

            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):
        """Updates weights by building the stretched component matrix
        `stretched_comps` with np.interp and solving a quadratic program
        for each signal."""
        sample_indices = np.arange(self.signal_length_)
        for signal in range(self.n_signals_):
            # Stretch factors for this signal across components:
            this_stretch = self.stretch_[:, signal]
            # Build stretched_comps[:, k] by interpolating component at frac.
            # pos. index / this_stretch[comp]
            stretched_comps = np.empty(
                (self.signal_length_, self.n_components_),
                dtype=self.components_.dtype,
            )
            for comp in range(self.n_components_):
                pos = sample_indices / this_stretch[comp]
                stretched_comps[:, comp] = np.interp(
                    pos,
                    sample_indices,
                    self.components_[:, comp],
                    left=self.components_[0, comp],
                    right=self.components_[-1, comp],
                )

            # Solve quadratic problem for a given signal and update its weight
            new_weight = self._solve_quadratic_program(
                t=stretched_comps, m=signal
            )
            self.weights_[:, signal] = new_weight

    def _stretch_residual_and_derivatives(self, stretch):
        stretched_components, d_stretch_comps, dd_stretch_comps = (
            self._compute_stretched_components(stretch=stretch)
        )
        intermediate = stretched_components.flatten(order="F").reshape(
            (self.signal_length_ * self.n_signals_, self.n_components_),
            order="F",
        )
        residuals = (
            intermediate.sum(axis=1).reshape(
                (self.signal_length_, self.n_signals_), order="F"
            )
            - self._source_matrix
        )
        return residuals, d_stretch_comps, dd_stretch_comps

    def _regularize_function(self, stretch=None):
        if stretch is None:
            stretch = self.stretch_

        residuals, d_stretch_comps, _ = self._stretch_residual_and_derivatives(
            stretch
        )

        fun = self._get_objective_function(residuals, stretch)

        tiled_res = np.tile(residuals, (1, self.n_components_))
        grad_flat = np.sum(d_stretch_comps * tiled_res, axis=0)
        gra = grad_flat.reshape(
            (self.n_signals_, self.n_components_), order="F"
        ).T
        gra += (
            self.rho
            * stretch
            @ (self._spline_smooth_operator.T @ self._spline_smooth_operator)
        )

        return fun, gra

    def _regularize_function_hessian(self, stretch):
        """Calculate the Hessian for the stretch optimization objective.

        The Hessian combines the Gauss-Newton curvature from the stretched
        component derivatives, the residual-weighted second derivatives of
        those stretched components, and the quadratic smoothing penalty on
        neighboring stretch factors.

        Parameters
        ----------
        stretch : ndarray of shape (n_components, n_signals)
            Stretching factors at which to evaluate the objective curvature.

        Returns
        -------
        ndarray of shape (n_components * n_signals, n_components * n_signals)
            Symmetric Hessian matrix for the flattened stretch variables.
        """
        residuals, d_stretch_comps, dd_stretch_comps = (
            self._stretch_residual_and_derivatives(stretch)
        )
        n_variables = self.n_components_ * self.n_signals_
        hessian = np.zeros((n_variables, n_variables), dtype=float)

        for signal in range(self.n_signals_):
            variable_indices = (
                np.arange(self.n_components_) * self.n_signals_ + signal
            )
            d_signal = d_stretch_comps[:, variable_indices]
            dd_signal = dd_stretch_comps[:, variable_indices]
            hessian[np.ix_(variable_indices, variable_indices)] += (
                d_signal.T @ d_signal
            )
            hessian[variable_indices, variable_indices] += np.sum(
                dd_signal * residuals[:, signal, None],
                axis=0,
            )

        smooth_hessian = (
            self._spline_smooth_operator.T @ self._spline_smooth_operator
        ).toarray()
        for comp in range(self.n_components_):
            component_slice = slice(
                comp * self.n_signals_,
                (comp + 1) * self.n_signals_,
            )
            hessian[component_slice, component_slice] += (
                self.rho * smooth_hessian
            )

        return 0.5 * (hessian + hessian.T)

    @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:
            print("Updating stretch factors...")

        # Flatten stretch for compatibility with the optimizer
        # (since SciPy expects 1D input)
        stretch_flat_initial = self.stretch_.flatten()

        # Define the optimization function
        def objective(stretch_vec):
            stretch_matrix = stretch_vec.reshape(
                self.stretch_.shape
            )  # Reshape back to matrix form
            fun, gra = self._regularize_function(stretch_matrix)
            gra = gra.flatten()
            return fun, gra

        def hessian(stretch_vec):
            stretch_matrix = stretch_vec.reshape(self.stretch_.shape)
            return self._regularize_function_hessian(stretch_matrix)

        unconstrained_result = minimize(
            fun=lambda stretch_vec: objective(stretch_vec)[0],
            x0=stretch_flat_initial,
            method="trust-exact",
            jac=lambda stretch_vec: objective(stretch_vec)[1],
            hess=hessian,
            options={"maxiter": 300},
        )
        unconstrained_stretch = unconstrained_result.x.reshape(
            self.stretch_.shape
        )
        if np.all(unconstrained_stretch >= 0.1):
            current_objective = self._regularize_function(self.stretch_)[0]
            candidate_objective = self._regularize_function(
                unconstrained_stretch
            )[0]
            if candidate_objective <= current_objective:
                self.stretch_ = unconstrained_stretch
                return

        # Optimization constraints: lower bound 0.1, no upper bound
        bounds = [
            (0.1, None)
        ] * stretch_flat_initial.size  # Equivalent to 0.1 * ones(K, M)

        # Solve optimization problem (equivalent to fmincon)
        result = minimize(
            fun=lambda stretch_vec: objective(stretch_vec)[0],
            x0=stretch_flat_initial,
            method="trust-constr",  # Substitute for 'trust-region-reflective'
            jac=lambda stretch_vec: objective(stretch_vec)[1],  # Gradient
            hess=hessian,
            bounds=bounds,
        )

        # 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(
        components, residuals, stretch, rho, eta, spline_smooth_operator
    ):
        r"""Computes the objective function used in stretched non-
        negative matrix factorization.

        Parameters
        ----------
        components : ndarray
            Non-negative matrix of component signals :math:`X`.
        residuals : ndarray
            Difference between reconstructed and observed data.
        stretch : ndarray
            Stretching factors :math:`A` applied to each component across
            samples.
        rho : float
            Regularization parameter enforcing smooth variation in :math:`A`.
        eta : float
            Sparsity-promoting regularization parameter applied to :math:`X`.
        spline_smooth_operator : ndarray
            Linear operator :math:`L` penalizing non-smooth changes
            in :math:`A`.

        Returns
        -------
        float
            Value of the stretched-NMF objective function.

        Notes
        -----
        The stretched-NMF objective function :math:`J` is

        .. math::

           J(X, Y, A) =
              \tfrac{1}{2} \lVert Z - Y\,S(A)X \rVert_F^2
              + \tfrac{\rho}{2} \lVert L A \rVert_F^2
              + \eta \sum_{i,j} \sqrt{X_{ij}} \,,

        where :math:`Z` is the data matrix, :math:`Y` contains the non-negative
        weights, :math:`S(A)` denotes the spline-interp. stretching operator,
        and :math:`\lVert \cdot \rVert_F` is the Frobenius norm.

        Special cases
        -------------
        - :math:`\rho = 0` — no smoothness regularization on stretch factors.
        - :math:`\eta = 0` — no sparsity promotion on components.
        - :math:`\rho = \eta = 0` — reduces to the classical NMF least-squares
          objective :math:`\tfrac{1}{2} \lVert Z - YX \rVert_F^2`.
        """
        residual_term = 0.5 * np.linalg.norm(residuals, "fro") ** 2
        regularization_term = (
            0.5
            * rho
            * np.linalg.norm(spline_smooth_operator @ stretch.T, "fro") ** 2
        )
        sparsity_term = eta * np.sum(np.sqrt(components))
        return residual_term + regularization_term + sparsity_term


def _cubic_largest_real_root(p, q):
    """Solves x^3 + p*x + q = 0 element-wise for matrices, returning the
    largest real root."""
    # 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

    # 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)
    a2 = (-q / 2 - d) ** (1 / 3)

    # Compute cube roots of unity
    w = (np.sqrt(3) * 1j - 1) / 2

    # Compute the three possible roots (element-wise)
    y1 = a1 + a2
    y2 = w * a1 + w**2 * a2
    y3 = w**2 * a1 + w * a2

    # 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)
    general_root = np.where(delta < 0, np.max(r_roots, axis=0), 0.0)

    return np.where(q_is_zero, zero_q_root, general_root)


def _reconstruct_matrix(components, weights, stretch):
    """Construct the approximation of the source matrix corresponding to
    the given components, weights, and stretch factors.

    Each component profile is stretched, interpolated to fractional positions,
    weighted per signal, and summed to form the reconstruction.

    Parameters
    ----------
    components : (signal_len, n_components) array
    weights    : (n_components, n_signals) array
    stretch    : (n_components, n_signals) array

    Returns
    -------
    reconstructed_matrix : (signal_len, n_signals) array
    """
    signal_len = components.shape[0]
    n_components = components.shape[1]
    n_signals = weights.shape[1]

    reconstructed_matrix = np.zeros((signal_len, n_signals))
    sample_indices = np.arange(signal_len)

    for comp in range(n_components):  # loop over components
        reconstructed_matrix += (
            np.interp(
                sample_indices[:, None]
                / stretch[comp][
                    None, :
                ],  # fractional positions (signal_len, n_signals)
                sample_indices,  # (signal_len,)
                components[:, comp],  # component profile (signal_len,)
                left=components[0, comp],
                right=components[-1, comp],
            )
            * weights[comp][None, :]  # broadcast (n_signals,) over rows
        )

    return reconstructed_matrix