updates for overhauled testsuite

Author: FFroehlich

.github/workflows/test_petab_sciml.yml

@@ -22,7 +22,7 @@ jobs:
 
     strategy:
       matrix:
-        python-version: ["3.11"]
+        python-version: ["3.12"]
 
     steps:
       - name: Set up Python ${{ matrix.python-version }}
@@ -56,13 +56,10 @@ jobs:
             && pip3 install wheel pytest shyaml pytest-cov
 
       # retrieve test models
-      - name: Download and install PEtab SciML test suite
+      - name: Download and install PEtab SciML
         run: |
-          git clone --depth 1 --branch main \
-            https://github.com/sebapersson/petab_sciml.git \
-            && export SCIML_TESTSUITE="$(pwd)/petab_sciml" \
-            && source venv/bin/activate \
-            && python -m pip install -e $SCIML_TESTSUITE/src/python
+          source ./venv/bin/activate \
+            && python -m pip install git+https://github.com/sebapersson/petab_sciml.git@unify_data#subdirectory=src/python \
 
 
       - name: Install petab

python/sdist/amici/jax/nn.py

@@ -30,22 +30,22 @@ def tanhshrink(x: jnp.ndarray) -> jnp.ndarray:
     return x - jnp.tanh(x)
 
 
-def generate_equinox(ml_model: "MLModel", filename: Path | str):  # noqa: F821
+def generate_equinox(nn_model: "NNModel", filename: Path | str):  # noqa: F821
     # TODO: move to top level import and replace forward type definitions
     from petab_sciml import Layer
 
     filename = Path(filename)
     layer_indent = 12
     node_indent = 8
 
-    layers = {layer.layer_id: layer for layer in ml_model.layers}
+    layers = {layer.layer_id: layer for layer in nn_model.layers}
 
     tpl_data = {
-        "MODEL_ID": ml_model.mlmodel_id,
+        "MODEL_ID": nn_model.nn_model_id,
         "LAYERS": ",\n".join(
             [
                 _generate_layer(layer, layer_indent, ilayer)
-                for ilayer, layer in enumerate(ml_model.layers)
+                for ilayer, layer in enumerate(nn_model.layers)
             ]
         )[layer_indent:],
         "FORWARD": "\n".join(
@@ -58,19 +58,19 @@ def generate_equinox(ml_model: "MLModel", filename: Path | str):  # noqa: F821
                         Layer(layer_id="dummy", layer_type="Linear"),
                     ).layer_type,
                 )
-                for node in ml_model.forward
+                for node in nn_model.forward
             ]
         )[node_indent:],
-        "INPUT": ", ".join([f"'{inp.input_id}'" for inp in ml_model.inputs]),
+        "INPUT": ", ".join([f"'{inp.input_id}'" for inp in nn_model.inputs]),
         "OUTPUT": ", ".join(
             [
                 f"'{arg}'"
                 for arg in next(
-                    node for node in ml_model.forward if node.op == "output"
+                    node for node in nn_model.forward if node.op == "output"
                 ).args
             ]
         ),
-        "N_LAYERS": len(ml_model.layers),
+        "N_LAYERS": len(nn_model.layers),
     }
 
     filename.parent.mkdir(parents=True, exist_ok=True)

python/sdist/amici/jax/ode_export.py

@@ -267,11 +267,10 @@ def _generate_jax_code(self) -> None:
 
     def _generate_nn_code(self) -> None:
         for net_name, net in self.hybridisation.items():
-            for model in net["model"]:
-                generate_equinox(
-                    model,
-                    self.model_path / f"{net_name}.py",
-                )
+            generate_equinox(
+                net["model"],
+                self.model_path / f"{net_name}.py",
+            )
 
     def set_paths(self, output_dir: str | Path | None = None) -> None:
         """

python/sdist/amici/jax/petab.py

@@ -6,6 +6,7 @@
 from collections.abc import Sized, Iterable
 from pathlib import Path
 from collections.abc import Callable
+import logging
 
 
 import diffrax
@@ -23,6 +24,7 @@
     ParameterMappingForCondition,
     create_parameter_mapping,
 )
+from amici.logging import get_logger
 from amici.jax.model import JAXModel, ReturnValue
 
 DEFAULT_CONTROLLER_SETTINGS = {
@@ -39,6 +41,8 @@
     petab.LOG10: 2,
 }
 
+logger = get_logger(__name__, logging.WARNING)
+
 
 def jax_unscale(
     parameter: jnp.float_,
@@ -512,28 +516,35 @@ def _get_nominal_parameter_values(
             }
             for net_id, nn in model.nns.items()
         }
+        # load nn parameters from file
+        par_arrays = {
+            array_id: h5py.File(file_spec["location"], "r")
+            for array_id, file_spec in self._petab_problem.extensions_config[
+                "array_files"
+            ].items()
+            # TODO: FIXME (https://github.com/sebapersson/petab_sciml_testsuite/issues/1)
+        }
+
         # extract nominal values from petab problem
         for pname, row in self._petab_problem.parameter_df.iterrows():
             if (net := pname.split(".")[0]) in model.nns:
                 to_set = []
                 nn = model_pars[net]
-                scalar = True
                 try:
                     value = float(row[petab.NOMINAL_VALUE])
                 except ValueError:
-                    value = h5py.File(row[petab.NOMINAL_VALUE], "r")
+                    value = par_arrays[row[petab.NOMINAL_VALUE]]
                     scalar = False
-
                 if len(pname.split(".")) > 1:
-                    layer = nn[pname.split(".")[1]]
+                    layer_name = pname.split(".")[1]
+                    layer = nn[layer_name]
                     if len(pname.split(".")) > 2:
-                        to_set.append(
-                            (pname.split(".")[1], pname.split(".")[2])
-                        )
+                        attribute_name = pname.split(".")[2]
+                        to_set.append((layer_name, attribute_name))
                     else:
                         to_set.extend(
                             [
-                                (pname.split(".")[1], attribute)
+                                (layer_name, attribute)
                                 for attribute in layer.keys()
                             ]
                         )
@@ -549,15 +560,20 @@ def _get_nominal_parameter_values(
                 for layer, attribute in to_set:
                     if scalar:
                         nn[layer][attribute] = value * jnp.ones_like(
-                            nn[layer][attribute]
+                            model.nns[net].layers[layer][attribute]
                         )
                     else:
-                        nn[layer][attribute] = value[layer][attribute]
+                        nn[layer][attribute] = jnp.array(
+                            value[layer][attribute]
+                        )
 
         # set values in model
         for net_id in model_pars:
             for layer_id in model_pars[net_id]:
                 for attribute in model_pars[net_id][layer_id]:
+                    logger.debug(
+                        f"Setting {attribute} of layer {layer_id} in network {net_id} to {model_pars[net_id][layer_id][attribute]}"
+                    )
                     model = eqx.tree_at(
                         lambda model: getattr(
                             model.nns[net_id].layers[layer_id], attribute

python/sdist/amici/petab/petab_import.py

@@ -148,7 +148,7 @@ def import_petab_problem(
         logger.info(f"Compiling model {model_name} to {model_output_dir}.")
 
         if "neural_nets" in petab_problem.extensions_config:  # TODO: fixme
-            from petab_sciml import PetabScimlStandard
+            from petab_sciml.standard import NNModelStandard
 
             config = petab_problem.extensions_config
             # TODO: only accept YAML format for now
@@ -169,9 +169,9 @@ def import_petab_problem(
             )
             hybridization = {
                 net_id: {
-                    "model": PetabScimlStandard.load_data(
+                    "model": NNModelStandard.load_data(
                         Path() / net_config["location"]
-                    ).models,
+                    ),
                     "input_vars": [
                         input_mapping[petab_id]
                         for petab_id, model_id in petab_problem.mapping_df.loc[

python/sdist/pyproject.toml

@@ -72,7 +72,6 @@ test = [
     "scipy",
     "pooch",
     "beartype",
-    ""
 ]
 vis = [
     "matplotlib",

tests/sciml/test_sciml.py

@@ -22,7 +22,7 @@
 import h5py
 from contextlib import contextmanager
 
-from petab_sciml import PetabScimlStandard
+from petab_sciml import NNModelStandard
 
 
 @contextmanager
@@ -78,7 +78,7 @@ def test_net(test):
         net_file = cases_dir / test.replace("_alt", "") / solutions["net_file"]
     else:
         net_file = test_dir / solutions["net_file"]
-    ml_models = PetabScimlStandard.load_data(net_file)
+    ml_models = NNModelStandard.load_data(net_file)
 
     nets = {}
     outdir = Path(__file__).parent / "models" / test
@@ -197,49 +197,8 @@ def test_ude(test):
             compile_=True,
             jax=True,
         )
-        # non_numeric = pd.to_numeric(petab_problem.parameter_df[petab.NOMINAL_VALUE], errors='coerce').isna()
-        # par_files = petab_problem.parameter_df.loc[non_numeric, petab.NOMINAL_VALUE].unique()
-        # par_values = {
-        #     par_file: h5py.File(par_file, "r")
-        #     for par_file in par_files
-        # }
-        # for par_id, row in petab_problem.parameter_df.iterrows():
-        #     if not non_numeric[par_id]:
-        #         continue
-        #     petab_problem.parameter_df.loc[par_id, petab.NOMINAL_VALUE] = \
-        #         (par_values[row[petab.NOMINAL_VALUE]],)
-        # petab_problem.parameter_df.loc[np.logical_not(non_numeric), petab.NOMINAL_VALUE] = pd.to_numeric(
-        #     petab_problem.parameter_df.loc[np.logical_not(non_numeric), petab.NOMINAL_VALUE]
-        # )
 
         jax_problem = JAXProblem(jax_model, petab_problem)
-        # for net, net_config in petab_problem.extensions_config.items(): # TODO: FIXME (https://github.com/sebapersson/petab_sciml_testsuite/issues/1)
-        #     pars = h5py.File(
-        #         net_config["net1_ps_file"]['location'], "r" # TODO: check format and actually use propoer petab nominal parameter infrastructure
-        #     )
-        #     for layer_name, layer in jax_problem.model.nns[net].layers.items():
-        #         for attribute in dir(layer):
-        #             if not isinstance(
-        #                 getattr(layer, attribute), jax.numpy.ndarray
-        #             ):
-        #                 continue
-        #             value = jnp.array(pars[layer_name][attribute])
-        #
-        #             if (
-        #                 isinstance(layer, eqx.nn.ConvTranspose)
-        #                 and attribute == "weight"
-        #             ):
-        #                 # see FAQ in https://docs.kidger.site/equinox/api/nn/conv/#equinox.nn.ConvTranspose
-        #                 value = jnp.flip(
-        #                     value, axis=tuple(range(2, value.ndim))
-        #                 ).swapaxes(0, 1)
-        #             jax_problem = eqx.tree_at(
-        #                 lambda x: getattr(
-        #                     x.model.nns[net].layers[layer_name], attribute
-        #                 ),
-        #                 jax_problem,
-        #                 value,
-        #             )
 
     # llh
     if test in (
@@ -281,7 +240,7 @@ def test_ude(test):
         controller=diffrax.PIDController(atol=1e-14, rtol=1e-14),
         max_steps=2**16,
     )
-    for component, file in solutions["grad_llh_files"].items():
+    for component, file in solutions["grad_files"].items():
         actual_dict = {}
         if component == "mech":
             expected = pd.read_csv(test_dir / file, sep="\t").set_index(