import copy
from typing import Dict, List, Sequence, Tuple
import numpy as np
from ...C import (
AMICI_SIGMAY,
AMICI_SY,
AMICI_Y,
FVAL,
GRAD,
HESS,
MODE_RES,
RDATAS,
RES,
SRES,
X_INNER_OPT,
)
from ...objective.amici.amici_calculator import (
AmiciCalculator,
AmiciModel,
AmiciSolver,
)
from ...objective.amici.amici_util import (
filter_return_dict,
init_return_values,
)
from .problem import SplineInnerProblem
from .solver import SplineInnerSolver
try:
import amici
from amici.parameter_mapping import ParameterMapping
except ImportError:
pass
[docs]class SplineAmiciCalculator(AmiciCalculator):
"""A calculator is passed as `calculator` to the pypesto.AmiciObjective.
The object is called by :func:`pypesto.AmiciObjective.call_unprocessed`
to calculate the current objective function values and gradient.
"""
[docs] def __init__(
self,
inner_problem: SplineInnerProblem,
inner_solver: SplineInnerSolver = None,
):
"""Initialize the calculator from the given problem.
Parameters
----------
inner_problem:
The optimal scaling inner problem.
inner_solver:
A solver to solve ``inner_problem``.
Defaults to ``SplineInnerSolver``.
"""
super().__init__()
self.inner_problem = inner_problem
if inner_solver is None:
inner_solver = SplineInnerSolver()
self.inner_solver = inner_solver
[docs] def initialize(self):
"""Initialize."""
self.inner_solver.initialize()
[docs] def __call__(
self,
x_dct: Dict,
sensi_orders: Tuple[int, ...],
mode: str,
amici_model: AmiciModel,
amici_solver: AmiciSolver,
edatas: List['amici.ExpData'],
n_threads: int,
x_ids: Sequence[str],
parameter_mapping: ParameterMapping,
fim_for_hess: bool,
rdatas: List['amici.ReturnData'] = None,
):
"""Perform the actual AMICI call.
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).
inner_rdatas:
AMICI simulation return data. In case the calculator is part of
the :class:`pypesto.objective.amici.InnerCalculatorCollector`,
it will already simulate the model and pass the results here.
Returns
-------
inner_result:
A dict containing the calculation results: FVAL, GRAD, RDATAS and X_INNER_OPT.
"""
if mode == MODE_RES:
raise ValueError(
"Spline approximation cannot be called with residual mode."
)
if 2 in sensi_orders:
raise ValueError(
"Hessian and FIM are not implemented for the spline calculator."
)
# get dimension of outer problem
dim = len(x_ids)
# initialize return values
nllh, snllh, s2nllh, chi2, res, sres = init_return_values(
sensi_orders, mode, dim
)
# set order in solver
sensi_order = 0
if sensi_orders:
sensi_order = max(sensi_orders)
# If AMICI ReturnData is not provided, we need to simulate the model
if rdatas is None:
amici_solver.setSensitivityOrder(sensi_order)
x_dct = copy.deepcopy(x_dct)
# fill in parameters
amici.parameter_mapping.fill_in_parameters(
edatas=edatas,
problem_parameters=x_dct,
scaled_parameters=True,
parameter_mapping=parameter_mapping,
amici_model=amici_model,
)
# run amici simulation
rdatas = amici.runAmiciSimulations(
amici_model,
amici_solver,
edatas,
num_threads=min(n_threads, len(edatas)),
)
inner_result = {
FVAL: nllh,
GRAD: snllh,
HESS: s2nllh,
RES: res,
SRES: sres,
RDATAS: rdatas,
X_INNER_OPT: self.inner_problem.get_inner_parameter_dictionary(),
}
# if any amici simulation failed, it's unlikely we can compute
# meaningful inner parameters, so we better just fail early.
if any(rdata.status != amici.AMICI_SUCCESS for rdata in rdatas):
inner_result[FVAL] = np.inf
# if the gradient was requested,
# we need to provide some value for it
if 1 in sensi_orders:
inner_result[GRAD] = np.full(
shape=len(x_ids), fill_value=np.nan
)
return filter_return_dict(inner_result)
sim = [rdata[AMICI_Y] for rdata in rdatas]
sigma = [rdata[AMICI_SIGMAY] for rdata in rdatas]
# Clip negative simulation values to zero, to avoid numerical issues.
for i in range(len(sim)):
sim[i] = sim[i].clip(min=0)
# compute optimal inner parameters
x_inner_opt = self.inner_solver.solve(self.inner_problem, sim, sigma)
inner_result[FVAL] = self.inner_solver.calculate_obj_function(
x_inner_opt
)
inner_result[
X_INNER_OPT
] = self.inner_problem.get_inner_parameter_dictionary()
# Calculate analytical gradients if requested
if sensi_order > 0:
sy = [rdata[AMICI_SY] for rdata in rdatas]
inner_result[GRAD] = self.inner_solver.calculate_gradients(
problem=self.inner_problem,
x_inner_opt=x_inner_opt,
sim=sim,
sigma=sigma,
sy=sy,
parameter_mapping=parameter_mapping,
par_opt_ids=x_ids,
par_sim_ids=amici_model.getParameterIds(),
snllh=snllh,
)
return filter_return_dict(inner_result)