fix(pipeline): correct state update logic during pruning

Author: xImoZA

rework_pysatl_mpest/estimators/iterative/pipeline.py

@@ -165,20 +165,16 @@ def fit(self, X: ArrayLike, mixture: MixtureModel[DType]) -> MixtureModel[DType]
 
         X = np.asarray(X, dtype=mixture.dtype)
         copied_mixture = copy(mixture)  # Copy to avoid modifying the original object
-        removed_indices: list[int] = []
         state = PipelineState(X, None, None, copied_mixture, None)
 
         while True:
             # Updating the state before starting an iteration
             state.prev_mixture = copy(state.curr_mixture)
-            # Update responsibility matrix H
-            if state.H is not None:
-                state.H = np.delete(state.H, removed_indices, axis=1)
 
             # Performing steps
             for step in self.steps:
                 result_state = step.run(state)
-                step.clear_after_prune(removed_indices)
+                # Log the error state before exiting
                 if result_state.error:
                     if len(self.history) > 0:
                         self.history[-1].error = result_state.error
@@ -219,11 +215,18 @@ def fit(self, X: ArrayLike, mixture: MixtureModel[DType]) -> MixtureModel[DType]
             # Pruning
             for pruner in self.pruners:
                 state, removed_components_indices = pruner.prune(state)
-                if len(removed_components_indices) != 0:
-                    removed_indices.extend(removed_components_indices)
-            # Save iteration record
-            self.history.save_record(
-                IterationRecord(self.history._counter, state.curr_mixture, state.X, state.H, self.pruners, state.error)
+                if removed_components_indices:
+                    # Update optimization blocks in steps to drop blocks associated with removed components
+                    for step in self.steps:
+                        step.clear_after_prune(removed_components_indices)
+
+                    # Update responsibility matrix H to match the new mixture size
+                    if state.H is not None:
+                        state.H = np.delete(state.H, removed_components_indices, axis=1)
+
+            # Log
+            self.logger.log(
+                IterationRecord(self.logger._counter, state.curr_mixture, state.X, state.H, self.pruners, state.error)
             )
 
             # Checking stopping criteria