# Neuromatch Academy: Week 2, Day 4, Tutorial 1
# Neuronal Network Dynamics: Neural Rate Models

__Content creators:__ Qinglong Gu, Songtin Li, Arvind Kumar, John Murray, Julijana Gjorgjieva

__Content reviewers:__  Maryam Vaziri-Pashkam, Ella Batty, Lorenzo Fontolan, Richard Gao, Spiros Chavlis, Michael Waskom

---
# Tutorial Objectives

The brain is a complex system, not because it is composed of a large number of diverse types of neurons, but mainly because of how neurons are connected to each other. The brain is indeed a network of highly specialized neuronal networks.

The activity of a neural network constantly evolves in time. For this reason, neurons can be modeled as dynamical systems. The dynamical system approach is only one of the many modeling approaches that computational neuroscientists have developed (other points of view include information processing,  statistical models, etc.).

How the dynamics of neuronal networks affect the representation and processing of information in the brain is an open question. However, signatures of altered brain dynamics present in many brain diseases (e.g., in epilepsy or Parkinson's disease) tell us that it is crucial to study network activity dynamics if we want to understand the brain.

In this tutorial, we will simulate and study one of the simplest models of biological neuronal networks. Instead of modeling and simulating individual excitatory neurons (e.g., LIF models that you implemented yesterday), we will treat them as a single homogeneous population and approximate their dynamics using a single one-dimensional equation describing the evolution of their average spiking rate in time.

In this tutorial, we will learn how to build a firing rate model of a single population of excitatory neurons.

**Steps:**
- Write the equation for the firing rate dynamics of a 1D excitatory population.
- Visualize the response of the population as a function of parameters such as threshold level and gain, using the frequency-current (F-I) curve.
- Numerically simulate the dynamics of the excitatory population and find the fixed points of the system.
- Investigate the stability of the fixed points by linearizing the dynamics around them.

---
# Setup

```python
# Imports
import matplotlib.pyplot as plt
import numpy as np
import scipy.optimize as opt       # root-finding algorithm
```

```python
# @title Figure Settings
import ipywidgets as widgets       # interactive display
%config InlineBackend.figure_format = 'retina'
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle")
```

```python
# @title Helper functions

def plot_fI(x, f):
  plt.figure(figsize=(6, 4))  # plot the figure
  plt.plot(x, f, 'k')
  plt.xlabel('x (a.u.)', fontsize=14)
  plt.ylabel('F(x)', fontsize=14)
  plt.show()

def plot_dr_r(r, drdt, x_fps=None):
  plt.figure()
  plt.plot(r, drdt, 'k')
  plt.plot(r, 0. * r, 'k--')
  if x_fps is not None:
    plt.plot(x_fps, np.zeros_like(x_fps), "ko", ms=12)
  plt.xlabel(r'$r$')
  plt.ylabel(r'$\frac{dr}{dt}$', fontsize=20)
  plt.ylim(-0.1, 0.1)

def plot_dFdt(x, dFdt):
  plt.figure()
  plt.plot(x, dFdt, 'r')
  plt.xlabel('x (a.u.)', fontsize=14)
  plt.ylabel('dF(x)', fontsize=14)
  plt.show()
```

---
# Section 1: Neuronal network dynamics

```python
# @title Video 1: Dynamic networks
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="p848349hPyw", width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

## Section 1.1: Dynamics of a single excitatory population

Individual neurons respond by spiking. When we average the spikes of neurons in a population, we can define the average firing activity of the population. In this model, we are interested in how the population-averaged firing varies as a function of time and network parameters. Mathematically, we can describe the firing rate dynamic as:

\begin{align}
\tau \frac{dr}{dt} &= -r + F(w\cdot r + I_{\text{ext}})  \quad\qquad (1)
\end{align}

$r(t)$ represents the average firing rate of the excitatory population at time $t$, $\tau$ controls the timescale of the evolution of the average firing rate, $w$ denotes the strength (synaptic weight) of the recurrent input to the population, $I_{\text{ext}}$ represents the external input, and the transfer function $F(\cdot)$ (which can be related to f-I curve of individual neurons described in the next sections) represents the population activation function in response to all received inputs.

To start building the model, please execute the cell below to initialize the simulation parameters.

```python
# @markdown *Execute this cell to set default parameters for a single excitatory population model*

def default_pars_single(**kwargs):
  pars = {}

  # Excitatory parameters
  pars['tau'] = 1.     # Timescale of the E population [ms]
  pars['a'] = 1.2      # Gain of the E population
  pars['theta'] = 2.8  # Threshold of the E population

  # Connection strength
  pars['w'] = 0.  # E to E, we first set it to 0

  # External input
  pars['I_ext'] = 0.

  # simulation parameters
  pars['T'] = 20.       # Total duration of simulation [ms]
  pars['dt'] = .1       # Simulation time step [ms]
  pars['r_init'] = 0.2  # Initial value of E

  # External parameters if any
  pars.update(kwargs)

  # Vector of discretized time points [ms]
  pars['range_t'] = np.arange(0, pars['T'], pars['dt'])

  return pars
```

You can now use:
- `pars = default_pars_single()` to get all the parameters, and then you can execute `print(pars)` to check these parameters.
- `pars = default_pars_single(T=T_sim, dt=time_step)` to set new simulation time and time step
- To update an existing parameter dictionary, use `pars['New_para'] = value`

Because `pars` is a dictionary, it can be passed to a function that requires individual parameters as arguments using `my_func(**pars)` syntax.

## Section 1.2: F-I curves
In electrophysiology, a neuron is often characterized by its spike rate output in response to input currents. This is often called the **F-I** curve, denoting the output spike frequency (**F**) in response to different injected currents (**I**). We estimated this for an LIF neuron in yesterday's tutorial.

The transfer function $F(\cdot)$ in Equation $1$ represents the gain of the population as a function of the total input. The gain is often modeled as a sigmoidal function, i.e., more input drive leads to a nonlinear increase in the population firing rate. The output firing rate will eventually saturate for high input values.

A sigmoidal $F(\cdot)$ is parameterized by its gain $a$ and threshold $\theta$.

$$ F(x;a,\theta) = \frac{1}{1+\text{e}^{-a(x-\theta)}} - \frac{1}{1+\text{e}^{a\theta}}  \quad(2)$$

The argument $x$ represents the input to the population. Note that the second term is chosen so that $F(0;a,\theta)=0$.

Many other transfer functions (generally monotonic) can be also used. Examples are the rectified linear function $ReLU(x)$ or the hyperbolic tangent $tanh(x)$.

### Exercise 1: Implement F-I curve

Let's first investigate the activation functions before simulating the dynamics of the entire population.

In this exercise, you will implement a sigmoidal **F-I** curve or transfer function $F(x)$, with gain $a$ and threshold level $\theta$ as parameters.

```python
def F(x, a, theta):
  """
  Population activation function.

  Args:
    x (float): the population input
    a (float): the gain of the function
    theta (float): the threshold of the function

  Returns:
    float: the population activation response F(x) for input x
  """
  #################################################
  ## TODO for students: compute f = F(x) ##
  # Fill out function and remove
  raise NotImplementedError("Student excercise: implement the f-I function")
  #################################################

  # Define the sigmoidal transfer function f = F(x)
  f = ...

  return f

pars = default_pars_single()  # get default parameters
x = np.arange(0, 10, .1)      # set the range of input

# Uncomment below to test your function
# f = F(x, pars['a'], pars['theta'])
# plot_fI(x, f)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D4_DynamicNetworks/solutions/W2D4_Tutorial1_Solution_45ddc05f.py)

*Example output:*

### Interactive Demo: Parameter exploration of F-I curve
Here's an interactive demo that shows how the F-I curve changes for different values of the gain and threshold parameters. How do the gain and threshold parameters affect the F-I curve?

```python
# @title

# @markdown Make sure you execute this cell to enable the widget!

def interactive_plot_FI(a, theta):
  """
  Population activation function.

  Expecxts:
    a     : the gain of the function
    theta : the threshold of the function

  Returns:
    plot the F-I curve with give parameters
  """

  # set the range of input
  x = np.arange(0, 10, .1)
  plt.figure()
  plt.plot(x, F(x, a, theta), 'k')
  plt.xlabel('x (a.u.)', fontsize=14)
  plt.ylabel('F(x)', fontsize=14)
  plt.show()

_ = widgets.interact(interactive_plot_FI, a=(0.3, 3, 0.3), theta=(2, 4, 0.2))
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D4_DynamicNetworks/solutions/W2D4_Tutorial1_Solution_1c0165d7.py)

## Section 1.3: Simulation scheme of E dynamics

Because $F(\cdot)$ is a nonlinear function, the exact solution of Equation $1$ can not be determined via analytical methods. Therefore, numerical methods must be used to find the solution. In practice, the derivative on the left-hand side of Equation $1$ can be approximated using the Euler method on a time-grid of stepsize $\Delta t$:

\begin{align}
&\frac{dr}{dt} \approx \frac{r[k+1]-r[k]}{\Delta t}
\end{align}
where $r[k] = r(k\Delta t)$.

Thus,

$$\Delta r[k] = \frac{\Delta t}{\tau}[-r[k] + F(w\cdot r[k] + I_{\text{ext}}[k];a,\theta)]$$

Hence, Equation (1) is updated at each time step by:

$$r[k+1] = r[k] + \Delta r[k]$$

```python
# @markdown *Execute this cell to enable the single population rate model simulator: `simulate_single`*

def simulate_single(pars):
  """
  Simulate an excitatory population of neurons

  Args:
    pars   : Parameter dictionary

  Returns:
    rE     : Activity of excitatory population (array)

  Example:
    pars = default_pars_single()
    r = simulate_single(pars)
  """

  # Set parameters
  tau, a, theta = pars['tau'], pars['a'], pars['theta']
  w = pars['w']
  I_ext = pars['I_ext']
  r_init = pars['r_init']
  dt, range_t = pars['dt'], pars['range_t']
  Lt = range_t.size

  # Initialize activity
  r = np.zeros(Lt)
  r[0] = r_init
  I_ext = I_ext * np.ones(Lt)

  # Update the E activity
  for k in range(Lt - 1):
      dr = dt / tau * (-r[k] + F(w * r[k] + I_ext[k], a, theta))
      r[k+1] = r[k] + dr

  return r

help(simulate_single)
```

### Interactive Demo: Parameter Exploration of single population dynamics

Note that $w=0$, as in the default setting, means no recurrent input to the neuron population in Equation (1). Hence, the dynamics are entirely determined by the external input $I_{\text{ext}}$. Explore these dynamics in this interactive demo.

How does $r_{\text{sim}}(t)$ change with different $I_{\text{ext}}$ values? How does it change with different $\tau$ values? Investigate the relationship between $F(I_{\text{ext}}; a, \theta)$ and the steady value of $r(t)$.

Note that, $r_{\rm ana}(t)$ denotes the analytical solution - you will learn how this is computed in the next section.

```python
# @title

# @markdown Make sure you execute this cell to enable the widget!

# get default parameters
pars = default_pars_single(T=20.)

def Myplot_E_diffI_difftau(I_ext, tau):
  # set external input and time constant
  pars['I_ext'] = I_ext
  pars['tau'] = tau

  # simulation
  r = simulate_single(pars)

  # Analytical Solution
  r_ana = (pars['r_init']
           + (F(I_ext, pars['a'], pars['theta'])
           - pars['r_init']) * (1. - np.exp(-pars['range_t'] / pars['tau'])))

  # plot
  plt.figure()
  plt.plot(pars['range_t'], r, 'b', label=r'$r_{\mathrm{sim}}$(t)', alpha=0.5,
           zorder=1)
  plt.plot(pars['range_t'], r_ana, 'b--', lw=5, dashes=(2, 2),
           label=r'$r_{\mathrm{ana}}$(t)', zorder=2)
  plt.plot(pars['range_t'],
           F(I_ext, pars['a'], pars['theta']) * np.ones(pars['range_t'].size),
           'k--', label=r'$F(I_{\mathrm{ext}})$')
  plt.xlabel('t (ms)', fontsize=16.)
  plt.ylabel('Activity r(t)', fontsize=16.)
  plt.legend(loc='best', fontsize=14.)
  plt.show()

_ = widgets.interact(Myplot_E_diffI_difftau, I_ext=(0.0, 10., 1.),
                     tau=(1., 5., 0.2))
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D4_DynamicNetworks/solutions/W2D4_Tutorial1_Solution_65dee3e7.py)

## Think!
Above, we have numerically solved a system driven by a positive input. Yet, $r_E(t)$ either decays to zero or reaches a fixed non-zero value.
- Why doesn't the solution of the system "explode" in a finite time? In other words, what guarantees that $r_E$(t) stays finite?
- Which parameter would you change in order to increase the maximum value of the response?

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D4_DynamicNetworks/solutions/W2D4_Tutorial1_Solution_5a95a98e.py)

---
# Section 2: Fixed points of the single population system

```python
# @title Video 2: Fixed point
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="Ox3ELd1UFyo", width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

As you varied the two parameters in the last Interactive Demo, you noticed that, while at first the system output quickly changes, with time, it reaches its maximum/minimum value and does not change anymore. The value eventually reached by the system is called the **steady state** of the system, or the **fixed point**. Essentially, in the steady states the derivative with respect to time of the activity ($r$) is zero, i.e. $\displaystyle \frac{dr}{dt}=0$.

We can find that the steady state of the Equation. (1) by setting $\displaystyle{\frac{dr}{dt}=0}$ and solve for $r$:

$$-r_{\text{steady}} + F(w\cdot r_{\text{steady}} + I_{\text{ext}};a,\theta) = 0, \qquad (3)$$

When it exists, the solution of Equation. (3) defines a **fixed point** of the dynamical system in Equation (1). Note that if $F(x)$ is nonlinear, it is not always possible to find an analytical solution, but the solution can be found via numerical simulations, as we will do later.

From the Interactive Demo, one could also notice that the value of $\tau$ influences how quickly the activity will converge to the steady state from its initial value.

In the specific case of $w=0$, we can also analytically compute  the solution of Equation (1) (i.e., the thick blue dashed line) and deduce the role of $\tau$ in determining the convergence to the fixed point:

$$\displaystyle{r(t) = \big{[}F(I_{\text{ext}};a,\theta) -r(t=0)\big{]} (1-\text{e}^{-\frac{t}{\tau}})} + r(t=0)$$ \\

We can now numerically calculate the fixed point with a root finding algorithm.

## Exercise 2: Visualization of the fixed points

When it is not possible to find the solution for Equation (3) analytically, a graphical approach can be taken. To that end, it is useful to plot $\displaystyle{\frac{dr}{dt}}$ as a function of $r$. The values of $r$ for which the plotted function crosses zero on the y axis correspond to fixed points.

Here, let us, for example, set $w=5.0$ and $I^{\text{ext}}=0.5$. From Equation (1), you can obtain

$$\frac{dr}{dt} = [-r + F(w\cdot r + I^{\text{ext}})]\,/\,\tau $$

Then, plot the $dr/dt$ as a function of $r$, and check for the presence of fixed points.

```python
def compute_drdt(r, I_ext, w, a, theta, tau, **other_pars):
  """Given parameters, compute dr/dt as a function of r.

  Args:
    r (1D array) : Average firing rate of the excitatory population
    I_ext, w, a, theta, tau (numbers): Simulation parameters to use
    other_pars : Other simulation parameters are unused by this function

  Returns
    drdt function for each value of r
  """
  #########################################################################
  # TODO compute drdt and disable the error
  raise NotImplementedError("Finish the compute_drdt function")
  #########################################################################

  # Calculate drdt
  drdt = ...

  return drdt

# Define a vector of r values and the simulation parameters
r = np.linspace(0, 1, 1000)
pars = default_pars_single(I_ext=0.5, w=5)

# Uncomment to test your function
# drdt = compute_drdt(r, **pars)
# plot_dr_r(r, drdt)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D4_DynamicNetworks/solutions/W2D4_Tutorial1_Solution_c5280901.py)

*Example output:*

## Exercise 3: Fixed point calculation

We will now find the fixed points numerically. To do so, we need to specif initial values ($r_{\text{guess}}$) for the root-finding algorithm to start from. From the line $\displaystyle{\frac{dr}{dt}}$ plotted above in Exercise 2, initial values can be chosen as a set of values close to where the line crosses zero on the y axis (real fixed point).

The next cell defines three helper functions that we will use:

- `my_fp_single(r_guess, **pars)` uses a root-finding algorithm to locate a fixed point near a given initial value
- `check_fp_single(x_fp, **pars)`, verifies that the values of $r_{\rm fp}$ for which $\displaystyle{\frac{dr}{dt}} = 0$ are the true fixed points
- `my_fp_finder(r_guess_vector, **pars)` accepts an array of initial values and finds the same number of fixed points, using the above two functions

```python
# @markdown *Execute this cell to enable the fixed point functions*

def my_fp_single(r_guess, a, theta, w, I_ext, **other_pars):
  """
  Calculate the fixed point through drE/dt=0

  Args:
    r_guess  : Initial value used for scipy.optimize function
    a, theta, w, I_ext : simulation parameters

  Returns:
    x_fp    : value of fixed point
  """
  # define the right hand of E dynamics
  def my_WCr(x):
    r = x
    drdt = (-r + F(w * r + I_ext, a, theta))
    y = np.array(drdt)

    return y

  x0 = np.array(r_guess)
  x_fp = opt.root(my_WCr, x0).x.item()

  return x_fp

def check_fp_single(x_fp, a, theta, w, I_ext, mytol=1e-4, **other_pars):
  """
   Verify |dr/dt| < mytol

  Args:
    fp      : value of fixed point
    a, theta, w, I_ext: simulation parameters
    mytol   : tolerance, default as 10^{-4}

  Returns :
    Whether it is a correct fixed point: True/False
  """
  # calculate Equation(3)
  y = x_fp - F(w * x_fp + I_ext, a, theta)

  # Here we set tolerance as 10^{-4}
  return np.abs(y) < mytol

def my_fp_finder(pars, r_guess_vector, mytol=1e-4):
  """
  Calculate the fixed point(s) through drE/dt=0

  Args:
    pars    : Parameter dictionary
    r_guess_vector  : Initial values used for scipy.optimize function
    mytol   : tolerance for checking fixed point, default as 10^{-4}

  Returns:
    x_fps   : values of fixed points

  """
  x_fps = []
  correct_fps = []
  for r_guess in r_guess_vector:
    x_fp = my_fp_single(r_guess, **pars)
    if check_fp_single(x_fp, **pars, mytol=mytol):
      x_fps.append(x_fp)

  return x_fps

help(my_fp_finder)
```

```python
r = np.linspace(0, 1, 1000)
pars = default_pars_single(I_ext=0.5, w=5)
drdt = compute_drdt(r, **pars)

#############################################################################
# TODO for students:
# Define initial values close to the intersections of drdt and y=0
# (How many initial values? Hint: How many times do the two lines intersect?)
# Calculate the fixed point with these initial values and plot them
#############################################################################
r_guess_vector = [...]

# Uncomment to test your values
# x_fps = my_fp_finder(pars, r_guess_vector)
# plot_dr_r(r, drdt, x_fps)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D4_DynamicNetworks/solutions/W2D4_Tutorial1_Solution_0637b6bf.py)

*Example output:*

## Interactive Demo: fixed points as a function of recurrent and external inputs.

You can now explore how the previous plot changes when the recurrent coupling $w$ and the external input $I_{\text{ext}}$ take different values. How does the number of fixed points change?

```python
# @title

# @markdown Make sure you execute this cell to enable the widget!

def plot_intersection_single(w, I_ext):
  # set your parameters
  pars = default_pars_single(w=w, I_ext=I_ext)

  # find fixed points
  r_init_vector = [0, .4, .9]
  x_fps = my_fp_finder(pars, r_init_vector)

  # plot
  r = np.linspace(0, 1., 1000)
  drdt = (-r + F(w * r + I_ext, pars['a'], pars['theta'])) / pars['tau']

  plot_dr_r(r, drdt, x_fps)

_ = widgets.interact(plot_intersection_single, w=(1, 7, 0.2),
                     I_ext=(0, 3, 0.1))
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D4_DynamicNetworks/solutions/W2D4_Tutorial1_Solution_20486792.py)

---
# Summary

In this tutorial, we have investigated the dynamics of a rate-based single population of neurons.

We learned about:
- The effect of the input parameters and the time constant of the network on the dynamics of the population.
- How to find the fixed point(s) of the system.

Next, we have two Bonus, but important concepts in dynamical system analysis and simulation. If you have time left, watch the next video and proceed to solve the exercises. You will learn:

- How to determine the stability of a fixed point by linearizing the system.
- How to add realistic inputs to our model.

---
# Bonus 1: Stability of a fixed point

```python
# @title Video 3: Stability of fixed points
from IPython.display import YouTubeVideo
video = YouTubeVideo(id="KKMlWWU83Jg", width=854, height=480, fs=1)
print("Video available at https://youtube.com/watch?v=" + video.id)
video
```

#### Initial values and trajectories

Here, let us first set $w=5.0$ and $I_{\text{ext}}=0.5$, and investigate the dynamics of $r(t)$ starting with different initial values $r(0) \equiv r_{\text{init}}$. We will plot the trajectories of $r(t)$ with $r_{\text{init}} = 0.0, 0.1, 0.2,..., 0.9$.

```python
# @markdown Execute this cell to see the trajectories!

pars = default_pars_single()
pars['w'] = 5.0
pars['I_ext'] = 0.5

plt.figure(figsize=(8, 5))
for ie in range(10):
  pars['r_init'] = 0.1 * ie  # set the initial value
  r = simulate_single(pars)  # run the simulation

  # plot the activity with given initial
  plt.plot(pars['range_t'], r, 'b', alpha=0.1 + 0.1 * ie,
           label=r'r$_{\mathrm{init}}$=%.1f' % (0.1 * ie))

plt.xlabel('t (ms)')
plt.title('Two steady states?')
plt.ylabel(r'$r$(t)')
plt.legend(loc=[1.01, -0.06], fontsize=14)
plt.show()
```

## Interactive Demo: dynamics as a function of the initial value

Let's now set $r_{\rm init}$ to a value of your choice in this demo. How does the solution change? What do you observe?

```python
# @title

# @markdown Make sure you execute this cell to enable the widget!

pars = default_pars_single(w=5.0, I_ext=0.5)

def plot_single_diffEinit(r_init):
  pars['r_init'] = r_init
  r = simulate_single(pars)

  plt.figure()
  plt.plot(pars['range_t'], r, 'b', zorder=1)
  plt.plot(0, r[0], 'bo', alpha=0.7, zorder=2)
  plt.xlabel('t (ms)', fontsize=16)
  plt.ylabel(r'$r(t)$', fontsize=16)
  plt.ylim(0, 1.0)
  plt.show()

_ = widgets.interact(plot_single_diffEinit, r_init=(0, 1, 0.02))
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D4_DynamicNetworks/solutions/W2D4_Tutorial1_Solution_4d2de6a0.py)

### Stability analysis via linearization of the dynamics

Just like Equation $1$ in the case ($w=0$) discussed above, a generic linear system
$$\frac{dx}{dt} = \lambda (x - b),$$
has a fixed point for $x=b$. The analytical solution of such a system can be found to be:
$$x(t) = b + \big{(} x(0) - b \big{)} \text{e}^{\lambda t}.$$
Now consider a small perturbation of the activity around the fixed point: $x(0) = b+ \epsilon$, where $|\epsilon| \ll 1$. Will the perturbation $\epsilon(t)$ grow with time or will it decay to the fixed point? The evolution of the perturbation with time can be written, using the analytical solution for $x(t)$, as:
 $$\epsilon (t) = x(t) - b = \epsilon \text{e}^{\lambda t}$$

- if $\lambda < 0$, $\epsilon(t)$ decays to zero, $x(t)$ will still converge to $b$ and the fixed point is "**stable**".

- if $\lambda > 0$, $\epsilon(t)$ grows with time, $x(t)$ will leave the fixed point $b$ exponentially, and the fixed point is, therefore, "**unstable**" .

### Compute the stability of Equation $1$

Similar to what we did in the linear system above, in order to determine the stability of a fixed point $r^{*}$ of the excitatory population dynamics, we perturb Equation (1) around $r^{*}$ by $\epsilon$, i.e. $r = r^{*} + \epsilon$. We can plug in Equation (1) and obtain the equation determining the time evolution of the perturbation $\epsilon(t)$:

\begin{align}
\tau \frac{d\epsilon}{dt} \approx -\epsilon + w F'(w\cdot r^{*} + I_{\text{ext}};a,\theta) \epsilon
\end{align}

where $F'(\cdot)$ is the derivative of the transfer function $F(\cdot)$. We can rewrite the above equation as:

\begin{align}
\frac{d\epsilon}{dt} \approx \frac{\epsilon}{\tau }[-1 + w F'(w\cdot r^* + I_{\text{ext}};a,\theta)]
\end{align}

That is, as in the linear system above, the value of

$$\lambda = [-1+ wF'(w\cdot r^* + I_{\text{ext}};a,\theta)]/\tau \qquad (4)$$

determines whether the perturbation will grow or decay to zero, i.e., $\lambda$ defines the stability of the fixed point. This value is called the **eigenvalue** of the dynamical system.

## Exercise 4: Compute $dF$

The derivative of the sigmoid transfer function is:
\begin{align}
\frac{dF}{dx} & = \frac{d}{dx} (1+\exp\{-a(x-\theta)\})^{-1}  \\
& = a\exp\{-a(x-\theta)\} (1+\exp\{-a(x-\theta)\})^{-2}. \qquad (5)
\end{align}

Let's now find the expression for the derivative $\displaystyle{\frac{dF}{dx}}$ in the following cell and plot it.

```python
def dF(x, a, theta):
  """
  Population activation function.

  Args:
    x     : the population input
    a     : the gain of the function
    theta : the threshold of the function

  Returns:
    dFdx  : the population activation response F(x) for input x
  """

  ###########################################################################
  # TODO for students: compute dFdx ##
  raise NotImplementedError("Student excercise: compute the deravitive of F")
  ###########################################################################

  # Calculate the population activation
  dFdx = ...

  return dFdx

pars = default_pars_single()  # get default parameters
x = np.arange(0, 10, .1)      # set the range of input

# Uncomment below to test your function
# df = dF(x, pars['a'], pars['theta'])
# plot_dFdt(x, df)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D4_DynamicNetworks/solutions/W2D4_Tutorial1_Solution_ce2e3bc5.py)

*Example output:*

## Exercise 5: Compute eigenvalues

As discussed above, for the case with $w=5.0$ and $I_{\text{ext}}=0.5$, the system displays **three** fixed points. However, when we simulated the dynamics and varied the initial conditions $r_{\rm init}$, we could only obtain **two** steady states. In this exercise, we will now check the stability of each of the three fixed points by calculating the corresponding eigenvalues with the function `eig_single`. Check the sign of each eigenvalue (i.e., stability of each fixed point). How many of the fixed points are stable?

Note that the expression of the eigenvalue at fixed point $r^*$
$$\lambda = [-1+ wF'(w\cdot r^* + I_{\text{ext}};a,\theta)]/\tau$$

```python
def eig_single(fp, tau, a, theta, w, I_ext, **other_pars):
  """
  Args:
    fp   : fixed point r_fp
    tau, a, theta, w, I_ext : Simulation parameters

  Returns:
    eig : eigevalue of the linearized system
  """
  #####################################################################
  ## TODO for students: compute eigenvalue and disable the error
  raise NotImplementedError("Student excercise: compute the eigenvalue")
  ######################################################################
  # Compute the eigenvalue
  eig = ...

  return eig

# Find the eigenvalues for all fixed points of Exercise 2
pars = default_pars_single(w=5, I_ext=.5)
r_guess_vector = [0, .4, .9]
x_fp = my_fp_finder(pars, r_guess_vector)

# Uncomment below lines after completing the eig_single function.

# for fp in x_fp:
#   eig_fp = eig_single(fp, **pars)
#   print(f'Fixed point1 at {fp:.3f} with Eigenvalue={eig_fp:.3f}')
```

**SAMPLE OUTPUT**

```
Fixed point1 at 0.042 with Eigenvalue=-0.583
Fixed point2 at 0.447 with Eigenvalue=0.498
Fixed point3 at 0.900 with Eigenvalue=-0.626
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D4_DynamicNetworks/solutions/W2D4_Tutorial1_Solution_e285f60d.py)

## Think!
Throughout the tutorial, we have assumed $w> 0 $, i.e., we considered a single population of **excitatory** neurons. What do you think will be the behavior of a population of inhibitory neurons, i.e., where $w> 0$ is replaced by $w< 0$?

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W2D4_DynamicNetworks/solutions/W2D4_Tutorial1_Solution_579bc9c9.py)

---
# Bonus 2: Noisy input drives the transition between two stable states

## Ornstein-Uhlenbeck (OU) process

As discussed in several previous tutorials, the OU process is usually used to generate a noisy input into the neuron. The OU input $\eta(t)$ follows:

$$\tau_\eta \frac{d}{dt}\eta(t) = -\eta (t) + \sigma_\eta\sqrt{2\tau_\eta}\xi(t)$$

Execute the following function `my_OU(pars, sig, myseed=False)` to generate an OU process.

```python
# @title OU process `my_OU(pars, sig, myseed=False)`

# @markdown Make sure you execute this cell to visualize the noise!

def my_OU(pars, sig, myseed=False):
  """
  A functions that generates Ornstein-Uhlenback process

  Args:
    pars       : parameter dictionary
    sig        : noise amplitute
    myseed     : random seed. int or boolean

  Returns:
    I          : Ornstein-Uhlenbeck input current
  """

  # Retrieve simulation parameters
  dt, range_t = pars['dt'], pars['range_t']
  Lt = range_t.size
  tau_ou = pars['tau_ou']  # [ms]

  # set random seed
  if myseed:
      np.random.seed(seed=myseed)
  else:
      np.random.seed()

  # Initialize
  noise = np.random.randn(Lt)
  I_ou = np.zeros(Lt)
  I_ou[0] = noise[0] * sig

  # generate OU
  for it in range(Lt - 1):
    I_ou[it + 1] = (I_ou[it]
                    + dt / tau_ou * (0. - I_ou[it])
                    + np.sqrt(2 * dt / tau_ou) * sig * noise[it + 1])

  return I_ou

pars = default_pars_single(T=100)
pars['tau_ou'] = 1.  # [ms]
sig_ou = 0.1
I_ou = my_OU(pars, sig=sig_ou, myseed=2020)
plt.figure(figsize=(10, 4))
plt.plot(pars['range_t'], I_ou, 'r')
plt.xlabel('t (ms)')
plt.ylabel(r'$I_{\mathrm{OU}}$')
plt.show()
```

## Example: Up-Down transition

In the presence of two or more fixed points, noisy inputs can drive a transition between the fixed points! Here, we stimulate an E population for 1,000 ms applying OU inputs.

```python
# @title Simulation of an E population with OU inputs

# @markdown Make sure you execute this cell to spot the Up-Down states!

pars = default_pars_single(T=1000)
pars['w'] = 5.0
sig_ou = 0.7
pars['tau_ou'] = 1.  # [ms]
pars['I_ext'] = 0.56 + my_OU(pars, sig=sig_ou, myseed=2020)

r = simulate_single(pars)

plt.figure(figsize=(10, 4))
plt.plot(pars['range_t'], r, 'b', alpha=0.8)
plt.xlabel('t (ms)')
plt.ylabel(r'$r(t)$')
plt.show()
```


# DCEGM Upper Envelope
## ["The endogenous grid method for discrete-continuous dynamic choice models with (or without) taste shocks"](https://onlinelibrary.wiley.com/doi/abs/10.3982/QE643)

For the following badges: GitHub does not allow click-through redirects; right-click to get the link, then paste into navigation bar

[](https://mybinder.org/v2/gh/econ-ark/DemARK/master?filepath=notebooks%2FDCEGM-Upper-Envelope.ipynb)

[](https://colab.research.google.com/github/econ-ark/DemARK/blob/master/notebooks/DCEGM-Upper-Envelope.ipynb)

This notebook provides a simple introduction to the upper envelope calculation in the "DCEGM" algorithm . It takes the EGM method proposed in , and extends it to the mixed choice (discrete and continuous) case. It handles various constraints. It works on a 1-dimensional problems.

The main challenge in the types of models considered in DCEGM is, that the first order conditions to the Bellman equations are no longer sufficient to find an optimum.  Though, they are still necessary in a broad class of models. This means that our EGM step will give us (resource, consumption) pairs that do fulfill the FOCs, but that are sub-optimal (there's another consumption choices for the same initial resources that gives a higher value).

Take a consumption model formulated as:
$$
\max_{\{c_t\}^T_{t=1}} \sum^T_{t=1}\beta^t\cdot u(c_t)
$$
given some initial condition on $x$ and some laws of motion for the states, though explicit references to states are omitted. Then, if we're in a class of models described in EGM
, we can show that
$$
c_t = {u_{c}}^{-1}[E_t(u_c(c_{t+1}))]
$$
uniquely determines an optimal consumption today given the expected marginal utility of consuming  tomorrow. However, if there is a another choice in the choice set, and that choice is discrete, we get
$$
\max_{\{c_t, d_t\}^T_{t=1}} \sum^T_{t=1}\beta^t\cdot u(c_t, d_t)
$$
again given initial conditions and the laws of motion. Then, we can show that
$$
c_t = {u_{c}}^{-1}[E_t(u_c(c_{t+1}))]
$$
will produce solutions that are necessary but not sufficient. Note, that there is no explicit mentioning of the discrete choices in the expectation, but they obviously vary over the realized states in general. For the optimal consumption, it doesn't matter what the choice is exactly, only what expected marginal utility is tomorrow. The algorithm presented in [1] is designed to take advantage of models with this structure.

To visualize the problem, consider the following pictures that show the output of an EGM step from the model in the REMARK [linkhere].

```python
# imports
import numpy as np
import matplotlib.pyplot as plt
```

```python
# here for now, should be
# from HARK import discontools or whatever name is chosen
from HARK.interpolation import LinearInterp
from HARK.dcegm import calcSegments, calcMultilineEnvelope
```

```python
m_common = np.linspace(0,1.0,100)
m_egm = np.array([0.0, 0.04, 0.25, 0.15, 0.1, 0.3, 0.6,0.5, 0.35, 0.6, 0.75,0.85])
c_egm = np.array([0.0, 0.03, 0.1, 0.07, 0.05, 0.36, 0.4, 0.6, 0.8, 0.9,0.9,0.9])
vt_egm = np.array( [0.0, 0.05, 0.1,0.04, 0.02,0.2, 0.7, 0.5, 0.2, 0.9, 1.0, 1.2])
```

```python
plt.plot(m_egm, vt_egm)
plt.xlabel("resources")
plt.ylabel("transformed values")
```

```python
plt.plot(m_egm, c_egm)
plt.xlabel("resources")
plt.ylabel("consumption")
plt.show()
```

The point of DCEGM is to realize, that the segments on the `(m, vt)` curve that are decreasing, cannot be optimal. This leaves us with a set of increasing line segments, as seen below (`dcegmSegments` is the function in HARK that calculates the breaks where the curve goes from increasing to decreasing).

```python
rise, fall = calcSegments(m_egm, vt_egm)
```

In `rise` we have all the starting indices for the segments that are "good", that is `(m, vt)` draws an increasing curve.

```python
rise
```

    array([0, 4, 8])

We see that `rise` has its first index at `0`, then again at `4`, and lastly at `8`. Let's look at `fall`.

```python
fall
```

    array([ 2,  6, 11])

We see that the last segment is increasing (as the last element of `rise` is larger than the last element of `fall`), and we see that `len(fall)` is one larger than number of problematic segments in the plot. The index of the last point in `m_egm`/`c_egm`/`vt_egm` is added for convenience when we do the upper envelope step (and is also convenient below for drawing the segments!).

We can use `fall` and `rise` to draw only the relevant segments that we will use to construct an upper envelope.

```python
for j in range(len(fall)):
    idx = range(rise[j],fall[j]+1)
    plt.plot(m_egm[idx], vt_egm[idx])
plt.xlabel("resources")
plt.ylabel("transformed values")
plt.show()
```

Let us now use the `calcMultilineEnvelope` function to do the full DCEGM step: find segments and calculate upper envelope in one sweep.

```python
m_upper, c_upper, v_upper = calcMultilineEnvelope(m_egm, c_egm, vt_egm, m_common)
```

```python
for j in range(len(fall)):
    idx = range(rise[j],fall[j]+1)
    plt.plot(m_egm[idx], vt_egm[idx])
plt.plot(m_upper, v_upper, 'k')
plt.xlabel("resources")
plt.ylabel("transformed values")
plt.show()
```

And there we have it! These functions are the building blocks for univariate discrete choice modeling in HARK, so hopefully this little demo helped better understand what goes on under the hood, or it was a help if you're extending some existing class with a discrete choice.

# An example: writing a will

We now present a basic example to illustrate the use of the previous tools in solving dynamic optimization problems with discrete and continuous decisions.

The model represents an agent that lives for three periods and decides how much of his resources to consume in each of them. On the second period, he must additionally decide whether to hire a lawyer to write a will. Having a will has the upside of allowing the agent to leave a bequest in his third and last period of life, which gives him utility, but has the downside that the lawyer will charge a fraction of his period 3 resources.

On each period, the agent receives a deterministic amount of resources $w$. The problem, therefore, is fully deterministic.

I now present the model formally, solving it backwards.

But first, some setup and calibration:

```python
# Import tools for linear interpolation and finding optimal
# discrete choices.
from HARK.interpolation import calcLogSumChoiceProbs

# Import CRRA utility (and related) functions from HARK
from HARK.utilities import CRRAutility, CRRAutilityP, CRRAutilityP_inv

# Solution method parameters
aGrid = np.linspace(0,8,400) # Savings grid for EGM.

# Model parameters

# Parameters that need to be fixed
# Relative risk aversion. This is fixed at 2 in order to mantain
# the analytical solution that we use, from Carroll (2000)
CRRA   = 2

# Parameters that can be changed.
w          = 1    # Deterministic wage per period.
willCstFac = 0.35 # Fraction of resources charged by lawyer for writing a will.
DiscFac    = 0.98 # Time-discount factor.

# Define utility (and related) functions
u     = lambda x: CRRAutility(x,CRRA)
uP    = lambda x: CRRAutilityP(x, CRRA)
uPinv = lambda x: CRRAutilityP_inv(x, CRRA)

# Create a grid for market resources
mGrid       = (aGrid-aGrid[0])*1.5
mGridPlots  = np.linspace(w,10*w,100)
mGridPlotsC = np.insert(mGridPlots,0,0)

# Transformations for value funtion interpolation
vTransf   = lambda x: np.exp(x)
vUntransf = lambda x: np.log(x)
```

# The third (last) period of life

In the last period of life, the agent's problem is determined by his total amount of resources $m_3$ and a state variable $W$ that indicates whether he wrote a will ($W=1$) or not ($W=0$).

### The agent without a will

An agent who does not have a will simply consumes all of his available resources. Therefore, his value and consumption functions will be:

\begin{equation}
V_3(m_3,W=0) = u(m_3)
\end{equation}

\begin{equation}
c_3(m_3, W=0) = m_3
\end{equation}

Where $u(\cdot)$ gives the utility from consumption. We assume a CRRA specification $u(c) = \frac{c^{1-\rho}}{1-\rho}$.

### The agent with a will

An agent who wrote a will decides how to allocate his available resources $m_3$ between his consumption and a bequest. We assume an additive specification for the utility of a given consumption-bequest combination that follows a particular case in [Carroll (2000)](http://www.econ2.jhu.edu/people/ccarroll/Why.pdf). The component of utility from leaving a bequest $x$ is assumed to be $\ln (x+1)$. Therefore, the agent's value function is

\begin{equation}
V_3(m_3, W=1) = \max_{0\leq c_3 \leq m_3} u(c_3) + \ln(m_3 - c_3 + 1)
\end{equation}

For ease of exposition we consider the case $\rho = 2$, where [Carroll (2000)](http://www.econ2.jhu.edu/people/ccarroll/Why.pdf) shows that the optimal consumption level is given by

\begin{equation}
c_3(m_3, W=1) = \min \left[m_3, \frac{-1 + \sqrt{1 + 4(m_3+1)}}{2} \right].
\end{equation}

The consumption function shows that $m_3=1$ is the level of resources at which an important change of behavior occurs: agents leave bequests only for $m_3 > 1$. Since an important change of behavior happens at this point, we call it a 'kink-point' and add it to our grids.

```python
# Agent without a will
mGrid3_no = mGrid
cGrid3_no = mGrid
vGrid3_no = u(cGrid3_no)

# Create functions
c3_no  = LinearInterp(mGrid3_no, cGrid3_no) # (0,0) is already here.
vT3_no = LinearInterp(mGrid3_no, vTransf(vGrid3_no), lower_extrap = True)
v3_no  = lambda x: vUntransf(vT3_no(x))

# Agent with a will

# Define an auxiliary function with the analytical consumption expression
c3will = lambda m: np.minimum(m, -0.5 + 0.5*np.sqrt(1+4*(m+1)))

# Find the kink point
mKink = 1.0
indBelw = mGrid < mKink
indAbve = mGrid > mKink

mGrid3_wi = np.concatenate([mGrid[indBelw],
                            np.array([mKink]),
                            mGrid[indAbve]])

cGrid3_wi = c3will(mGrid3_wi)

cAbve = c3will(mGrid[indAbve])
beqAbve = mGrid[indAbve] - c3will(mGrid[indAbve])
vGrid3_wi = np.concatenate([u(mGrid[indBelw]),
                             u(np.array([mKink])),
                             u(cAbve) + np.log(1+beqAbve)])

# Create functions
c3_wi  = LinearInterp(mGrid3_wi, cGrid3_wi) # (0,0) is already here
vT3_wi = LinearInterp(mGrid3_wi, vTransf(vGrid3_wi), lower_extrap = True)
v3_wi  = lambda x: vUntransf(vT3_wi(x))

plt.figure()

plt.plot(mGridPlots, v3_wi(mGridPlots), label = 'Will')
plt.plot(mGridPlots, v3_no(mGridPlots), label = 'No Will')
plt.title('Period 3: Value functions')
plt.xlabel('Market resources')
plt.legend()
plt.show()

plt.plot(mGridPlotsC, c3_wi(mGridPlotsC), label = 'Will')
plt.plot(mGridPlotsC, c3_no(mGridPlotsC), label = 'No Will')
plt.title('Period 3: Consumption Functions')
plt.xlabel('Market resources')
plt.legend()
plt.show()
```

# The second period

On the second period, the agent takes his resources as given (the only state variable) and makes two decisions:
- Whether to write a will or not.
- What fraction of his resources to consume.

These decisions can be seen as happening sequentially: the agent first decides whether to write a will or not, and then consumes optimally in accordance with his previous decision. Since we solve the model backwards in time, we first explore the consumption decision, conditional on the choice of writing a will or not.

## An agent who decides not to write a will

After deciding not to write a will, an agent solves the optimization problem expressed in the following conditional value function

\begin{equation}
\begin{split}
\nu (m_2|w=0) &= \max_{0\leq c \leq m_2} u(c) + \beta V_3(m_3,W=0)\\
s.t.&\\
m_3 &= m_2 - c + w
\end{split}
\end{equation}

We can approximate a solution to this problem through the method of endogenous gridpoints. This yields approximations to $\nu(\cdot|w=0)$ and $c_2(\cdot|w=0)$

```python
# Second period, not writing a will

# Compute market resources at 3 with and without a will
mGrid3_cond_nowi = aGrid + w
# Compute marginal value of assets in period 3 for each ammount of savings in 2
vPGrid3_no = uP(c3_no(mGrid3_cond_nowi))
# Get consumption through EGM inversion of the euler equation
cGrid2_cond_no = uPinv(DiscFac*vPGrid3_no)

# Get beginning-of-period market resources
mGrid2_cond_no = aGrid + cGrid2_cond_no

# Compute value function
vGrid2_cond_no = u(cGrid2_cond_no) + DiscFac*v3_no(mGrid3_cond_nowi)

# Create interpolating value and consumption functions
vT2_cond_no = LinearInterp(mGrid2_cond_no, vTransf(vGrid2_cond_no), lower_extrap = True)
v2_cond_no  = lambda x: vUntransf(vT2_cond_no(x))
c2_cond_no  = LinearInterp(np.insert(mGrid2_cond_no,0,0), np.insert(cGrid2_cond_no,0,0))
```

## An agent who decides to write a will

An agent who decides to write a will also solves for his consumption dinamically. We assume that the lawyer that helps the agent write his will takes some fraction $\tau$ of his total resources in period 3. Therefore, the evolution of resources is given by $m_3 = (1-\tau)(m_2 - c_2 + w)$. The conditional value function of the agent is therefore:

\begin{equation}
\begin{split}
\nu (m_2|w=1) &= \max_{0\leq c \leq m_2} u(c) + \beta V_3(m_3,W=1)\\
s.t.&\\
m_3 &= (1-\tau)(m_2 - c + w)
\end{split}
\end{equation}

We also approximate a solution to this problem using the EGM. This yields approximations to $\nu(\cdot|w=1)$ and $c_2(\cdot|w=1)$.

```python
# Second period, writing a will

# Compute market resources at 3 with and without a will
mGrid3_cond_will = (1-willCstFac)*(aGrid + w)
# Compute marginal value of assets in period 3 for each ammount of savings in 2
vPGrid3_wi = uP(c3_wi(mGrid3_cond_will))
# Get consumption through EGM inversion of the euler equation
cGrid2_cond_wi = uPinv(DiscFac*(1-willCstFac)*vPGrid3_wi)
# Get beginning-of-period market resources
mGrid2_cond_wi = aGrid + cGrid2_cond_wi

# Compute value function
vGrid2_cond_wi = u(cGrid2_cond_wi) + DiscFac*v3_wi(mGrid3_cond_will)

# Create interpolating value and consumption functions
vT2_cond_wi = LinearInterp(mGrid2_cond_wi, vTransf(vGrid2_cond_wi), lower_extrap = True)
v2_cond_wi  = lambda x: vUntransf(vT2_cond_wi(x))
c2_cond_wi  = LinearInterp(np.insert(mGrid2_cond_wi,0,0), np.insert(cGrid2_cond_wi,0,0))
```

## The decision whether to write a will or not

With the conditional value functions at hand, we can now express and solve the decision of whether to write a will or not, and obtain the unconditional value and consumption functions.

\begin{equation}
V_2(m_2) = \max \{ \nu (m_2|w=0), \nu (m_2|w=1) \}
\end{equation}

\begin{equation}
w^*(m_2) = \arg \max_{w \in \{0,1\}} \{ \nu (m_2|w=w) \}
\end{equation}

\begin{equation}
c_2(m_2) = c_2(m_2|w=w^*(m_2))
\end{equation}

We now construct these objects.

```python
# We use HARK's 'calcLogSumchoiceProbs' to compute the optimal
# will decision over our grid of market resources.
# The function also returns the unconditional value function
# Use transformed values since -given sigma=0- magnitudes are unimportant. This
# avoids NaNs at m \approx 0.
vTGrid2, willChoice2 = calcLogSumChoiceProbs(np.stack((vT2_cond_wi(mGrid),
                                                     vT2_cond_no(mGrid))),
                                           sigma = 0)
vGrid2 = vUntransf(vTGrid2)

# Plot the optimal decision rule
plt.plot(mGrid, willChoice2[0])
plt.title('$w^*(m)$')
plt.ylabel('Write will (1) or not (0)')
plt.xlabel('Market resources: m')
plt.show()

# With the decision rule we can get the unconditional consumption function
cGrid2 = (willChoice2*np.stack((c2_cond_wi(mGrid),c2_cond_no(mGrid)))).sum(axis=0)

vT2 = LinearInterp(mGrid, vTransf(vGrid2), lower_extrap = True)
v2  = lambda x: vUntransf(vT2(x))
c2  = LinearInterp(mGrid, cGrid2)

# Plot the conditional and unconditional value functions
plt.plot(mGridPlots, v2_cond_wi(mGridPlots), label = 'Cond. Will')
plt.plot(mGridPlots, v2_cond_no(mGridPlots), label = 'Cond. No will')
plt.plot(mGridPlots, v2(mGridPlots), 'k--',label = 'Uncond.')
plt.title('Period 2: Value Functions')
plt.xlabel('Market resources')
plt.legend()
plt.show()

# Plot the conditional and unconditiional consumption
# functions
plt.plot(mGridPlotsC, c2_cond_wi(mGridPlotsC), label = 'Cond. Will')
plt.plot(mGridPlotsC, c2_cond_no(mGridPlotsC), label = 'Cond. No will')
plt.plot(mGridPlotsC, c2(mGridPlotsC), 'k--',label = 'Uncond.')
plt.title('Period 2: Consumption Functions')
plt.xlabel('Market resources')
plt.legend()
plt.show()
```

# The first period

In the first period, the agent simply observes his market resources and decides what fraction of them to consume. His problem is represented by the following value function

\begin{equation}
\begin{split}
V (m_1) &= \max_{0\leq c \leq m_1} u(c) + \beta V_2(m_2)\\
s.t.&\\
m_2 &= m_1 - c + w.
\end{split}
\end{equation}

Although this looks like a simple problem, there are complications introduced by the kink in $V_2(\cdot)$, which is clearly visible in the plot from the previous block. Particularly, note that $V_2'(\cdot)$ and $c_2(\cdot)$ are not monotonic: there are now multiple points $m$ for which the slope of $V_2(m)$ is equal. Thus, the Euler equation becomes a necessary but not sufficient condition for optimality and the traditional EGM inversion step can generate non-monotonic endogenous $m$ gridpoints.

We now illustrate this phenomenon.

```python
# EGM step

# Period 2 resources implied by the exogenous savings grid
mGrid2 = aGrid + w
# Envelope condition
vPGrid2 = uP(c2(mGrid2))
# Inversion of the euler equation
cGrid1 = uPinv(DiscFac*vPGrid2)
# Endogenous gridpoints
mGrid1 = aGrid + cGrid1
vGrid1 = u(cGrid1) + DiscFac*v2(mGrid2)

plt.plot(mGrid1)
plt.title('Endogenous gridpoints')
plt.xlabel('Position: i')
plt.ylabel('Endogenous grid point: $m_i$')
plt.show()

plt.plot(mGrid1,vGrid1)
plt.title('Value function at grid points')
plt.xlabel('Market resources: m')
plt.ylabel('Value function')
plt.show()
```

The previous cell applies the endogenous gridpoints method to the first period problem. The plots illustrate that the sequence of resulting endogenous gridpoints $\{m_i\}_{i=1}^N$ is not monotonic. This results in intervals of market resources over which we have multiple candidate values for the value function. This is the point where we must apply the upper envelope function illustrated above.

We finally use the resulting consumption and value grid points to create the first period value and consumption functions.

```python
# Calculate envelope
vTGrid1 = vTransf(vGrid1) # The function operates with *transformed* value grids

rise, fall = calcSegments(mGrid1, vTGrid1)
mGrid1_up, cGrid1_up, vTGrid1_up = calcMultilineEnvelope(mGrid1, cGrid1,
                                                         vTGrid1, mGrid)

# Create functions
c1_up  = LinearInterp(mGrid1_up, cGrid1_up)
v1T_up = LinearInterp(mGrid1_up, vTGrid1_up)
v1_up  = lambda x: vUntransf(v1T_up(x))

# Show that there is a non-monothonicity and that the upper envelope fixes it
plt.plot(mGrid1,vGrid1, label = 'EGM Points')
plt.plot(mGridPlots, v1_up(mGridPlots), 'k--', label = 'Upper Envelope')
plt.title('Period 1: Value function')
plt.legend()
plt.show()

plt.plot(mGrid1,cGrid1, label = 'EGM Points')
plt.plot(mGridPlotsC,c1_up(mGridPlotsC),'k--', label = 'Upper Envelope')
plt.title('Period 1: Consumption function')
plt.legend()
plt.show()
```

# References
[1] Iskhakov, F. , Jørgensen, T. H., Rust, J. and Schjerning, B. (2017), The endogenous grid method for discrete‐continuous dynamic choice models with (or without) taste shocks. Quantitative Economics, 8: 317-365. doi:10.3982/QE643

[2] Carroll, C. D. (2006). The method of endogenous gridpoints for solving dynamic stochastic optimization problems. Economics letters, 91(3), 312-320.


## 7. Nonlinear ensemble filtering for the Lorenz-96 problem

In this notebook, we apply the stochastic map filter developed in Spantini et al. [5] to the Lorenz-96 problem.

References:

[1] Evensen, G., 1994. Sequential data assimilation with a nonlinear quasi‐geostrophic model using Monte Carlo methods to forecast error statistics. Journal of Geophysical Research: Oceans, 99(C5), pp.10143-10162.

[2] Asch, M., Bocquet, M. and Nodet, M., 2016. Data assimilation: methods, algorithms, and applications. Society for Industrial and Applied Mathematics.

[3] Bishop, C.H., Etherton, B.J. and Majumdar, S.J., 2001. Adaptive sampling with the ensemble transform Kalman filter. Part I: Theoretical aspects. Monthly weather review, 129(3), pp.420-436.

[4] Lorenz, E.N., 1963. Deterministic nonperiodic flow. Journal of atmospheric sciences, 20(2), pp.130-141.

[5] Spantini, A., Baptista, R. and Marzouk, Y., 2019. Coupling techniques for nonlinear ensemble filtering. arXiv preprint arXiv:1907.00389.

### The basic steps
To carry out sequential inference in `TransportMapBasedInference`, we need to carry out a few basic steps:
* **Specify the problem**: Define the state-space model: initial condition, dynamical and observation models (including process and observation noise)
* **Specify the inflation parameters**: Determine the levels of covariance inflation to properly balance the dynamical system and the observations from the truth system
* **Specify the filter**: Choose the ensemble filter to assimilate the observations in the state estimate
* **Perform the sequential inference**: Perform the sequential inference

We will go through all of these here.

```julia
using Revise
using LinearAlgebra
using TransportBasedInference
using Statistics
using Distributions
using OrdinaryDiffEq
```

    ┌ Info: Precompiling TransportBasedInference [bdf749b0-1400-4207-80d3-e689c0e3f03d]
    └ @ Base loading.jl:1342
    ┌ Warning: Type annotations on keyword arguments not currently supported in recipes. Type information has been discarded
    └ @ RecipesBase ~/.julia/packages/RecipesBase/qpxEX/src/RecipesBase.jl:123
    ┌ Warning: Type annotations on keyword arguments not currently supported in recipes. Type information has been discarded
    └ @ RecipesBase ~/.julia/packages/RecipesBase/qpxEX/src/RecipesBase.jl:123
    ┌ Warning: Type annotations on keyword arguments not currently supported in recipes. Type information has been discarded
    └ @ RecipesBase ~/.julia/packages/RecipesBase/qpxEX/src/RecipesBase.jl:123
    ┌ Warning: Type annotations on keyword arguments not currently supported in recipes. Type information has been discarded
    └ @ RecipesBase ~/.julia/packages/RecipesBase/qpxEX/src/RecipesBase.jl:123
    ┌ Warning: Type annotations on keyword arguments not currently supported in recipes. Type information has been discarded
    └ @ RecipesBase ~/.julia/packages/RecipesBase/qpxEX/src/RecipesBase.jl:123
    ┌ Warning: Type annotations on keyword arguments not currently supported in recipes. Type information has been discarded
    └ @ RecipesBase ~/.julia/packages/RecipesBase/qpxEX/src/RecipesBase.jl:123
    ┌ Warning: Type annotations on keyword arguments not currently supported in recipes. Type information has been discarded
    └ @ RecipesBase ~/.julia/packages/RecipesBase/qpxEX/src/RecipesBase.jl:123

Load some packages to make nice figures

```julia
using Plots
default(tickfont = font("CMU Serif", 18),
        titlefont = font("CMU Serif", 18),
        guidefont = font("CMU Serif", 18),
        legendfont = font("CMU Serif", 18),
        colorbar_tickfontsize = 18,
        colorbar_titlefontsize = 18,
        annotationfontsize = 18,
        annotationfontfamily = "CMU Serif",
        grid = false)
pyplot()

PyPlot.rc("text", usetex = "true")
PyPlot.rc("font", family = "CMU Serif")

using LaTeXStrings
using ColorSchemes
```

The Lorenz-63  model is a three dimensional system that models the atmospheric convection [4]. This system is a classical benchmark problem in data assimilation. The state $\boldsymbol{x} = (x_1, x_2, x_3)$ is governed by the following set of ordinary differential equations:

\begin{equation}
\begin{aligned}
&\frac{\mathrm{d} x_1}{\mathrm{d} t}=\sigma(x_2-x_1)\\
&\frac{\mathrm{d} x_2}{\mathrm{d} t}=x_1(\rho-x_2)-x_2\\
&\frac{\mathrm{d} x_3}{\mathrm{d} t}=x_1 x_2-\beta x_3,
\end{aligned}
\end{equation}

where $\sigma = 10, \beta = 8/3, \rho = 28$. For these values, the system is chaotic and behaves like a strange attractor. We integrate this system of ODEs with time step $\Delta t_{dyn} = 0.05$. The state is fully observed $h(t,\boldsymbol{x}) = \boldsymbol{x}$ with $\Delta t_{obs}=0.1$. The initial distribution $\pi_{\mathsf{X}_0}$ is the standard Gaussian. The process noise is Gaussian with zero mean and covariance $10^{-4}\boldsymbol{I}_3$. The measurement noise has a Gaussian distribution with zero mean and covariance $\theta^2\boldsymbol{I}_3$ where $\theta^2 = 4.0$.

### Simple twin-experiment

Define the dimension of the state and observation vectors

```julia
Nx = 40
Ny = 20
```

    20

Define the time steps $\Delta t_{dyn}, \Delta t_{obs}$  of the dynamical and observation models. Observations from the truth are assimilated every $\Delta t_{obs}$.

```julia
Δtdyn = 0.05
Δtobs = 0.4
```

    0.4

Define the time span of interest

```julia
t0 = 0.0
tf = 2000.
Tf = ceil(Int64, (tf-t0)/Δtobs)
```

    5000

 Define the distribution for the initial condition $\pi_{\mathsf{X}_0}$

```julia
π0 = MvNormal(zeros(Nx), Matrix(1.0*I, Nx, Nx))
```

    FullNormal(
    dim: 40
    μ: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0  …  0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    Σ: [1.0 0.0 … 0.0 0.0; 0.0 1.0 … 0.0 0.0; … ; 0.0 0.0 … 1.0 0.0; 0.0 0.0 … 0.0 1.0]
    )

We construct the state-space representation `F` of the system composed of the deterministic part of the dynamical and observation models.

The dynamical model is provided by the right hand side of the ODE to solve. For a system of ODEs, we will prefer an in-place syntax `f(du, u, p, t)`, where `p` are parameters of the model.
We rely on `OrdinaryDiffEq` to integrate the dynamical system with the Tsitouras 5/4 Runge-Kutta method adaptive time marching.

We assume that the state is observed every two components, i.e. $h(x, t) = x[1:2:Nx-1]$.

```julia
h(x, t) = x[1:2:end]
h(x, idx, t) = x[idx]

F = StateSpace(lorenz96!, h)
```

    StateSpace(TransportBasedInference.lorenz96!, h)

Define the additive inflation for the dynamical and observation models

```julia
### Process and observation noise
σx = 1e-4
σy = sqrt(0.5)

ϵx = AdditiveInflation(Nx, zeros(Nx), σx)
ϵy = AdditiveInflation(Ny, zeros(Ny), σy)
```

    AdditiveInflation(20, [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.5000000000000001 0.0 … 0.0 0.0; 0.0 0.5000000000000001 … 0.0 0.0; … ; 0.0 0.0 … 0.5000000000000001 0.0; 0.0 0.0 … 0.0 0.5000000000000001], [0.7071067811865476 0.0 … 0.0 0.0; 0.0 0.7071067811865476 … 0.0 0.0; … ; 0.0 0.0 … 0.7071067811865476 0.0; 0.0 0.0 … 0.0 0.7071067811865476])

```julia
model = Model(Nx, Ny, Δtdyn, Δtobs, ϵx, ϵy, π0, 0, 0, 0, F);
```

To perform the nonlinear ensemble filtering, we first need to estimate the transport map $\boldsymbol{S}^{\boldsymbol{\mathcal{X}}}$.

In this notebook, we are going to assume that the basis of features does not change over time, but solely the coefficients $c_{\boldsymbol{\alpha}}$ of the expansion.

To estimate the map, we generate joint samples $(\boldsymbol{y}^i, \boldsymbol{x}^i), \; i = 1, \ldots, N_e$ where $\{\boldsymbol{x}^i\}$ are i.i.d. samples from pushforward of the standard Gaussian distribution by the flow of the Lorenz-63 system.

```julia
# Time span
tspan = (0.0, tf)
```

    (0.0, 2000.0)

Set initial condition of the true system

```julia
x0 = rand(model.π0);
```

```julia
data = generate_lorenz96(model, x0, Tf);
```

Initialize the ensemble matrix `X` $\in \mathbb{R}^{(N_y + N_x) \times N_e}$.

```julia
# Ensemble size
Ne = 200

X0 = zeros(model.Ny + model.Nx, Ne)

# Generate the initial conditions for the state.
viewstate(X0, model.Ny, model.Nx) .= rand(model.π0, Ne);
```

We use the stochastic ensemble Kalman filter for the spin-up phase. There is no reason to use the stochastic map filter over the first cycles, as the performance of the inference is mostly determined by the quality of the ensemble, not by the quality of the filter.

```julia
enkf = StochEnKF(x->x, ϵy, Δtdyn, Δtobs)
```

    Stochastic EnKF  with filtered = false

```julia
Xsenkf = seqassim(F, data, Tf, model.ϵx, enkf, deepcopy(X0), model.Ny, model.Nx, t0);
```

    Progress: 100%|█████████████████████████████████████████| Time: 0:00:34

```julia
Δ = 2
yidx = 1:Δ:model.Nx
idx = vcat(collect(1:length(yidx))', collect(yidx)')
β = 1.0
Lrad = 6

senkfloc = SeqStochEnKF(x->x, F.h, MultiplicativeInflation(β), model.ϵy, Ny, Nx, model.Δtdyn, model.Δtobs,
                    idx, zeros(Nx+1, Ne), false, true)

# Create Localization structure
Gxy(i,j) = periodicmetric!(i,yidx[j], model.Nx)
Gyy(i,j) = periodicmetric!(yidx[i],yidx[j], model.Nx)
Loc = Localization(Lrad, Gxy, Gyy)
```

    Localization(6.0, Gxy, Gyy)

```julia
Xsenkfloc = seqassim(F, data, Tf, model.ϵx, senkfloc, deepcopy(X0), model.Ny, model.Nx, t0, Loc);
```

    Progress: 100%|█████████████████████████████████████████| Time: 0:01:34

```julia
tspin = 1000.0
Tspin = ceil(Int64, (tspin-t0)/Δtobs)

tsmf = 2000.0
Tsmf = ceil(Int64, (tsmf-tspin)/Δtobs);
```

Time average root-mean-squared error after the spin-up phase

```julia
rmse_senkf = mean(map(i->norm(data.xt[:,i]-mean(Xsenkf[i+1]; dims = 2))/sqrt(Nx), Tspin:Tf))
```

    0.9153862618614615

```julia
rmse_senkfloc = mean(map(i->norm(data.xt[:,i]-mean(Xsenkfloc[i+1]; dims = 2))/sqrt(Nx), Tspin:Tf))
```

    0.9058753605395535

```julia
# Estimate of the first component
nb = Tspin-200
ne = Tspin+Tsmf
Δ = 5
plt = plot(xlim = (-Inf, Inf), ylim = (-Inf, Inf), xlabel = L"t")
plot!(plt, data.tt[nb:Δ:ne], data.xt[1,nb:Δ:ne], linewidth =  3, color = :teal,
          ylabel = latexstring("x_"*string(1)), legend = :bottomleft, label = "True")
plot!(plt, data.tt[nb:Δ:ne], mean_hist(Xenkf)[1,1+nb:Δ:1+ne], linewidth = 3, grid = false,
          color = :orangered2, linestyle = :dash, label = "sEnKF")
scatter!(plt, data.tt[nb:Δ:ne], data.yt[1,nb:Δ:ne], linewidth = 3, color = :grey,
          markersize = 5, alpha = 0.5, label  = "Observation")
vline!(plt, [tspin], color = :grey2, linestyle = :dash, label = "")

plt
```

Initialize the ensemble matrix for the radial stochastic map filter

```julia
Xspin = vcat(zeros(Ny, Ne), deepcopy(Xenkf[Tspin+1]));
```

Initialize the structure of the map

```julia
p = 2
Lrad = 3
nonid_rad = 9

# parameters of the radial map
γ = 2.0
λ = 0.01
δ = 1e-8
κ = 4.0
β = 1.0

dist = Float64.(metric_lorenz(Nx))
order = setup_order(Nx, p, 0, p, Lrad, nonid_rad, dist)

# Define index of the measurement
idx = vcat(collect(1:Ny)',collect(1:2:Nx)')

smf = SparseRadialSMF(x->x, F.h, β, ϵy, order, γ, λ, δ, κ,
                      Ny, Nx, Ne,
                      Δtdyn, Δtobs,
                      dist, idx; islocalized = true);
```

```julia
# Not that the data assimilation can be accelerated by running the code in parallel
```

```julia
Xsmf = seqassim(F, data, Tsmf, model.ϵx, smf, deepcopy(Xspin), model.Ny, model.Nx, tspin);
```

    Progress: 100%|█████████████████████████████████████████| Time: 0:13:28

```julia
rmse_smf = mean(map(i->norm(data.xt[:,Tspin+i]-mean(Xsmf[i+1]; dims = 2))/sqrt(Nx), 1:Tsmf))
```

    0.7873191891500781

```julia
(rmse_enkf-rmse_smf)/rmse_enkf
```

    0.08884559890411144

```julia
nb = 1
ne = Tspin+Tsmf
Δ = 10
plt = plot( xlim = (-Inf, Inf), ylim = (-Inf, Inf), xlabel = L"t")
plot!(plt, data.tt[nb:Δ:ne], data.xt[1,nb:Δ:ne], linewidth =  3, color = :teal,
      ylabel = latexstring("x_"*string(1)), legend = :bottomleft, label = "True")
plot!(plt, data.tt[nb:Δ:ne], mean_hist(vcat(Xenkf[1:Tspin+1], Xsmf[2:end]))[1,1+nb:Δ:1+ne], linewidth = 3, grid = false,
      color = :orangered2, linestyle = :dash, label = "sEnKF for spin-up + Radial stochastic filter")
scatter!(plt, data.tt[nb:Δ:ne], data.yt[1,nb:Δ:ne], linewidth = 3, color = :grey,
      markersize = 5, alpha = 0.5, label  = "Observation")
vline!(plt, [tspin], color = :grey2, linestyle = :dash, label = "")

plt
```

```julia

```


# One Dimensional Lattice Models with QuTiP: Introduction

Saumya Biswas (saumyab@uoregon.edu)

For more information about QuTiP see [http://qutip.org](http://qutip.org)

We introduce the basic functionality of QuTiPs Lattice1d class of the lattice module.

# About

The qutip.lattice module enables defining tight binding/lattice models for bosons and fermions on lattices and calculating their fundamental properties specially features arising from the translational symmetry of the lattice. The lattice classes defined are compatible with the rest of the functionalities of QuTiP and can make use of them quite conveniently.

The module can be used to perform the widely used quantum optics calculations involving unitary/nonunitary evolutions, steadystates, correlation functions for bosons and fermions on a lattice. It also facilitates the traditional condensed matter problem calculations about topology and ground state properties.

The lattice1d class is well-suited for studying composite systems with multiple orbitals/excitations/spins per site of a crystal(with a unit cell possibly composed of multiple sites) with tensor structures of indices(cell number, site number, degrees of freedom) aligned with the spirit of QuTiP's tensor functionalities.

## Single-particle and Multiparticle physics

All the functionalities of the Lattice1d class are for sigle particle physics. The multi-particle physics calculations can be performed with a subclass of the Lattice1d class which is initiated with an instance of the Lattice1d class and inherits all the information about the lattice and basis.

## Unitcell structure in the Second Quantized notation
Defining an instance of the Lattice1d class requires formatting the second Quantized Hamiltonian in a unitcell based structure with nearest neighbor coupling only. Howewver, the functionality is limited to single particle physics only in Lattice1d class methods.

\begin{eqnarray}
H = \sum_i \psi_i^{\dagger} D \psi_i + \sum_{i} \left( \psi_i^{\dagger} T \psi_{i+1} +  \psi_{i+1}^{\dagger} T^{\dagger} \psi_i  \right) \label{eq:TB_block}
\end{eqnarray}

where $\psi_i$ is the annihilation operator for a unit cell at coordinate i,$D$ is the cell Hamiltonian of the unit cell, $T$ is the inter cell hopping. Any 1d lattice can be put in the form of the equation above by resolving it into unit cells with coupling limited to the nearest neighbors only.

The Lattice1d class is based on this unit cell and nearest neighbor hopping format. A unit cell can be comprised of one or more sites with one or more orbitals, spins, excitations or any other degrees of freedom. An 1d lattice with next nearest neighbor coupling can be equivalently represented as a 1d lattice with unit cells of larger size limiting the hopping terms to nearest neighbors only.

## How to Define a One Dimensional Lattice

```python
from qutip import *
import matplotlib.pyplot as plt
import numpy as np
```

### Declaring a tight binding chain with a single site unit cell
As a default the instance of Lattice1d class is initialized as an atomic chain with a unit
cell with one atom only. The user need only define the number of cells and the boundary
condition.

```python
boundary_condition = "periodic"
cells = 3
Periodic_Atom_Chain = Lattice1d(num_cell=cells, boundary = boundary_condition)
```

The user can call Periodic_Atom_Chain to print all its information.

```python
Periodic_Atom_Chain
```

    Lattice1d object: Number of cells = 3,
    Number of sites in the cell = 1,
    Degrees of freedom per site = [],
    Lattice tensor configuration = [3],
    basis_Hamiltonian = Quantum object: dims = [[1], [1]], shape = (1, 1), type = bra
    Qobj data =
    [[0.]],
    inter_hop = [Quantum object: dims = [[1], [1]], shape = (1, 1), type = bra
    Qobj data =
    [[-1.]]],
    cell_tensor_config = [1]
    Boundary Condition:  Periodic

To define a lattice with more than one site per unit cell and one or more degrees of freedom per site, the cell_num_site and cell_site_dof arguments must be used. In a case like this, specifying the intra and inter cell interactions would also be necessary (through the arguments cell_Hamiltonian and inter_hop) in most cases. However, Lattice1d() will initiate the instance with default cell_Hamiltonian and inter_hop
if the user does not input it.

### The use of display_unit_cell() and display_lattice()

The functions display_unit_cell() and display_lattice() can be used at any stage to produce visual symbolizations and elementwise information.

```python
H = Periodic_Atom_Chain.display_unit_cell(label_on = True)
T = Periodic_Atom_Chain.display_lattice()
```

The user can review the attribute values of H and T from the retuned Qobjs.

```python
print(H[0][0])
print(T)
```

    Quantum object: dims = [[1], [1]], shape = (1, 1), type = bra
    Qobj data =
    [[0.]]
    Quantum object: dims = [[1], [1]], shape = (1, 1), type = bra
    Qobj data =
    [[-1.]]

### Multiple site per unitcell and multiple degrees of freedom per site

Specifying cell_num_site enables choosing multiple sites for a unit call and any combination of degrees of freedom can be chosen for each site with cell_site_dof.

```python
boundary_condition = "periodic"
cells = 3
cell_num_site = 2
cell_site_dof = [2,3]    # It could be 2 orbitals and 3 spins per sites or
                         # any other combination of such degrees of freedom
lattice_3223 = Lattice1d(num_cell=cells, boundary = boundary_condition,
                           cell_num_site = cell_num_site, cell_site_dof = cell_site_dof)
```

###  The use of cell_structures()

There is an aide function that help the user form the cell_Hamiltonian and inter_hop
arguments namely cell_structures().

```python
val_s = ['site0', 'site1', 'site2']
val_t = [' orb0', ' orb1']
(cell_H_form,inter_cell_T_form,cell_H,inter_cell_T) = cell_structures( val_s, val_t)
```

The cell_structure() function returns two list of lists cell_H_form and inter_cell_T_form
that prints str s that can guide the user enter the nonzero elements at cell_H
and inter_cell_T which are np.zeros with the appropriate size. The procedure would
be to check a certain element in cell_H_form and insert the value for cell_H and so on.

```python
cell_H_form[0][5]
```

    '<site0 orb0 H site2 orb1>'

```python
cell_H[0][5] = -1-0.5j   # Calculated value from hand calculation
cell_H[5][0] = -1+0.5j   # keeping it Hermitian
```

Similarly, we set more elements to non-zero values.

```python
cell_H_form[2][5]
```

    '<site1 orb0 H site2 orb1>'

```python
cell_H[2][5] = -1+0.25j   # Calculated value from hand calculation
cell_H[5][2] = -1-0.25j   # keeping it Hermitian
```

The user would have to enter all the nonzero elements in cell_H and inter_cell_T
and then convert them into Qobjs and use them in declaring the instance of Lattice_1d.

```python
inter_cell_T_form[5][0]
```

    '<cell(i):site2 orb1 H site0 orb0:cell(i+1) >'

```python
inter_cell_T[5][0] = -0.5
inter_cell_T[0][5] = -0.5
```

```python
cell_H = Qobj(cell_H)
inter_cell_T = Qobj(inter_cell_T)
lattice_324 = Lattice1d(num_cell=3, boundary = "periodic", cell_num_site = 3, cell_site_dof = [2], Hamiltonian_of_cell = cell_H, inter_hop = inter_cell_T )
```

cell_cite_dof can take care of composite degrees of freedom such as orbits, spins and/or excitations. For example, if each site has 4 orbitals and 2 spins, we set cell_site_dof = [4,2] defines that lattice. With the aid of the Lattice1d.basis() operator we can access particles localized at specific cell,site,orbitals and spin.
Valid inputs for cell_site_dof are one int(e.g. 4) or a list of int's(e.g. [4,2]). A single dof can be entered either as an int or a list with that int. So cell_site_dof = [4] and cell_site_dof = 4 are the same.

```python
H = lattice_324.display_unit_cell(label_on = True)
T = lattice_324.display_lattice()
```

The labels of the diagrams can be read off from the returned H and T Qobjs. For example, $H_{12}$ can be read off as follows.

```python
H[1][2]
```

Quantum object: dims = [[2], [2]], shape = (2, 2), type = oper, isherm = False\begin{equation*}\left(\begin{array}{*{11}c}0.0 & (-1.0+0.250j)\\0.0 & 0.0\\\end{array}\right)\end{equation*}

### Basis function: ket vector initialized at specific cell, site, dof:

The basis() function enables the user to initialize a ket vector at a specific cell,
site and dof.

```python
lattice_3224 = Lattice1d(num_cell=3, boundary = "periodic", \
                           cell_num_site = 2, cell_site_dof = [4,2])
psi0 = lattice_3224.basis(1,0,[2,1])
print( psi0.dag() )  # Because plotting the dag() takes up less space
```

    Quantum object: dims = [[1, 1, 1, 1], [3, 2, 4, 2]], shape = (1, 48), type = bra
    Qobj data =
    [[0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0.
      0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]]

### Position Operator

Calling the position operator, x() returns an operator in matrix form that gives the
cell number for all the dofs present on the diagonal elements. The length of the unit
cell is always considered 1.

```python
lattice_412 = Lattice1d(num_cell=4, boundary = "periodic", cell_num_site = 1, cell_site_dof = [2])
lattice_412.x()
```

Quantum object: dims = [[4, 2], [4, 2]], shape = (8, 8), type = oper, isherm = True\begin{equation*}\left(\begin{array}{*{11}c}0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 1.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 1.0 & 0.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 2.0 & 0.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 2.0 & 0.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 3.0 & 0.0\\0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 0.0 & 3.0\\\end{array}\right)\end{equation*}

### The crystal momentum operator

The crystal momentum operator can be produced with the k() method.

```python
lattice_411 = Lattice1d(num_cell=4, boundary = "periodic", cell_num_site = 1, cell_site_dof = [1])
k = lattice_411.k()
print(k)
```

    Quantum object: dims = [[4], [4]], shape = (4, 4), type = oper, isherm = True
    Qobj data =
    [[ 0.78539816+0.j         -0.78539816-0.78539816j  0.78539816+0.j
      -0.78539816+0.78539816j]
     [-0.78539816+0.78539816j  0.78539816+0.j         -0.78539816-0.78539816j
       0.78539816+0.j        ]
     [ 0.78539816+0.j         -0.78539816+0.78539816j  0.78539816+0.j
      -0.78539816-0.78539816j]
     [-0.78539816-0.78539816j  0.78539816+0.j         -0.78539816+0.78539816j
       0.78539816+0.j        ]]

### Distribute A Operator
The operator_at_cells() function distributes a user input operator on cells specified
in a list and identity operator on the rest.
The operator distribute_operator() distributes it over all the cells indiscriinately.

```python
lattice_412 = Lattice1d(num_cell=4, boundary = "periodic", cell_num_site = 1, cell_site_dof = [2])
op = Qobj(np.array([[0,1],[1,0]]) )
op_sp = lattice_412.operator_at_cells(op, cells = [1,2])
op_all = lattice_412.distribute_operator(op)
print(op_sp)
print(op_all)
```

    Quantum object: dims = [[4, 2], [4, 2]], shape = (8, 8), type = oper, isherm = True
    Qobj data =
    [[0. 0. 0. 0. 0. 0. 0. 0.]
     [0. 0. 0. 0. 0. 0. 0. 0.]
     [0. 0. 0. 1. 0. 0. 0. 0.]
     [0. 0. 1. 0. 0. 0. 0. 0.]
     [0. 0. 0. 0. 0. 1. 0. 0.]
     [0. 0. 0. 0. 1. 0. 0. 0.]
     [0. 0. 0. 0. 0. 0. 0. 0.]
     [0. 0. 0. 0. 0. 0. 0. 0.]]
    Quantum object: dims = [[4, 2], [4, 2]], shape = (8, 8), type = oper, isherm = True
    Qobj data =
    [[0. 1. 0. 0. 0. 0. 0. 0.]
     [1. 0. 0. 0. 0. 0. 0. 0.]
     [0. 0. 0. 1. 0. 0. 0. 0.]
     [0. 0. 1. 0. 0. 0. 0. 0.]
     [0. 0. 0. 0. 0. 1. 0. 0.]
     [0. 0. 0. 0. 1. 0. 0. 0.]
     [0. 0. 0. 0. 0. 0. 0. 1.]
     [0. 0. 0. 0. 0. 0. 1. 0.]]

### Hamiltonian:
The Hamiltoian() function returns the Hamiltonian for the lattice.

```python
boundary_condition = "periodic"
cells = 8
Periodic_Atom_Chain = Lattice1d(num_cell=cells, boundary = boundary_condition)
Hamt = Periodic_Atom_Chain.Hamiltonian()
print(Hamt)
```

    Quantum object: dims = [[8], [8]], shape = (8, 8), type = oper, isherm = True
    Qobj data =
    [[ 0. -1.  0.  0.  0.  0.  0. -1.]
     [-1.  0. -1.  0.  0.  0.  0.  0.]
     [ 0. -1.  0. -1.  0.  0.  0.  0.]
     [ 0.  0. -1.  0. -1.  0.  0.  0.]
     [ 0.  0.  0. -1.  0. -1.  0.  0.]
     [ 0.  0.  0.  0. -1.  0. -1.  0.]
     [ 0.  0.  0.  0.  0. -1.  0. -1.]
     [-1.  0.  0.  0.  0.  0. -1.  0.]]

### Dispersion Relation:
plot_dispersion() plots the valid (same as the number of unit cells) points in k-space
over the dispersion relation of an infinite crystal.

get_dispersion() returns the tuple of two np.ndarrays (knxA,val_kns). knxA has the valid k-values in it and val_kns has the band energies at those k-values. The length of
the unit cell is always set to 1.

```python
Periodic_Atom_Chain.plot_dispersion()
[knxA,val_kns] = Periodic_Atom_Chain.get_dispersion()
print(knxA)
print(val_kns)
```

## Unitary evolution a Gaussian Wavepacket with mesolve

### With no initial momentum

```python
num_cellN = 51
discrete_space_periodic = Lattice1d(num_cell=num_cellN, boundary = "periodic", cell_num_site = 1,
                                    cell_site_dof = [1])
H0 = discrete_space_periodic.Hamiltonian()
xs = np.linspace(0, num_cellN-1, num_cellN)
```

```python
sig = 3     # A standard deviation of 3
xm = num_cellN //2 + 15
psi0 = 1/np.sqrt(2*np.pi*sig**2) * np.exp(-(xs-xm)**2/2/sig/sig)
psi0 = Qobj(np.sqrt(psi0))
```

```python
tlist = np.linspace(0,24,801)
options = Options(atol=1e-12)
options.store_states = True
states_Gauss_0 = mesolve(H0, psi0, tlist, [], [], options=options)
```

```python
t0 = 0
t1 = 180
t2 = 360
t3 = 540
t4 = 720
x_t0 = states_Gauss_0.states[t0]
x_t1 = states_Gauss_0.states[t1]
x_t2 = states_Gauss_0.states[t2]
x_t3 = states_Gauss_0.states[t3]
x_t4 = states_Gauss_0.states[t4]

plt.plot(xs, np.abs(x_t0))
plt.plot(xs, np.abs(x_t1))
plt.plot(xs, np.abs(x_t2))
plt.plot(xs, np.abs(x_t3))
plt.plot(xs, np.abs(x_t4))
plt.xlabel('space', fontsize=14)
plt.ylabel('Wavepacket shape', fontsize=14)
plt.legend(['t0', 't1', 't2', 't3', 't4'])
plt.show()
plt.close()
```

The wavepacket disperses over time keeping the periodic nature in space, since we picked a periodic boundary cndition for space.

### With initial momentum

We initialize a Gaussian state with initial crystal momentum $\pi/3$ and plot the momentum expectation value with the unitary Schrodinger evolution.

```python
sig = 3
xm = num_cellN //2 + 15
psi0 = 1/np.sqrt(2*np.pi*sig**2) * np.exp(-(xs-xm)**2/2/sig/sig)
psi0 = Qobj(np.sqrt(psi0) * np.exp(np.pi*1j*xs/3) )
```

```python
k = discrete_space_periodic.k()
tlist = np.linspace(0,24,801)
options = Options(atol=1e-12)
options.store_states = True
states_Gauss_k = mesolve(H0, psi0, tlist, [], [k], options=options)
```

```python
plt.plot(tlist, states_Gauss_k.expect[0])
plt.xlabel('Time', fontsize=14)
plt.ylabel(r'$\langle k \rangle$', fontsize=14)
plt.ylim([np.pi/3.01, np.pi/2.99])
plt.show()
plt.close()
```

```python
np.pi/3
```

    1.0471975511965976

The crystal momentum operator commutes with the Hamiltonian, so it is conserved in a Hamiltonian evolution, as expected.

```python
t0 = 0
t1 = 40
t2 = 80
t3 = 120
t4 = 160
x_t0 = states_Gauss_k.states[t0]
x_t1 = states_Gauss_k.states[t1]
x_t2 = states_Gauss_k.states[t2]
x_t3 = states_Gauss_k.states[t3]
x_t4 = states_Gauss_k.states[t4]

plt.plot(xs, np.abs(x_t0))
plt.plot(xs, np.abs(x_t1))
plt.plot(xs, np.abs(x_t2))
plt.plot(xs, np.abs(x_t3))
plt.plot(xs, np.abs(x_t4))
plt.xlabel('space', fontsize=14)
plt.ylabel('Wavepacket shape', fontsize=14)
plt.legend(['t0', 't1', 't2', 't3', 't4'])
plt.show()
plt.close()
```

Due to the initial momentum, the wave-packet moves to the right keeping the momentum as well as disperses.

### Simulating a hardwall

We can simulate a hardwall by putting in a large potential on the first and last site of the lattice as well as change the boundary condition to "aperiodic".

```python
sig = 3
xm = num_cellN //2 + 5
psi0 = 1/np.sqrt(2*np.pi*sig**2) * np.exp(-(xs-xm)**2/2/sig/sig)
psi0 = Qobj(np.sqrt(psi0) * np.exp(np.pi*1j*xs/3) )
```

```python
discrete_space_aperiodic = Lattice1d(num_cell=num_cellN, boundary = "aperiodic",
                                     cell_num_site = 1, cell_site_dof = [1])

psiL = discrete_space_aperiodic.basis(0,0,[0])
psiR = discrete_space_aperiodic.basis(num_cellN-1,0,[0])
Ha = discrete_space_aperiodic.Hamiltonian()
H_p = 1e4*(psiL * psiL.dag() + psiR * psiR.dag() )
tlist = np.linspace(0,30,5001)
options = Options(atol=1e-12)
options.store_states = True
states_Gauss_k_HW = mesolve(Ha+H_p, psi0, tlist, [], [k], options=options)
# Warning:  This calculation takes upto a minute
```

```python
t0 = 0
t1 = 1000
t2 = 2000
t3 = 3000
t4 = 4000
t5 = 5000
x_t0 = states_Gauss_k_HW.states[t0]
x_t1 = states_Gauss_k_HW.states[t1]
x_t2 = states_Gauss_k_HW.states[t2]
x_t3 = states_Gauss_k_HW.states[t3]
x_t4 = states_Gauss_k_HW.states[t4]
x_t5 = states_Gauss_k_HW.states[t5]

plt.plot(xs, np.abs(x_t0))
plt.plot(xs, np.abs(x_t1))
plt.plot(xs, np.abs(x_t2))
plt.plot(xs, np.abs(x_t3))
plt.plot(xs, np.abs(x_t4))
plt.plot(xs, np.abs(x_t5))
plt.xlabel('space', fontsize=14)
plt.ylabel('Wavepacket shape', fontsize=14)
plt.legend(['t0', 't1', 't2', 't3', 't4', 't5'])
plt.show()
plt.close()
```

```python
plt.plot(tlist, states_Gauss_k_HW.expect[0])
plt.xlabel('Time', fontsize=14)
plt.ylabel(r'$\langle k \rangle$', fontsize=14)
plt.show()
plt.close()
```

```python
kd = discrete_space_aperiodic.k()
psi_f = states_Gauss_k_HW.states[3200]
kex0 = psi0.dag() * kd * psi0
kexf = psi_f.dag() * kd * psi_f
print('Initital momentum: ', kex0)
print('Final momentum: ', kexf)
```

We confirm that, the final momentum is indeed exactly the opposite of the initial momentum.

### Dissipation induced translational motion

We can initialize a Gaussian state with no initial momentum and devise a dissipation induced translational motion scheme with forming a collapse operator that translates the wavepacket by one lattice point periodically.

```python
num_cellN = 51
discrete_space_periodic = Lattice1d(num_cell=num_cellN, boundary = "periodic", cell_num_site = 1,
                                    cell_site_dof = [1])
H0 = discrete_space_periodic.Hamiltonian()
xp = discrete_space_periodic.x()
kp = discrete_space_periodic.k()
xs = np.linspace(0, num_cellN-1, num_cellN)
```

```python
sig = 3     # A standard deviation of 3
xm = num_cellN //2
psi0 = 1/np.sqrt(2*np.pi*sig**2) * np.exp(-(xs-xm)**2/2/sig/sig)
psi0 = Qobj(np.sqrt(psi0))
```

```python
lat_trR = np.diag(np.zeros(num_cellN-1)+1, -1)
lat_trR[0, num_cellN-1] = 1    # translate right
lat_trL = np.diag(np.zeros(num_cellN-1)+1, 1)
lat_trL[num_cellN-1, 0] = 1     #  translate left
trR = Qobj(lat_trR)
trL = Qobj(lat_trL)
```

```python
gamma = 2
col_op = [np.sqrt(gamma) * trR ]
```

```python
tlistC = np.linspace(0,24,801)
options = Options(atol=1e-12)
options.store_states = True
rho0 = psi0 * psi0.dag()
states_Gauss_0 = mesolve(H0, rho0, tlistC, col_op, [kp], options=options)
```

```python
plt.plot(tlistC, states_Gauss_0.expect[0])
plt.xlabel('Time', fontsize=14)
plt.ylabel(r'$\langle k \rangle$', fontsize=14)
plt.ylim([-1e-8, 1e-8])
plt.show()
plt.close()
```

```python
t0 = 0
t1 = 140
t2 = 280
t3 = 420
t4 = 560

diag_x0 = np.diag(states_Gauss_0.states[t0])
diag_x1 = np.diag(states_Gauss_0.states[t1])
diag_x2 = np.diag(states_Gauss_0.states[t2])
diag_x3 = np.diag(states_Gauss_0.states[t3])
diag_x4 = np.diag(states_Gauss_0.states[t4])

plt.plot(xs, np.abs(diag_x0))
plt.plot(xs, np.abs(diag_x1))
plt.plot(xs, np.abs(diag_x2))
plt.plot(xs, np.abs(diag_x3))
plt.plot(xs, np.abs(diag_x4))
plt.xlabel('space', fontsize=14)
plt.ylabel('Wavepacket shape', fontsize=14)
plt.title('Nonunitary evolution')
plt.show()
plt.close()
```

The wave-packet disperses and trannslates to the right, but the momentum expectation remains zero, since the translation is induced by the dissipation.

# Example: A Coupled Resonator Optical Waveguide¶

We now demonstrate the basic functionality of QuTiPs Lattice1d class of the lattice module with the example of a Coupled Resonator Optical Waveguide(CROW)(ref. [2]).

\begin{eqnarray}
H_0 = \sum\limits_{n} \left(H_a + H_b + H_{ab} + H^{\dagger}_{ab}   \right) \\
H_a = \frac{J}{2} a_n^{\dagger} \left( e^{-i\eta} a_{n-1} + e^{i\eta} a_{n+1} \right) \\
H_b = \frac{J}{2} b_n^{\dagger} \left( e^{i\eta} b_{n-1} + e^{-i\eta} b_{n+1} \right) \\
H_{ab} = J a_n^{\dagger} \left( sin (\eta) b_n + \frac{1}{2} \left(b_{n-1} + b_{n+1}  \right)  \right)
\end{eqnarray}

For implementation with Lattice1d class, we resolve the Hamiltonian into unitcells.
\begin{equation}
H = \sum\limits_{n} H_n
\end{equation}

\begin{eqnarray}
H_{n}= \begin{bmatrix}
    a_{n}^{\dagger} & b_{n}^{\dagger}
\end{bmatrix}
\begin{bmatrix}
      o     & J sin(\eta) \\
    J sin(\eta)    & 0
\end{bmatrix}
\begin{bmatrix}
    a_{n}  \\
    b_{n}
\end{bmatrix}
\ \ \ \ \ \ \ \ \ \ \\
+ \left( \begin{bmatrix}
    a_{n}^{\dagger} & b_{n}^{\dagger}
\end{bmatrix}
\begin{bmatrix}
      e^{i\eta}     & 1 \\
    1    & e^{-i\eta}
\end{bmatrix}
\begin{bmatrix}
    a_{n+1}  \\
    b_{n+1}
\end{bmatrix} + H.C. \right)
\end{eqnarray}

In the present case, we have 1 site in every unit cell and 2 dofs per site. And

\begin{equation}
\text{cell_Hamiltonian} = \begin{bmatrix}
      o     & J sin(\eta) \\
    J sin(\eta)    & 0
\end{bmatrix}
\end{equation}
\begin{equation}
\text{inter_hop} = \begin{bmatrix}
      e^{i\eta}     & 1 \\
    1    & e^{-i\eta}
\end{bmatrix}
\end{equation}

```python
cells = 4
cell_num_site = 1
cell_site_dof = [2]
J = 2
```

```python
### For eta = 0
eta = 0
H_cell = Qobj(np.array([[0, J * np.sin(eta)], [J * np.sin(eta), 0]]))
inter_cell_T = (J/2) * Qobj(np.array([[np.exp(eta * 1j), 1], [1, np.exp(-eta*1j)]]))
CROW_lattice = Lattice1d(num_cell=cells, boundary = "periodic", cell_num_site = 1,
                         cell_site_dof = [2], Hamiltonian_of_cell = H_cell,
                         inter_hop = inter_cell_T )
CROW_lattice.plot_dispersion()
```

```python
### For eta = pi/4
eta = np.pi/4
H_cell = Qobj(np.array([[0, J * np.sin(eta)], [J * np.sin(eta), 0]]))
inter_cell_T = (J/2) * Qobj(np.array([[np.exp(eta * 1j), 1], [1, np.exp(-eta*1j)]]))
CROW_lattice = Lattice1d(num_cell=cells, boundary = "periodic", cell_num_site = 1,
                         cell_site_dof = [2], Hamiltonian_of_cell = H_cell,
                         inter_hop = inter_cell_T )
CROW_lattice.plot_dispersion()
```

```python
### For eta = pi/2
eta = np.pi/2
H_cell = Qobj(np.array([[0, J * np.sin(eta)], [J * np.sin(eta), 0]]))
inter_cell_T = (J/2) * Qobj(np.array([[np.exp(eta * 1j), 1], [1, np.exp(-eta*1j)]]))
CROW_lattice = Lattice1d(num_cell=cells, boundary = "periodic", cell_num_site = 1,
                         cell_site_dof = [2], Hamiltonian_of_cell = H_cell,
                         inter_hop = inter_cell_T )
CROW_lattice.plot_dispersion()
```

The three dispersion relationships for the three values of $\eta$ can be compared with the published results in Ref [2].

## Unitary dynamics example

### No initial momentum

```python
num_cell = 100
J = 2
eta = np.pi/2
H_cell = Qobj(np.array([[0, J * np.sin(eta)], [J * np.sin(eta), 0]]))
inter_cell_T = (J/2) * Qobj(np.array([[np.exp(eta * 1j), 1], [1, np.exp(-eta*1j)]]))
CROW_lattice = Lattice1d(num_cell=num_cell, boundary = "periodic", cell_num_site = 2,
                         cell_site_dof = [1], Hamiltonian_of_cell = H_cell,
                         inter_hop = inter_cell_T)
```

```python
HCROW = CROW_lattice.Hamiltonian()
kC = CROW_lattice.k()
```

```python
nx = 1
ne = 2
positions = np.kron(range(nx), [1/nx for i in range(ne)])
S = np.kron(np.ones(num_cell), positions)
R = np.kron(range(0, num_cell), np.ones(nx*ne))
xA = R+S
sig = 3     # A standard deviation of 3
xm = num_cell //2
psi0 = 1/np.sqrt(2*np.pi*sig**2) * np.exp(-(xA-xm)**2/2/sig/sig)
psi0 = Qobj(np.sqrt(psi0))
tlistW = np.linspace(0,30,5001)
options = Options(atol=1e-12)
options.store_states = True
states_CROW_u = mesolve(HCROW, psi0, tlistW, [], [kC], options=options)
```

```python
plt.plot(tlistW, states_CROW_u.expect[0])
plt.xlabel('Time', fontsize=14)
plt.ylabel(r'$\langle k \rangle$', fontsize=14)
plt.ylim([-1e-8, 1e-8])
plt.show()
plt.close()
```

```python
t0 = 0
t1 = 1000
t2 = 2000
t3 = 3000
t4 = 4000
t5 = 5000
x_t0 = states_CROW_u.states[t0]
x_t1 = states_CROW_u.states[t1]
x_t2 = states_CROW_u.states[t2]
x_t3 = states_CROW_u.states[t3]
x_t4 = states_CROW_u.states[t4]
x_t5 = states_CROW_u.states[t5]

plt.plot(xA, np.abs(x_t0))
plt.plot(xA, np.abs(x_t1))
plt.plot(xA, np.abs(x_t2))
plt.plot(xA, np.abs(x_t3))
plt.plot(xA, np.abs(x_t4))
plt.plot(xA, np.abs(x_t5))
plt.xlabel('space', fontsize=14)
plt.ylabel('Wavepacket shape', fontsize=14)
plt.legend(['t0', 't1', 't2', 't3', 't4', 't5'])
plt.show()
plt.close()
```

### With Initial momentum

```python
sig = 3
xm = num_cell //2 + 15
psi0 = 1/np.sqrt(2*np.pi*sig**2) * np.exp(-(xA-xm)**2/2/sig/sig)
psi0 = Qobj(np.sqrt(psi0) * np.exp(1*np.pi*1j*xA/3) )
tlistCk = np.linspace(0,30,5001)
options = Options(atol=1e-12)
options.store_states = True
states_CROW_uk = mesolve(HCROW, psi0, tlistCk, [], [kC], options=options)
```

```python
plt.plot(tlistCk, states_CROW_uk.expect[0])
plt.xlabel('Time', fontsize=14)
plt.ylabel(r'$\#langle k \rangle$', fontsize=14)
plt.ylim([1.046, 1.048])
plt.show()
plt.close()
```

```python
t0 = 0
t1 = 1000
t2 = 2000
t3 = 3000
t4 = 4000
t5 = 5000
x_t0 = states_CROW_u.states[t0]
x_t1 = states_CROW_u.states[t1]
x_t2 = states_CROW_u.states[t2]
x_t3 = states_CROW_u.states[t3]
x_t4 = states_CROW_u.states[t4]
x_t5 = states_CROW_u.states[t5]

plt.plot(xA, np.abs(x_t0))
plt.plot(xA, np.abs(x_t1))
plt.plot(xA, np.abs(x_t2))
plt.plot(xA, np.abs(x_t3))
plt.plot(xA, np.abs(x_t4))
plt.plot(xA, np.abs(x_t5))
plt.xlabel('space', fontsize=14)
plt.ylabel('Wavepacket shape', fontsize=14)
plt.legend(['t0', 't1', 't2', 't3', 't4', 't5'])
plt.show()
plt.close()
```

```python
t0 = 0
t1 = 1000
t2 = 2000
t3 = 3000
t4 = 4000
t5 = 5000
x_t0 = states_CROW_u.states[t0]
x_t1 = states_CROW_u.states[t1]
x_t2 = states_CROW_u.states[t2]
x_t3 = states_CROW_u.states[t3]
x_t4 = states_CROW_u.states[t4]
x_t5 = states_CROW_u.states[t5]

plt.plot(xA[range(0,200,2)], np.abs(x_t0.full()[range(0,200,2)]))
plt.plot(xA[range(0,200,2)], np.abs(x_t1.full()[range(0,200,2)]))
plt.plot(xA[range(0,200,2)], np.abs(x_t2.full()[range(0,200,2)]))
plt.plot(xA[range(0,200,2)], np.abs(x_t3.full()[range(0,200,2)]))
plt.plot(xA[range(0,200,2)], np.abs(x_t4.full()[range(0,200,2)]))
plt.plot(xA[range(0,200,2)], np.abs(x_t5.full()[range(0,200,2)]))
plt.xlabel('space(left sublattice)', fontsize=14)
plt.ylabel('Wavepacket shape', fontsize=14)
plt.legend(['t0', 't1', 't2', 't3', 't4', 't5'])
plt.show()
plt.close()
```

```python
t0 = 0
t1 = 1000
t2 = 2000
t3 = 3000
t4 = 4000
t5 = 5000
x_t0 = states_CROW_u.states[t0]
x_t1 = states_CROW_u.states[t1]
x_t2 = states_CROW_u.states[t2]
x_t3 = states_CROW_u.states[t3]
x_t4 = states_CROW_u.states[t4]
x_t5 = states_CROW_u.states[t5]

plt.plot(xA[range(1,200,2)], np.abs(x_t0.full()[range(1,200,2)]))
plt.plot(xA[range(1,200,2)], np.abs(x_t1.full()[range(1,200,2)]))
plt.plot(xA[range(1,200,2)], np.abs(x_t2.full()[range(1,200,2)]))
plt.plot(xA[range(1,200,2)], np.abs(x_t3.full()[range(1,200,2)]))
plt.plot(xA[range(1,200,2)], np.abs(x_t4.full()[range(1,200,2)]))
plt.plot(xA[range(1,200,2)], np.abs(x_t5.full()[range(1,200,2)]))
plt.xlabel('space(right sublattice)', fontsize=14)
plt.ylabel('Wavepacket shape', fontsize=14)
plt.legend(['t0', 't1', 't2', 't3', 't4', 't5'])
plt.show()
plt.close()
```

## translation by dissipation

```python
cells = 100
nx = 2
ne = 1
positions = np.kron(range(nx), [1/nx for i in range(ne)])
S = np.kron(np.ones(cells), positions)
R = np.kron(range(0, cells), np.ones(nx*ne))
xA = R+S
```

```python
eta = np.pi/2
H_cell = Qobj(np.array([[0, J * np.sin(eta)], [J * np.sin(eta), 0]]))
inter_cell_T = (J/2) * Qobj(np.array([[np.exp(eta * 1j), 1], [1, np.exp(-eta*1j)]]))
CROW_lattice = Lattice1d(num_cell=cells, boundary = "periodic", cell_num_site = 2,
                         cell_site_dof = [1], Hamiltonian_of_cell = H_cell,
                         inter_hop = inter_cell_T)
```

```python
HCROW = CROW_lattice.Hamiltonian()
kC = CROW_lattice.k()
```

```python
lat_trR = np.diag(np.zeros(cells-1)+1, -1)
lat_trR[0, cells-1] = 1   # translate to the right

lat_trL = np.diag(np.zeros(cells-1)+1, 1)
lat_trL[cells-1, 0] = 1   # translate to the left

trR = Qobj(lat_trR)
trL = Qobj(lat_trL)
```

```python
gamma = 0.5
col_op = [np.sqrt(gamma) * tensor(trL, qeye(2)) ]  # We could have used trR for translation to the right
```

```python
sig = 3
xm = cells //2 + 15
psi0 = 1/np.sqrt(2*np.pi*sig**2) * np.exp(-(xA-xm)**2/2/sig/sig)
psi0 = Qobj(np.sqrt(psi0))
tlistCN = np.linspace(0,30,601)
options = Options(atol=1e-12)
options.store_states = True
states_CROW_nu = mesolve(HCROW, psi0, tlistCN, col_op, [kC], options=options)
```

```python
plt.plot(tlistCN, states_CROW_nu.expect[0])
plt.xlabel('Time', fontsize=14)
plt.ylabel(r'$\#langle k \rangle$', fontsize=14)
plt.ylim([-1e-8, 1e-8])
plt.show()
plt.close()
```

```python
t0 = 0
t1 = 100
t2 = 200
t3 = 300
t4 = 400
t5 = 500
x_t0 = np.diag(states_CROW_nu.states[t0])
x_t1 = np.diag(states_CROW_nu.states[t1])
x_t2 = np.diag(states_CROW_nu.states[t2])
x_t3 = np.diag(states_CROW_nu.states[t3])
x_t4 = np.diag(states_CROW_nu.states[t4])
x_t5 = np.diag(states_CROW_nu.states[t5])

plt.plot(xA, np.abs(x_t0))
plt.plot(xA, np.abs(x_t1))
plt.plot(xA, np.abs(x_t2))
plt.plot(xA, np.abs(x_t3))
plt.plot(xA, np.abs(x_t4))
plt.plot(xA, np.abs(x_t5))
plt.xlabel('space', fontsize=14)
plt.ylabel('Wavepacket shape', fontsize=14)
plt.legend(['t0', 't1', 't2', 't3', 't4', 't5'])
plt.show()
plt.close()
```

### References
[1] J. R. Johansson, P. D. Nation, and F. Nori, Comp. Phys. Comm. 183, 1760 (2012). http://qutip.org

[2] Han, JungYun, Clemens Gneiting, and Daniel Leykam. "Helical transport in coupled resonator waveguides." Physical Review B 99.22 (2019): 224201.
https://journals.aps.org/prb/abstract/10.1103/PhysRevB.99.224201#

```python
qutip.about()
```

```python
qutip.cite()
```

qutip-notebooks-master/examples/lattice-Introduction to 1d Lattice models.ipynb