{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model import using the PEtab format"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we illustrate how to use [pyPESTO](https://github.com/icb-dcm/pypesto.git) together with [PEtab](https://github.com/petab-dev/petab.git) and [AMICI](https://github.com/AMICI-dev/AMICI). The notebook first details the individual steps of the import and the creation of the objective function. Note that those steps can be summarised, demonstrated at the end of the 'Import' section. After that, optimization and visualisation are showcased.\n",
    "We employ models from the [benchmark collection](https://github.com/benchmarking-initiative/benchmark-models-petab), which we first download:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# install if not done yet\n",
    "# !apt install libatlas-base-dev swig\n",
    "# %pip install pypesto[amici,petab] --quiet\n",
    "# %pip install git+https://github.com/Benchmarking-Initiative/Benchmark-Models-PEtab.git@master#subdirectory=src/python --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pprint import pprint\n",
    "\n",
    "import benchmark_models_petab as models\n",
    "import numpy as np\n",
    "import petab\n",
    "\n",
    "import pypesto.optimize as optimize\n",
    "import pypesto.petab\n",
    "import pypesto.visualize as visualize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Manage PEtab model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A PEtab problem comprises all the information on the model, the data and the parameters to perform parameter estimation. We import a model as a `petab.Problem`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# a collection of models that can be simulated\n",
    "\n",
    "# model_name = \"Zheng_PNAS2012\"\n",
    "model_name = \"Boehm_JProteomeRes2014\"\n",
    "# model_name = \"Fujita_SciSignal2010\"\n",
    "# model_name = \"Sneyd_PNAS2002\"\n",
    "# model_name = \"Borghans_BiophysChem1997\"\n",
    "# model_name = \"Elowitz_Nature2000\"\n",
    "# model_name = \"Crauste_CellSystems2017\"\n",
    "# model_name = \"Lucarelli_CellSystems2018\"\n",
    "# model_name = \"Schwen_PONE2014\"\n",
    "# model_name = \"Blasi_CellSystems2016\"\n",
    "\n",
    "# the yaml configuration file links to all needed files\n",
    "yaml_config = os.path.join(models.MODELS_DIR, model_name, model_name + \".yaml\")\n",
    "\n",
    "# create a petab problem\n",
    "petab_problem = petab.Problem.from_yaml(yaml_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Import model to AMICI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model must be imported to pyPESTO and AMICI. Therefore, we create a `pypesto.PetabImporter` from the problem, and create an AMICI model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "importer = pypesto.petab.PetabImporter(petab_problem)\n",
    "\n",
    "model = importer.create_model(verbose=False)\n",
    "\n",
    "# some model properties\n",
    "print(\"Model parameters:\", list(model.getParameterIds()), \"\\n\")\n",
    "print(\"Model const parameters:\", list(model.getFixedParameterIds()), \"\\n\")\n",
    "print(\"Model outputs:   \", list(model.getObservableIds()), \"\\n\")\n",
    "print(\"Model states:    \", list(model.getStateIds()), \"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create objective function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To perform parameter estimation, we need to define an objective function, which integrates the model, data, and noise model defined in the PEtab problem."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Creating the objective from PEtab with default settings can be done in as little as two lines."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "importer = pypesto.petab.PetabImporter.from_yaml(yaml_config)\n",
    "problem = importer.create_problem()  # creating the problem from the importer. The objective can be found at problem.objective"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you need more flexibility, e.g., to define whether you need residuals of the objective function, what sensitivities you want to use, or fix certain parameters, you can also create the problem from a customized objective:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import libsbml\n",
    "\n",
    "converter_config = libsbml.SBMLLocalParameterConverter().getDefaultProperties()\n",
    "petab_problem.sbml_document.convert(converter_config)\n",
    "\n",
    "obj = importer.create_objective()\n",
    "\n",
    "# for some models, hyperparameters need to be adjusted\n",
    "# obj.amici_solver.setMaxSteps(10000)\n",
    "# obj.amici_solver.setRelativeTolerance(1e-7)\n",
    "# obj.amici_solver.setAbsoluteTolerance(1e-7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can request variable derivatives via `sensi_orders`, function values or residuals as specified via `mode`. Passing `return_dict`, we obtain the direct result of the AMICI simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ret = obj(\n",
    "    petab_problem.x_nominal_scaled,\n",
    "    mode=\"mode_fun\",\n",
    "    sensi_orders=(0, 1),\n",
    "    return_dict=True,\n",
    ")\n",
    "pprint(ret)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The problem defined in PEtab also defines the fixed parameters and parameter bounds. This information is contained in a `pypesto.Problem`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "problem = importer.create_problem(obj)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In particular, the problem accounts for the fixation of parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"{problem.x_fixed_indices=}\")\n",
    "print(f\"{problem.x_free_indices=}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The problem creates a copy of the objective function that takes into account the fixed parameters. The objective function is able to calculate function values and derivatives. A finite difference checks whether the computed gradient is accurate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "objective = problem.objective\n",
    "fval, gradient = objective(\n",
    "    petab_problem.x_nominal_free_scaled, sensi_orders=(0, 1)\n",
    ")\n",
    "print(f\"{fval=}\\n{gradient=}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eps = 1e-4\n",
    "\n",
    "\n",
    "def fd(x):\n",
    "    grad = np.zeros_like(x)\n",
    "    j = 0\n",
    "    for i, _xi in enumerate(x):\n",
    "        mask = np.zeros_like(x)\n",
    "        mask[i] += eps\n",
    "        valinc, _ = objective(x + mask, sensi_orders=(0, 1))\n",
    "        valdec, _ = objective(x - mask, sensi_orders=(0, 1))\n",
    "        grad[j] = (valinc - valdec) / (2 * eps)\n",
    "        j += 1\n",
    "    return grad\n",
    "\n",
    "\n",
    "fdval = fd(petab_problem.x_nominal_free_scaled)\n",
    "for i, (g, f) in enumerate(zip(gradient, fdval)):\n",
    "    print(f\"{i=}: {g=:9f},\\t{f=:9f},\\t{g - f=:9f}\")\n",
    "print(f\"l2 difference: {np.linalg.norm(gradient - fdval):.2e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Given the problem, we can perform optimization. We can specify an optimizer to use, and a parallelization engine to speed things up."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "optimizer = optimize.ScipyOptimizer()\n",
    "\n",
    "# engine = pypesto.engine.SingleCoreEngine()\n",
    "engine = pypesto.engine.MultiProcessEngine()\n",
    "\n",
    "# do the optimization\n",
    "result = optimize.minimize(\n",
    "    problem=problem, optimizer=optimizer, n_starts=10, engine=engine\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "### Dealing with function evaluations at the initial point\n",
    "\n",
    "It is quite common in real applications that the objective function is evaluable at every point in parameter space. Therefore, some local optimizations may fail directly at their initial point. Such results are usually not very informative and would be discarded. To directly discard such initial points, we can select a startpoint method that will resample starting points if the objective function (`check_fval`) or its gradient (`check_grad`) are non-finite:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "problem = importer.create_problem(\n",
    "    startpoint_kwargs={\"check_fval\": True, \"check_grad\": True}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The results are contained in a `pypesto.Result` object. It contains e.g., the optimal function values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result.optimize_result.fval"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use the standard pyPESTO plotting routines to visualize and analyze the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ref = visualize.create_references(\n",
    "    x=petab_problem.x_nominal_scaled, fval=obj(petab_problem.x_nominal_scaled)\n",
    ")\n",
    "\n",
    "visualize.waterfall(result, reference=ref, scale_y=\"lin\")\n",
    "visualize.parameters(result, reference=ref);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We can also conveniently visualize the model fit. This plots the petab visualization using optimized parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# we need to explicitly import the method\n",
    "from pypesto.visualize.model_fit import visualize_optimized_model_fit\n",
    "\n",
    "visualize_optimized_model_fit(\n",
    "    petab_problem=petab_problem, result=result, pypesto_problem=problem\n",
    ");"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "vscode": {
   "interpreter": {
    "hash": "44a9cdcbdccbf05a880e90d2e6fe72470baab4d1b82472d890be0596ed887a6b"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}