Fix nondeterministic preprocessing extension

dingo/preprocess.py

@@ -1,6 +1,7 @@
 
 import cobra
 import cobra.manipulation
+import warnings
 from collections import Counter
 from dingo import MetabolicNetwork, PolytopeSampler
 from dingo.utils import correlated_reactions
@@ -165,7 +166,7 @@ def _remove_model_reactions(self):
         return self._model
 
 
-    def reduce(self, extend=False):
+    def reduce(self, extend=False, steady_states=None):
         """
         A function that calls the "remove_model_reactions" function
         and removes blocked, zero-flux and metabolically less efficient 
@@ -183,6 +184,15 @@ def reduce(self, extend=False):
         If this removal produces an infeasible solution (or a solution of 0) 
         to the objective function, these reactions are restored to their initial bounds.
 
+        steady_states -- Optional precomputed steady states matrix. If provided,
+                        internal MCMC sampling is skipped and this matrix is used
+                        directly for correlation estimation. Expected shape is
+                        (n_reactions, n_samples) where n_reactions matches the
+                        number of reactions in the reduced model (after initial
+                        blocked/zero-flux/mle removal) and rows correspond to the
+                        reaction ordering of that reduced model. Default is None,
+                        which triggers the standard internal sampling path.
+
         A dingo-type tuple is then created from the cobra model 
         using the "cobra_dingo_tuple" function.
 
@@ -221,8 +231,23 @@ def reduce(self, extend=False):
 
             reduced_dingo_model = MetabolicNetwork.from_cobra_model(self._model)
             reactions = reduced_dingo_model.reactions      
-            sampler = PolytopeSampler(reduced_dingo_model)
-            steady_states = sampler.generate_steady_states()
+
+            if steady_states is not None:
+                if steady_states.ndim != 2:
+                    raise ValueError("steady_states must be 2-dimensional")
+                if steady_states.shape[0] != len(reactions):
+                    raise ValueError(
+                        f"steady_states first dimension ({steady_states.shape[0]}) must match "
+                        f"the number of reactions in the reduced model ({len(reactions)})"
+                    )
+            else:
+                warnings.warn(
+                    "extend=True uses internal MCMC sampling, which may produce "
+                    "non-reproducible results. Provide steady_states for reproducibility.",
+                    UserWarning,
+                )
+                sampler = PolytopeSampler(reduced_dingo_model)
+                steady_states = sampler.generate_steady_states()
 
             # calculate correlation matrix with additional filtering from copula indicator       
             corr_matrix = correlated_reactions(

tests/preprocess.py

@@ -1,6 +1,6 @@
 
 from cobra.io import load_json_model
-from dingo import MetabolicNetwork
+from dingo import MetabolicNetwork, PolytopeSampler
 from dingo.preprocess import PreProcess
 import unittest
 import numpy as np
@@ -63,6 +63,62 @@ def test_preprocess(self):
         self.assertTrue(abs(final_fba_solution - initial_fba_solution) < 1e-03)   
 
 
+    def test_reduce_extend_with_steady_states(self):
+        """Test that supplying the same steady_states produces deterministic results."""
+
+        # load cobra model
+        cobra_model = load_json_model("ext_data/e_coli_core.json")
+
+        # Create a preprocessor and perform initial removal to get the model state
+        obj = PreProcess(cobra_model.copy(), tol=1e-6, open_exchanges=False, verbose=False)
+        obj.reduce(extend=False)
+
+        # Get the reduced model after initial removal and generate steady states
+        reduced_model = MetabolicNetwork.from_cobra_model(obj._model)
+        sampler = PolytopeSampler(reduced_model)
+        steady_states = sampler.generate_steady_states()
+
+        # Create two new preprocessors from fresh model copies
+        cobra_model2 = load_json_model("ext_data/e_coli_core.json")
+        cobra_model3 = load_json_model("ext_data/e_coli_core.json")
+
+        obj2 = PreProcess(cobra_model2, tol=1e-6, open_exchanges=False, verbose=False)
+        obj3 = PreProcess(cobra_model3, tol=1e-6, open_exchanges=False, verbose=False)
+
+        # Both should produce identical results with the same steady_states
+        removed2, _ = obj2.reduce(extend=True, steady_states=steady_states)
+        removed3, _ = obj3.reduce(extend=True, steady_states=steady_states)
+
+        # Verify identical reaction sets
+        self.assertEqual(set(removed2), set(removed3))
+
+    def test_reduce_extend_invalid_dimensionality(self):
+        """Test that invalid steady_states dimensionality is rejected."""
+
+        cobra_model = load_json_model("ext_data/e_coli_core.json")
+        obj = PreProcess(cobra_model, tol=1e-6, open_exchanges=False, verbose=False)
+
+        # Test with 1D array
+        with self.assertRaises(ValueError):
+            obj.reduce(extend=True, steady_states=np.array([1, 2, 3]))
+
+        # Test with 3D array
+        with self.assertRaises(ValueError):
+            obj.reduce(extend=True, steady_states=np.random.rand(10, 5, 3))
+
+    def test_reduce_extend_reaction_count_mismatch(self):
+        """Test that reaction count mismatch is rejected."""
+
+        cobra_model = load_json_model("ext_data/e_coli_core.json")
+        obj = PreProcess(cobra_model, tol=1e-6, open_exchanges=False, verbose=False)
+
+        # Create steady_states with wrong number of reactions
+        wrong_steady_states = np.random.rand(50, 100)
+
+        with self.assertRaises(ValueError):
+            obj.reduce(extend=True, steady_states=wrong_steady_states)
+
+
 if __name__ == "__main__":
     unittest.main()