Update/use ProblemConfig (#353)

* Update/use ProblemConfig

* add to_yaml
* refactor model, avoid extra elements in object tree
* use ProblemConfig in petab1to2

* review

* format_version_tuple

Author: dweindl

petab/v1/problem.py

@@ -11,7 +11,7 @@
 from warnings import warn
 
 import pandas as pd
-from pydantic import AnyUrl, BaseModel, Field, RootModel
+from pydantic import AnyUrl, BaseModel, Field
 
 from ..versions import get_major_version
 from . import (
@@ -1185,33 +1185,14 @@ def add_measurement(
         )
 
 
-class VersionNumber(RootModel):
-    root: str | int
-
-
-class ListOfFiles(RootModel):
-    """List of files."""
-
-    root: list[str | AnyUrl] = Field(..., description="List of files.")
-
-    def __iter__(self):
-        return iter(self.root)
-
-    def __len__(self):
-        return len(self.root)
-
-    def __getitem__(self, index):
-        return self.root[index]
-
-
 class SubProblem(BaseModel):
     """A `problems` object in the PEtab problem configuration."""
 
-    sbml_files: ListOfFiles = []
-    measurement_files: ListOfFiles = []
-    condition_files: ListOfFiles = []
-    observable_files: ListOfFiles = []
-    visualization_files: ListOfFiles = []
+    sbml_files: list[str | AnyUrl] = []
+    measurement_files: list[str | AnyUrl] = []
+    condition_files: list[str | AnyUrl] = []
+    observable_files: list[str | AnyUrl] = []
+    visualization_files: list[str | AnyUrl] = []
 
 
 class ProblemConfig(BaseModel):
@@ -1227,6 +1208,16 @@ class ProblemConfig(BaseModel):
         description="The base path to resolve relative paths.",
         exclude=True,
     )
-    format_version: VersionNumber = 1
+    format_version: str | int = 1
     parameter_file: str | AnyUrl | None = None
     problems: list[SubProblem] = []
+
+    def to_yaml(self, filename: str | Path):
+        """Write the configuration to a YAML file.
+
+        :param filename: Destination file name. The parent directory will be
+            created if necessary.
+        """
+        from .yaml import write_yaml
+
+        write_yaml(self.model_dump(), filename)

petab/v2/petab1to2.py

@@ -10,10 +10,10 @@
 from pandas.io.common import get_handle, is_url
 
 from .. import v1, v2
-from ..v1 import Problem as ProblemV1
-from ..v1.yaml import get_path_prefix, load_yaml, validate, write_yaml
+from ..v1.yaml import get_path_prefix, load_yaml, validate
 from ..versions import get_major_version
 from .models import MODEL_TYPE_SBML
+from .problem import ProblemConfig
 
 __all__ = ["petab1to2"]
 
@@ -59,7 +59,7 @@ def petab1to2(yaml_config: Path | str, output_dir: Path | str = None):
     validate(yaml_config, path_prefix=path_prefix)
     if get_major_version(yaml_config) != 1:
         raise ValueError("PEtab problem is not version 1.")
-    petab_problem = ProblemV1.from_yaml(yaml_file or yaml_config)
+    petab_problem = v1.Problem.from_yaml(yaml_file or yaml_config)
     # get rid of conditionName column if present (unsupported in v2)
     petab_problem.condition_df = petab_problem.condition_df.drop(
         columns=[v1.C.CONDITION_NAME], errors="ignore"
@@ -71,26 +71,24 @@ def petab1to2(yaml_config: Path | str, output_dir: Path | str = None):
 
     # Update YAML file
     new_yaml_config = _update_yaml(yaml_config)
+    new_yaml_config = ProblemConfig(**new_yaml_config)
 
     # Update tables
     # condition tables, observable tables, SBML files, parameter table:
     #  no changes - just copy
     file = yaml_config[v2.C.PARAMETER_FILE]
     _copy_file(get_src_path(file), Path(get_dest_path(file)))
 
-    for problem_config in yaml_config[v2.C.PROBLEMS]:
+    for problem_config in new_yaml_config.problems:
         for file in chain(
-            problem_config.get(v2.C.OBSERVABLE_FILES, []),
-            (
-                model[v2.C.MODEL_LOCATION]
-                for model in problem_config.get(v2.C.MODEL_FILES, {}).values()
-            ),
-            problem_config.get(v2.C.VISUALIZATION_FILES, []),
+            problem_config.observable_files,
+            (model.location for model in problem_config.model_files.values()),
+            problem_config.visualization_files,
         ):
             _copy_file(get_src_path(file), Path(get_dest_path(file)))
 
         # Update condition table
-        for condition_file in problem_config.get(v2.C.CONDITION_FILES, []):
+        for condition_file in problem_config.condition_files:
             condition_df = v1.get_condition_df(get_src_path(condition_file))
             condition_df = v1v2_condition_df(condition_df, petab_problem.model)
             v2.write_condition_df(condition_df, get_dest_path(condition_file))
@@ -159,12 +157,12 @@ def create_experiment_id(sim_cond_id: str, preeq_cond_id: str) -> str:
                 raise ValueError(
                     f"Experiment table file {exp_table_path} already exists."
                 )
-            problem_config[v2.C.EXPERIMENT_FILES] = [exp_table_path.name]
+            problem_config.experiment_files.append("experiments.tsv")
             v2.write_experiment_df(
                 v2.get_experiment_df(pd.DataFrame(experiments)), exp_table_path
             )
 
-        for measurement_file in problem_config.get(v2.C.MEASUREMENT_FILES, []):
+        for measurement_file in problem_config.measurement_files:
             measurement_df = v1.get_measurement_df(
                 get_src_path(measurement_file)
             )
@@ -221,7 +219,7 @@ def create_experiment_id(sim_cond_id: str, preeq_cond_id: str) -> str:
 
     # Write new YAML file
     new_yaml_file = output_dir / Path(yaml_file).name
-    write_yaml(new_yaml_config, new_yaml_file)
+    new_yaml_config.to_yaml(new_yaml_file)
 
     # validate updated Problem
     validation_issues = v2.lint_problem(new_yaml_file)

petab/v2/problem.py

@@ -25,7 +25,6 @@
     yaml,
 )
 from ..v1.models.model import Model, model_factory
-from ..v1.problem import ListOfFiles, VersionNumber
 from ..v1.yaml import get_path_prefix
 from ..v2.C import *  # noqa: F403
 from ..versions import parse_version
@@ -994,13 +993,14 @@ class ModelFile(BaseModel):
 class SubProblem(BaseModel):
     """A `problems` object in the PEtab problem configuration."""
 
+    # TODO: consider changing str to Path
     model_files: dict[str, ModelFile] | None = {}
-    measurement_files: ListOfFiles = []
-    condition_files: ListOfFiles = []
-    experiment_files: ListOfFiles = []
-    observable_files: ListOfFiles = []
-    visualization_files: ListOfFiles = []
-    mapping_files: ListOfFiles = []
+    measurement_files: list[str | AnyUrl] = []
+    condition_files: list[str | AnyUrl] = []
+    experiment_files: list[str | AnyUrl] = []
+    observable_files: list[str | AnyUrl] = []
+    visualization_files: list[str | AnyUrl] = []
+    mapping_files: list[str | AnyUrl] = []
 
 
 class ExtensionConfig(BaseModel):
@@ -1024,7 +1024,23 @@ class ProblemConfig(BaseModel):
         description="The base path to resolve relative paths.",
         exclude=True,
     )
-    format_version: VersionNumber = "2.0.0"
+    format_version: str = "2.0.0"
     parameter_file: str | AnyUrl | None = None
     problems: list[SubProblem] = []
     extensions: list[ExtensionConfig] = []
+
+    def to_yaml(self, filename: str | Path):
+        """Write the configuration to a YAML file.
+
+        :param filename: Destination file name. The parent directory will be
+            created if necessary.
+        """
+        from ..v1.yaml import write_yaml
+
+        write_yaml(self.model_dump(), filename)
+
+    @property
+    def format_version_tuple(self) -> tuple[int, int, int, str]:
+        """The format version as a tuple of major/minor/patch `int`s and a
+        suffix."""
+        return parse_version(self.format_version)