{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Animation of a family of complex functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import plotly.graph_objects as go\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Plotly version of the HSV colorscale, corresponding to S=1, V=1, where S is saturation and V is the value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pl_hsv = [[0.0, 'rgb(0, 255, 255)'],\n",
    " [0.0833, 'rgb(0, 127, 255)'],\n",
    " [0.1667, 'rgb(0, 0, 255)'],\n",
    " [0.25, 'rgb(127, 0, 255)'],\n",
    " [0.3333, 'rgb(255, 0, 255)'],\n",
    " [0.4167, 'rgb(255, 0, 127)'],\n",
    " [0.5, 'rgb(255, 0, 0)'],\n",
    " [0.5833, 'rgb(255, 127, 0)'],\n",
    " [0.6667, 'rgb(254, 255, 0)'],\n",
    " [0.75, 'rgb(127, 255, 0)'],\n",
    " [0.8333, 'rgb(0, 255, 0)'],\n",
    " [0.9167, 'rgb(0, 255, 127)'],\n",
    " [1.0, 'rgb(0, 255, 255)']]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_function(func,  re=(-1,1), im=(-1,1), N=100):\n",
    "    # func is the complex function to be ploted\n",
    "    # re, im are the interval ends on the real and imaginary axes, defining the rectangular region in the complex plane\n",
    "    # N gives the number of points in an interval of length 1\n",
    "    l = re[1]-re[0]\n",
    "    h = im[1]-im[0]\n",
    "    resL = int(N*l) #horizontal resolution\n",
    "    resH = int(N*h) #vertical resolution\n",
    "    X = np.linspace(re[0], re[1], resL)\n",
    "    Y = np.linspace(im[0], im[1], resH)\n",
    "    x, y = np.meshgrid(X,Y)\n",
    "    z = x+1j*y\n",
    "    return X, Y, z"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For a particular parameter,  the corresponding function of the given family is represented as heatmap of its argument, which illustrates \n",
    "its zeros and poles position."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define tickvals and ticktext for the heatmap colorbar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tickvals = [-np.pi, -2*np.pi/3, -np.pi/3, 0, np.pi/3, 2*np.pi/3, np.pi]\n",
    "ticktext=['-\\u03c0', '-2\\u03c0/3', '-\\u03c0/3', '0', '\\u03c0/3', '2\\u03c0/3', '\\u03c0']   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define a family of complex functions, depending on the parameter t. To animate the motion of function zeros and poles, as t varies,\n",
    "we use `functools.partial` to get the\n",
    "function corresponding to a particular parameter:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "def func(t, z):\n",
    "    return z**4+np.exp(2*t*1j)/z**2*np.exp(1j*t)\n",
    "f0 =  partial(func, 0) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y, z =  evaluate_function(f0, re=(-1.5, 1.5), im=(-1.5,1.5), N=50)\n",
    "w = f0(z)\n",
    "argument = np.angle(w)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = go.Figure(go.Heatmap(x=x, y=y, z=argument,  colorscale=pl_hsv, \n",
    "                    colorbar=dict(thickness=20, tickvals=tickvals, \n",
    "                    ticktext=ticktext,  \n",
    "                    title='arg(f(z))')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frames = []\n",
    "t = np.linspace(0, 3, 45)  #6, 85\n",
    "for s  in t:\n",
    "    g = partial(func, s)\n",
    "    w = g(z)\n",
    "    argument = np.angle(w)\n",
    "    frames.append(go.Frame(data=[go.Heatmap(z=argument)]))\n",
    "fig.update(frames=frames);    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig.update_layout(width=500, height=475,\n",
    "                  updatemenus=[dict(type='buttons', \n",
    "                                y=1,\n",
    "                                x=1.45,\n",
    "                                active=0,\n",
    "                                buttons=[dict(label='Play',\n",
    "                                              method='animate',\n",
    "                                              args=[None, \n",
    "                                                    dict(frame=dict(duration=10, \n",
    "                                                                    redraw=True),\n",
    "                                                         transition=dict(duration=0),\n",
    "                                                         fromcurrent=True,\n",
    "                                                         mode='immediate')])])]);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Anim-complex-f](anim-fcompl.gif)"
   ]
  },
  {
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}