Add Tracker class for monitoring simulation metrics

pyology/krebs_cycle.py

@@ -5,6 +5,7 @@
 from .data import Effector, Enzyme
 from .organelle import Organelle
 from .utils import allosteric_regulation, hill_equation, michaelis_menten
+from .tracker import Tracker, EnergyTracker, CO2Tracker
 
 logger = logging.getLogger(__name__)
 
@@ -63,6 +64,9 @@ def __init__(self):
             malate=(0, 100),
         )
 
+        self.energy_tracker = EnergyTracker()
+        self.co2_tracker = CO2Tracker()
+
     def is_metabolite_available(self, metabolite: str, amount: float) -> bool:
         """Check if a metabolite is available in sufficient quantity."""
         if metabolite in self.metabolites:
@@ -137,6 +141,8 @@ def step1_citrate_synthase(self):
             self.produce_metabolites(Citrate=reaction_rate)
             self.cofactors["Coenzyme-A"] += reaction_rate
             logger.info(f"Citrate synthase: Produced {reaction_rate} Citrate")
+            self.energy_tracker.log_energy(0, 0, 0)
+            self.co2_tracker.log_co2_production(0)
             return True
         else:
             logger.warning("Insufficient substrates for step 1")
@@ -152,6 +158,8 @@ def step2_aconitase(self):
 
         if self.consume_metabolites(Citrate=reaction_rate):
             self.produce_metabolites(Isocitrate=reaction_rate)
+            self.energy_tracker.log_energy(0, 0, 0)
+            self.co2_tracker.log_co2_production(0)
         else:
             logger.warning("Insufficient Citrate for step 2")
 
@@ -185,6 +193,8 @@ def step3_isocitrate_dehydrogenase(self):
                     "CO2": reaction_rate,
                 }
             )
+            self.energy_tracker.log_energy(0, reaction_rate, 0)
+            self.co2_tracker.log_co2_production(reaction_rate)
         else:
             logger.warning("Insufficient substrates or NAD⁺ for step 3")
 
@@ -217,6 +227,8 @@ def step4_alpha_ketoglutarate_dehydrogenase(self):
                     "CO2": reaction_rate,
                 }
             )
+            self.energy_tracker.log_energy(0, reaction_rate, 0)
+            self.co2_tracker.log_co2_production(reaction_rate)
         else:
             logger.warning("Insufficient substrates or NAD⁺ for step 4")
 
@@ -236,6 +248,8 @@ def step5_succinyl_coa_synthetase(self):
                 GTP=reaction_rate,
                 **{"Coenzyme-A": reaction_rate},
             )
+            self.energy_tracker.log_energy(reaction_rate, 0, 0)
+            self.co2_tracker.log_co2_production(0)
         else:
             logger.warning("Insufficient substrates or GDP for step 5")
 
@@ -249,6 +263,8 @@ def step6_succinate_dehydrogenase(self):
 
         if self.consume_metabolites(Succinate=reaction_rate, FAD=reaction_rate):
             self.produce_metabolites(Fumarate=reaction_rate, FADH2=reaction_rate)
+            self.energy_tracker.log_energy(0, 0, reaction_rate)
+            self.co2_tracker.log_co2_production(0)
         else:
             logger.warning("Insufficient substrates or FAD for step 6")
 
@@ -262,6 +278,8 @@ def step7_fumarase(self):
 
         if self.consume_metabolites(Fumarate=reaction_rate):
             self.produce_metabolites(Malate=reaction_rate)
+            self.energy_tracker.log_energy(0, 0, 0)
+            self.co2_tracker.log_co2_production(0)
         else:
             logger.warning("Insufficient Fumarate for step 7")
 
@@ -275,10 +293,14 @@ def step8_malate_dehydrogenase(self):
 
         if self.consume_metabolites(Malate=reaction_rate, NAD=reaction_rate):
             self.produce_metabolites(Oxaloacetate=reaction_rate, NADH=reaction_rate)
+            self.energy_tracker.log_energy(0, reaction_rate, 0)
+            self.co2_tracker.log_co2_production(reaction_rate)
         else:
             logger.warning("Insufficient substrates or NAD⁺ for step 8")
 
     def run_cycle(self):
+        self.energy_tracker.start_tracking()
+        self.co2_tracker.start_tracking()
         if self.metabolites["Acetyl-CoA"].quantity > 0:
             self.step1_citrate_synthase()
             self.step2_aconitase()
@@ -290,6 +312,10 @@ def run_cycle(self):
             self.step8_malate_dehydrogenase()
         else:
             logger.warning("Insufficient Acetyl-CoA to start Krebs cycle")
+        energy_report = self.energy_tracker.report()
+        co2_report = self.co2_tracker.report()
+        logger.info(f"Energy Report: {energy_report}")
+        logger.info(f"CO2 Report: {co2_report}")
 
     def krebs_cycle_iterator(self, num_cycles: int = None):
         """Generator that yields the state after each Krebs cycle."""
@@ -374,4 +400,3 @@ def run_cycle_with_generators(self):
             # - Pausing or modifying the cycle based on certain conditions
 
         logger.info("Krebs cycle complete")
-

pyology/tracker.py

@@ -0,0 +1,41 @@
+class Tracker:
+    def __init__(self):
+        self.tracked_metrics = {}
+        self.simulation_step = 0
+
+    def start_tracking(self):
+        self.tracked_metrics = {}
+        self.simulation_step = 0
+
+    def log_change(self, metric: str, value: float):
+        if metric not in self.tracked_metrics:
+            self.tracked_metrics[metric] = []
+        self.tracked_metrics[metric].append((self.simulation_step, value))
+        self.simulation_step += 1
+
+    def report(self):
+        report_data = {}
+        for metric, values in self.tracked_metrics.items():
+            report_data[metric] = sum(value for step, value in values) / len(values)
+        return report_data
+
+    def reset(self):
+        self.start_tracking()
+
+
+class EnergyTracker(Tracker):
+    def __init__(self):
+        super().__init__()
+
+    def log_energy(self, atp_yield: float, nadh_contribution: float, fadh2_contribution: float):
+        self.log_change("ATP Yield", atp_yield)
+        self.log_change("NADH Contribution", nadh_contribution)
+        self.log_change("FADH2 Contribution", fadh2_contribution)
+
+
+class CO2Tracker(Tracker):
+    def __init__(self):
+        super().__init__()
+
+    def log_co2_production(self, co2_production: float):
+        self.log_change("CO2 Production", co2_production)

tests/test_tracker.py

@@ -0,0 +1,62 @@
+import unittest
+from pyology.tracker import Tracker, EnergyTracker, CO2Tracker
+
+class TestTracker(unittest.TestCase):
+    def setUp(self):
+        self.tracker = Tracker()
+
+    def test_start_tracking(self):
+        self.tracker.start_tracking()
+        self.assertEqual(self.tracker.tracked_metrics, {})
+        self.assertEqual(self.tracker.simulation_step, 0)
+
+    def test_log_change(self):
+        self.tracker.start_tracking()
+        self.tracker.log_change("ATP Yield", 10.0)
+        self.assertIn("ATP Yield", self.tracker.tracked_metrics)
+        self.assertEqual(self.tracker.tracked_metrics["ATP Yield"], [(0, 10.0)])
+        self.assertEqual(self.tracker.simulation_step, 1)
+
+    def test_report(self):
+        self.tracker.start_tracking()
+        self.tracker.log_change("ATP Yield", 10.0)
+        self.tracker.log_change("ATP Yield", 20.0)
+        report = self.tracker.report()
+        self.assertIn("ATP Yield", report)
+        self.assertEqual(report["ATP Yield"], 15.0)
+
+    def test_reset(self):
+        self.tracker.start_tracking()
+        self.tracker.log_change("ATP Yield", 10.0)
+        self.tracker.reset()
+        self.assertEqual(self.tracker.tracked_metrics, {})
+        self.assertEqual(self.tracker.simulation_step, 0)
+
+class TestEnergyTracker(unittest.TestCase):
+    def setUp(self):
+        self.energy_tracker = EnergyTracker()
+
+    def test_log_energy(self):
+        self.energy_tracker.start_tracking()
+        self.energy_tracker.log_energy(10.0, 5.0, 2.0)
+        self.assertIn("ATP Yield", self.energy_tracker.tracked_metrics)
+        self.assertIn("NADH Contribution", self.energy_tracker.tracked_metrics)
+        self.assertIn("FADH2 Contribution", self.energy_tracker.tracked_metrics)
+        self.assertEqual(self.energy_tracker.tracked_metrics["ATP Yield"], [(0, 10.0)])
+        self.assertEqual(self.energy_tracker.tracked_metrics["NADH Contribution"], [(0, 5.0)])
+        self.assertEqual(self.energy_tracker.tracked_metrics["FADH2 Contribution"], [(0, 2.0)])
+        self.assertEqual(self.energy_tracker.simulation_step, 1)
+
+class TestCO2Tracker(unittest.TestCase):
+    def setUp(self):
+        self.co2_tracker = CO2Tracker()
+
+    def test_log_co2_production(self):
+        self.co2_tracker.start_tracking()
+        self.co2_tracker.log_co2_production(10.0)
+        self.assertIn("CO2 Production", self.co2_tracker.tracked_metrics)
+        self.assertEqual(self.co2_tracker.tracked_metrics["CO2 Production"], [(0, 10.0)])
+        self.assertEqual(self.co2_tracker.simulation_step, 1)
+
+if __name__ == '__main__':
+    unittest.main()