From dc54c1288122256d3a5d9c36d89f12bd9446189e Mon Sep 17 00:00:00 2001
From: Daniel Weindl <daniel.weindl@uni-bonn.de>
Date: Thu, 16 Oct 2025 21:42:15 +0200
Subject: [PATCH] Add `get_simulation_df`

Add `get_simulation_df` to load simulation `DataFrame`s.
---
 src/python/benchmark_models_petab/__init__.py |  2 +-
 src/python/benchmark_models_petab/base.py     | 20 +++++++++++++++++++
 src/python/test/test_base.py                  |  5 +++++
 3 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/python/benchmark_models_petab/__init__.py b/src/python/benchmark_models_petab/__init__.py
index 66be078f..d6a210d6 100644
--- a/src/python/benchmark_models_petab/__init__.py
+++ b/src/python/benchmark_models_petab/__init__.py
@@ -4,7 +4,7 @@
 Python tool to access the model collection.
 """
 
-from .base import get_problem, get_problem_yaml_path
+from .base import get_problem, get_problem_yaml_path, get_simulation_df
 from .C import MODEL_DIRS, MODELS, MODELS_DIR
 from .overview import get_overview_df
 from importlib.metadata import PackageNotFoundError, version
diff --git a/src/python/benchmark_models_petab/base.py b/src/python/benchmark_models_petab/base.py
index 65507534..0a4a3246 100644
--- a/src/python/benchmark_models_petab/base.py
+++ b/src/python/benchmark_models_petab/base.py
@@ -6,6 +6,8 @@
 
 from .C import MODELS_DIR
 
+import pandas as pd
+
 
 def get_problem_yaml_path(id_: str) -> Path:
     """Get the path to the PEtab problem YAML file.
@@ -40,3 +42,21 @@ def get_problem(id_: str) -> petab.Problem:
     yaml_file = get_problem_yaml_path(id_)
     petab_problem = petab.Problem.from_yaml(yaml_file)
     return petab_problem
+
+
+def get_simulation_df(id_: str) -> pd.DataFrame | None:
+    """Get the simulation dataframe for the benchmark collection problem with
+    the given name.
+
+    Parameters
+    ----------
+    id_: Problem name, as in `benchmark_models_petab.MODELS`.
+
+    Returns
+    -------
+    The simulation dataframe if it exists, else None.
+    """
+    path = Path(MODELS_DIR, id_, f"simulatedData_{id_}.tsv")
+    if path.is_file():
+        return petab.get_simulation_df(path)
+    return None
diff --git a/src/python/test/test_base.py b/src/python/test/test_base.py
index 38082f39..282a4f6e 100644
--- a/src/python/test/test_base.py
+++ b/src/python/test/test_base.py
@@ -14,3 +14,8 @@ def test_get_problem():
     """Test whether extracting a petab problem works."""
     problem = models.get_problem(models.MODELS[0])
     assert problem.measurement_df is not None
+
+
+def test_get_simulation_df():
+    assert models.get_simulation_df("Elowitz_Nature2000").empty is False
+    assert models.get_simulation_df("not a problem name") is None