{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quasistatic vertex models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Quasitatic vertex models are based on the minimisation of a potential energy depending on the tissue geometry.\n",
    "\n",
    "It is the original formalism defined by Honda et al., and in its most widely used form by Farhadifar et al..\n",
    "\n",
    "`tyssue` provides a way to easily define such a potential and find the corresponding energy minimum through gradient descent.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import json\n",
    "import matplotlib.pylab as plt\n",
    "%matplotlib inline\n",
    "\n",
    "from scipy import optimize\n",
    "\n",
    "from tyssue import Sheet, config\n",
    "from tyssue import SheetGeometry as geom\n",
    "from tyssue.dynamics import effectors, model_factory\n",
    "from tyssue.dynamics import SheetModel as model\n",
    "from tyssue.solvers.quasistatic import QSSolver\n",
    "from tyssue.draw import sheet_view\n",
    "from tyssue.draw.plt_draw import plot_forces\n",
    "from tyssue.io import hdf5\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h5store = 'data/small_hexagonal.hf5'\n",
    "#h5store = 'data/before_apoptosis.hf5'\n",
    "\n",
    "datasets = hdf5.load_datasets(h5store,\n",
    "                              data_names=['face', 'vert', 'edge'])\n",
    "specs = config.geometry.cylindrical_sheet()\n",
    "sheet = Sheet('emin', datasets, specs)\n",
    "sheet.sanitize(trim_borders=True, order_edges=True)\n",
    "\n",
    "geom.update_all(sheet)\n",
    "\n",
    "\n",
    "fig, ax = sheet_view(sheet, ['z', 'x'], mode='quick')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Effector classes\n",
    "\n",
    "An effector designates a dynamical term in the epithelium governing equation. For quasistatic models, we need to provide  a method to compute the energy associated with this effector, and the corresponding gradient. \n",
    "\n",
    "For example, we can consider a line tension effector. The energy is $E_t = \\sum_{ij} \\Lambda\\ell_{ij}$ where the sum is over all the edges. For each half-edge, the gradient of $E_t$ is defined by two terms, one for the gradient term associated with the half-edge ${ij}$ source, the over for it's target. For the $x$ coordinate:\n",
    "$$\n",
    "\\frac{\\partial E_t}{\\partial x_i}\n",
    "= \\Lambda\\left(\\sum_k \\frac{\\partial \\ell_{ik}}{\\partial x_i} + \\sum_m \\frac{\\partial \\ell_{mi}}{\\partial x_i}\\right)\n",
    "= \\Lambda\\left(\\sum_k \\frac{x_i}{\\ell_{ik}}  - \\sum_m \\frac{x_i}{\\ell_{mi}}\\right)\n",
    "$$\n",
    "\n",
    "Where $\\sum_k$ is are over all the edges which vertex $i$ is a source, and $\\sum_m$ over all the edges of which vertex i is a target. \n",
    "\n",
    "\n",
    "Here is the definition of the line tension effector:\n",
    "```python\n",
    "class LineTension(AbstractEffector):\n",
    "\n",
    "    dimensions = units.line_tension\n",
    "    magnitude = \"line_tension\"\n",
    "    label = \"Line tension\"\n",
    "    element = \"edge\"\n",
    "    specs = {\"edge\": {\"is_active\": 1, \"line_tension\": 1e-2}}\n",
    "\n",
    "    spatial_ref = \"mean_length\", units.length\n",
    "\n",
    "    @staticmethod\n",
    "    def energy(eptm):\n",
    "        return eptm.edge_df.eval(\n",
    "            \"line_tension\" \"* is_active\" \"* length / 2\"\n",
    "        )  # accounts for half edges\n",
    "\n",
    "    @staticmethod\n",
    "    def gradient(eptm):\n",
    "        grad_srce = -eptm.edge_df[eptm.ucoords] * to_nd(\n",
    "            eptm.edge_df.eval(\"line_tension * is_active/2\"), len(eptm.coords)\n",
    "        )\n",
    "        grad_srce.columns = [\"g\" + u for u in eptm.coords]\n",
    "        grad_trgt = -grad_srce\n",
    "        return grad_srce, grad_trgt\n",
    "```\n",
    "\n",
    "For quasistatic methods, this effectors is agregated with others to define a `model` object. This object will have two methods `compute_energy` and `compute_gradient` that take an `Epithelium` object as single argument.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_specs = {\n",
    "    \"edge\": {\n",
    "        \"line_tension\": 1.,\n",
    "    },\n",
    "    \"face\": {\n",
    "        \"contractility\": 1.,\n",
    "        \"vol_elasticity\": 1.,\n",
    "        \"prefered_vol\": sheet.face_df[\"vol\"].mean(),\n",
    "    },\n",
    "}\n",
    "\n",
    "sheet.update_specs(model_specs, reset=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Et, Ec, Ev = model.compute_energy(sheet, full_output=True)\n",
    "energy = model.compute_energy(sheet, full_output=False)\n",
    "print('Total energy: {}'.format(energy))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grad_E = model.compute_gradient(sheet)\n",
    "\n",
    "grad_E.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "gradients = model.compute_gradient(sheet, components=True)\n",
    "gradients = {label: (srce, trgt) for label, (srce, trgt)\n",
    "             in zip(model.labels, gradients)}\n",
    "gradients['Line tension'][0].head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "solver = QSSolver()\n",
    "\n",
    "scale = 0.2\n",
    "fig, ax = plot_forces(sheet, geom, model, ['z', 'x'], scale)\n",
    "fig.set_size_inches(10, 12)\n",
    "for n, (vx, vy, vz) in sheet.vert_df[sheet.coords].iterrows():\n",
    "    shift = 0.6 * np.sign(vy)\n",
    "    ax.text(vz+shift-0.3, vx, str(n))\n",
    "\n",
    "app_grad_specs = config.draw.sheet_spec()['grad']\n",
    "app_grad_specs.update({'color':'r'})\n",
    "    \n",
    "def draw_approx_force(ax=None):\n",
    "    fig, ax = plot_forces(sheet, geom, model,\n",
    "                          ['z', 'x'], scaling=scale, ax=ax,\n",
    "                          approx_grad=solver.approx_grad, **{'grad':app_grad_specs})\n",
    "    fig.set_size_inches(10, 12)\n",
    "    return fig, ax\n",
    "\n",
    "## Uncomment bellow to recompute\n",
    "fig, ax = draw_approx_force(ax=ax)\n",
    "#fig"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "http://scipy.github.io/devdocs/generated/scipy.optimize.check_grad.html#scipy.optimize.check_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\n",
    "grad_err = solver.check_grad(sheet, geom, model)\n",
    "grad_err /= sheet.vert_df.size\n",
    "\n",
    "\n",
    "print(\"Error on the gradient (non-dim, per vertex): {:.3e}\".format(grad_err))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "options ={'disp':False,\n",
    "          'ftol':1e-5,\n",
    "          'gtol':1e-5}\n",
    "\n",
    "\n",
    "res = solver.find_energy_min(sheet, geom, model, options=options)\n",
    "print(res['success'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res['message']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res['fun']/sheet.face_df.is_alive.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plot_forces(sheet, geom, model, ['z', 'y'], 1)\n",
    "fig.set_size_inches(10, 12)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}