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
[docs]
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``.
"""
[docs]
def __init__(
self,
n_components=None,
max_iter=500,
min_iter=20,
tol=5e-7,
rho=0,
eta=0,
random_state=None,
show_plots=False,
verbose=False,
):
"""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.
"""
if n_components is not None and n_components < 1:
raise ValueError("n_components must be a positive 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._rng = np.random.default_rng(self.random_state)
self._plotter = SNMFPlotter() if self.show_plots else 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()
# 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,
)
[docs]
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 _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.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
for outiter in range(self.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 : Δ "
f"({self.objective_difference_:.2e})"
f" < Threshold ({convergence_threshold:.2e})\n"
)
if (
self.objective_difference_ < convergence_threshold
and outiter >= 7
):
break
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_
reconstructed_matrix = _reconstruct_matrix(
components, weights, stretch
)
residuals = reconstructed_matrix - self._source_matrix
return residuals
def _get_objective_function(self, residuals=None, stretch=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_.
Returns
-------
float
Current objective function value.
"""
return SNMFOptimizer._compute_objective_function(
components=self.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 = np.floor(t).astype(np.int64) # prior original index
alpha = t - i0.astype(float) # fractional distance between left/right
# Clip indices to range (0, signal_len - 1) to maintain original size
max_idx = signal_len - 1
i0 = np.clip(i0, 0, max_idx)
i1 = np.clip(i0 + 1, 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
# 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 _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
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
self.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_)
< 0
)
self.components_ = mask * self.components_
objective_improvement = self.objective_log[-1][
"objective"
] - self._get_objective_function(
residuals=self._get_residual_matrix()
)
# Check if objective function improves
if objective_improvement > 0:
break
# If not, increase step_size (step size)
step_size *= 2
if np.isinf(step_size):
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 _regularize_function(self, stretch=None):
if stretch is None:
stretch = 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
)
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)
)
# Hessian would go here
return fun, gra
def _update_stretch(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
# 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
bounds=bounds,
)
# Update stretch with the optimized values
self.stretch_ = result.x.reshape(self.stretch_.shape)
@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."""
# Handle special case where q == 0
y = np.where(
q == 0, np.maximum(0, -p) ** 0.5, np.zeros_like(p)
) # q=0 case
# 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
# 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)
y = np.max(r_roots, axis=0) * (
delta < 0
) # Keep only real roots when delta < 0
return y
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