diff --git a/news/logic-fix.rst b/news/logic-fix.rst new file mode 100644 index 0000000..2917e2a --- /dev/null +++ b/news/logic-fix.rst @@ -0,0 +1,24 @@ +**Added:** + +* Add test for recovery from failed component update +* Add a regression test for the cubic root solver + +**Changed:** + +* + +**Deprecated:** + +* + +**Removed:** + +* + +**Fixed:** + +* Maintain valid state through the algorithm + +**Security:** + +* diff --git a/src/diffpy/stretched_nmf/snmf_class.py b/src/diffpy/stretched_nmf/snmf_class.py index 0fa169a..a6d554d 100644 --- a/src/diffpy/stretched_nmf/snmf_class.py +++ b/src/diffpy/stretched_nmf/snmf_class.py @@ -467,7 +467,6 @@ def _normalize_results(self): self._prev_grad_components = np.zeros_like( self.components_ ) # Previous gradient of X (zeros for now) - self._fill_tail_zero = True try: self.residuals_ = self._get_residual_matrix() @@ -476,7 +475,6 @@ def _normalize_results(self): self._objective_history = [self.objective_function_] self._outer_iter = 0 self._inner_iter = 0 - normalization_max_iter = max(self.max_iter, 100) for outiter in range(normalization_max_iter): self._outer_iter = outiter @@ -686,7 +684,9 @@ def _reconstruct_from_stretched_components( order="F", ) - def _get_objective_function(self, residuals=None, stretch=None): + def _get_objective_function( + self, residuals=None, stretch=None, components=None + ): """Return the objective value, passing stored attributes or overrides to _compute_objective_function(). @@ -696,6 +696,8 @@ def _get_objective_function(self, residuals=None, stretch=None): Residual matrix to use instead of self.residuals_. stretch : ndarray, optional Stretch matrix to use instead of self.stretch_. + components : ndarray, optional + Component matrix to use instead of self.components_. Returns ------- @@ -703,7 +705,7 @@ def _get_objective_function(self, residuals=None, stretch=None): Current objective function value. """ return SNMFOptimizer._compute_objective_function( - components=self.components_, + components=self.components_ if components is None else components, residuals=self.residuals_ if residuals is None else residuals, stretch=self.stretch_ if stretch is None else stretch, rho=self.rho, @@ -1032,33 +1034,40 @@ def _update_components(self): self._prev_components - self._grad_components / step_size ) # Solve x^3 + p*x + q = 0 for the largest real root - self.components_ = np.square( + candidate_components = np.square( _cubic_largest_real_root( -components_step, self.eta / (2 * step_size) ) ) # Mask values that should be set to zero mask = ( - self.components_**2 * step_size / 2 - - step_size * self.components_ * components_step - + self.eta * np.sqrt(self.components_) + candidate_components**2 * step_size / 2 + - step_size * candidate_components * components_step + + self.eta * np.sqrt(candidate_components) < 0 ) - self.components_ = mask * self.components_ + candidate_components = mask * candidate_components objective_improvement = ( self.objective_function_ - self._get_objective_function( - residuals=self._get_residual_matrix() + components=candidate_components, + residuals=self._get_residual_matrix( + components=candidate_components + ), ) ) - # Check if objective function improves - if objective_improvement > 0: + if ( + np.isfinite(objective_improvement) + and objective_improvement > 0 + ): + self.components_ = candidate_components break # If not, increase step_size (step size) step_size *= 2 if np.isinf(step_size): + self.components_ = self._prev_components break def _update_weights(self): @@ -1383,10 +1392,12 @@ def _compute_objective_function( def _cubic_largest_real_root(p, q): """Solves x^3 + p*x + q = 0 element-wise for matrices, returning the largest real root.""" - # Handle special case where q == 0 - y = np.where( - q == 0, np.maximum(0, -p) ** 0.5, np.zeros_like(p) - ) # q=0 case + # For q == 0, the non-negative solution is available directly. Keep + # this branch separate: the general complex-root calculation below is + # numerically unstable at this degenerate cubic and previously overwrote + # the exact result. + q_is_zero = q == 0 + zero_q_root = np.maximum(0, -p) ** 0.5 # Compute discriminant delta = (q / 2) ** 2 + (p / 3) ** 3 @@ -1409,9 +1420,9 @@ def _cubic_largest_real_root(p, q): # Take the largest real root element-wise when delta < 0 r_roots = np.stack([np.real(y1), np.real(y2), np.real(y3)], axis=0) - y = np.where(delta < 0, np.max(r_roots, axis=0), 0.0) + general_root = np.where(delta < 0, np.max(r_roots, axis=0), 0.0) - return y + return np.where(q_is_zero, zero_q_root, general_root) def _reconstruct_matrix(components, weights, stretch): diff --git a/tests/test_snmf_optimizer.py b/tests/test_snmf_optimizer.py index 0d2bedf..f32bdad 100644 --- a/tests/test_snmf_optimizer.py +++ b/tests/test_snmf_optimizer.py @@ -2,7 +2,10 @@ import pytest from scipy.sparse import csr_matrix -from diffpy.stretched_nmf.snmf_class import SNMFOptimizer +from diffpy.stretched_nmf.snmf_class import ( + SNMFOptimizer, + _cubic_largest_real_root, +) def test_fit_recovers_rank_one_factors(): @@ -40,6 +43,40 @@ def test_fit_recovers_rank_one_factors(): assert np.allclose(model.weights_, expected_weights, rtol=0.2, atol=0.1) +def test_cubic_largest_real_root_preserves_tiny_zero_q_root(): + root = _cubic_largest_real_root(np.array([[-1e-300]]), np.zeros((1, 1))) + + np.testing.assert_allclose(root, [[1e-150]], rtol=1e-12, atol=0) + + +def test_failed_component_update_restores_previous_components(): + model = SNMFOptimizer(n_components=1, eta=0.0) + model.signal_length_ = model.n_signals_ = model.n_components_ = 1 + model.components_ = np.array([[1.0]]) + max_float = np.finfo(float).max + model.weights_ = np.array([[np.sqrt(max_float)]]) + model.stretch_ = np.ones((1, 1)) + model._source_matrix = np.zeros((1, 1)) + model._fill_tail_zero = True + model._outer_iter = model._inner_iter = 0 + model.objective_function_ = 0.0 + model._compute_stretched_components = lambda: ( + np.zeros((1, 1)), + None, + None, + ) + model._compute_component_gradient_zero_tail = lambda residuals: np.array( + [[np.finfo(float).max]] + ) + model._get_residual_matrix = lambda **kwargs: np.zeros((1, 1)) + model._get_objective_function = lambda **kwargs: 1.0 + + with np.errstate(over="ignore"): + model._update_components() + + np.testing.assert_array_equal(model.components_, [[1.0]]) + + @pytest.mark.parametrize( "inputs, expected", # inputs tuple: