1{
2 "cells": [
3  {
4   "cell_type": "markdown",
5   "id": "9a50cf41-28ba-4416-bf29-25621f95ef4f",
6   "metadata": {},
7   "source": [
8    "# Computing event sensitivity\n",
9    "\n",
10    "In this example, we will show how it is possible to use heyoka.py's [event detection](<./Event detection.ipynb>) system to compute the first-order *sensitivity* of an event's trigger time. For the sake of simplicity, we will consider here a dynamical system with a single parameter and we will consider only the sensitivity with respect to this parameter. The approach can easily be generalised to the computation of the sensitivities with respect to multiple parameters and/or the initial conditions.\n",
11    "\n",
12    "In order to illustrate the method, we will be focusing on an analytically-solvable system of ODEs, so that we will be able to determine an explicit expression for the sensitivity which we will then use to validate the numerical computation.\n",
13    "\n",
14    "## The analytical approach\n",
15    "\n",
16    "We consider the ODE system for the [harmonic oscillator](https://en.wikipedia.org/wiki/Simple_harmonic_motion):\n",
17    "\n",
18    "$$\n",
19    "   \\begin{cases}\n",
20    "   x^\\prime &= v \\\\\n",
21    "   v^\\prime &= -kx\n",
22    "   \\end{cases},\n",
23    "$$\n",
24    "\n",
25    "where $k>0$ is the spring constant. We fix the following initial conditions:\n",
26    "\n",
27    "$$\n",
28    "   \\begin{cases}\n",
29    "   x\\left( 0 \\right) &= 0 \\\\\n",
30    "   v\\left( 0 \\right) &= 1\n",
31    "   \\end{cases}.\n",
32    "$$\n",
33    "\n",
34    "The analytical solution for this simple initial-value problem is:\n",
35    "\n",
36    "$$\n",
37    "   \\begin{cases}\n",
38    "   x\\left( t \\right) &= \\frac{1}{\\sqrt{k}}\\sin\\left(\\sqrt{k}t\\right) \\\\\n",
39    "   v\\left( t \\right) &= \\cos\\left(\\sqrt{k}t\\right)\n",
40    "   \\end{cases}.\n",
41    "$$\n",
42    "\n",
43    "Let us now suppose that we want to detect when the oscillation reaches the maximum amplitude. We can do so by defining the simple event equation\n",
44    "\n",
45    "$$\n",
46    "g\\left( x, v \\right) \\equiv v = 0.\n",
47    "$$\n",
48    "\n",
49    "I.e., the maximum amplitude in the harmonic oscillator is reached when the velocity is zero. We can substitute the solution $v\\left( t \\right)$ into the event equation, yielding the time-dependent equation\n",
50    "\n",
51    "$$\n",
52    "\\cos\\left( \\sqrt{k} t \\right) = 0.\n",
53    "$$\n",
54    "\n",
55    "Solving this equation for $t$ gives us an analytical expression for the trigger time of the event, i.e., the time of maximum amplitude $t_M$:\n",
56    "\n",
57    "$$\n",
58    "t_M = \\frac{\\pi}{2\\sqrt{k}}.\n",
59    "$$\n",
60    "\n",
61    "The first-order sensitivity of $t_M$ with respect to $k$ is easily computed:\n",
62    "\n",
63    "$$\n",
64    "\\frac{d t_M}{d k} = -\\frac{\\pi}{4k^\\frac{3}{2}}.\n",
65    "$$\n",
66    "\n",
67    "\n",
68    "## The numerical approach\n",
69    "\n",
70    "If we do not have the analytical solution of the ODE system, it is not possible to compute an explicit expression for the event equation as a function of $t$ and $k$, like we did above. By extension, we cannot compute an explicit expression for $t_M$ and its sensitivity either.\n",
71    "We can however compute numerically the sensitivity with the help of the [variational equations](<./The variational equations.ipynb>).\n",
72    "\n",
73    "In the absence of an analytical solution, the left-hand side of the event equation can be seen as an unknown function of $t$ and $k$:\n",
74    "\n",
75    "$$\n",
76    "g\\left(t, k\\right) = 0.\n",
77    "$$\n",
78    "\n",
79    "We cannot write an explicit expression for $t_M$ using this equation, but via the formula for the derivative of an [implicit function](https://en.wikipedia.org/wiki/Implicit_function#Implicit_differentiation), we can write the sensitivity $d t_M / d k$ as\n",
80    "\n",
81    "$$\n",
82    "\\frac{d t_M}{d k} = -\\frac{\\frac{\\partial g}{\\partial k}}{\\frac{\\partial g}{\\partial t_M}}.\n",
83    "$$\n",
84    "\n",
85    "$\\frac{\\partial g}{\\partial t_M}$ is the time derivative of $g$ calculated at the trigger time $t_M$, which can be computed directly from the original definition of the event equation in terms of the state variables:\n",
86    "\n",
87    "$$\n",
88    "\\frac{\\partial g}{\\partial t_M} = -kx\\left(t_M\\right).\n",
89    "$$\n",
90    "\n",
91    "In order to compute $\\frac{\\partial g}{\\partial k}$, we need to augment the original ODE system with the variational equations for $x$, $v$ and $g$ with respect to $k$:\n",
92    "\n",
93    "$$\n",
94    "   \\begin{cases}\n",
95    "   x^\\prime & = v \\\\\n",
96    "   v^\\prime & = -kx \\\\\n",
97    "   \\left( \\frac{\\partial x}{\\partial k} \\right)^\\prime & = \\frac{\\partial v}{\\partial k} \\\\\n",
98    "   \\left( \\frac{\\partial v}{\\partial k} \\right)^\\prime & = -x-k\\frac{\\partial x}{\\partial k} \\\\\n",
99    "   \\left( \\frac{\\partial g}{\\partial k} \\right)^\\prime & = -x-k\\frac{\\partial x}{\\partial k}\n",
100    "   \\end{cases}.\n",
101    "$$\n",
102    "\n",
103    "We can now proceed to the definition of the heyoka.py integrator:"
104   ]
105  },
106  {
107   "cell_type": "code",
108   "execution_count": 1,
109   "id": "844df82e-9a63-4e8f-b9c5-1edaa4c09da0",
110   "metadata": {},
111   "outputs": [],
112   "source": [
113    "import heyoka as hy\n",
114    "import numpy as np\n",
115    "\n",
116    "# The dynamical variables (including the\n",
117    "# variational variables).\n",
118    "x, v, x_k, v_k, g_k = hy.make_vars(\"x\", \"v\", \"x_k\", \"v_k\", \"g_k\")\n",
119    "\n",
120    "# The spring constant.\n",
121    "k = hy.par[0]\n",
122    "\n",
123    "# The ODEs.\n",
124    "x_t = v\n",
125    "v_t = -k*x\n",
126    "x_k_t = v_k\n",
127    "v_k_t = -x-k*x_k\n",
128    "g_k_t = -x-k*x_k\n",
129    "\n",
130    "# The initial conditions.\n",
131    "ic = [0, 1., 0, 0, 0]\n",
132    "\n",
133    "# Event to detect the maximum amplitude.\n",
134    "ev = hy.t_event(v)\n",
135    "\n",
136    "# Definition of the integrator.\n",
137    "ta = hy.taylor_adaptive([(x, x_t),\n",
138    "                         (v, v_t),\n",
139    "                         (x_k, x_k_t),\n",
140    "                         (v_k, v_k_t),\n",
141    "                         (g_k, g_k_t)],\n",
142    "                        ic, t_events = [ev])\n",
143    "\n",
144    "# Pick a concrete value for the spring constant.\n",
145    "ta.pars[0] = .456"
146   ]
147  },
148  {
149   "cell_type": "markdown",
150   "id": "960f2a83-4444-4d2c-813c-69f47439a21f",
151   "metadata": {},
152   "source": [
153    "Let us propagate up to a large time coordinate. The integration will anyway be stopped almost immediately due to the event triggering when the maximum amplitude is reached:"
154   ]
155  },
156  {
157   "cell_type": "code",
158   "execution_count": 2,
159   "id": "c788f5df-58af-4f97-aab3-531bda9ca644",
160   "metadata": {},
161   "outputs": [
162    {
163     "data": {
164      "text/plain": [
165       "(<taylor_outcome.???: -1>, 1.2823207799363494, 1.2823207799363494, 2, None)"
166      ]
167     },
168     "execution_count": 2,
169     "metadata": {},
170     "output_type": "execute_result"
171    }
172   ],
173   "source": [
174    "ta.propagate_until(1e9)"
175   ]
176  },
177  {
178   "cell_type": "markdown",
179   "id": "f29afb61-d5e7-48ef-b213-93f3488916fc",
180   "metadata": {},
181   "source": [
182    "Recall that, analytically, we expect the sensitivity value to be $-\\frac{\\pi}{4k^\\frac{3}{2}}$, i.e.,"
183   ]
184  },
185  {
186   "cell_type": "code",
187   "execution_count": 3,
188   "id": "ab5e0e04-9022-465d-aacd-709438e43b24",
189   "metadata": {},
190   "outputs": [
191    {
192     "data": {
193      "text/plain": [
194       "-2.550601538829664"
195      ]
196     },
197     "execution_count": 3,
198     "metadata": {},
199     "output_type": "execute_result"
200    }
201   ],
202   "source": [
203    "-np.pi/(4*ta.pars[0]**(3./2))"
204   ]
205  },
206  {
207   "cell_type": "markdown",
208   "id": "d8fe6f26-ca02-4515-b6cf-3171df87f822",
209   "metadata": {},
210   "source": [
211    "From the numerical integration, we can compute the sensitivity value, as explained above, as $-\\frac{\\frac{\\partial g}{\\partial k}}{\\frac{\\partial g}{\\partial t_M}}$. The value of $\\frac{\\partial g}{\\partial k}$ can be read directly from the state vector (at index 4), while $\\frac{\\partial g}{\\partial t_M} = -kx\\left(t_M\\right)$:"
212   ]
213  },
214  {
215   "cell_type": "code",
216   "execution_count": 4,
217   "id": "1c3a3032-317a-43c9-ad22-e7f11e526719",
218   "metadata": {},
219   "outputs": [
220    {
221     "data": {
222      "text/plain": [
223       "-2.5506015388296643"
224      ]
225     },
226     "execution_count": 4,
227     "metadata": {},
228     "output_type": "execute_result"
229    }
230   ],
231   "source": [
232    "-ta.state[4]/(-ta.pars[0]*ta.state[0])"
233   ]
234  },
235  {
236   "cell_type": "markdown",
237   "id": "3a1439e0-97c7-47a5-8e69-da04bb44788e",
238   "metadata": {},
239   "source": [
240    "Indeed, the numerical value matches the analytical result to machine precision.\n",
241    "\n",
242    "## Application to optimisation problems\n",
243    "\n",
244    "Suppose that we want to determine what value the spring constant $k$ must assume in order for the maximum amplitude of the oscillator to be $A$ (a fixed constant). We can formulate this problem as the minimisation of the function\n",
245    "\n",
246    "$$\n",
247    "f\\left( k \\right) = \\left[ x\\left(t_M\\left(k\\right), k\\right) - A\\right]^2,\n",
248    "$$\n",
249    "\n",
250    "where $x\\left(t_M\\left(k\\right), k\\right)$ is the value assumed by the coordinates $x$ at the event trigger time $t_M$. Local optimisation algorithms can greatly benefit from the availability of the gradient of $f$ with respect to the optimisation variable $k$. For the harmonic oscillator system considered here, $df/dk$ is easily computed analytically as\n",
251    "\n",
252    "$$\n",
253    "\\frac{df}{dk} = -\\left(\\frac{1}{\\sqrt{k}} -A \\right)k^{-\\frac{3}{2}}.\n",
254    "$$\n",
255    "\n",
256    "If we assume that an analytical solution is not available (as it is generally the case), we can compute $df/dk$ numerically with the help of the [variational equations](<./The variational equations.ipynb>). Specifically, we can write:\n",
257    "\n",
258    "$$\n",
259    "\\frac{df}{dk} = 2 \\left[ x\\left(t_M\\left(k\\right), k\\right) - A\\right]\\left( \\frac{\\partial x}{\\partial t_M}\\frac{\\partial t_M}{\\partial k} + \\frac{\\partial x}{\\partial k} \\right),\n",
260    "$$\n",
261    "\n",
262    "where:\n",
263    "\n",
264    "* $\\partial x/\\partial t_M$ is the time derivative of $x$ (i.e., the velocity $v$) at the trigger time $t_M$,\n",
265    "* $\\partial t_M/\\partial k$ is the sensitivity of $t_M$ (as computed earlier),\n",
266    "* $\\partial x/\\partial k$ is the sensitivity of $x$ with respect to $k$ (which appears as an extra state variable in the augmented ODE system defined earlier).\n",
267    "\n",
268    "Let us now define the objective function and let us implement its gradient using the numerical solution of the augmented ODE:"
269   ]
270  },
271  {
272   "cell_type": "code",
273   "execution_count": 5,
274   "id": "76873c03-33c6-4281-9d71-cf203037b6f0",
275   "metadata": {},
276   "outputs": [],
277   "source": [
278    "# Pick a concrete value for the A constant.\n",
279    "A = 10.123\n",
280    "\n",
281    "# Objective function.\n",
282    "def fun(x):\n",
283    "    # Reset the integrator state.\n",
284    "    ta.reset_cooldowns()\n",
285    "    ta.time = 0\n",
286    "    ta.state[:] = ic\n",
287    "    ta.pars[0] = x[0]\n",
288    "    \n",
289    "    # Propagate until the event triggers.\n",
290    "    oc, _, _, _ = ta.propagate_until(1e9)\n",
291    "    if int(oc) != -1:\n",
292    "        raise\n",
293    "    \n",
294    "    return (ta.state[0] - A)**2\n",
295    "\n",
296    "# Derivative of 'fun' with respect to k.\n",
297    "def jac(x):\n",
298    "    # Reset the integrator state.\n",
299    "    ta.reset_cooldowns()\n",
300    "    ta.time = 0\n",
301    "    ta.state[:] = ic\n",
302    "    ta.pars[0] = x[0]\n",
303    "    \n",
304    "    # Propagate until the event triggers.\n",
305    "    oc, _, _, _, _ = ta.propagate_until(1e9)\n",
306    "    if int(oc) != -1:\n",
307    "        raise\n",
308    "\n",
309    "    # Compute the sensitivity of t_M.\n",
310    "    tM_k = -ta.state[4]/(-ta.pars[0]*ta.state[0])\n",
311    "\n",
312    "    return [2*(ta.state[0] - A)*(ta.state[1]*tM_k + ta.state[2])]"
313   ]
314  },
315  {
316   "cell_type": "markdown",
317   "id": "a5ca32c3-8c94-4245-8d29-24bc9f08dbef",
318   "metadata": {},
319   "source": [
320    "Let us now compute $df/dk$ via the analytical formula for a specific value of $k$:"
321   ]
322  },
323  {
324   "cell_type": "code",
325   "execution_count": 6,
326   "id": "af69ba9b-d32f-4d26-8542-973e4fce2d36",
327   "metadata": {},
328   "outputs": [
329    {
330     "data": {
331      "text/plain": [
332       "0.5015866697490922"
333      ]
334     },
335     "execution_count": 6,
336     "metadata": {},
337     "output_type": "execute_result"
338    }
339   ],
340   "source": [
341    "# Pick a value for k.\n",
342    "k_val = 7.23\n",
343    "\n",
344    "-(1/np.sqrt(k_val)-A)*k_val**(-3./2)"
345   ]
346  },
347  {
348   "cell_type": "markdown",
349   "id": "d9ad3a00-6ff1-401b-9874-7d162ecf18a0",
350   "metadata": {},
351   "source": [
352    "And here's the value computed via the numerical integration of the variational equations:"
353   ]
354  },
355  {
356   "cell_type": "code",
357   "execution_count": 7,
358   "id": "408b98cd-5c97-446e-958b-1c782f3c64bd",
359   "metadata": {},
360   "outputs": [
361    {
362     "data": {
363      "text/plain": [
364       "[0.5015866697490922]"
365      ]
366     },
367     "execution_count": 7,
368     "metadata": {},
369     "output_type": "execute_result"
370    }
371   ],
372   "source": [
373    "jac([k_val])"
374   ]
375  },
376  {
377   "cell_type": "markdown",
378   "id": "64f775f0-4056-4484-ac52-e116875c19d5",
379   "metadata": {},
380   "source": [
381    "Indeed, the two values agree to machine precision."
382   ]
383  }
384 ],
385 "metadata": {
386  "kernelspec": {
387   "display_name": "Python 3 (ipykernel)",
388   "language": "python",
389   "name": "python3"
390  },
391  "language_info": {
392   "codemirror_mode": {
393    "name": "ipython",
394    "version": 3
395   },
396   "file_extension": ".py",
397   "mimetype": "text/x-python",
398   "name": "python",
399   "nbconvert_exporter": "python",
400   "pygments_lexer": "ipython3",
401   "version": "3.8.10"
402  }
403 },
404 "nbformat": 4,
405 "nbformat_minor": 5
406}
407