refactor sciml specifics into more helpers

Author: BSnelling

petab/v2/core.py

@@ -312,7 +312,6 @@ def __iadd__(self, other: T) -> BaseTable[T]:
 # SciML extension classes — imported after BaseTable is defined to avoid
 # circular imports (sciml.py does not import from core.py).
 from .extensions.sciml import (  # noqa: E402
-    HybridizationTable,
     SciMLConfig,
     SciMLExt,
 )
@@ -1350,44 +1349,12 @@ def from_yaml(
             else None
         )
 
-        extensions = ProblemExtensions()
-        if config.extensions and config.extensions.get(C.EXT_ID_SCIML):
-            from petab_sciml import ArrayDataStandard, NNModel, NNModelStandard
-
-            # Neural network classes are constructed via pytorch for now to get
-            # the proper inputs
-            neural_networks = [
-                NNModel.from_pytorch_module(
-                    NNModelStandard.load_data(
-                        _generate_path(
-                            file_path=nn_config.location,
-                            base_path=base_path,
-                        )
-                    ).to_pytorch_module(),
-                    nn_model_id=nn_id,
-                )
-                for nn_id, nn_config in (
-                    config.extensions[C.EXT_ID_SCIML].neural_networks or {}
-                ).items()
-            ]
-
-            hybridization_tables = [
-                HybridizationTable.from_tsv(f, base_path)
-                for f in config.extensions[C.EXT_ID_SCIML].hybridization_files
-            ]
-
-            array_data_files = [
-                ArrayDataStandard.load_data(_generate_path(f, base_path))
-                for f in config.extensions[C.EXT_ID_SCIML].array_files
-            ]
-
-            extensions = ProblemExtensions(
-                sciml=SciMLExt(
-                    neural_networks=neural_networks,
-                    hybridization_tables=hybridization_tables,
-                    array_data_files=array_data_files,
-                )
-            )
+        sciml = (
+            SciMLExt.from_config(config, base_path)
+            if config.extensions and config.extensions.get(C.EXT_ID_SCIML)
+            else None
+        )
+        extensions = ProblemExtensions(sciml=sciml)
 
         return Problem(
             config=config,
@@ -2647,15 +2614,9 @@ def to_yaml(self, filename: str | Path):
             # The schema requires a valid id or no id field at all.
             del data["id"]
 
-        for ext_id, d_ext in data[C.EXTENSIONS].items():
+        for ext_id in list(data[C.EXTENSIONS]):
             if ext_id == C.EXT_ID_SCIML:
-                # convert Paths to strings
-                for key in ("array_files", "hybridization_files"):
-                    d_ext[key] = list(map(str, d_ext[key]))
-                for nn in d_ext["neural_networks"]:
-                    d_ext["neural_networks"][nn][C.MODEL_LOCATION] = str(
-                        d_ext["neural_networks"][nn][C.MODEL_LOCATION]
-                    )
+                data[C.EXTENSIONS][ext_id] = self.extensions[ext_id].to_yaml()
 
         write_yaml(data, filename)
 

petab/v2/extensions/sciml.py

@@ -153,6 +153,19 @@ class SciMLConfig(BaseModel):
         validate_assignment=True,
     )
 
+    def to_yaml(self) -> dict:
+        """Return a YAML-serializable dict with Paths converted to strings."""
+        from . import C
+
+        d = self.model_dump(by_alias=True)
+        for key in ("array_files", "hybridization_files"):
+            d[key] = list(map(str, d[key]))
+        for nn in d["neural_networks"] or {}:
+            d["neural_networks"][nn][C.MODEL_LOCATION] = str(
+                d["neural_networks"][nn][C.MODEL_LOCATION]
+            )
+        return d
+
 
 class SciMLExt:
     """SciML extension runtime state.
@@ -241,3 +254,52 @@ def add_array_data_from_hdf5(
         self.array_data_files.append(
             ArrayDataStandard.load_data(_generate_path(file_path, base_path))
         )
+
+    @staticmethod
+    def from_config(
+        config,
+        base_path: str | Path | None = None,
+    ) -> SciMLExt:
+        """Construct a SciMLExt from a ProblemConfig.
+
+        Arguments:
+            config: A ProblemConfig whose ``extensions[EXT_ID_SCIML]`` entry
+                is a :class:`SciMLConfig`.
+            base_path: Base path used to resolve relative file paths.
+        """
+        from petab_sciml import ArrayDataStandard, NNModel, NNModelStandard
+
+        sciml_config: SciMLConfig = config.extensions[C.EXT_ID_SCIML]
+
+        # Neural network classes are constructed via pytorch for now to get
+        # the proper inputs
+        neural_networks = [
+            NNModel.from_pytorch_module(
+                NNModelStandard.load_data(
+                    _generate_path(
+                        file_path=nn_config.location,
+                        base_path=base_path,
+                    )
+                ).to_pytorch_module(),
+                nn_model_id=nn_id,
+            )
+            for nn_id, nn_config in (
+                sciml_config.neural_networks or {}
+            ).items()
+        ]
+
+        hybridization_tables = [
+            HybridizationTable.from_tsv(f, base_path)
+            for f in sciml_config.hybridization_files
+        ]
+
+        array_data_files = [
+            ArrayDataStandard.load_data(_generate_path(f, base_path))
+            for f in sciml_config.array_files
+        ]
+
+        return SciMLExt(
+            neural_networks=neural_networks,
+            hybridization_tables=hybridization_tables,
+            array_data_files=array_data_files,
+        )

tests/v2/test_sciml.py

@@ -4,37 +4,36 @@
 from petab.v2.core import *
 from petab.v2.core import ModelFile
 from petab.v2.extensions.sciml import NeuralNetConfig, SciMLConfig
-from petab.v2.lint import sciml_validation_tasks
 from petab.v2.models.sbml_model import SbmlModel
 
 
 def _get_test_problem():
-    problem = Problem()
-    problem.validation_tasks = sciml_validation_tasks
-    problem.config = ProblemConfig(
-        format_version="2.0.0",
-        model_files=ConfigDict(
-            {"lv": ModelFile(location="lv.xml", language="sbml")}
-        ),
-        parameter_files=["parameters.tsv"],
-        measurement_files=["measurements.tsv"],
-        observable_files=["observables.tsv"],
-        experiment_files=["experiments.tsv"],
-        mapping_files=["mappings.tsv"],
-        extensions={
-            "sciml": SciMLConfig(
-                version="0.1.0",
-                array_files=["net1_ps.hdf5"],
-                hybridization_files=["hybridizations.tsv"],
-                neural_networks={
-                    "net1": NeuralNetConfig(
-                        location="net1.yaml",
-                        pre_initialization=False,
-                        format="YAML",
-                    )
-                },
-            )
-        },
+    problem = Problem(
+        config=ProblemConfig(
+            format_version="2.0.0",
+            model_files=ConfigDict(
+                {"lv": ModelFile(location="lv.xml", language="sbml")}
+            ),
+            parameter_files=["parameters.tsv"],
+            measurement_files=["measurements.tsv"],
+            observable_files=["observables.tsv"],
+            experiment_files=["experiments.tsv"],
+            mapping_files=["mappings.tsv"],
+            extensions={
+                "sciml": SciMLConfig(
+                    version="0.1.0",
+                    array_files=["net1_ps.hdf5"],
+                    hybridization_files=["hybridizations.tsv"],
+                    neural_networks={
+                        "net1": NeuralNetConfig(
+                            location="net1.yaml",
+                            pre_initialization=False,
+                            format="YAML",
+                        )
+                    },
+                )
+            },
+        )
     )
     problem.model = SbmlModel.from_antimony("""
     model lv