Benchtop 0.1.0__py3-none-any.whl

benchtop/__init__.py

@@ -0,0 +1,6 @@
+"""Benchtop: parallel in-silico experiments for biological models."""
+
+from .experiment import Experiment
+
+__version__ = "0.1.0"
+__all__ = ["Experiment"]

benchtop/_abstract_simulator.py

@@ -0,0 +1,24 @@
+"""Abstract base class for simulator wrappers (Tellurium, AMICI, etc.)."""
+
+from abc import ABC, abstractmethod
+from types import ModuleType
+
+
+class AbstractSimulator(ABC):
+    """Interface: load model, modify state, simulate over a time grid."""
+
+    def __init__(self, *args, **kwargs):
+        self.tool = type("Tool", (), {})()
+        self.load(*args, **kwargs)
+
+    @abstractmethod
+    def load(self, *args, **kwargs) -> ModuleType:
+        """Load or compile the model from constructor arguments."""
+
+    @abstractmethod
+    def modify(self, component: str, value: float) -> None:
+        """Set a parameter or species value before simulation."""
+
+    @abstractmethod
+    def simulate(self, start: float, stop: float, step: float):
+        """Integrate from start to stop with given step; return trajectory."""

benchtop/_observable_calculator.py

@@ -0,0 +1,190 @@
+"""Evaluate PEtab observable formulas and align simulation with measurements."""
+
+import math
+import re
+import gc
+from typing import List
+
+import numpy as np
+import pandas as pd
+
+
+class ObservableCalculator:
+    """Downsample trajectories, evaluate formulas, pair with experimental data."""
+
+    def __init__(self, parent):
+        self.parent = parent
+        self.problem_name = parent.record.current_problem_name
+        self.results_dict = {
+            key: entry
+            for key, entry in parent.record.cache.results_dict.items()
+            if key != parent.record.cache.PROBLEMS_META_KEY
+            and entry.get("problem") == self.problem_name
+        }
+        self.cache = parent.record.cache
+        self.observable_df = parent.record.problem.observable_files[0]
+        self.measurement_df = parent.record.problem.measurement_files[0]
+        self.data_groups = self._group_conditions_and_observables()
+        self.observable_results = self._build_observable_results_dict()
+
+    def _group_conditions_and_observables(self) -> pd.core.groupby.generic.DataFrameGroupBy:
+        if self.measurement_df.empty or self.observable_df.empty:
+            raise ValueError("PEtab DataFrame is empty; cannot group.")
+
+        return self.measurement_df.groupby(["simulationConditionId", "observableId"])
+
+    def _build_observable_results_dict(self) -> dict:
+        return {
+            entry: {
+                "problem": self.problem_name,
+                "conditionId": self.results_dict[entry]["conditionId"],
+                "cell": self.results_dict[entry]["cell"],
+            }
+            for entry in self.results_dict
+        }
+
+    def run(self) -> dict:
+        """Compute observables for every cached simulation entry in the current problem."""
+        for entry in self.results_dict:
+            condition_id = self.results_dict[entry]["conditionId"]
+            matched_formulas = self._get_entry_formulas(condition_id)
+            needed_columns = self._columns_for_formulas(matched_formulas)
+
+            dataset = self.cache.load_columns(entry, needed_columns)
+
+            for observable_key, formula in matched_formulas.items():
+                self.observable_results[entry][observable_key] = {}
+                group = self.data_groups.get_group((condition_id, observable_key))
+
+                self.observable_results[entry][observable_key]["experiment"] = (
+                    self._get_experimental_data(group)
+                )
+                self.observable_results[entry][observable_key]["simulation"] = (
+                    self._calculate_formula(dataset, formula, group)
+                )
+                self.observable_results[entry][observable_key]["time"] = (
+                    self._downsample_timepoints(dataset, group)
+                )
+            
+            # Dump unnecessary data after calculate
+            del dataset
+            gc.collect()
+        
+        return self.observable_results
+
+    def _columns_for_formulas(self, matched_formulas: dict) -> List[str]:
+        species: Set[str] = set()
+        for formula in matched_formulas.values():
+            null_like = {"", None, 0, "0", float("nan"), np.nan, "nan"}
+            if formula in null_like or (
+                isinstance(formula, float) and math.isnan(formula)
+            ):
+                continue
+            species.update(self._get_valid_species(formula))
+
+        return sorted(species | {"time"})
+
+    def _get_entry_formulas(self, condition_id: str) -> dict:
+        matched_obs_ids = self._get_condition_observables(condition_id)
+        return {
+            obs_id: self.observable_df["observableFormula"][
+                self.observable_df["observableId"] == obs_id
+            ].iloc[0]
+            for obs_id in matched_obs_ids
+        }
+
+    def _get_condition_observables(self, condition_id: str) -> list:
+        return [
+            obs
+            for (cond, obs) in self.data_groups.groups
+            if cond == condition_id
+        ]
+
+    def _get_experimental_data(self, group) -> np.ndarray:
+        return np.array(group["measurement"])
+
+    def _calculate_formula(self, dataset: pd.DataFrame, formula: str, group) -> np.ndarray:
+        """Substitute species arrays into formula and eval; downsample to exp times."""
+        null_like = {"", None, 0, "0", float("nan"), np.nan, "nan"}
+        if formula in null_like or (isinstance(formula, float) and math.isnan(formula)):
+            return None
+
+        species_names = self._get_valid_species(formula)
+        namespace = self._formula_namespace(species_names, dataset)
+
+        try:
+            formula_answer = eval(formula, {"np": np}, namespace)
+        except Exception:
+            raise RuntimeError(f"Failed to evaluate observable formula: {formula}")
+
+        return self._downsample_results(formula_answer, dataset, group)
+
+    @staticmethod
+    def _get_valid_species(formula: str) -> List[str]:
+        """Extract PEtab-compliant species identifiers from a formula string."""
+        if not isinstance(formula, str):
+            raise TypeError("Observable formula must be a string.")
+
+        valid_species = re.findall(
+            r"(?:@[A-Za-z_]+::[A-Za-z_]\w*|[A-Za-z_]\w*)(?:\(\))?",
+            formula,
+        )
+        if not valid_species:
+            raise ValueError("No valid species found in the observable formula.")
+
+        return valid_species
+
+    def _formula_namespace(self, species_names: list, dataset: pd.DataFrame) -> dict:
+        """Map species names to numpy arrays for safe eval."""
+        return {
+            name: np.asarray(self._safe_retrieve_array(dataset, name))
+            for name in species_names
+        }
+
+    @staticmethod
+    def _safe_retrieve_array(dataset, species_name) -> np.ndarray:
+        species_arr = dataset[species_name]
+        if species_arr is None:
+            raise KeyError(f"Species '{species_name}' not found in dataset.")
+
+        if isinstance(species_arr, (pd.Series, pd.DataFrame)):
+            species_arr = species_arr.to_numpy().ravel()
+        if not isinstance(species_arr, np.ndarray):
+            raise ValueError(
+                f"Replacement value for species '{species_name}' is not a valid array."
+            )
+
+        return species_arr
+
+    def _downsample_results(
+        self, observable_answer: np.ndarray, dataset: pd.DataFrame, group
+    ) -> np.ndarray:
+        """Slice simulation values to experimental measurement timepoints."""
+        valid_rows = group.dropna(subset=["measurement"])
+        if valid_rows.empty:
+            return observable_answer
+
+        exp_time = np.sort(
+            valid_rows["time"].dropna().astype(float).unique()
+        )
+        sim_indices = self._get_exp_time_indices(exp_time, dataset["time"])
+        return observable_answer[np.sort(sim_indices)]
+
+    @staticmethod
+    def _get_exp_time_indices(exp_time: np.ndarray, sim_time: np.ndarray) -> list:
+        """Index of closest simulation time for each experimental timepoint."""
+        return [
+            int(np.argmin(np.abs(sim_time - t)))
+            for t in np.sort(exp_time)
+        ]
+
+    def _downsample_timepoints(self, dataset: pd.DataFrame, group) -> np.ndarray:
+        valid_rows = group.dropna(subset=["measurement"])
+        if valid_rows.empty:
+            return dataset["time"]
+
+        exp_time = np.sort(
+            valid_rows["time"].dropna().astype(float).unique()
+        )
+        indices = self._get_exp_time_indices(exp_time, dataset["time"])
+        return np.unique(np.array(dataset["time"][indices]))

benchtop/_organizer.py

@@ -0,0 +1,147 @@
+"""Task scheduling: topological ordering, cell replication, and round-robin assignment."""
+
+import os
+from collections import defaultdict, deque
+
+import pandas as pd
+
+
+class Organizer:
+    """Build and assign simulation tasks respecting preequilibration dependencies."""
+
+    def __init__(self, workers: int = os.cpu_count()):
+        self.workers = workers
+
+    def task_organization(
+        self,
+        measurement_df: pd.DataFrame,
+        cell_count: int,
+    ) -> tuple[int, dict]:
+        """Return (num_rounds, rank_jobs_directory) for the full experiment."""
+        size = self.workers
+
+        topological_list = self.topologic_sort(measurements_df=measurement_df)
+        total_tasks = self.total_tasks(tasks=topological_list, cell_count=cell_count)
+        delayed_list = self.delay_secondary_conditions(
+            measurements_df=measurement_df,
+            task_list=total_tasks,
+            cell_count=cell_count,
+        )
+
+        rank_jobs_directory = {}
+        for i in range(size):
+            rank_ids = self.assign_tasks(i, len(delayed_list))
+            rank_jobs_directory[i] = [delayed_list[job] for job in rank_ids]
+
+        rounds_to_complete = -(-len(delayed_list) // size)  # ceiling division
+        return rounds_to_complete, rank_jobs_directory
+
+    def topologic_sort(self, measurements_df: pd.DataFrame) -> list:
+        """Kahn topological sort of conditions; preequilibration edges first."""
+        if "preequilibrationConditionId" not in measurements_df.columns:
+            return measurements_df["simulationConditionId"].dropna().unique().tolist()
+
+        sim_nodes = measurements_df["simulationConditionId"].dropna().unique().tolist()
+        pre_nodes = measurements_df["preequilibrationConditionId"].dropna().unique().tolist()
+        nodes = sorted(set(sim_nodes) | set(pre_nodes))
+
+        succs = defaultdict(list)
+        indegree = {n: 0 for n in nodes}
+
+        for _, row in measurements_df.dropna(
+            subset=["preequilibrationConditionId"]
+        ).iterrows():
+            pre = row["preequilibrationConditionId"]
+            sim = row["simulationConditionId"]
+            succs[pre].append(sim)
+            indegree[sim] += 1
+
+        for k in succs:
+            succs[k].sort()
+
+        queue = deque(sorted(n for n, d in indegree.items() if d == 0))
+        ordered = []
+
+        while queue:
+            n = queue.popleft()
+            ordered.append(n)
+
+            for m in succs[n]:
+                indegree[m] -= 1
+                if indegree[m] == 0:
+                    queue.append(m)
+
+            queue = deque(sorted(queue))
+
+        if len(ordered) != len(nodes):
+            raise RuntimeError("Circular dependency detected among conditions!")
+
+        return ordered
+
+    def delay_secondary_conditions(
+        self,
+        measurements_df: pd.DataFrame,
+        task_list: list,
+        cell_count: int,
+    ) -> list:
+        """Insert None padding so preequilibration completes before dependents."""
+        if "preequilibrationConditionId" not in measurements_df.columns:
+            return task_list
+
+        pre_conds = (
+            measurements_df["preequilibrationConditionId"]
+            .drop_duplicates()
+            .dropna()
+            .to_list()
+        )
+
+        for idx, job in enumerate(task_list):
+            if job is None:
+                continue
+
+            cond_id = job.split("+")[0]
+            if cond_id not in pre_conds:
+                continue
+
+            pause_ranks = max(self.workers - cell_count, 0)
+            while pause_ranks:
+                task_list.insert(idx + cell_count, None)
+                pause_ranks -= 1
+
+            pre_conds.pop(pre_conds.index(cond_id))
+
+        return task_list
+
+    def total_tasks(self, tasks: list, cell_count: int) -> list:
+        """Expand conditions into ``conditionId+cell`` task strings."""
+        return [
+            f"{cond}+{cell}"
+            for cond in tasks
+            for cell in range(1, cell_count + 1)
+        ]
+
+    def assign_tasks(self, rank: int, total_jobs: int) -> list:
+        """Round-robin job indices assigned to a single worker rank."""
+        size = self.workers
+        num_rounds = -(-total_jobs // size)
+
+        rank_jobs = []
+        for round_index in range(num_rounds):
+            job_id = rank + round_index * size
+            if job_id < total_jobs:
+                rank_jobs.append(job_id)
+
+        return rank_jobs
+
+    def task_assignment(self, rank_jobs_directory: dict, round_i: int) -> list:
+        """Collect one task per worker for the given round."""
+        round_i_tasks = []
+
+        for i in range(self.workers):
+            rank_jobs = rank_jobs_directory[i]
+            if round_i < len(rank_jobs):
+                round_i_tasks.append(rank_jobs[round_i])
+            else:
+                round_i_tasks.append(None)
+
+        return round_i_tasks

benchtop/_record.py

@@ -0,0 +1,204 @@
+"""Results index: maps condition/cell pairs to cache keys and lookup helpers."""
+
+import os
+import logging
+import uuid
+from types import SimpleNamespace
+from typing import List, Union
+
+import pandas as pd
+
+from benchtop._results_cacher import ResultCache, DEFAULT_CACHE
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+
+class Record:
+    """Shared results dictionary accessed by workers across processes."""
+
+    def __init__(
+        self,
+        problems: Union[List[SimpleNamespace], SimpleNamespace],
+        cache_dir: Union[os.PathLike, str] = DEFAULT_CACHE,
+        load_index: bool = False,
+    ) -> None:
+        if isinstance(problems, list):
+            self.problems = problems
+        else:
+            self.problems = [problems]
+
+        self.problem = self.problems[0]
+        self.current_problem_name = self.problem.name
+
+        if not load_index:
+            results_dict = self._results_dictionary()
+        else:
+            results_dict = None
+
+        self.cache = ResultCache(
+            results_dict=results_dict,
+            cache_dir=cache_dir,
+            load_index=load_index,
+            problem_names=[p.name for p in self.problems],
+        )
+
+        if load_index:
+            merged = self._merge_loaded_index(self.cache.results_dict)
+            self.cache.results_dict = merged
+            self.cache._write_cache_index()
+
+    def set_current_problem(
+        self, problem: SimpleNamespace, problem_name: str | None = None
+    ) -> None:
+        """Point workers and lookups at the problem currently being simulated."""
+        self.problem = problem
+        self.current_problem_name = problem_name or problem.name
+
+    def _results_dictionary(self) -> dict:
+        """Build initial index: one entry per problem × condition × cell replicate."""
+        jobs = self._expected_jobs()
+        return {key: entry for key, entry in jobs.items()}
+
+    def _expected_jobs(self) -> dict:
+        """Return job entries keyed by stable cache identifiers."""
+        results = {}
+
+        for problem in self.problems:
+            conditions_df = problem.condition_files[0]
+            measurement_df = problem.measurement_files[0]
+
+            for _, condition in conditions_df.iterrows():
+                condition_id = condition["conditionId"]
+
+                for cell in range(1, problem.cell_count + 1):
+                    if "datasetId" in measurement_df.columns:
+                        identifier = measurement_df["datasetId"][
+                            measurement_df["simulationConditionId"] == condition_id
+                        ].values[0]
+                    else:
+                        identifier = self._identifier_generator()
+
+                    results[identifier] = {
+                        "problem": problem.name,
+                        "conditionId": condition_id,
+                        "cell": cell,
+                        "complete": False,
+                    }
+
+        return results
+
+    def _job_lookup_key(self, entry: dict) -> tuple:
+        return (
+            entry.get("problem"),
+            str(entry["conditionId"]),
+            str(entry["cell"]),
+        )
+
+    def _merge_loaded_index(self, loaded: dict) -> dict:
+        """Align a loaded cache index with the current benchmark configuration."""
+        expected_jobs = self._expected_jobs()
+        problem_names = [p.name for p in self.problems]
+
+        meta = loaded.get(ResultCache.PROBLEMS_META_KEY, {})
+        for name in problem_names:
+            meta.setdefault(name, {"complete": False})
+
+        loaded_by_job = {}
+        for key in loaded:
+            if key == ResultCache.PROBLEMS_META_KEY:
+                continue
+            entry = loaded[key]
+            loaded_by_job[self._job_lookup_key(entry)] = (key, entry)
+
+        merged = {ResultCache.PROBLEMS_META_KEY: meta}
+        used_keys = set()
+
+        for _, expected_entry in expected_jobs.items():
+            job_key = self._job_lookup_key(expected_entry)
+            if job_key in loaded_by_job:
+                cache_key, loaded_entry = loaded_by_job[job_key]
+                merged[cache_key] = loaded_entry
+                used_keys.add(cache_key)
+            else:
+                identifier = self._identifier_generator()
+                merged[identifier] = expected_entry
+
+        for job_key, (cache_key, loaded_entry) in loaded_by_job.items():
+            if cache_key not in used_keys:
+                merged[cache_key] = loaded_entry
+
+        return merged
+
+    def incomplete_tasks_for_problem(self, problem_name: str) -> list[str]:
+        """Return ``conditionId+cell`` task strings not yet marked complete."""
+        incomplete = []
+        for key in self.cache.job_keys():
+            entry = self.cache.results_dict[key]
+            if not self._job_belongs_to_problem(entry, problem_name):
+                continue
+            if not entry["complete"]:
+                incomplete.append(f"{entry['conditionId']}+{entry['cell']}")
+        return incomplete
+
+    @staticmethod
+    def _job_belongs_to_problem(entry: dict, problem_name: str) -> bool:
+        entry_problem = entry.get("problem")
+        if entry_problem is None:
+            return True
+        return entry_problem == problem_name
+
+    def find_job_key(
+        self, condition_id: str, cell: int, problem_name: str | None = None
+    ) -> str | None:
+        """Return cache key for a condition/cell pair within a problem."""
+        problem_name = problem_name or self.current_problem_name
+
+        for key in self.cache.job_keys():
+            entry = self.cache.results_dict[key]
+            entry_problem = entry.get("problem")
+            if entry_problem is not None and entry_problem != problem_name:
+                continue
+            if (
+                str(entry["conditionId"]) == str(condition_id)
+                and str(entry["cell"]) == str(cell)
+            ):
+                return key
+
+        return None
+
+    def results_lookup(
+        self,
+        condition_id: str,
+        cell: int,
+        problem_name: str | None = None,
+    ) -> pd.DataFrame | None:
+        """Load cached trajectory for a condition/cell pair."""
+        key = self.find_job_key(condition_id, cell, problem_name)
+        if key is None:
+            logger.error(
+                "No prior results found for %s at cell %s (problem=%s)",
+                condition_id,
+                cell,
+                problem_name or self.current_problem_name,
+            )
+            return None
+
+        logger.debug("Results found for %s, cell %s", condition_id, cell)
+        return self.cache.load(key)
+
+    def condition_cell_id(self, rank_task: str, conditions_df) -> tuple:
+        """Parse ``conditionId+cell`` task string into condition row and IDs."""
+        condition_id, cell = rank_task.split("+")
+
+        matches = conditions_df.loc[conditions_df["conditionId"] == condition_id]
+        if matches.empty:
+            raise ValueError(f"Condition ID '{condition_id}' not found in conditions_df")
+
+        return matches.iloc[0], cell, condition_id
+
+    def _identifier_generator(self) -> str:
+        return str(uuid.uuid4())