Add get_problem_yaml_path (#223)

Sometimes we want to have just the path instead of a petab.Problem.


Co-authored-by: Dilan Pathirana <59329744+dilpath@users.noreply.github.com>

Author: dweindl

src/python/benchmark_models_petab/__init__.py

@@ -4,6 +4,6 @@
 Python tool to access the model collection.
 """
 
-from .base import get_problem
+from .base import get_problem, get_problem_yaml_path
 from .C import MODEL_DIRS, MODELS, MODELS_DIR
 from .version import __version__

src/python/benchmark_models_petab/base.py

@@ -1,12 +1,31 @@
 """Get a petab problem from the collection."""
 
-import os
+from pathlib import Path
 
 import petab
 
 from .C import MODELS_DIR
 
 
+def get_problem_yaml_path(id_: str) -> Path:
+    """Get the path to the PEtab problem YAML file.
+
+    Parameters
+    ----------
+    id_: Problem name, as in `benchmark_models_petab.MODELS`.
+
+    Returns
+    -------
+    The path to the PEtab problem YAML file.
+    """
+    yaml_path = Path(MODELS_DIR, id_, id_ + ".yaml")
+    if not yaml_path.exists():
+        yaml_path = Path(MODELS_DIR, id_, "problem.yaml")
+    if not yaml_path.exists():
+        raise ValueError(f"Could not find YAML for problem with ID `{id_}`.")
+    return yaml_path
+
+
 def get_problem(id_: str) -> petab.Problem:
     """Read PEtab problem from benchmark collection by name.
 
@@ -18,6 +37,6 @@ def get_problem(id_: str) -> petab.Problem:
     -------
     The PEtab problem.
     """
-    yaml_file = os.path.join(MODELS_DIR, id_, id_ + ".yaml")
+    yaml_file = get_problem_yaml_path(id_)
     petab_problem = petab.Problem.from_yaml(yaml_file)
     return petab_problem