```python
%matplotlib inline
%run ../setup/nb_setup
```

# Orbits 2: Orbits in Axisymmetric Potentials

Author(s): Adrian Price-Whelan

## Learning goals

In this tutorial, we will introduce axisymmetric potential models, and explore differences between orbits in non-spherical potentials with what we learned about orbits in spherical systems.

## Introduction

As we saw in the previous tutorial, while spherical potential models can have complex radial density profiles, orbits in spherical potential models are planar and can be characterized by their radial and azimuthal frequencies because these orbits have at least four isolating integrals of motion (energy and the three components of angular momentum). As we will see in this tutorial, as the symmetries of a potential model are relaxed (from spherical to axisymmetric, and then in the next tutorial, from axisymmetric to triaxial), the number of isolating integrals of motion for a generic orbit decreases to three, and they become difficult to compute. The implication of this is that generic orbits in non-spherical potential models are no longer confined to a plane: They fill a three-dimensional volume instead of a two-dimensional region. In addition, these more complex models have regions of *chaotic* or irregular orbits, which differ from regular orbits in many important ways, as we will see in the next tutorial.

In this tutorial, we will introduce some commonly-used axisymmetric gravitational potential models, compute the orbits of particles in some of these models, and analyze the properties of orbits in non-spherical models.

## Terminology and Notation

- (See Orbits tutorial 1)
- Cylindrical radius: $R = \sqrt{x^2 + y^2}$
- Maximum $z$ excursion: $z_\textrm{max} = \textrm{max}_\textrm{time}(z)$

### Notebook Setup and Package Imports

```python
from astropy.constants import G
import astropy.units as u

import matplotlib as mpl
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import numpy as np

import gala.dynamics as gd
import gala.integrate as gi
import gala.potential as gp
from gala.units import galactic
```

## Axisymmetric Potential Models

Rather than having spherical symmetry, axisymmetric potential models are symmetric to rotations around a particular axis. A common use case for these potential models is for representing flattened and disky systems (like the stellar and gas disks of the Milky Way). A common way of constructing axisymmetric potentials is to start with a spherical potential model, but replace the spherical radius with an elliptical radius that breaks the symmetry between $x, y$ and $z$. For example:

$$
r = \sqrt{x^2 + y^2 + z^2} \rightarrow \sqrt{x^2 + y^2 + (|z| + a)^2}
$$

By convention, in many contexts we will take the $z$ axis to be the symmetry axis, so we often use cylindrical coordinates $R, z$ when working with axisymmetric models.

One commonly-used axisymmetric model is the Miyamoto–Nagai (MN) potential, which is given by

$$
\Phi_{\textrm{MN}}(R, z) = - \frac{G\, M}{\sqrt{R^2 + (a + \sqrt{z^2 + b^2})^2}}
$$

In the limit that $a \rightarrow 0$, the MN potential reduces to the spherical Plummer model introduced in the previous tutorial. In the limit that $b \rightarrow 0$, the MN potential reduces to the potential generated by an infinitely-thin disk. Depending on the setting of $a$ and $b$, we can represent a variety of flattened density distributions, like galactic stellar or gas disks.

To get a feel for how orbits behave in axisymmetric potentials, we will compute some orbits in the MN potential using [Gala](http://gala.adrian.pw/). We will start by defining a model that has parameters similar to the stellar mass distribution in the local part of the Galactic disk (parameter values taken from [this paper](https://ui.adsabs.harvard.edu/abs/2021ApJ...910...17P/abstract)):

```python
mn_disk = gp.MiyamotoNagaiPotential(
    m=6.98e10 * u.Msun, a=3 * u.kpc, b=0.28 * u.kpc, units=galactic
)
```

To start off, we will first use Gala to plot equipotential and isodensity contours to visualize the structure of this potential–density pair. To do this, we will visualize the lines of equal potential/density in 2D slices of the 3D models. We therefore need to specify which axis to "slice," and what value in that coordinate to slice at. In the other two coordinates, we need to specify grids over which to compute the potential or density (so we can use matplotlib's `contourf()` function to visualize the curves). Here, we will make plots of x-y (z=0) and x-z (y=0) slices. We will compute the potential and density on grids of 256 by 256 points between (-10, 10) kpc:

```python
grid = np.linspace(-10, 10, 256) * u.kpc
```

We will first plot the isopotential contours using the `.plot_contours()` method on any Gala potential object:

```python
fig, axes = plt.subplots(
    1, 2, figsize=(10, 5), sharex=True, sharey=True, constrained_layout=True
)
mn_disk.plot_contours(grid=(grid, grid, 0), ax=axes[0])
mn_disk.plot_contours(grid=(grid, 0, grid), ax=axes[1])

for ax in axes:
    ax.set_xlabel("$x$")
axes[0].set_ylabel("$y$")
axes[1].set_ylabel("$z$")

for ax in axes:
    ax.set_aspect("equal")
fig.suptitle("Iso-potential contours", fontsize=22)
```

For comparison, we will now plot the same slices, but visualizing the isodensity contours using the `.plot_density_contours()` method of any Gala potential object:

```python
fig, axes = plt.subplots(
    1, 2, figsize=(10, 5), sharex=True, sharey=True, constrained_layout=True
)
mn_disk.plot_density_contours(grid=(grid, grid, 0), ax=axes[0])
mn_disk.plot_density_contours(grid=(grid, 0, grid), ax=axes[1])

for ax in axes:
    ax.set_xlabel("$x$")
axes[0].set_ylabel("$y$")
axes[1].set_ylabel("$z$")

for ax in axes:
    ax.set_aspect("equal")
fig.suptitle("Iso-density contours", fontsize=22)
```

Note that the density contours are *much* more "disky" than the potential contours (i.e. the axis ratio comparing $z$ to $R$ is much smaller in the density than in the potential). This is generally true: Any flattening or axis ratio in density leads to less flattening / rounder potential contours. Keep this in mind for later! Setting a very small flattening parameter in the potential can therefore lead to very strange (or unphysical) density distributions.

Below, we will compute some orbits in the MN model. However, we would like to compare these orbits to orbits computed in an equivalent spherical, Plummer model with the same mass and scale radius as our disk. We will therefore first define a plummer model with the same mass and scale radius to use below.

### Exercise: Defining an comparison Plummer model

As mentioned above, in a particular limit, the MN potential becomes the Plummer potential. What should we set the Plummer scale length to so that, in the midplane (z=0) it has the same profile as our MN potential? I.e. what combination of the MN scale parameters $a_\textrm{MN}$ and $b_\textrm{MN}$ should we set the Plummer scale length to?

Write the answer here:

...

Define the comparison Plummer model using the mass from the MN potential, and the correct combination of the MN scale parameters (you can use `mn_disk.parameters['m']`, `mn_disk.parameters['a']`, and `mn_disk.parameters['b']` to retrieve the parameter values so you don't have to re-define them).

```python
# plummer = gp.PlummerPotential(...)
```

Demonstrate that the mass enclosed (computed assuming that both potentials are spherical) are equivalent in the midplane at $(x, y, z) = (8, 0, 0)~\textrm{kpc}$

```python
# Menc_MN = ...
# Menc_plummer =
```

---

## Orbits in axisymmetric potentials

In spherical potential models, we saw that orbits are confined to a plane, and can either be circular or have some radial oscillations that cause orbits to form a rosette pattern.

In axisymmetric potential models, new types of orbital shapes are allowed, and circular orbits only exist in the symmetry plane (here, the $x$-$y$ plane). In particular, as we will see, generic orbits are no longer confined to a plane and instead can oscillate in the $z$ direction with a non-commensurate frequency.

Even though the concept of the circular velocity only makes sense in the symmetry plane, we can still use the value of the circular velocity (computed assuming a spherical mass enclosed) as a way of initializing orbits, because the value of $v_{\rm circ}$ will have the right order of magnitude to stay bound and remain *close* to circular in many regimes. To demonstrate the types of orbits that we see in axisymmetric potentials, we will therefore compute three orbits in both the MN and comparison Plummer potentials using the circular velocity to inform our initial conditions. In the example below, we will start one orbit in the $x,y$ plane, one started slightly above the $x,y$ plane, and a third started far above the $x,y$ plane.

```python
# We first define the positions: These are [x, y, z] values
# for the three orbital initial conditions
mn_xyz = ([[8, 0, 0.0], [8, 0, 1.0], [8, 0, 10.0]] * u.kpc).T

# We compute the "circular velocity" (assuming this is a
# spherical potential) at each of the locations
mn_vcirc = mn_disk.circular_velocity(mn_xyz)

# We then use the circular velocity to set the scale of our
# initial velocities: We set the vy equal to the circular velocity,
# and vx equal to 10% of the circular velocity.
# The line below uses Numpy array broadcasting:
# https://numpy.org/doc/stable/user/basics.broadcasting.html
mn_vxyz = mn_vcirc[np.newaxis] * np.array([0.1, 1, 0])[:, np.newaxis]

mn_w0 = gd.PhaseSpacePosition(pos=mn_xyz, vel=mn_vxyz)
```

We will use these same intial conditions to compute orbits in both potential models and compare them below:

```python
mn_dt = 1.0 * u.Myr
mn_steps = 4000

mn_orbits = mn_disk.integrate_orbit(mn_w0, dt=mn_dt, n_steps=mn_steps)

plummer_orbits = plummer.integrate_orbit(mn_w0, dt=mn_dt, n_steps=mn_steps)
```

Let's plot the 3D configurations of the orbits computed in each of the potentials: We can use the `.plot_3d()` method on any `Orbit` object to make these plots for us, and in each panel the different colors will correspond to the different initial conditions:

```python
fig, axes = plt.subplots(1, 2, figsize=(16, 8), subplot_kw=dict(projection="3d"))

_ = mn_orbits.plot_3d(ax=axes[0])
_ = plummer_orbits.plot_3d(ax=axes[1])

for ax in axes:
    ax.azim = 30
    ax.elev = 15

axes[0].set_title("Miyamoto–Nagai")
axes[1].set_title("Plummer")
```

We could also instead all 2D projections of the orbits (xy, xz, yz) using the `.plot()` method:

```python
for orbits, name in zip([mn_orbits, plummer_orbits], ["Miyamoto–Nagai", "Plummer"]):
    fig = orbits.plot()
    for ax in fig.axes:
        ax.set_xlim(-15, 15)
        ax.set_ylim(-15, 15)
    fig.suptitle(name, fontsize=22)
```

Visually, it the orbits in the Plummer potential are planar (2D; as we expect), but in the Miyamoto–Nagai potential, the orbits seem to fill a 3D volume (except for the orbit started in the symmetry plane). In fact, in $x-y$ projections these orbits look a lot like their spherical analogs, however these have vertical deviations as well that give vertical thickness to the orbits: These orbits are called "tube" orbits because the surface they cover looks like a hollowed-out tube.

Because of the azimuthal symmetry of axisymmetric potentials, orbits in these models are sometimes also plotted in the *meridional plane*, which plots the cylindrical radius $R$ vs. $z$. We can plot cylindrical coordinates by using the `.cylindrical` attribute of `Orbit` objects, and then specify that we want to plot just $R$ (called `rho` in Gala) and $z$ by passing these in with the `components=` keyword argument:

```python
fig = mn_orbits.cylindrical.plot(["rho", "z"], labels=["$R$ [kpc]", "$z$ [kpc]"])
fig.axes[0].set_title("Miyamoto–Nagai", fontsize=22)
```

### Exercise: Why do the orbits in the Miyamoto–Nagai potential look different from the Plummer orbits?

Compute and plot the three components of the angular momentum for all orbits in both potentials. Do you see any differences? Given what we discussed about integrals of motion in the last tutorial, what do you think the connection is between the angular momentum components and the phase-space properties of the orbit?

```python

```

## The epicyclic approximation: Building intuition for close-to-planar, nearly-circular orbits

Because of the azimuthal symmetry of axisymmetric potentials (so $L_z$ is conserved), the full Hamiltonian for any orbit in an axisymmetric potential (in terms of cylindrical position and conjugate momentum coordinates $p_R, p_\phi, p_z$)

$$
H(R, \phi, z, p_R, p_\phi, p_z) = \frac{1}{2}(p_R^2 + \frac{p_\phi^2}{R^2} + p_z^2) + \Phi(R, \phi, z)
$$

can be reduced to a 2D Hamiltonian that governs the motion in $R$ and $z$ (noting that $p_\phi = L_z$)

$$
H(R, z, p_R, p_z; L_z) = \frac{1}{2}(p_R^2 + p_z^2) + \Phi(R, z) + \frac{L_z^2}{2\,R^2}
$$

where now $L_z$ can be thought of as a parameter that labels an orbit, not as a coordinate. Because the terms in a Hamiltonian are often grouped into "terms that depend on the momentum coordinates" and "terms that depend on the position coordinates," the dependence on $\frac{L_z^2}{2\,R^2}$ is sometimes absorbed into the expression of the potential and referred to as the *effective potential* $\Phi_{\rm eff}(R, z)$:

$$
\Phi_{\rm eff}(R, z) = \Phi(R, z) + \frac{L_z^2}{2\,R^2}
$$

The equations of motion for $R$ and $z$ are therefore

$$
\begin{align}
    \dot{p_R} &= - \frac{\partial H}{\partial R} = - \frac{\partial \Phi_{\rm eff}}{\partial R}\\
    \dot{R} &= p_R\\
    \ddot{R} &= - \frac{\partial \Phi_{\rm eff}}{\partial R}\\
\end{align}
$$

and

$$
\begin{align}
    \dot{p_z} &= - \frac{\partial H}{\partial z} = - \frac{\partial \Phi}{\partial z}\\
    \dot{z} &= p_z\\
    \ddot{z} &= - \frac{\partial \Phi}{\partial z}\\
\end{align}
$$

In general, for relevant axisymmetric potentials used in galactic dynamics, the potential expressions are complex enough that the partial derivative expressions still contain terms that mix $R$ and $z$ so that these are still coupled differential equations.

In disk galaxies, however, most stars are on orbits such that their maximum excursions in $z$, sometimes called $z_\textrm{max}$, are much smaller than the mean cylindrical radius of the orbit, i.e. $z_\textrm{max} \ll \textrm{mean}(R)$. In this limit, it is often conceptually useful (and sometimes quantitatively reasonable) to treat the motion as if it were decoupled in the radial $R$ and vertical $z$ dimensions. In reality, the motion *is* coupled for any non-planar orbit, as we saw with the equations of motion above. We can also see this geometrically using the numerical orbits we computed above: A truly uncoupled, non-resonant orbit would fill a rectangular area in positional coordinates, whereas instead orbits in the meridional plane have a slope to their upper and lower $z$, and have curvature at maximum and minimum $R$. For example, for the non-planar, intermediate orbit we computed above, compare the area filled by this orbit to the rectangular frame of the meridional plane plot:

```python
fig = mn_orbits[:, 1].cylindrical.plot(["rho", "z"], labels=["$R$", "$z$"])
```

However, as an approximation to gain intuition, we can make the simplifying assumption that the motion is decoupled. Making this assumption is equivalent to assuming that the potential model is separable such that

$$
\Phi(R, z) \approx \Phi(R) + \Phi(z) \quad .
$$

With the assumption of decoupled motion, and from observing that the orbital trajectory in the meridional plane oscillates in both radius $R$ and height $z$, we can use the effective potential to find the central points in each coordinate about which the orbit oscillates. We do this by taking the derivative of the effective potential and setting it equal to zero. For the vertical direction, this is easy: We assume that any potential is symmetric about $z=0$ (a condition for it to be axisymmetric), and so

$$
0 = \frac{\partial \Phi(z)}{\partial z}
$$

must occur anywhere $z=0$.

For radius,

$$
\begin{align}
0 &= \frac{\partial \Phi_\textrm{eff}}{\partial R} \\
\frac{\partial \Phi(R)}{\partial R} &= \frac{L_z^2}{R^3}
\end{align}
$$

By convention, the radius at which this expression is valid is called the *guiding-center radius*, $R_g$:

$$
\left.\frac{\partial \Phi(R)}{\partial R}\right|_{R_g} = \frac{L_z^2}{R_g^3}
$$

The guiding center radius is an important conceptual quantity: The "guiding center" is an implicit component of the orbit that revolves around the center of the potential on a circular orbit with a constant frequency. Therefore, given our approximations so far, the only reason an orbit in an axisymmetric potential appears non-circular is because a given orbit may make radial $R$ and vertical $z$ oscillations away from the (circular, planar) guiding-center orbit.

### Exercise: Estimate the guiding center radius for an orbit in the Miyamoto–Nagai potential

Estimate the guiding center radius of a planar orbit in the MN potential (the `mn_disk` we defined above) with the initial conditions:

$$
(x,y,z) = (8.5, 0, 0.02)~\textrm{kpc}\\
(v_x,v_y,v_z) = (0, 168, 0)~\textrm{km}~\textrm{s}^{-1}\\
$$

Hint: you might find the root finder `scipy.optimize.root` useful!

Compute an orbit from these initial conditions, plot it in the meridional plane, and draw a vertical line on the plot at the location of the guiding center radius.

```python
from scipy.optimize import root
```

```python

```

---

For the orbit you computed in the exercise above, we can get a better understanding of the geometry of epicyclic motion by plotting the orbit in a coordinate frame that rotates with the azimuthal frequency of the guiding center orbit. This frame rotates with a constant angular speed around the $z$ axis with a frequency (from dimensional analysis):

$$
\Omega_{\phi}^* = \frac{L_z}{R_g^2}
$$

With Gala, we can define a rotating reference frame using the `ConstantRotatingFrame()` frame class:

```python
guiding_center_Omega_phi = trial_Lz / trial_Rg ** 2
guiding_center_frame = gp.ConstantRotatingFrame(
    [0, 0, 1] * guiding_center_Omega_phi, units=galactic
)
```

We can then transform from the default, static frame to the rotating frame using the `Orbit.to_frame()` method:

```python
guiding_center_orbit = trial_orbit.to_frame(guiding_center_frame)
guiding_center_orbit
```

The returned object is still an `Orbit` instance, so we can plot its trajectory in the x-y plane using `.plot()` as we have done above:

```python
fig, ax = plt.subplots(figsize=(6, 6))

guiding_center_orbit.plot(["x", "y"], auto_aspect=False, axes=ax)
ax.scatter(trial_Rg.value, 0, color="tab:red")

ax.set_xlim(7, 10)
ax.set_ylim(-1.5, 1.5)
```

In the figure above, the red dot shows the location of the guiding center, and the black smeared-out ellipse is the orbit: This is the radial epicycle! The fact that it does not close on itself (to form a perfect ellipse) is because our assumptions are approximate: the orbit also oscillates in $z$.

For even less eccentric orbits, a crude way of approximating an epicycle orbit is to combine two circular orbits: one orbit around the origin of the coordinate system (i.e. the guiding center), and the other around the guiding radius

$$
Z(t) = e^{i\,\Omega_\phi t} \, \times \, A_R\,e^{-i\,\Omega_R t}\\
x(t) = \textrm{Re}(Z)\\
y(t) = \textrm{Im}(Z)
$$

where the minus sign is because the epicycle rotates in the opposite sense, as we learned in the lectures, and $A_R$ is the amplitude of the radial epicycle (which is related to the eccentricity).

Here is an interactive plot that lets us vary the R amplitude and the ratio of $\Omega_\phi/\Omega_R$. If you recall the previous tutorial, there are two limiting cases for the frequency ratio: When $\Omega_\phi=\Omega_R$ we get a Keplerian orbit, and when $\Omega_\phi = \frac{1}{2}\Omega_R$ we get an elliptical orbit centered on the origin. Try playing with the parameter values (using the sliders under the plot below). Can you find any resonant orbits? What resonances do they correspond to?

```python
from ipywidgets import interact, widgets
from IPython.display import display

t = np.arange(0, 32, 1e-2)
Omega_phi = 2 * np.pi

fig, ax = plt.subplots(figsize=(6, 6))

(l,) = ax.plot([], [], marker="")

ax.set_xlim(-2, 2)
ax.set_ylim(-2, 2)

ax.set_xlabel("$x$")
ax.set_ylabel("$y$")

def plot_func(Omega_phi_over_Omega_R, amp_R):
    Omega_R = Omega_phi / Omega_phi_over_Omega_R
    zz = amp_R * np.exp(-1j * Omega_R * t) + 1
    zz = zz * np.exp(1j * Omega_phi * t)

    l.set_data(zz.real, zz.imag)

    fig.canvas.draw()
    display(fig)

plt.close()
```

```python
omega_slider = widgets.FloatSlider(min=0.5, max=1.0, step=0.02, value=1.74)
amp_R_slider = widgets.FloatSlider(min=0, max=1, step=0.1, value=0.3)

interact(plot_func, Omega_phi_over_Omega_R=omega_slider, amp_R=amp_R_slider);
```

### One more frequency: azimuthal, radial, and vertical frequencies for orbits in axisymmetric models

From the previous tutorial and the previous section here, you should now be familiar with the azimuthal and radial frequencies, $\Omega_\phi, \Omega_R$. Generic orbits in axisymmetric potentials also have a third frequency: The $z$ or *vertical* frequency $\Omega_z$. Under the epicycle approximation, and with an assumed separable (in $R$, $z$) Hamiltonian, the equations of motion for $R$ and $z$ of an object are:

$$
\ddot{R} = -\frac{\partial \Phi_\textrm{eff}}{\partial R}\\
\ddot{z} = -\frac{\partial \Phi_\textrm{eff}}{\partial z}
$$

As noted previously, generically we have to solve these expressions numerically. However, to gain some intuition about the expected orbital properties, like the radial and vertical frequencies, it is useful to make one more approximation related to the epicycle assumption. If we are dealing with near-circular and near-coplanar orbits, we can get an expression for the effective potential close to an orbit's guiding center position $(R, z) \sim (R_g, 0)$ by Taylor expanding the effective potential: This will allow us to solve the orbital equations analytically.

Expanding the effective potential $\Phi_\textrm{eff}$ around the guiding center, we get:

$$
\Phi_{\textrm{eff}}(R, z) \approx \Phi_{\textrm{eff}}(R_g, 0) +
    \frac{1}{2}\left.\frac{\partial^2\Phi_\textrm{eff}}{\partial R^2}\right|_{(R_g, 0)} \, (R-R_g)^2 +
    \frac{1}{2}\left.\frac{\partial^2\Phi_\textrm{eff}}{\partial z^2}\right|_{(R_g, 0)} \, z^2 +
    \mathcal{O}((R-R_g)\,z^2)
$$

Note that in this Taylor series approximation, the potential is separable up to mixed terms (like $(R-R_g)\,z^2$)! With this approximation, the equations of motion are (introducing the variable $X = R-R_g$ for convenience)

$$
\ddot{X} = - \left.\frac{\partial^2\Phi_\textrm{eff}}{\partial R^2}\right|_{(R_g, 0)} \, X\\
\ddot{z} = - \left.\frac{\partial^2\Phi}{\partial z^2}\right|_{(R_g, 0)} \, z
$$

which you may recognize as equations for two independent simple harmonic oscillators: One in radius, that oscillates around the guiding center, and one in vertical position that oscillates around the midplane of the potential. From these expressions, we can read off the expected frequencies of oscillation (for orbits started at the midplane, at the guiding center):

$$
\Omega_R^2 = \left.\frac{\partial^2\Phi_\textrm{eff}}{\partial R^2}\right|_{(R_g, 0)}\\
\Omega_z^2 = \left.\frac{\partial^2\Phi}{\partial z^2}\right|_{(R_g, 0)}
$$

### Exercise: Estimate (analytically) the radial, vertical, and azimuthal frequencies for an orbit at the Solar radius

The Sun is approximately at a radius of $R_\odot \approx 8.1~\textrm{kpc}$ in the Milky Way's disk. This region of the disk is still dominated (in mass) by the gravitational potential of the stars and gas, so we can neglect the dark matter halo for first approximation.

Assuming $R_g = R_\odot$, estimate the azimuthal frequency of a circular orbit at the Solar circle using the `mn_disk` potential we defined above.

```python

```

Recall that the expression for the effective potential for a Miyamoto–Nagai disk is:

$$
\Phi_\textrm{eff}(R, z) = - \frac{G \, M}{\sqrt{R^2 + (a + \sqrt{b^2 + z^2})^2}} + \frac{L_z^2}{2\,R^2}
$$

Estimate the radial and vertical frequencies for an orbit near the solar circle. How do the frequency and period values compare (radial to azimuthal to vertical)?

*Hint: you either want to take second derivatives of the expression above and evaluate this manually, or you can use the `.hessian()` method on any Gala potential object (but note that this computes the Cartesian 2nd derivative matrix at a specified position)*

```python

```

## Other Axisymmetric and Flattened Potential Models

We have so far worked a lot with the Miyamoto–Nagai potential model, however like with spherical potential models, there are many options for potential–density pairs for flattened, axisymmetric potential models. In fact, as hinted at above, we can take any spherical potential model, replace the radius value with an ellipsoidal radius, and this gets us a flattened potential. However, it is worth noting that doing this for a spherical potential–density pair does not always lead to an analytic potential–density pair in axisymmetric coordinates.

One other example of a flattened potential that can be related to a density distribution is the flattened logarithmic potential:

$$
\Phi_L(R, z) = \frac{1}{2} v_0^2 \ln\left( R^2 + \frac{z^2}{q^2} + r_c^2 \right)\\
\rho_L(R, z) = \frac{v_0^2}{4\pi \, G \, q^2} \, \frac{(2\,q^2 +1)\,r_c^2 + R^2 + (2 - \frac{1}{q^2})\,z^2}{(R^2 + \frac{z^2}{q^2} + r_c^2)^2}
$$

where $q$ sets the amount of flattening (when $q = 1$, this is a spherical model).

Like the spherical logarithmic potential, this flattened model has the feature that when $R \gg r_c$, the circular velocity curve is close to constant with a value $v_c \approx v_0$, and is therefore useful for constructing simple mass models for computing orbits in the combined stars + dark matter potential of a Galaxy. However, as this model was defined by turning a spherical logarithmic *potential* model into a flattened potential model, the density does not have to be physically meaningful at all positions for all parameter values.

### Exercise: Are there values for the parameters (q, r_c) that lead to unphysical density values?

What values of the parameters?

```python

```

## Integrals of motion for axisymmetric orbits: Connection to action-angle coordinates

We have talked a lot about approximating orbits in axisymmetric potentials because most stars in the Milky Way are on orbits that can be described using the terminology and quantities we have discussed so far. We have focused a lot on the different frequencies that an orbit can have in different cylindrical coordinates, assuming that the motion is separable. However another important aspect of the geometry of axisymmetric orbits is the *amplitude* of oscillation in the different coordinates.

As we discussed early on in this tutorial, in axisymmetric models, a generic orbit has three integrals of motion, which could be energy, $z$-component of angular momentum, and the "third integral." It turns out that any orbit also has a different set of three integrals of motion that has a closer connection to the orbital geometry we discussed: something close to the amplitude of oscillations in the radial direction, something close to the amplitude of oscillations in the vertical direction, and the $z$-component of angular momentum (which is the analog of the others for the azimuthal direction) — these are the "radial action" $J_R$, the "vertical action" $J_z$, and the "azimuthal action" $J_\phi = L_z$. I added the phrase "something close to" in the previous sentence because the actions are not exactly related to the amplitudes in $R$ and $z$: Most of the expressions and intuition-building approximations we have made above have been under the assumption that the potential is separable in these coordinates. But we also saw that this is not true in detail (there is curvature to the orbital boundaries when orbits are plotted in the meridional $R,Z$ plane). The actions are like a generalization of the conceptual / approximate arguments we made above, but which take into account the fact that the orbits are not separable in detail.

Actions are special and useful for a number of reasons. For one, they connect to the intuitions we built about orbits above: For example, orbits with larger radial action have larger oscillations in $R$, and the same for $z$. But, as mentioned briefly in the previous tutorial, they also help define a dynamically-useful coordinate system known as action-angle coordinates. In this coordinate system, any point in ordinary phase-space $(\boldsymbol{x}, \boldsymbol{v})$ can be represented instead by a location in action-angle coordinates $(\boldsymbol{\theta}, \boldsymbol{J})$, where for axisymmetric orbits the actions and angles are:

$$
\boldsymbol{J} = (J_R, J_\phi, J_z)\\
\boldsymbol{\theta} = (\theta_R, \theta_\phi, \theta_z)
$$

where $J_\phi = L_z$, and in the limit $J_R, J_z \rightarrow 0$, $\theta_\phi \rightarrow \phi$.

The actions, which are the momentum coordinates, are integrals of motion, so $\dot{\boldsymbol{J}} = 0$. The Hamiltonian equations of motion are therefore:

$$
\dot{J}_i = -\frac{\partial H}{\partial \theta_i} = 0\\
\dot{\theta}_i = \frac{\partial H}{\partial J_i}
$$

This implies that the Hamiltonian $H$ must depend only on the actions, and so the expressions $\frac{\partial H}{\partial J_i}$ must be a function only of the actions. We can therefore integrate the equations of motion for the angle variables:

$$
\dot{\theta}_i = \frac{\partial H}{\partial J_i}\\
\theta_i(t) = \frac{\partial H}{\partial J_i} \, t + \theta_{i}(0)
$$

The angle variables increase linearly with time with a constant rate set by $\frac{\partial H}{\partial J_i}$, which must have units of frequency: We therefore define the *fundamental frequencies* of an orbit to be:

$$
\Omega_i = \frac{\partial H}{\partial J_i}
$$

which, in vector form, will be $\boldsymbol{\Omega} = (\Omega_R, \Omega_\phi, \Omega_z)$!

Action-angle coordinates are conceptually and dynamically very useful, and there is a lot more to be said about them. However, a subtlety of working with these coordinates is that it is often challenging to transform from $(x, v) \rightarrow (J, \theta)$ and even more challenging to transform back $(J, \theta) \rightarrow (x, v)$. The transformation often requires either numerically integrating an orbit for a long time, or making approximations to the potential (that often are not valid for all types of orbits). We will not talk much about the different methods for numerically estimating actions and angles (unless you ask!), but we recommend checking out [Sanders et al. 2016](https://ui.adsabs.harvard.edu/abs/2016MNRAS.457.2107S/abstract) if you are interested.


# Computational and Numerical Methods
## Group 16
### Set 11 (08-10-2018): Nonlinear Systems and the Newton Method
#### Vidhin Parmar   201601003
#### Parth Shah      201601086

<form action="javascript:code_toggle()"><input type="submit" id="toggleButton" value="Show Code"></form>

```python
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import scipy
import scipy.linalg as la
from sympy import symbols, diff
from sympy import *
import mpmath
```

$x^2+4y^2-9=0$

$18y−14x^2+ 45 = 0$

```python
th = np.arange(0, 2*np.pi + np.pi/50, np.pi/50)
plt.plot(3*np.cos(th), 1.5*np.sin(th), label="x^2+y^2=9")
t = np.arange(-5, 5 + 0.01, 0.01)
plt.plot((9/14)*t, (-2.5 + ((9/28)*(t**2))), label="18y - 14x^2 +45 =0")
plt.legend(loc="best")
plt.grid(True)
plt.show()
```

```python
x,y = symbols('x y')
f = lambda a,b: a**2 + 4*b**2 - 9
g = lambda a,b: -14*a**2 + 18*b + 45

fx = lambdify((x,y), diff(x**2 + 4*y**2 - 9, x), "numpy")
fy = lambdify((x,y), diff(x**2 + 4*y**2 - 9, y), "numpy")
gx = lambdify((x,y), diff(-14*x**2 + 18*y + 45, x), "numpy")
gy = lambdify((x,y), diff(-14*x**2 + 18*y + 45, y), "numpy")

x0, y0 = 1, 1

for i in range(25):
    mat1 = np.array([[fx(x0,y0), fy(x0,y0)], [gx(x0,y0), gy(x0,y0)]], dtype=np.float32)
    mat2 = np.array([[-f(x0,y0)], [-g(x0,y0)]], dtype=np.float32)
    ans = la.solve(mat1, mat2)
    x0, y0 = x0+ans[0,0], y0+ans[1,0]

"Root1 : ", x0, y0
```

    ('Root1 : ', 2.1372167369636323, 1.0526519628113704)

```python
x,y = symbols('x y')
f = lambda a,b: a**2 + 4*b**2 - 9
g = lambda a,b: -14*a**2 + 18*b + 45

fx = lambdify((x,y), diff(x**2 + 4*y**2 - 9, x), "numpy")
fy = lambdify((x,y), diff(x**2 + 4*y**2 - 9, y), "numpy")
gx = lambdify((x,y), diff(-14*x**2 + 18*y + 45, x), "numpy")
gy = lambdify((x,y), diff(-14*x**2 + 18*y + 45, y), "numpy")

x0, y0 = -1, 1

for i in range(25):
    mat1 = np.array([[fx(x0,y0), fy(x0,y0)], [gx(x0,y0), gy(x0,y0)]], dtype=np.float32)
    mat2 = np.array([[-f(x0,y0)], [-g(x0,y0)]], dtype=np.float32)
    ans = la.solve(mat1, mat2)
    x0, y0 = x0+ans[0,0], y0+ans[1,0]

"Root2 : ", x0, y0
```

    ('Root2 : ', -2.1372167369636323, 1.0526519628113704)

```python
x,y = symbols('x y')
f = lambda a,b: a**2 + 4*b**2 - 9
g = lambda a,b: -14*a**2 + 18*b + 45

fx = lambdify((x,y), diff(x**2 + 4*y**2 - 9, x), "numpy")
fy = lambdify((x,y), diff(x**2 + 4*y**2 - 9, y), "numpy")
gx = lambdify((x,y), diff(-14*x**2 + 18*y + 45, x), "numpy")
gy = lambdify((x,y), diff(-14*x**2 + 18*y + 45, y), "numpy")

x0, y0 = -1, -1

for i in range(25):
    mat1 = np.array([[fx(x0,y0), fy(x0,y0)], [gx(x0,y0), gy(x0,y0)]], dtype=np.float32)
    mat2 = np.array([[-f(x0,y0)], [-g(x0,y0)]], dtype=np.float32)
    ans = la.solve(mat1, mat2)
    x0, y0 = x0+ans[0,0], y0+ans[1,0]

"Root3 : ", x0, y0
```

    ('Root3 : ', -1.2031669633477737, -1.3740805342399418)

```python
x,y = symbols('x y')
f = lambda a,b: a**2 + 4*b**2 - 9
g = lambda a,b: -14*a**2 + 18*b + 45

fx = lambdify((x,y), diff(x**2 + 4*y**2 - 9, x), "numpy")
fy = lambdify((x,y), diff(x**2 + 4*y**2 - 9, y), "numpy")
gx = lambdify((x,y), diff(-14*x**2 + 18*y + 45, x), "numpy")
gy = lambdify((x,y), diff(-14*x**2 + 18*y + 45, y), "numpy")

x0, y0 = 1, -1

for i in range(25):
    mat1 = np.array([[fx(x0,y0), fy(x0,y0)], [gx(x0,y0), gy(x0,y0)]], dtype=np.float32)
    mat2 = np.array([[-f(x0,y0)], [-g(x0,y0)]], dtype=np.float32)
    ans = la.solve(mat1, mat2)
    x0, y0 = x0+ans[0,0], y0+ans[1,0]

"Root4 : ", x0, y0
```

    ('Root4 : ', 1.2031669633477737, -1.3740805342399418)


# Reproduce and extend the A. Destexhe *Neural. Comp* 2009 paper on sustained asynchronous firing in networks of AdEx neurons

```python
%pylab inline
```

    Populating the interactive namespace from numpy and matplotlib

```python
from brian2 import *
```

```python
import input_factory as inpf
```

```python
def beautify_spikes(statemon,spikemon,neuron_id):
    vm = statemon[neuron_id].v[:]
    offset = statemon.t[0]#/defaultclock.dt
    dt = statemon.t[1]-statemon.t[0]
    spike_times = spikemon.t[spikemon.i == neuron_id]
    for t in spike_times:
        i = int((t-offset) / dt)
        vm[i] = 20*mV
    return vm
```

```python
rc('axes',grid=True)
```

```python
def raster_spikes(spikemon, offset=0, ax=None,color='k'):
    if ax is None:
        f,ax = subplots(1,1)
    ax.plot(spikemon.t/ms, offset+spikemon.i,',',c=color)

```

## AdEx model in Destexhe's formulation

\begin{align}
C_m\frac{dV}{dt} & =   -g_l(V-E_l) + g_l\Delta V\exp(\frac{V-V_T}{\Delta V}) - w/S\\
\tau_w\frac{dw}{dt} & =  a(V-E_l) - w
\end{align}

Здесь можно писать всевозможные объяснения. В том числе **полужирным** и *курсивом*

```python

```

### Parameters

```python
Cm = 1*uF/cm2
gl = 0.05 *mS/cm2
El = -60 *mV
VT = -50 *mV
Sarea = 20e3*um2
tau_w = 600 *ms
#tau_w = 100 *ms
deltaT = 2.5*mV
Vcut = -20*mV
#Vreset = -55*mV
Vreset = -55*mV

# Synapses
E_e = 0*mV     # Excitatory synaptic reversal potential (AMPA and NMDA receptors)
E_i = -80*mV   # Inhibitory synaptic reversal potential (GABAA receptors)
tau_e = 5*ms   # time scale of excitatory synaptic conductance
tau_i = 10*ms  # time scale of excitatory synaptic conductance

```

```python
AdEx_equations = Equations('''
dv/dt = (-gl*(v-El) + gl*deltaT*exp((v-VT)/deltaT) - u/Sarea + Isyn/Sarea + Ibias/Sarea + Iapp/Sarea)/Cm : volt (unless refractory)
du/dt = (a*(v-El) - u)/tau_w: amp  # adaptation variable
stim_amp : 1
Ibias : amp
Iapp = stim_amp*input_current(t,i): amp
''')

# Synaptic input
synaptic_equations = Equations("""
Isyn =  -g_e*(v-E_e) - g_i*(v-E_i) : amp
dg_e/dt = -g_e/tau_e : siemens
dg_i/dt = -g_i/tau_i : siemens
""")
```

```python
PY_pars = dict(a=0.001*uS, b=0.04*nA)
FS_pars = dict(a=0.001*uS, b=0.00*nA)
LTS_pars = dict(a=0.02*uS, b=0.0*nA)
TC_pars = dict(a=0.04*uS, b=0.0*nA)
RE_pars = dict(a=0.03*uS, b=0.08*nA)
RS1_pars = dict(a=0.001*uS, b=0.04*nA)
RS2_pars = dict(a=0.001*uS, b=0.005*nA)
```

```python
RS2_pars
```

    {'a': 1. * nsiemens, 'b': 5. * pamp}

## Reproducing Fig. 1

Below we learn to create groups of neurons with specified parameters and run a simulation.

```python
start_scope()

Nneurons = 3

defaultclock.dt = 0.1*ms

G_LTS = NeuronGroup(Nneurons,
                AdEx_equations+synaptic_equations,
                refractory=2.5*ms,
                namespace=LTS_pars,
                threshold='v>Vcut',
                reset='v=Vreset; u += b',
                method='exponential_euler')

G_LTS.set_states(dict(v=El,u=0))

G_LTS.stim_amp = (-0.25,0,0.25)
M_LTS = StateMonitor(G_LTS, ['v','u'], record=True)
S_LTS = SpikeMonitor(G_LTS,)

G_RS1 = NeuronGroup(Nneurons,
                AdEx_equations+synaptic_equations,
                refractory=2.5*ms,
                namespace=RS1_pars,
                threshold='v>Vcut',
                reset='v=Vreset; u += b',
                method='exponential_euler')

G_RS1.set_states(dict(v=El,u=0))

G_RS1.stim_amp = (-0.25,0,0.25)
M_RS1 = StateMonitor(G_RS1, ['v','u'], record=True)
S_RS1 = SpikeMonitor(G_RS1,)

G_RS2 = NeuronGroup(Nneurons,
                AdEx_equations+synaptic_equations,
                namespace=RS2_pars,
                threshold='v>Vcut',
                refractory=2.5*ms,
                reset='v=Vreset; u += b',
                method='exponential_euler')

G_RS2.set_states(dict(v=El,u=0))

G_RS2.stim_amp = (-0.25,0,0.25)
M_RS2 = StateMonitor(G_RS2, ['v','u'], record=True)
S_RS2 = SpikeMonitor(G_RS2,)

G_FS = NeuronGroup(Nneurons,
                AdEx_equations+synaptic_equations,
                namespace=FS_pars,
                refractory=2.5*ms,
                threshold='v>Vcut',
                reset='v=Vreset; u += b',
                method='exponential_euler')

G_FS.set_states(dict(v=El,u=0))

G_FS.stim_amp = (-0.25,0,0.25)
M_FS = StateMonitor(G_FS, ['v','u'], record=True)
S_FS = SpikeMonitor(G_FS,)

G_TC = NeuronGroup(Nneurons,
                AdEx_equations+synaptic_equations,
                namespace=TC_pars,
                threshold='v>Vcut',
                reset='v=Vreset; u += b',
                refractory=2.5*ms,
                method='exponential_euler')

G_TC.set_states(dict(v=El,u=0))

G_TC.stim_amp = (-0.25,0,0.25)
M_TC = StateMonitor(G_TC, ['v','u'], record=True)
S_TC = SpikeMonitor(G_TC,)

G_RE = NeuronGroup(Nneurons,
                AdEx_equations+synaptic_equations,
                namespace=RE_pars,
                threshold='v>Vcut',
                reset='v=Vreset; u += b',
                refractory=2.5*ms,
                method='exponential_euler')

G_RE.set_states(dict(v=El,u=0))

G_RE.stim_amp = (-0.25,0,0.25)
M_RE = StateMonitor(G_RE, ['v','u'], record=True)
S_RE = SpikeMonitor(G_RE,)

```

```python
input_current = inpf.get_step_current(200, 700, 1*ms, 1.0*nA,Nneurons=Nneurons)
```

```python
store()
```

```python
restore()
```

```python
%time run(1.5*second)
```

    CPU times: user 5.76 s, sys: 83.9 ms, total: 5.84 s
    Wall time: 15.8 s

```python
f, axs = subplots(2,3,sharex=True,sharey=True,figsize=(15,10))

axs[0,0].plot(M_RS1.t/ms, beautify_spikes(M_RS1,S_RS1,2)/mV,label='+0.25nA')
axs[0,0].plot(M_RS1.t/ms, beautify_spikes(M_RS1,S_RS1,0)/mV,label='-0.25nA')
axs[0,0].set_title("RS neuron, type1")
axs[0,0].legend()

axs[0,1].plot(M_RS2.t/ms, beautify_spikes(M_RS2,S_RS2,2)/mV,label='+0.25nA')
axs[0,1].plot(M_RS2.t/ms, beautify_spikes(M_RS2,S_RS2,0)/mV,label='-0.25nA')
axs[0,1].set_title("RS neuron, type2")
axs[0,1].legend()

axs[0,2].plot(M_FS.t/ms, beautify_spikes(M_FS,S_FS,2)/mV,label='+0.25nA')
axs[0,2].plot(M_FS.t/ms, beautify_spikes(M_FS,S_FS,0)/mV,label='-0.25nA')
axs[0,2].set_title("FS neuron")
axs[0,2].legend()

axs[1,0].plot(M_LTS.t/ms, beautify_spikes(M_LTS,S_LTS,2)/mV,label='+0.25nA')
axs[1,0].plot(M_LTS.t/ms, beautify_spikes(M_LTS,S_LTS,0)/mV,label='-0.25nA')
axs[1,0].set_title("LTS neuron")
axs[1,0].legend()

axs[1,1].plot(M_TC.t/ms, beautify_spikes(M_TC,S_TC,2)/mV,label='+0.25nA')
axs[1,1].plot(M_TC.t/ms, beautify_spikes(M_TC,S_TC,0)/mV,label='-0.25nA')
axs[1,1].set_title("TC neuron")
axs[1,1].legend()

axs[1,2].plot(M_RE.t/ms, beautify_spikes(M_RE,S_RE,2)/mV,label='+0.25nA')
axs[1,2].plot(M_RE.t/ms, beautify_spikes(M_RE,S_RE,0)/mV,label='-0.25nA')
axs[1,2].set_title("RE neuron")
axs[1,2].legend()

#title('Regular spiking (pyramidal) neuron')
xlabel('time, ms'); axs[0,0].set_ylabel('Vm, mV');

```

Итак, нейроны классов LTS, TC и RE демонстрируют спайки в ответ на снятие гиперполяризующего тока, потому что...

## Networking

Для начала определим ряд параметров и уравнений для описания событий, связанных со спайками в пресинаптическом нейроне, вероятностью выброса кванта нейромедиатора и пластичности.

```python
tau_ps = 0.8*second   # facilitation timescale
tau_ns = 1.5*second   # replenishing timescale
tau_stdp = 20*ms     # STDP time constant

p_s0 = 1.0            # Start with no facilitation
p_s0 = 0.8            # Start with some facilitation
#p_s0 = 0.6           # ground-state probability of release
depletion_amount = 0.0  # no depression; change in 0...1 interval
stdp_amount = 0 # change between 0 and 1

epsilon = 0.02        # sparsity synaptic connectivity

Apre0 = 0.01
Apost0 = -Apre0*1.05

w_e = 6*nS
w_i = 67*nS

w_ext = 3*nS

tm_plasticity_model = Equations('''
dp_s/dt = (p_s0-p_s)/tau_ps : 1 (event-driven)    # release probability
dn_s/dt = (1-n_s)/tau_ns   : 1    (event-driven)  # fraction of resources available
''')

stdp_model=Equations('''
dApre/dt = -Apre/tau_stdp : 1 (event-driven)    # STDP
dApost/dt = -Apost/tau_stdp : 1(event-driven)   # STDP
w_syn: 1
''')

tm_action_pre='''
p_s += p_s0*(1-p_s) # facilitation
r_s = p_s*n_s       # probability of release
n_s -= depletion_amount*r_s          # depletion

'''

stdp_action_pre='''
Apre += Apre0
'''

stdp_action_post='''
Apost += Apost0
'''

pre_actions_e = '''
w_syn = clip(w_syn+Apost, (1-stdp_amount)*w_e/nS, w_e/nS)
g_e_post += w_syn*r_s*nS
'''
pre_actions_i = '''
w_syn = clip(w_syn+Apost, (1-stdp_amount)*w_i/nS, w_i/nS)
g_i_post += w_syn*r_s*nS
'''

post_actions_e='''
w_syn = clip(w_syn+Apre,(1-stdp_amount)*w_e/nS,w_e/nS)
'''

post_actions_i='''
w_syn = clip(w_syn+Apre,(1-stdp_amount)*w_i/nS,w_i/nS)
'''
```

```python

```

## Reproducing Fig. 7

Попробуем собрать сеть из RS, LTS и FS нейронов. Первые два класса будут образовывать возбуждающие синапсы, последний класс будет описывать интернейроны и давать тормозящие синапсы.

Но сначала нам нужно создать группу "как-бы нейронов", которые будут выдавать ПД в случайные моменты времени в течение заданного периода. Этим мы будем моделировать внешние возбуждающие связи, необходимые для затравки собственной сетевой активности.

```python
def transient_PoissonGroup(N, rate, start_ms, stop_ms):
    thr = 'rand()<rates*dt and t<{stop}*ms and t >={start}*ms'.format(start=start_ms, stop=stop_ms)
    g = NeuronGroup(N, 'rates : Hz', threshold=thr)
    g.rates = rate
    return g

```

```python

```

```python
shared_kwargs = dict(
                 threshold='v>Vcut',
                 refractory=2.5*ms,
                 reset='v=Vreset; u += b',
                 method='exponential_euler')
```

```python
start_scope()

defaultclock.dt = 0.1*ms

exc_fraction = 0.8
lts_fraction = 0.1

Ntotal = 1000
Nin = int((1-exc_fraction)*Ntotal)             # inhibitory neurons (interneurons, aka fast spiking)
Nrs = int(exc_fraction*(1-lts_fraction)*Ntotal)        # regular spiking pyramidal neurons
Nlts = int(exc_fraction*lts_fraction*Ntotal)        # low-threshold spiking pyramidal neurons

prob_scale = 2000/Ntotal

P = transient_PoissonGroup(100, 200/second, 50, 250)
Pspikes = SpikeMonitor(P,record=True)

Gin = NeuronGroup(Nin, AdEx_equations+synaptic_equations,
                  namespace = FS_pars, **shared_kwargs)
Gin.set_states(dict(v=El, u=0))

Min = StateMonitor(Gin,['v','g_e'],record=True,dt=0.5*ms)
Sin = SpikeMonitor(Gin, record=True)

Glts = NeuronGroup(Nlts, AdEx_equations+synaptic_equations,
                   namespace=LTS_pars,
                   **shared_kwargs)
Glts.set_states(dict(v=El, u=0))

Mlts = StateMonitor(Glts,['v','g_e','g_i'],record=True,dt=0.5*ms)
Slts = SpikeMonitor(Glts, record=True)

Grs = NeuronGroup(Nrs, AdEx_equations+synaptic_equations,
                 namespace=RS2_pars,
                 **shared_kwargs)
Grs.set_states(dict(v=El, u=0))
Mrs = StateMonitor(Grs,['v','g_e','g_i'],record=True, dt=0.5*ms)
Srs = SpikeMonitor(Grs, record=True)
```

```python
Nrs,Nlts,Nin
```

    (720, 80, 199)

```python
## Excitatory synapses

# External input

S_ext = Synapses(P, Grs, on_pre='g_e_post += w_ext')

exc_syn_args = dict(model=tm_plasticity_model+stdp_model,
                    on_pre=tm_action_pre+stdp_action_pre+pre_actions_e,
                    on_post=stdp_action_post+post_actions_e)

inh_syn_args = dict(model=tm_plasticity_model+stdp_model,
                    on_pre=tm_action_pre+stdp_action_pre+pre_actions_i,
                    on_post=stdp_action_post+post_actions_i)
# From RS
S_rs_rs = Synapses(Grs, Grs, **exc_syn_args)
S_rs_lts = Synapses(Grs, Glts,**exc_syn_args)
S_rs_in = Synapses(Grs, Gin, **exc_syn_args)

# From LTS
S_lts_rs = Synapses(Glts, Grs, **exc_syn_args)
S_lts_lts = Synapses(Glts, Glts, **exc_syn_args)
S_lts_in = Synapses(Glts, Gin, **exc_syn_args)

# From FS
S_in_rs = Synapses(Gin, Grs, **inh_syn_args)
S_in_lts = Synapses(Gin, Glts, **inh_syn_args)
S_in_in = Synapses(Gin, Gin, **inh_syn_args)
```

```python
epsilon*prob_scale
```

    0.04

```python
%%time

S_ext.connect(p=0.005)

S_rs_rs.connect('i!=j',p=epsilon*prob_scale)
S_rs_lts.connect(p=epsilon*prob_scale)
S_rs_in.connect(p=epsilon*prob_scale)

S_lts_rs.connect(p=epsilon*prob_scale)
S_lts_lts.connect('i!=j',p=epsilon*prob_scale)
S_lts_in.connect(p=epsilon*prob_scale)

excitatory_synapses = [S_rs_rs,S_rs_lts,S_rs_in,
                       S_lts_rs,S_lts_lts,S_lts_in]

S_in_rs.connect(p=epsilon*prob_scale)
S_in_lts.connect(p=epsilon*prob_scale)
S_in_in.connect('i!=j',p=epsilon*prob_scale)

inhibitory_synapses = [S_in_rs, S_in_lts, S_in_in]
cortical_connections = excitatory_synapses+inhibitory_synapses
```

    CPU times: user 147 ms, sys: 3.97 ms, total: 151 ms
    Wall time: 150 ms

```python
cortical_net = Network(P, Gin, Grs, Glts, Mrs, Srs, Mlts, Slts, Min,Sin, Pspikes,
                       S_ext, *cortical_connections)
```

```python
input_current = inpf.get_step_current(200, 700, 1*ms, 1.0*nA,Nneurons=Ntotal)
```

```python
cortical_net.store()
```

`>>>>> Start over simulation from here <<<<<`

К этой ячейке нужно возвращаться, чтобы поменять такие параметры как задержка между нейронами, наличие фасилитации или депрессии в пресинаптической мембране и STDP.

```python
cortical_net.restore() # "вспоминаем сохраненное состояние"
```

```python
p_s0 = 1.0               # ground-state probability of release; change in 0...1 interval to allow for facilitation
depletion_amount = 0.00  # depression amount; change in 0...1 interval; increase for depressing synapses
stdp_amount = 0.0        # STDP amount; change between 0 and 1
```

```python
average_delay = 3*ms             # средняя синаптическая задержка (фактически, определяется длиной аксона)
delay_jitter = 0.5*average_delay # разброс задержек. Здесь, полностью случайный
```

```python
for s in cortical_connections:
    s.delay = "average_delay-0.5*delay_jitter + delay_jitter*rand()"  # delays in synapses.
    s.p_s = p_s0
    s.n_s = 1.0
```

```python
%%time

cortical_net.run(3*second, report='text') # три секунды модельного времени. У меня считает около 2 минут.
```

    Starting simulation at t=0. s for a duration of 3. s
    2.3465 s (78%) simulated in 10s, estimated 3s remaining.
    3. s (100%) simulated in 12s
    CPU times: user 1min 6s, sys: 471 ms, total: 1min 7s
    Wall time: 1min 37s

```python
#mean([Mrs.v[i].max() for i in range(Nrs)])
```

```python
k=randint(Nrs)
```

```python
f, axs = subplots(3,1, sharex=True, figsize=(14,5))

axs[0].plot(Mrs.t/ms, beautify_spikes(Mrs,Srs,k)/mV)

axs[0].set_ylabel('membrane potential, mV')
axs[0].set_title('активность произвольно выбранной RS-клетки')

axs[1].plot(Mrs.t/ms, Mrs.g_e[k]/nS)
axs[1].set_ylabel('g_e, nS')
axs[1].set_title('Суммарная проводимость возбуждающих синапсов в этой же клетке')

axs[2].plot(Mrs.t/ms, Mrs.g_i[k]/nS)
axs[2].set_ylabel('g_i, nS')
axs[2].set_title('Суммарная проводимость тормозящих синапсов в этой же клетке')
#xlim(200,400)
xlabel('time, ms');
xlim(0,Mrs.t[-1]/ms)

```

```python

```

```python
f,ax = subplots(1,1,figsize=(14,5))

raster_spikes(Pspikes,ax=ax,color='gray',offset=-100)
raster_spikes(Srs,ax=ax,color='red',offset=0)
raster_spikes(Slts,ax=ax,color='orange',offset=0+Nrs)
raster_spikes(Sin,ax=ax,color='blue',offset=0+Nrs+Nlts)

title('Raster of network activity (blue: IN, orange:LTS, red:PY. gray:external)')
xlabel('time, ms')
ylabel('neuron #')
xlim(0,Mrs.t[-1]/ms)
ylim(-100,Ntotal)
```

**Вопросы:**
   1. Как зависит наличие самоподдерживающейся активности от средней синаптической задержки? Оцените длину аксона, соответствующую задержке в 10 мс.
   1. Как зависит наличие самоподдерживающейся активности от разброса синаптических задержек?
   2. Как зависит наличие самоподдерживающейся активности от наличия синаптической фасилитации ($p_s0 < 1$)?
   2. Как зависит наличие самоподдерживающейся активности от наличия синаптической депрессии?
   3. Как активность сети от вклада механизма STDP?

```python
def get_mean_CV(spikemon):
    active_cells = unique(spikemon.i)
    acc = []
    for i in active_cells:
        spike_times = spikemon.t[spikemon.i==i]
        if len(spike_times)>10:
            isi_v = diff(spike_times/ms)
            acc.append(std(isi_v)/mean(isi_v))
    return mean(acc)
```

```python
get_mean_CV(Srs), get_mean_CV(Sin)
```

    (1.8012195180619281, 1.6800724809030774)

```python
def get_spike_counts(spikemon, neuronid,start_time=0*ms,stop_time=3*second, bin_size=5*ms):
    spike_times = spikemon.t[spikemon.i==neuronid]
    acc = [((spike_times>=start)*(spike_times<start+bin_size)).sum() for
          start in arange(start_time,stop_time-bin_size,bin_size)]
    return array(acc)
```

```python
def spike_count_entropy(spike_counts):
    N = len(spike_counts)
    probs = [sum(spike_counts==i)/N for i in range(0,max(spike_counts))]
    return -sum(p*log2(p) for p in probs)

```

```python
x=get_spike_counts(Srs,10)
plot(x)
```

```python
spike_count_entropy(x)
```

    0.04502804609154204

```python

```


# Import

```python
#source
#region

#import
#region
import math
from sympy import *
import matplotlib.pyplot as plt
from numpy import linspace
import numpy as np
from sympy.codegen.cfunctions import log10
from sympy.abc import x,t,y
from sympy.plotting import plot
#endregion

#symbol declaration
#region
x, t = symbols('x t')
f = symbols('f', cls=Function)
#endregion
```

# Read Input

```python
#input, output
#region
def ReadInput(file):
    f = file.readline()
    (lowT, upT) = map(lambda s: N(s), file.readline().split(","))
    (lowX, upX) = map(lambda s: N(s), file.readline().split(","))
    (t0, x0) = map(lambda s: N(s), file.readline().split(","))
    epsilon = N(file.readline())
    return (f, lowT, upT, lowX, upX, t0, x0, epsilon)
#endregion
```

# Main Function

```python
def Pica1(f, deltaT, deltaX, t0, x0, M, L, epsilon, mode = ""):

    N = GetN(M, L, deltaT, deltaX, epsilon)
    xn = SymbolicIntegrate(f, t0, x0, N, mode)

    return xn

def Pica2(f, deltaT, t0, x0, M, L, epsilon, length = 69, mode = ""):

    xn = []
    segmentLength = 2 * deltaT / length
    n = (int)(length / 2)
    for i in range(-n, n + 1):
        xn.append([t0 + i * segmentLength, x0])

    xn = NumericIntegrate(f, xn, x0, segmentLength, epsilon, mode)

    return xn

def Pica(filename, length = None, M = None, L = None, deltaT = None, mode = ""):
    try:
        file = open(filename, "r")
        (f, lowT, upT, lowX, upX, t0, x0, epsilon) = ReadInput(file)
        f = sympify(f)
    except:
        raise ValueError("invalid Pica input")
    file.close()

    if not lowX< x0 <upX or not lowT< t0< upT:
        raise ValueError("invalid Pica input")

    if M is None:
        M = GetM(x, lowT, upT, lowX, upX)
    else:
        if M <= 0:
            raise ValueError("invalid Pica input")

    if L is None:
        L = GetL(x, lowT, upT, lowX, upX)
    else:
        if L < 0:
            raise ValueError("invalid Pica input")
    if L == 0:
        SymbolicIntergrate(f, t0, x0, 1, mode)

    deltaX = min(x0 - lowX, upX - x0)
    if deltaT is None:
        deltaT = min(deltaX / M, 1 / (2 * L), t0 - lowT, upT - t0)
    interval = (float(t0-deltaT), float(t0+deltaT))
    if length is None:
        return (Pica1(f, deltaT, deltaX, t0, x0, M, L, epsilon, mode), interval)
    return Pica2(f, deltaT, t0, x0, M, L, epsilon, length, mode)
```

# Main loop (integrate)

```python

#region
def NumericIntegrate(f, xn, x0, segmentLength, epsilon, mode = ""):
    n = (int) (len(xn)/2)
    segmentLength /=2
    maxError = -math.inf
    loop = 0
    while abs(maxError) > epsilon:
        if mode == "test":
            dx = []
        loop += 1
        maxError = -math.inf

        integral = 0
        for i in range(n, 0, -1):
            integral = integral - segmentLength * (f.subs([(t, xn[i][0]), (x, xn[i][1])]) + f.subs([(t, xn[i - 1][0]), (x, xn[i - 1][1])]))
            newValue = x0 + integral
            error = abs(xn[i - 1][1] - newValue)
            xn[i - 1][1] = newValue
            if(error > maxError): maxError = error
            if mode == 'test':
                dx.append((xn[i][0], error))

        integral = 0
        for i in range(n, 2 * n):
            integral = integral + segmentLength * (f.subs([(t, xn[i][0]), (x, xn[i][1])]) + f.subs([(t, xn[i + 1][0]), (x, xn[i + 1][1])]))
            newValue = x0 + integral
            error = abs(xn[i + 1][1] - newValue)
            xn[i + 1][1] = x0 + integral
            if(error > maxError): maxError = error
            if mode == 'test':
                dx.append((xn[i][0], error))

        if mode == "test":
            print("Lặp lần ", loop, " với max error = ", maxError)
            PlotPairs(dx)
            plt.show()
    return xn

def GetN(M, L, deltaT, deltaX, epsilon, mode = ""):
    h = deltaT * L
    N = 1
    error = M * deltaT
    while error > epsilon:
        N+=1
        error = error * h / N
    return N

def SymbolicIntegrate(f, t0, x0, N, mode = ''):
     xn = x0
     for i in range(0,N):
        if mode == 'test':
            print(xn.evalf(2))
        xn = x0 + integrate(f.replace(x,xn), (t,t0,t))
     return xn

#endregion

```

# Not implemented supremum finder

```python
# sup
#region
def GetM(f, lowT, upT, lowX, upX):
    #not implemented
    return 10
def GetL(f, lowT, upT, lowX, upX):
    #not implemented
    return 10

#endregion
```

# Plot

```python
#plot
#region
def PlotPairs(pairList):
    t,x = zip(*pairList)
    plt.scatter(t,x)

def PlotSymbol(symbolOutput):
    func, interval = symbolOutput
    #t = linspace(interval[0], interval[1], 1000)
    #func = t**3/3 + t**7/67
    plot((func, (t, interval[0], interval[1])))

def PlotBoth(symbolOutput, pairList):
    t1, x1 = zip(*pairList)
    plt.scatter(t1,x1)
    func, interval = symbolOutput
    t_vals = linspace(interval[0], interval[1], 1000)
    lam_x = lambdify(t, func, modules=['numpy'])
    x_vals = lam_x(t_vals)
    plt.plot(t_vals, x_vals)

def Plot(f, interval):
    t_vals = linspace(interval[0], interval[1], 1000)
    lam_x = lambdify(t, f, modules = ['numpy'])
    x_vals = lam_x(t_vals)
    plt.plot(t_vals, x_vals)
    #plt.show()
#endregion

```

# Test

```python
filename = "input2.txt"

result = Pica(filename, M = 2.5, L = 1)
result1 = Pica(filename, M = 2.5, L = 1, length = 31, mode = 'test')

print(result[0].evalf(2))
print("Khoảng hội tụ:", result[1])
PlotBoth(result, result1)
```

```python
filename = "input3.txt"

result = Pica(filename, M = 5, L = 10)
result1 = Pica(filename, length = 31, M = 12, L = 10)

print(result[0].evalf(2))
print("Khoảng hội tụ:", result[1])
PlotBoth(result, result1)
print(result1[0][0])

```

```python
filename = "input4.txt"

#result = Pica(filename, M = 15, L = 1.5)
result1 = Pica(filename, length = 31, M = 15, L = 1.5)

PlotPairs(result1)
#print(result)
#PlotBoth(result, result1)
```

```python
filename = "input1.txt"
result = Pica(filename, M = 2, L = 1)
result1 = Pica(filename, length = 31, M = 2, L = 1, mode = 'test')

print(result[0].evalf(2))
print("Khoảng hội tụ:", result[1])
PlotBoth(result, result1)
#Plot(sin(10*t) + cos(10*t), result[1])
#plt.show()
```

```python
filename = "input5.txt"

result1 = Pica(filename, length = 222, M = 50, L = 1, mode = 'test')
PlotPairs(result1)
interval = (float(result1[0][0]), float(result1[len(result1)-1][0]))
Plot(cos(300*t), interval)
plt.show()
```

```python
filename = "input6.txt"

result = Pica(filename, M = 250, L = 100)
result1 = Pica(filename, length = 31, M = 250, L = 100, mode = 'test')

print(result[0].evalf(2))
print("Khoảng hội tụ:", result[1])
PlotBoth(result, result1)
#interval = (float(result1[0][0]), float(result1[len(result1)-1][0]))
#Plot(sin(100*t), interval)
```

```python

```

```python

```

Topic 5 - Solving Differential Equations/26.1.Pica/Pica.ipynb


```python
from devito import *
from examples.seismic.source import WaveletSource, TimeAxis
from examples.seismic import plot_image
import numpy as np

from sympy import init_printing, latex
init_printing(use_latex=True)
```

```python
# Initial grid: 1km x 1km, with spacing 100m
extent = (2000., 2000.)
shape = (81, 81)
x = SpaceDimension(name='x', spacing=Constant(name='h_x', value=extent[0]/(shape[0]-1)))
z = SpaceDimension(name='z', spacing=Constant(name='h_z', value=extent[1]/(shape[1]-1)))
grid = Grid(extent=extent, shape=shape, dimensions=(x, z))
```

```python
class DGaussSource(WaveletSource):

    def wavelet(self, f0, t):
        a = 0.004
        return -2.*a*(t - 1./f0) * np.exp(-a * (t - 1./f0)**2)

# Timestep size from Eq. 7 with V_p=6000. and dx=100
t0, tn = 0., 200.
dt = 1e2*(1. / np.sqrt(2.)) / 60.
time_range = TimeAxis(start=t0, stop=tn, step=dt)

src = DGaussSource(name='src', grid=grid, f0=0.01, time_range=time_range)
src.coordinates.data[:] = [1000., 1000.]
src.show()
```

```python
# Now we create the velocity and pressure fields
p = TimeFunction(name='p', grid=grid, space_order=4, time_order=1)
vx= TimeFunction(name='vx', grid=grid, staggered=(0, 1, 0), space_order=4, time_order=1)
vz = TimeFunction(name='vz', grid=grid, staggered=(0, 0, 1), space_order=4, time_order=1)
```

```python
def derivative(f, dim, order, stagger=centered):
    """
    Utility function to generate staggered derivatives
    """
    diff = dim.spacing
    if stagger == left:
        off = -.5
    elif stagger == right:
        off = .5
    else:
        off = 0.
    idx = [(dim + int(i+.5+off)*diff) for i in range(-int(order / 2), int(order / 2))]
    return -f.diff(dim).as_finite_difference(idx, x0=dim + off*dim.spacing)
```

```python
t = grid.stepping_dim
time = grid.time_dim

# We need some initial conditions
V_p = 4.0
#V_s = 1.0
density = 3.

dx = 100.
rox = 1/density * dt
roz = 1/density * dt
l2m = V_p*V_p*density * dt

c1 = 9.0/8.0;
c2 = -1.0/24.0;

# The source injection term
src_p = src.inject(field=p.forward, expr=src)

# 2nd order acoustic according to fdelmoc
u_vx_2 = Eq(vx.forward, vx + rox * derivative(p, x, 2, stagger=left))
u_vz_2 = Eq(vz.forward, vz + roz * derivative(p, z, 2, stagger=left))
u_p_2 = Eq(p.forward, p + l2m * (derivative(vx.forward,x,2,stagger=right) + derivative(vz.forward,z,2,stagger=right)))

# u_vx_2 = Eq(vx.forward, vx + rox * (-P[t, x, z] + P[t, x-1, z]))

#vz[ix*n1+iz] -= roz[ix*n1+iz]*(p[ix*n1+iz] - p[ix*n1+iz-1]);
# u_vz_2 = Eq(vz.forward, vz + roz * (-P[t, x, z] + P[t, x, z-1]))

# u_p_2 = Eq(p.forward, p + l2m * ( (-VX[t+1, x+1, z] + VX[t+1, x, z]) + (-VZ[t+1, x, z+1] + VZ[t+1, x, z])))

op_2 = Operator([u_vx_2, u_vz_2, u_p_2] + src_p)

# 4th order acoustic according to fdelmoc
u_vx_4 = Eq(vx.forward, vx + rox * derivative(p, x, 4, stagger=left))
u_vz_4 = Eq(vz.forward, vz + roz * derivative(p, z, 4, stagger=left))
u_p_4 = Eq(p.forward, p + l2m * (derivative(vx.forward,x,4,stagger=right) + derivative(vz.forward,z,4,stagger=right)))

op_4 = Operator([u_vx_4, u_vz_4, u_p_4] + src_p)
```

```python
### 2nd order acoustic

# Reset the fields
vx.data[:] = 0.
vz.data[:] = 0.
p.data[:] = 0.

op_2(time=src.time_range.num-1)

# Let's see what we got....
plot_image(vx.data[0])
plot_image(vz.data[0])
plot_image(p.data[0])
```

```python
### 4th order acoustic

# Reset the fields
vx.data[:] = 0.
vz.data[:] = 0.
p.data[:] = 0.

op_4(time=src.time_range.num-1)

# Let's see what we got....
plot_image(vx.data[0])
plot_image(vz.data[0])
plot_image(p.data[0])
```

```python

```


# Using Optimization in Accelerator Physics

by Kelly Anderson

Optimization in accelerator physics comes in many different forms depending on what effects you are looking for. There are various multipole magnets and cavities found along an accelerator for the purpose of bending and accelerating the particles. A collection of components in an accelerator is referred to as a lattice. One of the most common lattices used is called a FODO lattice. This lattice consists of a focusing quadrupole, drift space (open), a defocusing quadrupole, and another drift space. This lattice has the ultimate effect of focusing a group of charged particles.

Using a defocusing quadrupole may seem counter productive but it is needed. When we talk about the position in an accelerator we refer to the translational components, x and y, and the longitudinal component, s. The longitudinal component is the direction along the accelerator. The reason this defocusing quadrupole is needed is because a quadrupole will only focus one translational component and defocus the other. This will still have a focusing effect because the particles farther from the center of the magnet will be bent more. A component that doesn't have this defocusing effect, like a dipole, does not have strong enough focusing for higher energy.

Even this simple lattice can be optimized. If you assume the length of the magnets is relatively short you can actually treat the quadrupoles as optical lenses. If the accelerator is a ring then you might want the particles to be in the same position at the beginning and end of the ring so they form a closed orbit. This could be for storage of particles or light sources that use synchrotron radiation. If we assume a one dimensional system and a series of FODO cells as our entire ring (we will ignore the small error effects of dipoles in this case) we can treat each component as a matrix as shown.

\begin{equation}
    M = M_n M_{n-1}...M_2 M_1
\end{equation}

\begin{equation}\begin{pmatrix}
    x \\
    x'
\end{pmatrix}_{s_o+l} = M(s_o+l|s_o)
\begin{pmatrix}
    x \\
    x'
\end{pmatrix}_{s_o}\end{equation}

The matrix for each component can be multiplied together to get a single matrix for a whole ring. Because our system has constant energy, phase space area is conserved which will mean the determinant of the matrix is 1. By looking at the eigenvalues and making sure our results stay finite as you do infinite passes through the ring, we can also determine that the trace of the matrix must be less than 2.

\begin{equation}    \begin{pmatrix}
    x \\
    x'
    \end{pmatrix}_{s_o+nl}
    = [M(s_o+l|s_o)]^n
    \begin{pmatrix}
    x \\
    x'
    \end{pmatrix}_{s_o}
    = A_1 \lambda_1^n \mathbf{v}_1+A_2 \lambda_2^n \mathbf{v}_2
\end{equation}

\begin{equation}
\det(M)=1
\end{equation}
\begin{equation}
\det(M)=1=\lambda_1 \lambda_2
\end{equation}

\begin{equation}
    \lambda_2 = 1/\lambda_1
\end{equation}

\begin{equation}
    |\lambda_1| = |\lambda_2|=1
\end{equation}
\begin{equation}
    \lambda_1 = e^{i\mu}
\end{equation}
\begin{equation}
    \lambda_2 = e^{-i\mu}
\end{equation}

\begin{equation}
    TrM = \lambda_1 + \lambda_2 = e^{i\mu}+e^{-i\mu}=2\cos{\mu}
\end{equation}
\begin{equation}
    |TrM|\leq2
\end{equation}

Below is the matrix for a FODO cell where L is the drift length and f is the focal length. There is also a common transformation of variables used in accelerators called twiss parameters. Without too much detail, the beta function corresponds to the amplitude of the position argument and alpha is related to the change in beta. Shown is how to convert the (x, x’) matrix into a twiss one.

\begin{equation}
    M =
    \begin{pmatrix}
    1 && 0 \\ -\frac{1}{2f} && 1
    \end{pmatrix}
    \begin{pmatrix}
    1 && L \\ 0 && 1
    \end{pmatrix}
    \begin{pmatrix}
    1 && 0 \\ \frac{1}{f} && 1
    \end{pmatrix}
    \begin{pmatrix}
    1 && L \\ 0 && 1
    \end{pmatrix}
    \begin{pmatrix}
    1 && 0 \\ -\frac{1}{2f} && 1
    \end{pmatrix}
\end{equation}

\begin{equation}
    =
    \begin{pmatrix}
    1 - \frac{L^2}{2f^2} && 2L(1+\frac{L}{2f}) \\ -\frac{L}{2f}(1-\frac{L}{2f}) && 1 - \frac{L^2}{2f^2}
    \end{pmatrix}
\end{equation}

\begin{equation}
    \begin{pmatrix}
    \beta \\ \alpha \\ \gamma
    \end{pmatrix}_{s_2}
    =
    \begin{pmatrix}
    M_{11}^2 && -2 M_{11} M_{12} && M_{12}^2 \\
    -M_{11} M_{21} && M_{11} M_{22} + M_{12} M_{21} && -M_{12} M_{22} \\
    M_{21}^2 && -2 M_{21} M_{22} && M_{22}^2
    \end{pmatrix}
    \begin{pmatrix}
    \beta \\ \alpha \\ \gamma
    \end{pmatrix}_{s_1}
\end{equation}

Using the previous condition we can optimize the length and focal distance to ensure a stable orbit.

\begin{equation}
    \frac{1}{2}|TrM| = |1-\frac{L^2}{2f^2}| \leq 1
\end{equation}

\begin{equation}
    |\frac{L^2}{f^2}| \leq 4
\end{equation}

Here is a program that shows what happens if that condition is not met.

```python
import numpy as np
import matplotlib.pyplot as plt
from numpy import linalg as LA
from matplotlib.pyplot import figure
```

```python
class FODO(object):
    def __init__(self, drift, focal):
        self.L = drift
        self.f = focal

        self.M = np.zeros((2,2))

        self.M[0,0] = 1 - self.L**2/(2*(self.f**2))
        self.M[1,0] = 2*self.L*(1+self.L/(2*self.f))
        self.M[0,1] = (-self.L/(2*self.f**2))*(1-self.L/(2*self.f))
        self.M[1,1] = 1 - self.L**2/(2*(self.f**2))

        self.FODO=np.zeros((3,3))
        self.FODO[0,0] = self.M[0,0]**2
        self.FODO[1,0] =-2*self.M[0,0]*self.M[1,0]
        self.FODO[2,0] = self.M[1,0]**2
        self.FODO[0,1] = -self.M[0,0]*self.M[0,1]
        self.FODO[1,1] = self.M[0,0]*self.M[1,1]+self.M[1,0]*self.M[0,1]
        self.FODO[2,1] = -self.M[1,0]*self.M[1,1]
        self.FODO[0,2] = self.M[0,1]**2
        self.FODO[1,2] = -2*self.M[0,1]*self.M[1,1]
        self.FODO[2,2] = self.M[1,1]**2

    def __call__(self, x0, xp0, N):
        X=np.zeros((3,1))
        X_o=np.zeros((3,1))
        X_o[0,0] = x0
        X_o[1,0] = xp0
        X_o[2,0] = (1 + X_o[1,0]**2)/X_o[0,0]
        beta_o = X_o[0,0]
        alpha_o = X_o[1,0]

        beta_array = [X_o[0]]
        S = [0]

        for i in range(1,N):
            X = np.matmul(self.FODO,X_o)
            beta_array.append(X[0,0])
            X_o = X
            S.append(2*self.L*i)

        return(beta_o, alpha_o, S,beta_array)
```

```python
N_s = 250
L_s = 1
f_s = 15*L_s
N_u = 5
L_u = 1
f_u = np.sqrt(L_u**2/(4+1E-9))
x0 = 0.25E-2
xp0 = 0

FODO_unstable = FODO(L_u, f_u)
FODO_stable = FODO(L_s, f_s)

beta0_u, alpha0_u, S_u, beta_u = FODO_unstable(x0, xp0, N_u)
beta0_s, alpha0_s, S_s, beta_s = FODO_stable(x0, xp0, N_s)

print("Initial Conditions: beta_0 = " + str(beta0_s) + "[m], alpha_o = " + str(alpha0_s))
print("L = " + str(L_s) + "[m] f = " + str(f_s) + "[m]" + " N = " + str(N_s))
print("|TrM| = " + str(abs(FODO_stable.M[0,0] + FODO_stable.M[1,1])))
fig = plt.figure(figsize=(20,4))
grph = fig.add_axes([0.1,0.1,0.8,0.8])
grph.plot(S_s,beta_s)
plt.title("Stable Orbit")
plt.xlabel("s [m]")
plt.ylabel("beta [m]")
plt.show()
```

```python
print("Initial Conditions: beta_0 = " + str(beta0_u) + "[m], alpha_o = " + str(alpha0_u))
print("L = " + str(L_u) + "[m] f = " + str(f_u) + "[m]" + " N = " + str(N_u))
print("|TrM| = " + str(abs(FODO_unstable.M[0,0] + FODO_unstable.M[1,1])))
fig = plt.figure(figsize=(20,4))
grph = fig.add_axes([0.1,0.1,0.8,0.8])
grph.plot(S_u,beta_u)
plt.title("Unstable Orbit")
plt.xlabel("s [m]")
plt.ylabel("beta [m]")
plt.show()
```

Another example of optimization is with a more sophisticated lattice code called madx. This program allows one to put in components without the approximations we used before. It is used in accelerator physics to design and test lattices. Below is an example of a series of FODO cells in a ring. Small random errors are introduced to the alignment of the quadrupoles, like in real accelerators. Using dipoles placed after them call kickers, the program will automatically optimize the strength of these components to try and stabilize the x and y oscillations

```
circum=1000.0;
// circumfrance of ring;
lquad=3.0;
// length of quad;
ldi=5.0;
// length of dipole;
lbpm=0.0001;
// length of beam position monitor;
lcor = lbpm;
// lenght of corrector magnet;
ncell=20;
// number of cells in a ring;
lcell=circum/ncell;
// length of every cell;
ndc = 2;
// number of dipoles per cell;
nfqc = 1;
//number of focusing quads per cell;
ndqc = 1:
//number of defocusing quads per cell;
ndr1c = nqc;
// number of drift ones per cell;
ndr2c = ndr1c;
// numver of drift twos per cell;
ldrifthalf = (lcell - (nfqc+ndqc)*lquad - ndc*ldi - nfqc*(lbpm + lcor) - ndqc*(lbpm+lcor))/ndr1c;
// length of drift space per cell divided by the number of drift cells to account for 2 per quadropole;
fdrift1 = 0.5;
// % of drift space that is drift1;
ldrift1 = fdrift1*ldrifthalf;
// lenght of drift one;
ldrift2 = ldrifthalf - ldrift1;
// length of drift two;
fq=14.957;
// focal enght of quads;
k_1 = 1/(lquad*fq);
// k_1 for quads;

QF: QUADRUPOLE, L=lquad, K1=k_1, THICK=true;
QD: QUADRUPOLE, L=lquad, K1=-k_1, THICK=true;
DP: SBEND, L=ldi, ANGLE=(2*pi)/(ndc*ncell);
Dr1: DRIFT, L=ldrift1;
Dr2: DRIFT, L=ldrift2;
BPM: MONITOR, L=lbpm;
CH: HKICKER, L=lcor;
CV: VKICKER, L=lcor;

FODO: line=(QF,CH,BPM,Dr1,DP,Dr2,QD,CV,BPM,Dr2,DP,Dr1);
RING: line=20*(FODO);

return;
```

```
call, file='./Exercise_2/exercise_2.lat';
option, -echo;

Beam, particle = proton, energy=20.0;

use, sequence=RING;
eoption,seed=123456789,add=false;
select,flag=error,range=#S/#E,class=quadrupole;
ealign, dx:=tgauss(3.0)*1.0e-4, dy:=tgauss(3.0)*2.0e-4;

select, flag=twiss,column=name,s,x,y, dx, dy;

twiss,centre;

plot, haxis=s,vaxis=x,y, title="Position Pre Correction", file="./Exercise_2/ex2_xy_plot", colour=100.0;

coption,print=3;
correct,sequence=RING,flag=ring,mode=micado,plane=x,error=1e-6,ncorr=20,clist="./Exercise_2/ex2_xcor.txt",mlist="./Exercise_2/ex2_xbpm.txt";
correct,sequence=RING,flag=ring,mode=micado,plane=y,error=1e-6,ncorr=20,clist="./Exercise_2/ex2_ycor.txt",mlist="./Exercise_2/ex2_ybpm.txt";

twiss, centre;
plot, haxis=s,vaxis=x,y, title="Position Post Correction", colour=100.0;
stop;
```

---
# References

S. Y. Lee, Accelerator Physics, 4th Edition

D. A. Edwards, M. J. Syphers, An Introduction to the Physics of High Energy Accelerators

The beta funciton code was adapted from a code I wrote for a demonstration in a class Fall 2019.

The madx code is from what I wrote as an exercise working for Dr. Yue Hao.

```python

```