modified cli inputs and config

Author: AdishAssain

configs/sample-sir-ode.yaml

@@ -4,13 +4,17 @@ NetworkFile: data/networks/sample-network-static.csv
 SeedFile: data/seeds/sample-sir-ode-patchA-2.csv
 Logging: False
 
+# Model configuration
+ModelName: sample-sir-ode
+
 # Simulation parameters
-Beta: 0.3
-Gamma: 0.1
 TMax: 50
 Tolerance: 1e-8
 MaxIter: 10000
 StartDate: 2020-01-01
 EndDate: 2022-12-31
 OutputDir: output/sample-sir-ode
+compartments: ["S", "I", "R"]
+Parameters: {"beta": 0.5, "gamma": 0.1}
+Transitions: [{"from": "S", "to": "I", "rate": "beta * S * I / (S + I + R)"}, {"from": "I", "to": "R", "rate": "gamma * I"}]
 

src/patchsim/cli.py

@@ -1,9 +1,7 @@
 import argparse
 import logging
-import os
-from datetime import datetime
 
-from patchsim.models.sample_sir_ode import run_simulation
+from patchsim.core.simulation import load_config, run_simulation, setup_simulation
 
 
 def main():
@@ -18,9 +16,6 @@ def main():
 Examples:
   # Run simulation with sample SIR model
   patchsim --config configs/sample-sir-ode.yaml
-
-  # Run simulation with custom model
-  patchsim --config path/to/your/config.yaml
         """
     )
     parser.add_argument(
@@ -31,25 +26,13 @@ def main():
     )
     args = parser.parse_args()
 
-    # Create output directories
-    os.makedirs("output/logs", exist_ok=True)
-
-    # Set up logging with timestamp
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    log_file = f"output/logs/cli_{timestamp}.log"
-
-    logging.basicConfig(
-        level=logging.INFO,
-        format="%(asctime)s - %(levelname)s - %(message)s",
-        handlers=[
-            logging.FileHandler(log_file),
-            logging.StreamHandler(),
-        ],
-    )
-
     try:
         logging.info("Starting PatchSim simulation...")
-        run_simulation(args.config)
+        config = load_config(args.config)
+
+        # Set up simulation
+        net, y0, patches, num_patches = setup_simulation(config)
+        run_simulation(config, config['ModelName'], net, y0, patches, num_patches)
         logging.info("Simulation completed successfully.")
     except Exception as e:
         logging.error(f"Simulation failed: {e}")

src/patchsim/core/model.py

@@ -2,9 +2,8 @@
 Core model implementation for compartmental models.
 """
 
-from typing import Any, Callable
+from typing import Any, Callable, Dict
 
-import numpy as np
 from scipy.integrate import odeint
 
 
@@ -21,8 +20,8 @@ def compute_rates(self, state: dict[str, float]) -> dict[str, float]:
         """Compute transition rates for each compartment."""
         rates = {}
         for transition in self.transitions:
-            source = transition['source']
-            target = transition['target']
+            source = transition['from']
+            target = transition['to']
             rate = transition['rate']
             # Handle rate expressions
             if isinstance(rate, str):
@@ -45,89 +44,93 @@ def __init__(self, base_model: CompartmentalModel, num_patches: int, network_mat
         self.network = network_matrix
         self.all_compartments = [f"{c}_{i}" for i in range(num_patches) for c in base_model.compartments]
 
-    def compute_force_of_infection(self, full_state: dict[str, float]) -> list[float]:
-        """Compute force of infection for each patch."""
+    def get_patch_state(self, full_state: Dict[str, float], patch_idx: int) -> Dict[str, float]:
+        """Get state for a specific patch."""
+        return {c: full_state[f"{c}_{patch_idx}"] for c in self.base_model.compartments}
+
+    def get_patch_population(self, state: Dict[str, float]) -> float:
+        """Get total population for a patch."""
+        return sum(state[c] for c in self.base_model.compartments)
+
+    def compute_force_of_infection(self, full_state: dict[str, float], infected_compartment: str = "I") -> list[float]:
+        """Compute force of infection for each patch.
+
+        Args:
+            full_state: Current state of all compartments
+            infected_compartment: Name of the compartment representing infected individuals
+        """
         lambdas = []
         for i in range(self.num_patches):
-            force = 0
-            for j in range(self.num_patches):
-                infected_j = full_state[f"I_{j}"]
-                pop_j = sum(full_state[f"{c}_{j}"] for c in self.base_model.compartments)
-                force += self.network[i][j] * infected_j / pop_j
+            if self.num_patches == 1:
+                # Single patch case: local force of infection
+                patch_state = self.get_patch_state(full_state, 0)
+                infected = patch_state[infected_compartment]
+                total_pop = self.get_patch_population(patch_state)
+                force = self.base_model.parameters['beta'] * infected / total_pop
+            else:
+                # Multi-patch case: network-based force of infection
+                force = 0
+                for j in range(self.num_patches):
+                    patch_state_j = self.get_patch_state(full_state, j)
+                    infected_j = patch_state_j[infected_compartment]
+                    pop_j = self.get_patch_population(patch_state_j)
+                    force += self.network[i][j] * infected_j / pop_j
             lambdas.append(force)
         return lambdas
 
+    def compute_derivatives(self, state: dict[str, float]) -> dict[str, float]:
+        """Compute derivatives for all compartments based on transitions."""
+        derivatives = {c: 0.0 for c in self.all_compartments}
+
+        # Process each patch
+        for i in range(self.num_patches):
+            # Get state for this patch
+            patch_state = self.get_patch_state(state, i)
+
+            # Get rates for this patch
+            rates = self.base_model.compute_rates(patch_state)
+
+            # Update derivatives based on rates
+            for transition in self.base_model.transitions:
+                source = transition['from']
+                target = transition['to']
+                rate = rates[f"{source}_to_{target}"]
+
+                # Decrease source compartment
+                derivatives[f"{source}_{i}"] -= rate
+                # Increase target compartment
+                derivatives[f"{target}_{i}"] += rate
+
+        return derivatives
+
     def simulate_discrete(self, y0_dict: dict[str, float], t_range: list[float]) -> dict[str, list[float]]:
         """Run discrete-time simulation."""
         state = y0_dict.copy()
         history = {c: [state[c]] for c in self.all_compartments}
+
         for _ in t_range[1:]:
-            new_state = state.copy()
-            lambdas = self.compute_force_of_infection(state)
-            for i in range(self.num_patches):
-                for c in self.base_model.compartments:
-                    comp_key = f"{c}_{i}"
-                    if c == 'S':
-                        new_state[comp_key] -= lambdas[i] * state[comp_key]
-                    elif c == 'I':
-                        new_state[comp_key] += (
-                            lambdas[i] * state[f"S_{i}"] - 
-                            self.base_model.parameters['gamma'] * state[comp_key]
-                        )
-                    elif c == 'R':
-                        new_state[comp_key] += self.base_model.parameters['gamma'] * state[f"I_{i}"]
+            derivatives = self.compute_derivatives(state)
+            new_state = {
+                c: state[c] + derivatives[c]
+                for c in self.all_compartments
+            }
             state = new_state
             for c in self.all_compartments:
                 history[c].append(state[c])
+
         return history
 
     def simulate_ode(
         self, y0_dict: dict[str, float], t_range: list[float], integrator: Callable = odeint
     ) -> tuple[list[float], dict[str, list[float]]]:
         """Run ODE simulation."""
-        if (
-            hasattr(self, 'all_compartments') and 
-            len(self.all_compartments) > 0 and 
-            all('_' in c for c in self.all_compartments)
-        ):
-            y0 = [y0_dict[c] for c in self.all_compartments]
-
-            def rhs(y, t):
-                state = {c: y[i] for i, c in enumerate(self.all_compartments)}
-                lambdas = self.compute_force_of_infection(state)
-                dydt = np.zeros_like(y)
-                for i in range(self.num_patches):
-                    for c in self.base_model.compartments:
-                        idx = self.all_compartments.index(f"{c}_{i}")
-                        if c == 'S':
-                            dydt[idx] = -lambdas[i] * state[f"S_{i}"]
-                        elif c == 'I':
-                            dydt[idx] = (
-                                lambdas[i] * state[f"S_{i}"] - 
-                                self.base_model.parameters['gamma'] * state[f"I_{i}"]
-                            )
-                        elif c == 'R':
-                            dydt[idx] = self.base_model.parameters['gamma'] * state[f"I_{i}"]
-                return dydt
-
-            sol = integrator(rhs, y0, t_range)
-            out = {c: sol[:, i] for i, c in enumerate(self.all_compartments)}
-            return t_range, out
-
-        y0 = [y0_dict[c] for c in self.base_model.compartments]
+        y0 = [y0_dict[c] for c in self.all_compartments]
 
         def rhs(y, t):
-            state = {c: y[i] for i, c in enumerate(self.base_model.compartments)}
-            rates = self.base_model.compute_rates(state)
-            dydt = np.zeros_like(y)
-            for i, c in enumerate(self.base_model.compartments):
-                for transition in self.base_model.transitions:
-                    if transition['source'] == c:
-                        dydt[i] -= rates[f"{c}_to_{transition['target']}"]
-                    if transition['target'] == c:
-                        dydt[i] += rates[f"{transition['source']}_to_{c}"]
-            return dydt
+            state = {c: y[i] for i, c in enumerate(self.all_compartments)}
+            derivatives = self.compute_derivatives(state)
+            return [derivatives[c] for c in self.all_compartments]
 
         sol = integrator(rhs, y0, t_range)
-        out = {c: sol[:, i] for i, c in enumerate(self.base_model.compartments)}
+        out = {c: sol[:, i] for i, c in enumerate(self.all_compartments)}
         return t_range, out

src/patchsim/models/sample_sir_ode.py

@@ -16,10 +16,10 @@
 
 def read_config(config_path: str) -> dict[str, str]:
     """Read and parse a YAML configuration file.
-    
+
     Args:
         config_path: Path to the YAML configuration file
-    
+
     Returns:
         Dictionary containing configuration parameters
     """