PEtab v2 import

Support for PEtab v2.

**WIP**

Author: dweindl

pypesto/objective/amici/amici.py

@@ -40,6 +40,8 @@
 
     try:
         import amici
+        import amici.petab.petab_importer
+        import pandas as pd
         from amici.petab.parameter_mapping import ParameterMapping
     except ImportError:
         pass
@@ -723,3 +725,67 @@ def update_from_problem(
             ) in condition_mapping.map_preeq_fix.items():
                 if (val := id_to_val.get(mapped_to_par)) is not None:
                     condition_mapping.map_preeq_fix[model_par] = val
+
+
+class AmiciPetabV2Objective(AmiciObjective):
+    """An AMICI objective constructed from a PEtab v2 problem."""
+
+    def __init__(
+        self,
+        petab_importer: amici.petab.petab_importer.PetabImporter,
+        **kwargs,
+    ) -> None:
+        from .amici_calculator import AmiciCalculatorPetabV2
+
+        self._petab_simulator: amici.petab.petab_importer.PetabSimulator = (
+            petab_importer.create_simulator()
+        )
+        self.petab_problem = petab_importer.petab_problem
+        amici_model = self._petab_simulator.model
+        amici_solver = self._petab_simulator.solver
+        edatas = self._petab_simulator.exp_man.create_edatas()
+
+        super().__init__(
+            amici_model=amici_model,
+            amici_solver=amici_solver,
+            edatas=edatas,
+            calculator=AmiciCalculatorPetabV2(self._petab_simulator),
+            **kwargs,
+        )
+
+    def __deepcopy__(self, memo=None):
+        """Override AmiciObjective.__deepcopy__."""
+        if memo is None:
+            memo = {}
+        cls = self.__class__
+        result = cls.__new__(cls)
+        memo[id(self)] = result
+        for k, v in self.__dict__.items():
+            setattr(result, k, copy.deepcopy(v, memo))
+        return result
+
+    def __getstate__(self) -> dict:
+        """Use Python's default pickling semantics (shallow copy of instance dict)."""
+        return dict(self.__dict__)
+
+    def __setstate__(self, state: dict) -> None:
+        """Restore state using the instance dict (default unpickling behaviour)."""
+        self.__dict__.update(state)
+
+    def rdatas_to_simulation_df(
+        self,
+        rdatas: Sequence[amici.ReturnData],
+    ) -> pd.DataFrame:
+        """
+        See :meth:`rdatas_to_measurement_df`.
+
+        Except a petab simulation dataframe is created, i.e. the measurement
+        column label is adjusted.
+        """
+        from amici.petab.petab_importer import rdatas_to_simulation_df
+
+        return rdatas_to_simulation_df(
+            rdatas,
+            self._petab_simulator._model,
+            self._petab_simulator._petab_problem,
+        )

pypesto/objective/amici/amici_calculator.py

@@ -149,6 +149,136 @@ def __call__(
         )
 
 
+class AmiciCalculatorPetabV2(AmiciCalculator):
+    """Class to perform the AMICI call and obtain objective function values."""
+
+    def __init__(
+        self,
+        petab_simulator: amici.petab.petab_importer.PetabSimulator,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.petab_simulator = petab_simulator
+
+    def __call__(
+        self,
+        x_dct: dict,
+        sensi_orders: tuple[int],
+        mode: ModeType,
+        amici_model: AmiciModel,
+        amici_solver: AmiciSolver,
+        edatas: list[amici.ExpData],
+        n_threads: int,
+        x_ids: Sequence[str],
+        parameter_mapping: ParameterMapping,
+        fim_for_hess: bool,
+    ):
+        """Perform the actual AMICI call.
+
+        Called within the :func:`AmiciObjective.__call__` method.
+
+        Parameters
+        ----------
+        x_dct:
+            Parameters for which to compute function value and derivatives.
+        sensi_orders:
+            Tuple of requested sensitivity orders.
+        mode:
+            Call mode (function value or residual based).
+        amici_model:
+            The AMICI model.
+        amici_solver:
+            The AMICI solver.
+        edatas:
+            The experimental data.
+        n_threads:
+            Number of threads for AMICI call.
+        x_ids:
+            Ids of optimization parameters.
+        parameter_mapping:
+            Mapping of optimization to simulation parameters.
+        fim_for_hess:
+            Whether to use the FIM (if available) instead of the Hessian (if
+            requested).
+        """
+        amici_solver = self.petab_simulator._solver
+
+        if 2 in sensi_orders:
+            raise NotImplementedError(
+                "Second order sensitivities are not yet supported for "
+                "PEtab v2."
+            )
+
+        if mode != MODE_FUN:
+            raise NotImplementedError(
+                "Only function value mode is currently supported for "
+                f"PEtab v2. Got mode {mode}."
+            )
+
+        # TODO: -> method
+        # set order in solver
+        sensi_order = 0
+        if sensi_orders:
+            sensi_order = max(sensi_orders)
+
+        if sensi_order == 2 and fim_for_hess:
+            # we use the FIM
+            amici_solver.set_sensitivity_order(sensi_order - 1)
+        else:
+            amici_solver.set_sensitivity_order(sensi_order)
+
+        # run amici simulation
+        res = self.petab_simulator.simulate(x_dct)
+        rdatas = res[RDATAS]
+
+        if (
+            not self._known_least_squares_safe
+            and mode == MODE_RES
+            and 1 in sensi_orders
+        ):
+            if not amici_model.get_add_sigma_residuals() and any(
+                (
+                    (r["ssigmay"] is not None and np.any(r["ssigmay"]))
+                    or (r["ssigmaz"] is not None and np.any(r["ssigmaz"]))
+                )
+                for r in rdatas
+            ):
+                raise RuntimeError(
+                    "Cannot use least squares solver with"
+                    "parameter dependent sigma! Support can be "
+                    "enabled via "
+                    "amici_model.setAddSigmaResiduals()."
+                )
+            self._known_least_squares_safe = True  # don't check this again
+
+        grad = None
+        if 1 in sensi_orders:
+            if res["sllh"] is None and np.isnan(res["llh"]):
+                # TODO: to amici -- set sllh even if llh is nan
+                grad = np.full(len(x_ids), np.nan)
+            else:
+                # llh to nllh, dict to array
+                grad = -np.array(
+                    [
+                        res["sllh"][x_id]  # if x_id in res["sllh"] else 0.0
+                        for x_id in x_ids
+                        if x_id in x_dct.keys()
+                    ]
+                )
+
+        ret = {
+            FVAL: -res["llh"],
+            GRAD: grad,
+            # TODO
+            # HESS: s2nllh,
+            # RES: res,
+            # SRES: sres,
+            RDATAS: rdatas,
+        }
+
+        return filter_return_dict(ret)
+
+
 def calculate_function_values(
     rdatas,
     sensi_orders: tuple[int, ...],

pypesto/petab/importer.py

@@ -14,6 +14,7 @@
 
 import pandas as pd
 import petab.v1 as petab
+from petab import v2
 
 try:
     import roadrunner
@@ -39,6 +40,7 @@
 from ..startpoint import StartpointMethod
 from .objective_creator import (
     AmiciObjectiveCreator,
+    AmiciPetabV2ObjectiveCreator,
     ObjectiveCreator,
     PetabSimulatorObjectiveCreator,
     RoadRunnerObjectiveCreator,
@@ -70,7 +72,7 @@ class PetabImporter:
 
     def __init__(
         self,
-        petab_problem: petab.Problem,
+        petab_problem: petab.Problem | v2.Problem,
         output_folder: str | None = None,
         model_name: str | None = None,
         validate_petab: bool = True,
@@ -149,8 +151,18 @@ def __init__(
 
         self.validate_petab = validate_petab
         if self.validate_petab:
-            if petab.lint_problem(petab_problem):
+            if isinstance(petab_problem, petab.Problem) and petab.lint_problem(
+                petab_problem
+            ):
                 raise ValueError("Invalid PEtab problem.")
+            if (
+                isinstance(petab_problem, v2.Problem)
+                and (
+                    validation_result := petab_problem.validate()
+                ).has_errors()
+            ):
+                validation_result.log(logger=logger)
+                raise ValueError("Invalid PEtab v2 problem.")
         if self._hierarchical and validate_petab_hierarchical:
             from ..hierarchical.petab import (
                 validate_hierarchical_petab_problem,
@@ -186,7 +198,15 @@ def from_yaml(
         simulator_type: str = AMICI,
     ) -> PetabImporter:
         """Simplified constructor using a petab yaml file."""
-        petab_problem = petab.Problem.from_yaml(yaml_config)
+        from petab.versions import get_major_version
+
+        match get_major_version(yaml_config):
+            case 1:
+                petab_problem = petab.Problem.from_yaml(yaml_config)
+            case 2:
+                petab_problem = v2.Problem.from_yaml(yaml_config)
+            case _:
+                raise ValueError("Only PEtab v1 and v2 are supported.")
 
         return PetabImporter(
             petab_problem=petab_problem,
@@ -278,6 +298,42 @@ def create_prior(self) -> NegLogParameterPriors | None:
         else:
             return None
 
+    def _create_prior_v2(self) -> NegLogParameterPriors | None:
+        """Create prior for PEtab v2 problem."""
+        import petab.v2.C as petab_c
+
+        import pypesto.C as pypesto_c
+
+        petab_to_pypesto = {
+            petab_c.LAPLACE: pypesto_c.LAPLACE,
+            petab_c.LOG_LAPLACE: pypesto_c.LOG_LAPLACE,
+            petab_c.LOG_NORMAL: pypesto_c.LOG_NORMAL,
+            petab_c.LOG_UNIFORM: pypesto_c.LOG_UNIFORM,
+            petab_c.NORMAL: pypesto_c.NORMAL,
+            petab_c.UNIFORM: pypesto_c.UNIFORM,
+        }
+
+        prior_list = []
+        for parameter in self.petab_problem.parameters:
+            if not parameter.estimate or parameter.prior_distribution is None:
+                continue
+
+            prior_list.append(
+                get_parameter_prior_dict(
+                    index=len(prior_list),
+                    # TODO map names
+                    prior_type=petab_to_pypesto.get(
+                        str(parameter.prior_distribution),
+                        str(parameter.prior_distribution),
+                    ),
+                    prior_parameters=parameter.prior_parameters,
+                    parameter_scale="lin",
+                )
+            )
+        if prior_list:
+            return NegLogParameterPriors(prior_list)
+        return None
+
     def create_startpoint_method(self, **kwargs) -> StartpointMethod:
         """Create a startpoint method.
 
@@ -307,6 +363,22 @@ def create_objective_creator(
             has to be provided. Otherwise the argument is not used.
 
         """
+        if isinstance(self.petab_problem, v2.Problem):
+            if simulator_type != AMICI:
+                raise ValueError(
+                    "Only 'amici' simulator type is supported for PEtab v2 "
+                    "problems."
+                )
+            return AmiciPetabV2ObjectiveCreator(
+                petab_problem=self.petab_problem,
+                output_folder=self.output_folder,
+                model_name=self.model_name,
+                hierarchical=self._hierarchical,
+                inner_options=self.inner_options,
+                non_quantitative_data_types=self._non_quantitative_data_types,
+                validate_petab=self.validate_petab,
+            )
+
         if simulator_type == AMICI:
             return AmiciObjectiveCreator(
                 petab_problem=self.petab_problem,
@@ -365,10 +437,26 @@ def create_problem(
             objective = self.objective_constructor.create_objective(**kwargs)
 
         x_fixed_indices = self.petab_problem.x_fixed_indices
-        x_fixed_vals = self.petab_problem.x_nominal_fixed_scaled
         x_ids = self.petab_problem.x_ids
-        lb = self.petab_problem.lb_scaled
-        ub = self.petab_problem.ub_scaled
+        if isinstance(self.petab_problem, petab.Problem):
+            # PEtab v1
+            x_fixed_vals = self.petab_problem.x_nominal_fixed_scaled
+            lb = self.petab_problem.lb_scaled
+            ub = self.petab_problem.ub_scaled
+            x_scales = [
+                self.petab_problem.parameter_df.loc[
+                    x_id, petab.PARAMETER_SCALE
+                ]
+                for x_id in x_ids
+            ]
+            prior = self.create_prior()
+        else:
+            # PEtab v2 -- no parameter scaling
+            x_fixed_vals = self.petab_problem.x_nominal_fixed
+            lb = self.petab_problem.lb
+            ub = self.petab_problem.ub
+            x_scales = [petab.LIN for x_id in x_ids]
+            prior = self._create_prior_v2()
 
         # Raise error if the correct calculator is not used.
         if self._hierarchical:
@@ -388,19 +476,12 @@ def create_problem(
                 map(x_ids.index, self.petab_problem.x_fixed_ids)
             )
 
-        x_scales = [
-            self.petab_problem.parameter_df.loc[x_id, petab.PARAMETER_SCALE]
-            for x_id in x_ids
-        ]
-
         if problem_kwargs is None:
             problem_kwargs = {}
 
         if startpoint_kwargs is None:
             startpoint_kwargs = {}
 
-        prior = self.create_prior()
-
         if prior is not None:
             if self._hierarchical:
                 raise NotImplementedError(