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