# Imports and Simulation Parameters

```python
import numpy as np
import math
import cmath
import scipy
import scipy.integrate
import sys

import matplotlib
import matplotlib.pyplot as plt
%matplotlib inline

hbar = 1.0 / (2.0 * np.pi)

ZERO_TOLERANCE = 10**-6
```

    /Users/joseph/anaconda/lib/python2.7/site-packages/matplotlib/font_manager.py:273: UserWarning: Matplotlib is building the font cache using fc-list. This may take a moment.
      warnings.warn('Matplotlib is building the font cache using fc-list. This may take a moment.')

```python
MAX_VIBRATIONAL_STATES = 200

STARTING_GROUND_STATES = 5
STARTING_EXCITED_STATES = 5

time_scale_set = 10  #will divide the highest energy to give us the time step
low_frequency_cycles = 25.0 #will multiply the frequency of the lowest frequency mode to get the max time
```

```python
#See if a factorial_Storage dictionary exists already and if not, create one
try:
    a = factorial_storage
except:
    factorial_storage = {}
```

# Defining Parameters of the System

```python
energy_g = 0
energy_gamma = .1

energy_e = 0
energy_epsilon = .31

Huang_Rhys_Parameter = .80
S = Huang_Rhys_Parameter

#bookkeeping
overlap_storage = {}
electronic_energy_gap = energy_e + .5*energy_epsilon - (energy_g  + .5 * energy_gamma)
min_energy = energy_g + energy_gamma * .5

```

```python
mu_0 = 1.0
```

If we set the central frequency of a pulse at the 0->0 transition, and we decide we care about the ratio of the 0->1 transition to the 0->0 transition and set that to be $\tau$ then the desired pulse width will be
\begin{align}
    \sigma &= \sqrt{-\frac{2 \ln (\tau)}{\omega_{\epsilon}^2}}
\end{align}

```python
def blank_wavefunction(number_ground_states, number_excited_states):
    return np.zeros((number_ground_states + number_excited_states))

def perturbing_function(time):
#     stdev = 30000.0 * dt #very specific to 0->0 transition!
    stdev = 3000.0 * dt #clearly has a small amount of amplitude on 0->1 transition
    center = 6 * stdev
    return np.cos(electronic_energy_gap*(time - center) / hbar)*np.exp( - (time - center)**2 / (2 * stdev**2)) / stdev

def time_function_handle_from_tau(tau_proportion):
    stdev = np.sqrt( -2.0 * np.log(tau_proportion) / (energy_epsilon/hbar)**2)
    center = 6 * stdev
    return center, stdev, lambda t: np.cos(electronic_energy_gap*(t - center) / hbar)*np.exp( - (t - center)**2 / (2 * stdev**2)) / stdev

def time_function_handle_from_tau_and_kappa(tau_proportion, kappa_proportion):
    stdev = np.sqrt( -2.0 * np.log(tau_proportion) / (energy_epsilon/hbar)**2)
    center = 6 * stdev
    return center, stdev, lambda t: kappa_proportion * energy_gamma * np.cos(electronic_energy_gap*(t - center) / hbar)*np.exp( - (t - center)**2 / (2 * stdev**2)) / stdev

def perturbing_function_define_tau(time, tau_proportion):
    center, stdev, f = time_function_handle_from_tau(tau_proportion)
    return f(time)
```

```python

```

```python

```

# Defining Useful functions

$ O_{m}^{n} = \left(-1\right)^{n} \sqrt{\frac{e^{-S}S^{m+n}}{m!n!}} \sum_{j=0}^{\min \left( m,n \right)} \frac{m!n!}{j!(m-j)!(n-j)!}(-1)^j S^{-j} $

```python
def factorial(i):
    if i in factorial_storage:
        return factorial_storage[i]
    if i <= 1:
        return 1.0
    else:
        out =  factorial(i - 1) * i
        factorial_storage[i] = out
        return out

def ndarray_factorial(i_array):
    return np.array([factorial(i) for i in i_array])
```

```python
def overlap_function(ground_quantum_number, excited_quantum_number):
    m = ground_quantum_number
    n = excited_quantum_number
    if (m,n) in overlap_storage:
        return overlap_storage[(m,n)]
    output = (-1)**n
    output *= math.sqrt(math.exp(-S) * S**(m + n) /(factorial(m) * factorial(n)) )
    j_indeces = np.array(range(0, min(m,n) + 1))
    j_summation = factorial(m) * factorial(n) * np.power(-1.0, j_indeces) * np.power(S, -j_indeces)
    j_summation = j_summation / (ndarray_factorial(j_indeces) * ndarray_factorial( m - j_indeces) * ndarray_factorial(n - j_indeces) )
    output *= np.sum(j_summation)
    overlap_storage[(m,n)] = output
    return output
```

# Solving the Differential Equation

\begin{align*}
	\left(\frac{d G_a(t)}{dt} + \frac{i}{\hbar}\Omega_{(a)}\right) &=-E(t)\frac{i}{\hbar}  \sum_{b} E_b(t) \mu_{a}^{b}\\
	\left(\frac{d E_b(t)}{dt} + \frac{i}{\hbar} \Omega^{(b)} \right) &=-E(t)\frac{i}{\hbar} \sum_{a} G_a(t) \mu_{a}^{b}
\end{align*}
Or in a more compact form:
\begin{align*}
	\frac{d}{dt}\begin{bmatrix}
		G_a(t) \\
		E_b(t)
	\end{bmatrix}
	= -\frac{i}{\hbar}
	\begin{bmatrix}
		\Omega_{(a)} & E(t) \mu_{a}^{b} \\
		E(t) \mu_{a}^{b} & \Omega^{b}
	\end{bmatrix}
	\cdot
	\begin{bmatrix}
		G_a(t) \\
		E_b(t)
	\end{bmatrix}
\end{align*}

```python
def ode_diagonal_matrix(number_ground_states, number_excited_states):
    #Define the Matrix on the RHS of the above equation
    ODE_DIAGONAL_MATRIX = np.zeros((number_ground_states + number_excited_states, number_ground_states + number_excited_states), dtype=np.complex)
    #set the diagonals
    for ground_i in range(number_ground_states):
        ODE_DIAGONAL_MATRIX[ground_i, ground_i] = -1.0j * (energy_g + energy_gamma * (ground_i + .5)) / hbar

    for excited_i in range(number_excited_states):
        excited_index = excited_i + number_ground_states #the offset since the excited states comes later
        ODE_DIAGONAL_MATRIX[excited_index, excited_index] = -1.0j * (energy_e + energy_epsilon * (excited_i + .5)) / hbar

    return ODE_DIAGONAL_MATRIX

#now for the off-diagonals
def mu_matrix(c, number_ground_states, number_excited_states):
    MU_MATRIX = np.zeros((number_ground_states, number_excited_states), dtype = np.complex)
    for ground_a in range(number_ground_states):
        for excited_b in range(number_excited_states):
            new_mu_entry = overlap_function(ground_a, excited_b)
            if ground_a >0:
                new_mu_entry += c * math.sqrt(ground_a) * overlap_function(ground_a - 1, excited_b)

            new_mu_entry += c * math.sqrt(ground_a+1) * overlap_function(ground_a + 1, excited_b)

            MU_MATRIX[ground_a, excited_b] = new_mu_entry
    return MU_MATRIX

def ode_off_diagonal_matrix(c_value, number_ground_states, number_excited_states):
    output  = np.zeros((number_ground_states + number_excited_states, number_ground_states + number_excited_states), dtype=np.complex)
    MU_MATRIX = mu_matrix(c_value, number_ground_states, number_excited_states)
    output[0:number_ground_states, number_ground_states:] = -1.0j * mu_0 * MU_MATRIX  / hbar
    output[number_ground_states:, 0:number_ground_states] = -1.0j * mu_0 * MU_MATRIX.T  / hbar
    return output

def IR_transition_dipoles(number_ground_states, number_excited_states):
    "outputs matrices to calculate ground and excited state IR emission spectra.  Can be combined for total"
    output_g = np.zeros((number_ground_states + number_excited_states, number_ground_states + number_excited_states), dtype=np.complex)
    output_e = np.zeros((number_ground_states + number_excited_states, number_ground_states + number_excited_states), dtype=np.complex)

    for ground_a in range(number_ground_states):
        try:
            output_g[ground_a, ground_a + 1] = math.sqrt(ground_a + 1)
            output_g[ground_a + 1, ground_a] = math.sqrt(ground_a + 1)
        except:
            pass

        try:
            output_g[ground_a, ground_a - 1] = math.sqrt(ground_a)
            output_g[ground_a - 1, ground_a] = math.sqrt(ground_a)
        except:
            pass

    for excited_a in range(number_excited_states):
        matrix_index_e = number_ground_states + excited_a -1 #because of how 'number_ground_states' is defined
        try:
            output_e[matrix_index_e, matrix_index_e + 1] = math.sqrt(excited_a + 1)
            output_e[matrix_index_e + 1, matrix_index_e] = math.sqrt(excited_a + 1)
        except:
            pass

        try:
            output_e[matrix_index_e, matrix_index_e - 1] = math.sqrt(excited_a)
            output_e[matrix_index_e - 1, matrix_index_e] = math.sqrt(excited_a)
        except:
            pass

    return output_g, output_e
```

\begin{align*}
	\mu(x) &= \mu_0 \left(1 + \lambda x \right) \\
	&= \mu_0 \left(1 + c\left(a + a^{\dagger} \right) \right) \\
	\mu_{a}^{b} &= \mu_0\left(O_{a}^{b}  + c\left(\sqrt{a}O_{a-1}^{b} + \sqrt{a+1}O_{a+1}^{b}\right) \right)
\end{align*}

```python

```

```python
class VibrationalStateOverFlowException(Exception):
    def __init__(self):
        pass
```

```python

```

```python
def propagate_amplitude_to_end_of_perturbation(c_value, ratio_01_00, kappa=1, starting_ground_states=STARTING_GROUND_STATES, starting_excited_states=STARTING_EXCITED_STATES):
    center_time, stdev, time_function = time_function_handle_from_tau_and_kappa(ratio_01_00, kappa)
    ending_time = center_time + 8.0 * stdev

    number_ground_states = starting_ground_states
    number_excited_states = starting_excited_states

    while number_excited_states + number_ground_states < MAX_VIBRATIONAL_STATES:

        #define time scales
        max_energy = energy_e + energy_epsilon * (.5 + number_excited_states) + kappa * energy_gamma * mu_0
        dt = 1.0 / (time_scale_set *  max_energy)

        ODE_DIAGONAL = ode_diagonal_matrix(number_ground_states, number_excited_states)
        ODE_OFF_DIAGONAL = ode_off_diagonal_matrix(c_value, number_ground_states, number_excited_states)

        def ODE_integrable_function(time, coefficient_vector):
            ODE_TOTAL_MATRIX = ODE_OFF_DIAGONAL * time_function(time) + ODE_DIAGONAL
            return np.dot(ODE_TOTAL_MATRIX, coefficient_vector)

        #define the starting wavefuntion
        initial_conditions = blank_wavefunction(number_ground_states, number_excited_states)
        initial_conditions[0] = 1

        #create ode solver
        current_time = 0.0
        ode_solver = scipy.integrate.complex_ode(ODE_integrable_function)
        ode_solver.set_initial_value(initial_conditions, current_time)

        #Run it
        results = []
        try:  #this block catches an overflow into the highest ground or excited vibrational state
            while current_time < ending_time:
#                 print(current_time, ZERO_TOLERANCE)
                #update time, perform solution
                current_time = ode_solver.t+dt
                new_result = ode_solver.integrate(current_time)
                results.append(new_result)
                #make sure solver was successful
                if not ode_solver.successful():
                    raise Exception("ODE Solve Failed!")
                #make sure that there hasn't been substantial leakage to the highest excited states
                re_start_calculation = False
                if abs(new_result[number_ground_states - 1])**2 >= ZERO_TOLERANCE:
                    number_ground_states +=1
#                     print("Increasing Number of Ground vibrational states to %i " % number_ground_states)
                    re_start_calculation = True
                if abs(new_result[-1])**2 >= ZERO_TOLERANCE:
                    number_excited_states +=1
#                     print("Increasing Number of excited vibrational states to %i " % number_excited_states)
                    re_start_calculation = True
                if re_start_calculation:
                    raise VibrationalStateOverFlowException()
        except VibrationalStateOverFlowException:
            #Move on and re-start the calculation
            continue

        #Finish calculating
        results = np.array(results)
        return results, number_ground_states, number_excited_states
    raise Exception("NEEDED TOO MANY VIBRATIONAL STATES!  RE-RUN WITH DIFFERENT PARAMETERS!")
```

```python
def get_average_quantum_number_time_series(c_value, ratio_01_00, kappa=1, starting_ground_states=STARTING_GROUND_STATES, starting_excited_states=STARTING_EXCITED_STATES):
    results, number_ground_states, number_excited_states = propagate_amplitude_to_end_of_perturbation(c_value, ratio_01_00, kappa, starting_ground_states, starting_excited_states)
    probabilities = np.abs(results)**2
    #calculate the average_vibrational_quantum_number series
    average_ground_quantum_number = probabilities[:,0:number_ground_states].dot(np.array(range(number_ground_states)) )
    average_excited_quantum_number = probabilities[:,number_ground_states:].dot(np.array(range(number_excited_states)))
    return average_ground_quantum_number, average_excited_quantum_number, results, number_ground_states, number_excited_states
```

```python
def IR_emission_spectrum_after_excitation(c_value, ratio_01_00, kappa=1, starting_ground_states=STARTING_GROUND_STATES, starting_excited_states=STARTING_EXCITED_STATES):
    center_time, stdev, time_function = time_function_handle_from_tau_and_kappa(ratio_01_00, kappa)

    perturbation_ending_time = center_time + 8.0 * stdev
    simulation_ending_time = perturbation_ending_time + low_frequency_cycles * hbar/min_energy

    number_ground_states = starting_ground_states
    number_excited_states = starting_excited_states

    while number_excited_states + number_ground_states < MAX_VIBRATIONAL_STATES:

        ir_transDipole_g, ir_transDipole_e = IR_transition_dipoles(number_ground_states, number_excited_states)

        time_emission_g = [0]
        time_emission_e = [0]

        #define time scales
        e = energy_e + energy_epsilon * (.5 + number_excited_states)
        g = energy_g + energy_gamma* (.5 + number_ground_states)
        plus = e + g
        minus = e - g
        J = kappa * energy_gamma * mu_0
        max_split_energy = plus + math.sqrt(minus**2 + 4 * J**2)
        max_energy = max_split_energy * .5
        dt = 1.0 / (time_scale_set *  max_energy)

        time_values = np.arange(0, simulation_ending_time, dt)

        ODE_DIAGONAL = ode_diagonal_matrix(number_ground_states, number_excited_states)
        ODE_OFF_DIAGONAL = ode_off_diagonal_matrix(c_value, number_ground_states, number_excited_states)

        def ODE_integrable_function(time, coefficient_vector):
            ODE_TOTAL_MATRIX = ODE_OFF_DIAGONAL * time_function(time) + ODE_DIAGONAL
            return np.dot(ODE_TOTAL_MATRIX, coefficient_vector)

        def ODE_jacobean(time, coefficient_vector):
            ODE_TOTAL_MATRIX = ODE_OFF_DIAGONAL * time_function(time) + ODE_DIAGONAL
            return ODE_TOTAL_MATRIX

        #define the starting wavefuntion
        initial_conditions = blank_wavefunction(number_ground_states, number_excited_states)
        initial_conditions[0] = 1

        #create ode solver
        current_time = 0.0
        try:
            del ode_solver
        except:
            pass
#         ode_solver = scipy.integrate.complex_ode(ODE_integrable_function)
        ode_solver = scipy.integrate.complex_ode(ODE_integrable_function, jac = ODE_jacobean)
#         ode_solver.set_integrator("lsoda")
        ode_solver.set_integrator("vode", with_jacobian=True)

        ode_solver.set_initial_value(initial_conditions, current_time)

        #Run it
        results = []
        try:  #this block catches an overflow into the highest ground or excited vibrational state
            while current_time < simulation_ending_time:
#                 print(current_time, ZERO_TOLERANCE)
                #update time, perform solution
                current_time = ode_solver.t+dt
                new_result = ode_solver.integrate(current_time)
                results.append(new_result)
                #make sure solver was successful
                if not ode_solver.successful():
                    raise Exception("ODE Solve Failed!")

                if current_time < perturbation_ending_time:
                    #make sure that there hasn't been substantial leakage to the highest excited states
                    re_start_calculation = False
                    if abs(new_result[number_ground_states - 1])**2 >= ZERO_TOLERANCE:
                        number_ground_states +=1
#                         print("Increasing Number of Ground vibrational states to %i " % number_ground_states)
                        re_start_calculation = True
                    if abs(new_result[-1])**2 >= ZERO_TOLERANCE:
                        number_excited_states +=1
#                         print("Increasing Number of excited vibrational states to %i " % number_excited_states)
                        re_start_calculation = True
                    if re_start_calculation:
                        raise VibrationalStateOverFlowException()
                #calculate IR emission
                time_emission_g.append(np.conj(new_result).T.dot(ir_transDipole_g.dot(new_result)))
                time_emission_e.append(np.conj(new_result).T.dot(ir_transDipole_e.dot(new_result)))
                #on to next time value...

        except VibrationalStateOverFlowException:
            #Move on and re-start the calculation
            continue

        #Finish calculating
        results = np.array(results)
        n_t = len(time_emission_e)

        time_emission_g = np.array(time_emission_g)
        time_emission_e = np.array(time_emission_e)

        filter_x = np.array(range(n_t))
        filter_center = n_t / 2.0
        filter_sigma = n_t / 10.0
        filter_values = np.exp(-(filter_x - filter_center)**2 / (2 * filter_sigma**2))

        frequencies = np.fft.fftshift(np.fft.fftfreq(time_emission_g.shape[0], d= dt))
        frequency_emission_g = dt * np.fft.fftshift(np.fft.fft(time_emission_g * filter_values))
        frequency_emission_e = dt * np.fft.fftshift(np.fft.fft(time_emission_e * filter_values))

        return results, frequencies, frequency_emission_g, frequency_emission_e, number_ground_states, number_excited_states
    raise Exception("NEEDED TOO MANY VIBRATIONAL STATES!  RE-RUN WITH DIFFERENT PARAMETERS!")
```

```python
c_values = np.logspace(-4, np.log10(1), 5)
tau_values = np.logspace(-4, np.log10(.95), 10)
kappa_values = np.logspace(-2, np.log10(5), 20)

number_calcs = c_values.shape[0] * tau_values.shape[0] * kappa_values.shape[0]

heating_results_ground = np.zeros((kappa_values.shape[0],  tau_values.shape[0], c_values.shape[0]))
ir_amplitudes = np.zeros(heating_results_ground.shape)
heating_results_excited = np.zeros(heating_results_ground.shape)

# Keep track of the IR Spectrum

n_g = STARTING_GROUND_STATES
n_e = STARTING_EXCITED_STATES

counter = 1
# we will use the value of c as a bellweather for how many starting states to work with.
c_to_ng = {}
c_to_ne = {}
for i_kappa, kappa in enumerate(kappa_values):
    # as we increase in both tau and
    for i_tau, tau in enumerate(tau_values):

        for i_c, c in enumerate(c_values):
            try:
                n_g = c_to_ng[c]
                n_e = c_to_ne[c]
            except:
                n_g = STARTING_GROUND_STATES
                n_e = STARTING_EXCITED_STATES
                c_to_ng[c] = n_g
                c_to_ne[c] = n_e

            sys.stdout.flush()
            sys.stdout.write("\r%i / %i Calculating kappa=%f, c=%f, tau=%f at n_g = %i and n_e=%i..." %(counter, number_calcs, kappa, c, tau, n_g, n_e))
#             print("\r%i / %i Calculating kappa=%f, c=%f, tau=%f at n_g = %i and n_e=%i..." %(counter, number_calcs, kappa, c, tau, n_g, n_e))
#             n_bar_g, n_bar_e, results, num_g, num_e = get_average_quantum_number_time_series(c,
#                                                                                tau,
#                                                                                kappa,
#                                                                                starting_ground_states = n_g,
#                                                                                starting_excited_states = n_e)

#             heating_results_ground[i_kappa, i_tau, i_c] = n_bar_g[-1]
#             heating_results_excited[i_kappa, i_tau, i_c] = n_bar_e[-1]

            _, frequencies, emission_g, emission_e, num_g, num_e = IR_emission_spectrum_after_excitation(c,
                                                                                tau,
                                                                                kappa,
                                                                                starting_ground_states = n_g,
                                                                                starting_excited_states = n_e)
            if num_g > c_to_ng[c]:
                c_to_ng[c] = num_g
            if num_e > c_to_ne[c]:
                c_to_ne[c] = num_e
            vibrational_frequency_index = np.argmin(np.abs(energy_gamma - frequencies))
            ir_power = np.abs(emission_g[vibrational_frequency_index])**2
            ir_amplitudes[i_kappa, i_tau, i_c] = ir_power
            counter +=1
#             plt.figure()
#             plt.title(r"$\kappa{}, \tau={}, c={}".format(kappa, tau, c))
#             plt.plot(frequencies, emission_g)
#             plt.xlim(0, 2)

```

    780 / 1000 Calculating kappa=1.351344, c=1.000000, tau=0.016212 at n_g = 17 and n_e=14...

```python

```

```python
# decreased dt, does that keep these parameters from failing? last one is 600 / 800 Calculating kappa=8.858668, c=0.750000, tau=0.850000
for i_kappa, kappa in enumerate(kappa_values):
    for i_tau, tau in enumerate(tau_values):
        ir_power = ir_amplitudes[i_kappa, i_tau, :]
        plt.loglog(c_values, ir_power, "*-")
plt.xlabel(r"$c$")
plt.figure()
for i_kappa, kappa in enumerate(kappa_values):
    for i_c, c in enumerate(c_values):
        ir_power = ir_amplitudes[i_kappa, :, i_c]
        plt.loglog(tau_values, ir_power, "*-")
plt.xlabel(r"$\tau$")
plt.figure()
for i_tau, tau in enumerate(tau_values):
    for i_c, c in enumerate(c_values):
#     for i_c in [0,-1]:
        ir_power = ir_amplitudes[:, i_tau, i_c]
#         plt.loglog(kappa_values, ir_power, ["blue", "red"][i_c])
        plt.loglog(kappa_values, ir_power)
plt.xlabel(r"$\kappa$")
# plt.xlim(-.1, 1.1)
# plt.ylim(-10, 500)
```

```python
c_log = np.log10(c_values)
tau_log = np.log10(tau_values)
kappa_log = np.log10(kappa_values)
log_ir_amplitudes = np.log(ir_amplitudes)
num_levels = 100
contours = np.linspace(np.min(log_ir_amplitudes), np.max(log_ir_amplitudes), num_levels)
```

```python

# for i_kappa, kappa in enumerate(kappa_values):
#     ir_power = log_ir_amplitudes[i_kappa, :, :]
#     plt.figure()
#     plt.contourf(c_log, tau_log, ir_power, contours)
#     plt.title(r"$\kappa = {}$".format(kappa))
#     plt.ylabel(r"$c$")
#     plt.xlabel(r"$\tau$")
#     plt.colorbar()

for i_tau, tau in enumerate(tau_values):
    ir_power = log_ir_amplitudes[:, i_tau, :]
    plt.figure()
    plt.contourf(c_log, kappa_log, ir_power, contours)
    plt.title(r"$\tau = {}$".format(tau))
    plt.ylabel(r"$\kappa$")
    plt.xlabel(r"$c$")
    plt.colorbar()

```

```python
for i_c, c in enumerate(c_values):
    ir_power = log_ir_amplitudes[:, :, i_c]
    plt.figure()
    plt.contourf(tau_log, kappa_log, ir_power, contours)
    plt.title(r"$c = {}$".format(c))
    plt.ylabel(r"$\kappa$")
    plt.xlabel(r"$\tau$")
    plt.colorbar()
```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python

```

```python
np.linspace()
```


Calculate Numerical Diffusivity
============================

(For global files, output glued from miltiprocessor runs)

This ipython notebook calculates the numerical diffusivity experienced by a tracer, associated to a specific configuration of the MITgcm. In particular, it was developed to calculate the equivalent diffusivity $\kappa$, defined (here) as $\kappa = \kappa_{pres}+\kappa_{num}$, where $\kappa_{pres}$
is the prescibed or explicit tracer diffusivity one imposes on the model and $k_{num}$ is the additional diffusivity due to numerical truncation errors. Note that there are two $\kappa_{pres}$ and therefore two $\kappa$, one for the horizontal dimensions and one for the vertical one.

These calculations try to reproduce the method used by [1] Abernathy et al. 2010, [2] Hill et al. 2011, and [3] Leibensperger and Plumb, 2013 to determine the numerical diffusivity MITgcm Southern Ocean configurations [1,2] and a baroclinic flow simulation simulation [3].

The method goes as follows:

From the evolution equation for the variance of the tracer concentration in the model output

\begin{equation}
\frac{1}{2}\frac{\partial{\overline{q^{2}}}}{\partial{t}}=-\kappa_{h} \overline{|\nabla_h q|^2}-\kappa_{v} \overline{(\frac{\partial{q}}{\partial {z}})^{2}}
\end{equation}

one can fit by a least squares regression, suitable values of $\kappa_h$ and $\kappa_v$ that satisfy the equation.

```python
#KRM
import os

import numpy as np

import matplotlib.pyplot as plt

from scipy.stats import nanmean

from math import *

%matplotlib inline

import numpy as np

import matplotlib.colors as mcolors

from netCDF4 import Dataset

import pylab as pl

from decimal import getcontext, Decimal, Inexact
```

```python
#'''
#NAME
#    Custom Colormaps for Matplotlib
#PURPOSE
#    This program shows how to implement make_cmap which is a function that
#    generates a colorbar.  If you want to look at different color schemes,
#    check out https://kuler.adobe.com/create.
#PROGRAMMER(S)
#    Chris Slocum
#REVISION HISTORY
#    20130411 -- Initial version created
#    20140313 -- Small changes made and code posted online
#    20140320 -- Added the ability to set the position of each color
#'''

def make_cmap(colors, position=None, bit=False):
    #'''
    #make_cmap takes a list of tuples which contain RGB values. The RGB
    #values may either be in 8-bit [0 to 255] (in which bit must be set to
    #rue when called) or arithmetic [0 to 1] (default). make_cmap returns
    #a cmap with equally spaced colors.
    #Arrange your tuples so that the first color is the lowest value for the
    #colorbar and the last is the highest.
    #position contains values from 0 to 1 to dictate the location of each color.
    #'''
    import matplotlib as mpl
    import numpy as np
    bit_rgb = np.linspace(0,1,256)
    if position == None:
        position = np.linspace(0,1,len(colors))
    else:
        if len(position) != len(colors):
            sys.exit("position length must be the same as colors")
        elif position[0] != 0 or position[-1] != 1:
            sys.exit("position must start with 0 and end with 1")
    if bit:
        for i in range(len(colors)):
            colors[i] = (bit_rgb[colors[i][0]],
                         bit_rgb[colors[i][1]],
                         bit_rgb[colors[i][2]])
    cdict = {'red':[], 'green':[], 'blue':[]}
    for pos, color in zip(position, colors):
        cdict['red'].append((pos, color[0], color[0]))
        cdict['green'].append((pos, color[1], color[1]))
        cdict['blue'].append((pos, color[2], color[2]))

    cmap = mpl.colors.LinearSegmentedColormap('my_colormap',cdict,256)
    return cmap
```

```python
# Get field from MITgcm netCDF output
#
# statefile : string with /path/to/state.0000000000.t001.nc
# fieldname : string with the variable name as written on the netCDF file ('Temp', 'S','Eta', etc.)
#
def getField(statefile, fieldname):

    StateOut = Dataset(statefile)

    Fld = StateOut.variables[fieldname][:]

    shFld = np.shape(Fld)

    if len(shFld) == 1:

        return Fld

    if len(shFld) == 2:

        Fld2 = np.reshape(Fld,(shFld[0],shFld[1])) # reshape to pcolor order
        return Fld2

    elif len(shFld) == 3:

        Fld2 = np.zeros((shFld[0],shFld[1],shFld[2]))
        Fld2 = np.reshape(Fld,(shFld[0],shFld[1],shFld[2])) # reshape to pcolor order
        return Fld2

    elif len(shFld) == 4:

        Fld2 = np.zeros((shFld[0],shFld[1],shFld[2],shFld[3]))
        Fld2 = np.reshape(Fld,(shFld[0],shFld[1],shFld[2],shFld[3])) # reshape to pcolor order
        return Fld2

    else:

        print (' Check size of field ')

```

```python

def dsum(iterable):

    "Full precision summation using Decimal objects for intermediate values"
    # Transform (exactly) a float to m * 2 ** e where m and e are integers.
    # Convert (mant, exp) to a Decimal and add to the cumulative sum.
    # If the precision is too small for exact conversion and addition,
    # then retry with a larger precision.
    from decimal import getcontext, Decimal, Inexact
    getcontext().traps[Inexact] = True

    total = Decimal(0)
    for x in iterable:
        mant, exp = frexp(x)
        mant, exp = int(mant * 2.0 ** 53), exp-53
        while True:
            try:
                total += mant * Decimal(2) ** exp
                break
            except Inexact:
                getcontext().prec += 1
    return float(total)

```

Important input parameters:
------------------------

```python
#nx = 120  # number of x points in domain
#ny = 90  # number of y points in domain
#nz = 35   # number of z levels in domain
#nt = 10   # number of output files (n time)

nx = 360  # number of x points in domain
ny = 360  # number of y points in domain
nz = 90   # number of z levels in domain
nt = 10  # number of output files (n time)
```

Load output filenames:

```python
### Load output filenames ###

filename = '/ocean/kramosmu/MITgcm/CanyonUpwelling/360x360x90_BodyForcing_6Tr_LinProfiles/run02/tave.0000000000.glob.nc'#'/ocean/kramosmu/MITgcm/CanyonUpwelling/120x90x35_UpwLinRampWind_1Tr_LinProfiles/run30/output_0001/state.0000000000.t001.nc'
StateOut = Dataset(filename)
#for dimobj in StateOut.variables.values():
#    print dimobj

filename2 = '/ocean/kramosmu/MITgcm/CanyonUpwelling/360x360x90_BodyForcing_1Tr/run01/mnc_0001/grid.t001.nc'#'/ocean/kramosmu/MITgcm/CanyonUpwelling/120x90x35_UpwLinRampWind_1Tr_LinProfiles/run30/output_0001/grid.t001.nc'
GridOut = Dataset(filename2)
#for dimobj in StateOut.variables.values():
#    print dimobj
filename3 = '/ocean/kramosmu/MITgcm/CanyonUpwelling/360x360x90_BodyForcing_6Tr_LinProfiles/run02/ptr_tave.0000000000.glob.nc'#'/ocean/kramosmu/MITgcm/CanyonUpwelling/120x90x35_UpwLinRampWind_1Tr_LinProfiles/run30/output_0001/ptracers.0000000000.t001.nc'
PtracersOut = Dataset(filename3)

```

Make some plots

```python
tracer1 = getField(filename3,'Tr5')

tracer1msk = np.ma.masked_values(tracer1, 0)

MASK = np.ma.getmask(tracer1msk)

### PLOT TRACER 1 at zlev and time###

zlev = 20
time = 0

plt.rcParams.update({'font.size':13})
plt.figure(figsize=(10,7))

ax=plt.gca()
ax.set_axis_bgcolor((205/255.0, 201/255.0, 201/255.0))

plt.contourf(tracer1msk[time,zlev,:,:],30, cmap='rainbow')
cb = plt.colorbar()
cb.set_label('C', position=(0, 1),rotation=0)

z=StateOut.variables['Z']

plt.title("Tracer at z=%d m" % -z[zlev])

plt.figure(figsize=(10,7))

ax=plt.gca()
ax.set_axis_bgcolor((205/255.0, 201/255.0, 201/255.0))

plt.contourf(tracer1msk[time,:,:,60],30, cmap='rainbow')
plt.gca().invert_yaxis()
cb = plt.colorbar()
cb.set_label('C', position=(0, 1),rotation=0)

z=StateOut.variables['Z']

plt.title("Tracer at x=60 nx")

```

Calculate the volume of the domain
-------------------------------
The volume of a tracer cell (remember we have an Arakawa C grid, so this changes depending on which kind of cell we are thinking about) is given by

$V(i,j,k)=depth \times area = (hfacC(i,j,k)\times dRf(k)) \times rA(i,j) = (hfacC(i,j,k)\times dRf(k)) \times dXg(i,j) \times dYg(i,j)$,

where hfacC is the fraction of the cell that is open (not occupied with land). So, the total volume of the domain is

$\sum\limits_{i=1}^{nx}{\sum\limits_{j=1}^{ny}{\sum\limits_{k=1}^{nz}{(hfacC(i,j,k)\times dRf(k)) \times rA(i,j)}}}$

```python
Dom=0.0
DomVolume=0.0

hFacC = getField(filename2,'HFacC')
drF   = getField(filename2,'drF')
rAC   = getField(filename2,'rA')

for ii in np.arange(nx):

    for jj in np.arange(ny):

        Dom = Dom + hFacC[:,jj,ii]*drF[:]*rAC[jj,ii]

DomVolume=sum(Dom)
print('\n The total volume is %e m^3 \n' %DomVolume)

```

     The total volume is 6.102853e+12 m^3

1st Term: The volume-weighted average of the squared concentration
-------------------------------------------------------------
The first term in the variance evolution equation is $\frac{1}{2}\frac{\partial{\overline{q^{2}}}}{\partial{t}}$. Note that we care about the time derivative of the variance, so that the mean concentration that usually appears in the definition of variance will not play a role here, since it is constant in time (we are not putting in or letting out any tracer).

We are going to calculate $\overline{q^2}$, the volume-weighted average of the squared concentration, and then the time derivative of that using a centered difference scheme.

```python
### Calculate the volume-weighted mean of q^2

q2mean = np.zeros((nt,1))

for tt in np.arange(nt):
    q2sum = 0.0
    for kk in np.arange(nz):

        for jj in np.arange(ny):

            q2sum = q2sum + np.sum(((hFacC[kk,jj,:]*drF[kk]*rAC[jj,:])/DomVolume)*(tracer1[tt,kk,jj,:])**2)

    q2mean[tt] = q2sum

    print('volume-weighted mean of q^2 nt=%d is %e  \n' %(tt, q2mean[tt]))

```

    volume-weighted mean of q^2 nt=0 is 5.102666e+02

    volume-weighted mean of q^2 nt=1 is 4.641340e+02

    volume-weighted mean of q^2 nt=2 is 4.501812e+02

    volume-weighted mean of q^2 nt=3 is 4.452379e+02

    volume-weighted mean of q^2 nt=4 is 4.432554e+02

    volume-weighted mean of q^2 nt=5 is 4.422911e+02

    volume-weighted mean of q^2 nt=6 is 4.417853e+02

    volume-weighted mean of q^2 nt=7 is 4.415508e+02

    volume-weighted mean of q^2 nt=8 is 4.413793e+02

    volume-weighted mean of q^2 nt=9 is 4.412380e+02

```python
### Calculate time derivative
dt = (3600.*12.) # half day

dqdt = np.zeros(nt-2)

for tt in np.arange(nt-2):

    dqdt[tt] = (q2mean[tt+2]-q2mean[tt])/(2.*dt) # approximation of time derivative of v-w mean q^2

```

2nd Term: The volume-weighted average of the squared horizontal gradient
------------------------------------------------------------------
The second term in the variance evolution equation is $-\kappa_{h} \overline{|\nabla_h q|^2}$. Next, we calculate the square of the horizontal gradient $|\nabla_h q|^2=(\frac{\partial{q}}{\partial{x}})^2+(\frac{\partial{q}}{\partial{y}})^2$.

Spatial derivatives are approximated using a centered-difference scheme.

```python

### Calculate mean of (dq/dx)^2 + (dq/dy)^2)

dxG = getField(filename2,'dxG')
dyG = getField(filename2,'dyG')

qmeanDh = np.zeros((nt,1))
qmeanDx = np.zeros((nt,1))
qmeanDy = np.zeros((nt,1))

for tt in np.arange(nt):

    sumDx = 0.0

    tracer2dx = tracer1[tt,:,:,2:nx]
    tracer0dx = tracer1[tt,:,:,0:nx-2]
    tracerDerX = ((tracer2dx-tracer0dx)/(2.0*dxG[:-1,1:nx-1]))**2

    for kk in np.arange(nz):
        for jj in np.arange(ny):
            sumDx = sumDx + np.sum((hFacC[kk,jj,1:nx-1]*drF[kk]*rAC[jj,1:nx-1])*(tracerDerX[kk,jj,:]))

        #print(sumDx)

    qmeanDx[tt] = sumDx/DomVolume

    print('The v-w mean of (dq/dx)^2 at nt=%d is %e \n' % (tt,qmeanDx[tt]))

for tt in np.arange(nt):

    sumDy = 0.0

    tracer2dy = tracer1[tt,:,2:ny,:]
    tracer0dy = tracer1[tt,:,0:ny-2,:]
    tracerDerY = ((tracer2dy-tracer0dy)/(2.0*dyG[1:ny-1,:-1]))**2

    for kk in np.arange(nz):
        for ii in np.arange(nx):
            sumDy = sumDy + np.sum((hFacC[kk,1:ny-1,ii]*drF[kk]*rAC[1:ny-1,ii])*(tracerDerY[kk,:,ii]))

    qmeanDy[tt] = sumDy/DomVolume

    print('The v-w mean of (dq/dy)^2 at nt=%d is %e \n' % (tt,qmeanDy[tt]))

qmeanDh = qmeanDx + qmeanDy

```

    The v-w mean of (dq/dx)^2 at nt=0 is 2.023410e-06

    The v-w mean of (dq/dx)^2 at nt=1 is 1.413322e-06

    The v-w mean of (dq/dx)^2 at nt=2 is 1.127260e-06

    The v-w mean of (dq/dx)^2 at nt=3 is 8.201572e-07

    The v-w mean of (dq/dx)^2 at nt=4 is 6.599211e-07

    The v-w mean of (dq/dx)^2 at nt=5 is 5.649565e-07

    The v-w mean of (dq/dx)^2 at nt=6 is 5.330503e-07

    The v-w mean of (dq/dx)^2 at nt=7 is 5.513010e-07

    The v-w mean of (dq/dx)^2 at nt=8 is 5.330182e-07

    The v-w mean of (dq/dx)^2 at nt=9 is 5.016216e-07

    The v-w mean of (dq/dy)^2 at nt=0 is 6.427223e-06

    The v-w mean of (dq/dy)^2 at nt=1 is 4.561694e-06

    The v-w mean of (dq/dy)^2 at nt=2 is 3.857274e-06

    The v-w mean of (dq/dy)^2 at nt=3 is 3.483634e-06

    The v-w mean of (dq/dy)^2 at nt=4 is 3.283760e-06

    The v-w mean of (dq/dy)^2 at nt=5 is 3.199735e-06

    The v-w mean of (dq/dy)^2 at nt=6 is 3.126418e-06

    The v-w mean of (dq/dy)^2 at nt=7 is 3.089124e-06

    The v-w mean of (dq/dy)^2 at nt=8 is 3.087190e-06

    The v-w mean of (dq/dy)^2 at nt=9 is 3.081658e-06

```python

```

3rd Term: The volume-weighted average of the squared vertical derivative
------------------------------------------------------------------
The third term in the variance evolution equation is $-\kappa_{v} \overline{(\frac{\partial{q}}{\partial{z}})^2}$. Next, we calculate the square of the vertical gradient $(\frac{\partial{q}}{\partial{z}})^2$.

The vertical derivative is approximated using a centered-difference scheme.

```python
### Calculate mean of (dq/dz)^2

qmeanDz = np.zeros((nt,1))

for tt in np.arange(nt):
    qsumDz = 0.0
    for ii in np.arange(nx):

        for jj in np.arange(ny):

            for kk in np.arange(nz-2):

                if (hFacC[kk,jj,ii]*drF[kk])+(hFacC[kk+1,jj,ii]*drF[kk+1]) == 0:

                    continue

                else :

                   qsumDz =qsumDz +(hFacC[kk,jj,ii]*drF[kk]*rAC[jj,ii])*(((tracer1[tt,kk+2,jj,ii]-tracer1[tt,kk,jj,ii])/((hFacC[kk+1,jj,ii]*drF[kk+1])+(hFacC[kk,jj,ii]*drF[kk])))**2)

    qmeanDz[tt] = qsumDz/DomVolume

    print('Mean (dq/dz)^2 at nt=%d is %e \n' %(tt, qmeanDz[tt]))

```

    Mean (dq/dz)^2 at nt=0 is 3.810315e-02

    Mean (dq/dz)^2 at nt=1 is 2.703776e-02

    Mean (dq/dz)^2 at nt=2 is 2.742958e-02

    Mean (dq/dz)^2 at nt=3 is 2.840859e-02

    Mean (dq/dz)^2 at nt=4 is 2.950646e-02

    Mean (dq/dz)^2 at nt=5 is 3.048540e-02

    Mean (dq/dz)^2 at nt=6 is 3.028547e-02

    Mean (dq/dz)^2 at nt=7 is 3.007359e-02

    Mean (dq/dz)^2 at nt=8 is 3.032591e-02

    Mean (dq/dz)^2 at nt=9 is 3.049262e-02

Tracer variance plots:
--------------------

```python

fig1=plt.figure(figsize=(21,3))

plt.subplot(1,3,1)
ax=plt.gca()
ax.set_axis_bgcolor((205/255.0, 201/255.0, 201/255.0))

plt.plot(qmeanDh[1:len(qmeanDh)-1],dqdt)

plt.xlabel('$<|\partial_{h}q|^2>$')

plt.ylabel('$\partial{<q^2>}/\partial{t}$')

plt.subplot(1,3,2)
ax=plt.gca()
ax.set_axis_bgcolor((205/255.0, 201/255.0, 201/255.0))

plt.plot(qmeanDz[1:len(qmeanDz)-1],dqdt)

plt.xlabel('$<\partial{q^2}/\partial{z}>$')

plt.ylabel('$\partial{<q^2>}/\partial{t}$')

plt.subplot(1,3,3)
ax=plt.gca()
ax.set_axis_bgcolor((205/255.0, 201/255.0, 201/255.0))

plt.plot(np.arange(0,nt*2,2.0),q2mean)

plt.xlabel('hours')

plt.ylabel('$<q^2>$')

plt.tight_layout()
```

Linear regression for $\kappa_h$ and $\kappa_v$
------------------------------------------

In the regression results box, x1 corresponds to $\kappa_h$ and x2 to $\kappa_v$.

```python
import pandas as pd
import statsmodels.api as sm

dt = 1*3600.0*24.0 # 2hrs

x = np.column_stack((qmeanDh[4:len(qmeanDh)-1], qmeanDz[4:len(qmeanDz)-1])) #variance of space derivatives

y = 0.5*dqdt[3:len(dqdt)]

## fit a OLS model with intercept
#x = sm.add_constant(x)
#est = sm.OLS(y, x).fit()
est = sm.GLS(y,x).fit()
est.summary()

```

GLS Regression Results

  Dep. Variable:	            y	          R-squared:         	    0.984

  Model:	                   GLS	         Adj. R-squared:    	    0.973

  Method:	             Least Squares	    F-statistic:       	    91.26

  Date:	             Fri, 05 Dec 2014	   Prob (F-statistic):	  0.00206

  Time:	                 14:08:46	       Log-Likelihood:    	   61.400

  No. Observations:	           5	        AIC:               	   -118.8

  Df Residuals:	               3	        BIC:               	   -119.6

  Df Model:	                   2

   	     coef	     std err	      t	      P>|t|	 [95.0% Conf. Int.]

  x1	   -36.4510	     4.235	    -8.608	  0.003	   -49.928   -22.974

  x2	     0.0043	     0.001	     8.173	  0.004	     0.003     0.006

  Omnibus:	          nan	   Durbin-Watson:     	    1.979

  Prob(Omnibus):	    nan	   Jarque-Bera (JB):  	    0.202

  Skew:	          -0.223	   Prob(JB):          	    0.904

  Kurtosis:	       2.122	   Cond. No.          	 1.97e+05

```python

dt = 1*3600.0*24.0 # 2hrs

#x = np.column_stack((qmeanDh[1:len(qmeanDh)-1], qmeanDz[1:len(qmeanDz)-1])) #variance of space derivatives

y = 0.5*dqdt[0:len(dqdt)]

x1=qmeanDh[1:len(qmeanDh)-1]
x2=qmeanDz[1:len(qmeanDz)-1]
print(np.shape(dqdt))
print(np.shape(qmeanDh))

## fit a OLS model with intercept
#x = sm.add_constant(x)
#est = sm.OLS(y, x).fit()
est = sm.GLS(y,x1).fit()
est.summary()

```

    (8,)
    (10, 1)

GLS Regression Results

  Dep. Variable:	            y	          R-squared:         	    0.427

  Model:	                   GLS	         Adj. R-squared:    	    0.345

  Method:	             Least Squares	    F-statistic:       	    5.220

  Date:	             Fri, 05 Dec 2014	   Prob (F-statistic):	  0.0563

  Time:	                 14:08:51	       Log-Likelihood:    	   62.471

  No. Observations:	           8	        AIC:               	   -122.9

  Df Residuals:	               7	        BIC:               	   -122.9

  Df Model:	                   1

   	     coef	     std err	      t	      P>|t|	 [95.0% Conf. Int.]

  x1	   -19.6927	     8.620	    -2.285	  0.056	   -40.075     0.689

  Omnibus:	       18.036	   Durbin-Watson:     	    0.664

  Prob(Omnibus):	  0.000	   Jarque-Bera (JB):  	    7.159

  Skew:	          -1.995	   Prob(JB):          	   0.0279

  Kurtosis:	       5.356	   Cond. No.          	     1.00

```python
dt = 1*3600.0*24.0 # 2hrs

x = np.column_stack((qmeanDh[1:len(qmeanDh)-1], qmeanDz[1:len(qmeanDz)-1])) #variance of space derivatives

y = 0.5*dqdt[0:len(dqdt)]

x1=qmeanDh[1:len(qmeanDh)-1]
x2=qmeanDz[1:len(qmeanDz)-1]
## fit a OLS model with intercept
#x = sm.add_constant(x)
#est = sm.OLS(y, x).fit()
est = sm.GLS(y,x2).fit()
est.summary()

```

GLS Regression Results

  Dep. Variable:	            y	          R-squared:         	    0.230

  Model:	                   GLS	         Adj. R-squared:    	    0.120

  Method:	             Least Squares	    F-statistic:       	    2.091

  Date:	             Fri, 05 Dec 2014	   Prob (F-statistic):	   0.191

  Time:	                 14:08:54	       Log-Likelihood:    	   61.288

  No. Observations:	           8	        AIC:               	   -120.6

  Df Residuals:	               7	        BIC:               	   -120.5

  Df Model:	                   1

   	     coef	     std err	      t	      P>|t|	 [95.0% Conf. Int.]

  x1	    -0.0021	     0.001	    -1.446	  0.191	    -0.006     0.001

  Omnibus:	       16.062	   Durbin-Watson:     	    0.607

  Prob(Omnibus):	  0.000	   Jarque-Bera (JB):  	    6.114

  Skew:	          -1.885	   Prob(JB):          	   0.0470

  Kurtosis:	       5.032	   Cond. No.          	     1.00

```python

```

```python

```


# Nearest neighbors
[Index](0-index.ipynb)

Here we compute matrices the matrix of distances between communities, $[d_{ab}]$.

## Imports and global variables

```python
%load_ext autoreload
%autoreload 2
%matplotlib inline
```

```python
from pathlib import Path
import sys
import numpy as np
import pandas as pd
import geopandas as gpd
import datetime

import matplotlib.pyplot as plt
import matplotlib.colors as mco
import matplotlib.gridspec as mgs
import matplotlib.ticker as ticker
import matplotlib.cm as cm
from matplotlib import animation
plt.rcParams['svg.fonttype'] = 'none'

from IPython.display import HTML
from IPython.display import Image
```

```python
sys.path.append('../code')
from functions import show_image, geo_dist
```

```python
resdir = Path('../results/')
if not resdir.is_dir():
    raise ValueError('No results directory!')
```

```python
resfile = resdir / 'safegraph_analysis.hdf5'
complevel=7
complib='zlib'
with pd.HDFStore(resfile, complevel=complevel, complib=complib) as store:
    print(f"File {resfile.stem} has {len(store.keys())} entries.")
```

    File safegraph_analysis has 831 entries.

```python
tfmt = '%Y-%m-%d'
```

```python
figdir = Path('../figures') / '4-distances'
if not figdir.is_dir():
    figdir.mkdir(parents=True, exist_ok=True)
exts = ['.png', '.svg']
```

## Functions
We use the [formula](https://en.wikipedia.org/wiki/Geographical_distance) of the Spherical Earth projected to a plane. for the 'Spherical Earth projected to a plane':
\begin{equation}
d(M_1,M_2) = R \sqrt{\Delta \phi^2 + \cos{\bar{\phi}} \Delta \lambda^2},
\end{equation}
where $\phi$ is the latitude, $\lambda$ is the longitude, and $\bar{\phi} = (\phi_1 + \phi_2)/2$. $R = 6, 371.009 \, \mathrm{km}$ is the radius of the Earth.

Here `X` is the longitude and `Y` is the latitude (see this [notebook](06-cross_compare_cssegi.ipynb)).

```python
R_earth = 6.371009e6
```

## Clusters

### Load clusters

```python
key = "/clustering/clusters"
with pd.HDFStore(resfile, complevel=complevel, complib=complib) as store:
    clusters = store[key]
# clusters = pd.read_hdf(resfile, key)
N = len(clusters)
print(f"N = {N}")
clusters
```

    N = 1024

      leaves
      X
      Y
      cbg_clusters
      population
      mobile_count

      index

      0
      746
      -65.794825
      18.277415
      [720531504003, 720531501044, 720531505001, 720...
      390728
      8705.198547

      1
      458
      -66.083320
      18.131576
      [720572702011, 720572703004, 720572705001, 720...
      488287
      10942.351090

      2
      873
      -66.079340
      18.396328
      [720210307001, 720210310113, 720210316223, 720...
      1011530
      28389.915254

      3
      8
      -66.344566
      18.363775
      [720399559001, 720399556001, 720399557001, 720...
      425126
      9988.154964

      4
      641
      -66.582628
      18.044742
      [721130724005, 721130727011, 721130727012, 721...
      354598
      12057.116223

      ...
      ...
      ...
      ...
      ...
      ...
      ...

      1019
      911
      -158.026269
      21.423245
      [150030078092, 150030078102, 150030099044, 150...
      459209
      22780.496368

      1020
      7
      -157.813184
      21.329035
      [150090318011, 150099900000, 150030047001, 150...
      529047
      21868.060533

      1021
      392
      -156.517054
      20.867763
      [150090311032, 150090304042, 150090309034, 150...
      161290
      7878.004843

      1022
      808
      -155.865732
      19.692965
      [150010216012, 150010214023, 150010216041, 150...
      88670
      3758.949153

      1023
      134
      -155.089968
      19.680449
      [150010210053, 150010210052, 150019904000, 150...
      105010
      4298.900726

1024 rows × 6 columns

## Compute distance matrix using Spherical Earth projected to a plane formula

```python
XY = clusters.loc[:, ['X','Y']].astype('float64').to_numpy()

Long, Lat = XY.T * np.pi/180.

LLat1, LLat2 = np.meshgrid(Lat, Lat, indexing='ij')
dLLat = LLat1 - LLat2
LLat_bar = 0.5*(LLat1 + LLat2)

LLong1, LLong2 = np.meshgrid(Long, Long, indexing='ij')
dLLong = LLong1 - LLong2

D = R_earth * np.sqrt(dLLat**2 + np.cos(LLat_bar)*dLLong**2)
```

```python
df_dist = pd.DataFrame(data=D, index=clusters.index, columns=clusters.index)

key = "/clustering/distances"
with pd.HDFStore(resfile, complevel=complevel, complib=complib) as store:
    store[key] = df_dist
```

Compare with [equirectangular projection](https://en.wikipedia.org/wiki/Equirectangular_projection) distances

```python
P = clusters.loc[:, ['Y', 'X']].to_numpy()
PP = np.einsum('ik,jk->ij',P,P)
U = np.ones((N,1))
Delta = PP.diagonal().reshape(-1,1)
Dequi = (U*Delta.T) + (Delta*U.T) - 2*PP
Dequi = R_earth * np.pi/180.* np.sqrt(Dequi)
```

```python
err = np.linalg.norm(D-Dequi)/np.linalg.norm(D)
print("Equirectangular projection error = {:.1f} %".format(err*100))
```

    Equirectangular projection error = 15.1 %

## Some statistics on the actual separation

```python
key = "/clustering/distances"
with pd.HDFStore(resfile, complevel=complevel, complib=complib) as store:
    df_dist = store[key]

df_dist
```

      index
      0
      1
      2
      3
      4
      5
      6
      7
      8
      9
      ...
      1014
      1015
      1016
      1017
      1018
      1019
      1020
      1021
      1022
      1023

      index

      0
      0.000000e+00
      3.522132e+04
      3.353937e+04
      6.032823e+04
      8.922338e+04
      1.020716e+05
      1.423116e+05
      1.367488e+05
      3.191928e+06
      3.076757e+06
      ...
      1.056508e+07
      1.084609e+07
      1.163205e+07
      1.178416e+07
      1.010751e+07
      9.952485e+06
      9.930626e+06
      9.796315e+06
      9.740499e+06
      9.656786e+06

      1
      3.522132e+04
      0.000000e+00
      2.944231e+04
      3.831536e+04
      5.498573e+04
      7.612572e+04
      1.141486e+05
      1.041719e+05
      3.206031e+06
      3.090694e+06
      ...
      1.054908e+07
      1.082948e+07
      1.161449e+07
      1.175815e+07
      1.007947e+07
      9.924250e+06
      9.902359e+06
      9.767885e+06
      9.711705e+06
      9.627974e+06

      2
      3.353937e+04
      2.944231e+04
      0.000000e+00
      2.895685e+04
      6.710617e+04
      7.027471e+04
      1.109507e+05
      1.094806e+05
      3.176673e+06
      3.061348e+06
      ...
      1.053037e+07
      1.081124e+07
      1.159690e+07
      1.175017e+07
      1.007443e+07
      9.919505e+06
      9.897663e+06
      9.763436e+06
      9.707797e+06
      9.624100e+06

      3
      6.032823e+04
      3.831536e+04
      2.895685e+04
      0.000000e+00
      4.386493e+04
      4.177766e+04
      8.226182e+04
      8.112963e+04
      3.178601e+06
      3.063137e+06
      ...
      1.050835e+07
      1.078883e+07
      1.157388e+07
      1.172304e+07
      1.004655e+07
      9.891557e+06
      9.869703e+06
      9.735421e+06
      9.679656e+06
      9.595954e+06

      4
      8.922338e+04
      5.498573e+04
      6.710617e+04
      4.386493e+04
      0.000000e+00
      4.288828e+04
      6.809197e+04
      5.042405e+04
      3.212721e+06
      3.097135e+06
      ...
      1.050948e+07
      1.078911e+07
      1.157288e+07
      1.170787e+07
      1.002752e+07
      9.872132e+06
      9.850216e+06
      9.715614e+06
      9.659143e+06
      9.575399e+06

      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...
      ...

      1019
      9.952485e+06
      9.924250e+06
      9.919505e+06
      9.891557e+06
      9.872132e+06
      9.849459e+06
      9.809312e+06
      9.821268e+06
      9.521824e+06
      9.507770e+06
      ...
      3.765796e+06
      3.699263e+06
      3.810609e+06
      2.018228e+06
      1.732468e+05
      0.000000e+00
      2.514994e+04
      1.734402e+05
      3.017573e+05
      3.706394e+05

      1020
      9.930626e+06
      9.902359e+06
      9.897663e+06
      9.869703e+06
      9.850216e+06
      9.827605e+06
      9.787448e+06
      9.799353e+06
      9.506791e+06
      9.492439e+06
      ...
      3.779319e+06
      3.714520e+06
      3.829800e+06
      2.043401e+06
      1.983721e+05
      2.514994e+04
      0.000000e+00
      1.483572e+05
      2.775173e+05
      3.456700e+05

      1021
      9.796315e+06
      9.767885e+06
      9.763436e+06
      9.735421e+06
      9.715614e+06
      9.693324e+06
      9.653116e+06
      9.664762e+06
      9.408967e+06
      9.392986e+06
      ...
      3.851454e+06
      3.797212e+06
      3.936568e+06
      2.191844e+06
      3.466731e+05
      1.734402e+05
      1.483572e+05
      0.000000e+00
      1.482724e+05
      2.026106e+05

      1022
      9.740499e+06
      9.711705e+06
      9.707797e+06
      9.679656e+06
      9.659143e+06
      9.637554e+06
      9.597227e+06
      9.608305e+06
      9.416180e+06
      9.397826e+06
      ...
      3.992413e+06
      3.941662e+06
      4.085833e+06
      2.309319e+06
      4.698623e+05
      3.017573e+05
      2.775173e+05
      1.482724e+05
      0.000000e+00
      8.371375e+04

      1023
      9.656786e+06
      9.627974e+06
      9.624100e+06
      9.595954e+06
      9.575399e+06
      9.553856e+06
      9.513527e+06
      9.524567e+06
      9.342283e+06
      9.323377e+06
      ...
      4.009333e+06
      3.965071e+06
      4.124629e+06
      2.385999e+06
      5.423150e+05
      3.706394e+05
      3.456700e+05
      2.026106e+05
      8.371375e+04
      0.000000e+00

1024 rows × 1024 columns

```python
Ns = df_dist.index.to_numpy()
dists = df_dist.to_numpy()
dists_sorted = np.sort(dists, axis=-1)

# median distance of 1st, 2nd, 3rd...etc neighbor
dist_med = np.array([np.median(dists_sorted[:,n]) for n in Ns[1:]])
```

```python
# parameters
figsize = (6,4.5)
dpi = 300
ms=2
lw=1

fig = plt.figure(facecolor='w', figsize=figsize)
ax = fig.gca()

ax.plot(Ns[1:], dist_med/1.0e3, '-', ms=ms, color='darkblue')

# ax.set_xscale('log')
# ax.set_yscale('log')
ax.set_xlim(Ns[1],None)
ax.set_xlabel("kth nearest neighbor", fontsize="medium")
ax.set_ylabel("median distance (km)", fontsize="medium")
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.tick_params(left=True, labelleft=True, bottom=True, labelbottom=True)
ax.tick_params(axis='both', length=4)
fig.tight_layout()

fname = 'median_distance_nearest_neighbors'
for ext in exts:
    filepath = figdir / (fname + ext)
    fig.savefig(filepath, bbox_inches='tight', pad_inches=0, dpi=dpi)
    print("Written file: {:s}".format(str(filepath)))
fig.clf()
plt.close('all')
```

    Written file: ../figures/4-distances/median_distance_nearest_neighbors.png
    Written file: ../figures/4-distances/median_distance_nearest_neighbors.svg

```python
filepath = figdir / (fname + '.png')
Image(filename=filepath, width=4./3*360)
```


# Integrating the moment equations - Protein

(c) 2017 Manuel Razo. This work is licensed under a [Creative Commons Attribution License CC-BY 4.0](https://creativecommons.org/licenses/by/4.0/). All code contained herein is licensed under an [MIT license](https://opensource.org/licenses/MIT)

---

```python
import cloudpickle

# Our numerical workhorses
import numpy as np
import scipy as sp
import pandas as pd

# Import matplotlib stuff for plotting
import matplotlib.pyplot as plt
import matplotlib.cm as cm
import matplotlib as mpl

# Seaborn, useful for graphics
import seaborn as sns

# Import the utils for this project
import chann_cap_utils as chann_cap

# Set PBoC plotting style
chann_cap.set_plotting_style()

# Magic function to make matplotlib inline; other style specs must come AFTER
%matplotlib inline

# This enables SVG graphics inline (only use with static plots (non-Bokeh))
%config InlineBackend.figure_format = 'svg'

tmpdir = '../../tmp/'
figdir = '../../fig/moment_dynamics_numeric/'
```

### $\LaTeX$ macros
$\newcommand{kpon}{k^{(p)}_{\text{on}}}$
$\newcommand{kpoff}{k^{(p)}_{\text{off}}}$
$\newcommand{kron}{k^{(r)}_{\text{on}}}$
$\newcommand{kroff}{k^{(r)}_{\text{off}}}$
$\newcommand{rm}{r _m}$
$\newcommand{gm}{\gamma _m}$
$\newcommand{mm}{\left\langle m \right\rangle}$
$\newcommand{ee}[1]{\left\langle #1 \right\rangle}$
$\newcommand{bb}[1]{\mathbf{#1}}$

## Numerically integrating the moment equations.

In this notebook we will numerically integrate the differential equations for both, the mRNA and protein distribution moments for the two- and three-state promoter.

Up to this stage we wrote the chemical master equation in matrix notation and did some tricks when it came to compute the $n^{\text{th}}$ moment. Now we will use the equations we obtained for the time derivative of the mRNA distribution moments and try to solve them.

### Defining the moment equations.

We will define a single function that caputes all moments that lead to computing the protein third moment $\ee{\bb{p}^3}$. That is the first three moments of the mRNA distribution, the cross correlation terms between protein and mRNA and the three moments of the protein distribution. To have it as a reference let's write down all equations.

### mRNA moments
\begin{align}
% zeroth moment
{d \ee{\bb{m}^0} \over dt}  &= \mathbf{K}_m \left\langle \mathbf{m}^0 \right\rangle.\\
% first moment
{d \ee{\bb{m}^1}\over dt} &=
\left( \mathbf{K}_m - \mathbf{\Gamma}_m  \right)\ee{\bb{m}^1}
+ \mathbf{R}_m \left\langle \mathbf{m}^0 \right\rangle.\\
% second moment
{d\ee{\bb{m}^2} \over dt} &= \left( \mathbf{K}_m - 2 \mathbf{\Gamma}_m \right)
             \left\langle \mathbf{m}^2 \right\rangle
             + \left( 2 \mathbf{R}_m + \mathbf{\Gamma}_m \right)
             \left\langle \mathbf{m}^1 \right\rangle
             + \mathbf{R}_m
             \left\langle \mathbf{m}^0 \right\rangle\\
% third moment
{d\ee{\bb{m}^3} \over dt} &= \left( \mathbf{K}_m - 3 \mathbf{\Gamma}_m \right)
             \left\langle \mathbf{m}^3 \right\rangle
             + \left( 3 \mathbf{R}_m + 3 \mathbf{\Gamma}_m \right)
             \left\langle \mathbf{m}^2 \right\rangle
             + \left( 3 \mathbf{R}_m - \mathbf{\Gamma}_m \right)
             \left\langle \mathbf{m}^1 \right\rangle
             + \mathbf{R}_m
             \left\langle \mathbf{m}^0 \right\rangle
\end{align}

### Protein moments (and cross correlations)
\begin{align}
% first moment
{d \ee{\bb{p}^1}\over dt} &=
\left( \mathbf{K} - \mathbf{\Gamma}_p  \right) \ee{\bb{p}^1}
+ \mathbf{R}_p \left\langle \mathbf{m} \right\rangle.\\
% <mp>
{d \ee{\bb{mp}}\over dt} &=
\left( \mathbf{K} - \mathbf{\Gamma}_m -  \mathbf{\Gamma}_p  \right)
\left\langle \mathbf{mp} \right\rangle
+ \mathbf{R}_m \left\langle \mathbf{p} \right\rangle
+ \mathbf{R}_p \left\langle \mathbf{m}^2 \right\rangle \\
% second moment
{d \ee{\bb{p}^2}\over dt} &=
\left( \mathbf{K} - 2 \mathbf{\Gamma}_p \right)
\left\langle \mathbf{p}^2 \right\rangle
+ \mathbf{\Gamma}_p \left\langle \mathbf{p} \right\rangle
+ \mathbf{R}_p \mm
+ 2 \mathbf{R}_p \left\langle \mathbf{mp} \right\rangle \\
% <m^2p>
{d \ee{\bb{m}^2\bb{p}}\over dt} &=
\left( \bb{K} - 2 \bb{\Gamma}_m - \bb{\Gamma}_p \right)
\ee{\bb{m}^2\bb{p}} +
\bb{R}_m \ee{\bb{p}} +
\left( \bb{R}_m + \bb{\Gamma}_m \right) \ee{\bb{mp}} +
\bb{R}_p \ee{\bb{m}^3}.\\
% <mp^2>
{d \ee{\bb{m}\bb{p}^2}\over dt} &=
\left( \bb{K} - \bb{\Gamma}_m - 2 \bb{\Gamma}_p \right)
\ee{\bb{mp}^2} +
\bb{R}_m \ee{\bb{p}^2} +
2 \bb{R}_p \ee{\bb{m}^2\bb{p}} +
\bb{R}_p \ee{\bb{m}^2} +
\bb{\Gamma}_p \ee{\bb{mp}}.\\
% third moment
{d \ee{\bb{p}^3}\over dt} &=
\left( \bb{K} - 3 \bb{\Gamma}_p \right) \ee{\bb{p}^3} +
3 \bb{\Gamma}_p \ee{\bb{p}^2} -
\bb{\Gamma}_p \ee{\bb{p}} +
3 \bb{R}_p \ee{\bb{mp}^2} +
3 \bb{R}_p \ee{\bb{mp}} +
\bb{R}_p \ee{\bb{m}}.
\end{align}

Let's now implemenet these equations into the function

```python
def dpdt(mp, t, Kmat, Rm, Gm, Rp, Gp):
    '''
    function to integrate all mRNA and protein moment dynamics
    using scipy.integrate.odeint
    Parameters
    ----------
    m : array-like.
        Array containing all moments (mRNA, protein and cross correlations)
        Unregulated
        mp[0] = m0_P (RNAP bound)
        mp[1] = m0_E (Empty promoter)
        mp[2] = m1_P (RNAP bound)
        mp[3] = m1_P (Empty promoter)
        mp[4] = m2_P (RNAP bound)
        mp[5] = m2_P (Empty promoter)
        mp[6] = m3_P (RNAP bound)
        mp[7] = m3_P (Empty promoter)
        mp[8] = p1_P (RNAP bound)
        mp[9] = p1_P (Empty promoter)
        mp[10] = mp_P (RNAP bound)
        mp[11] = mp_P (Empty promoter)
        mp[12] = p2_P (RNAP bound)
        mp[13] = p2_P (Empty promoter)
        mp[14] = m2p_P (RNAP bound)
        mp[15] = m2p_P (Empty promoter)
        mp[16] = mp2_P (RNAP bound)
        mp[17] = mp2_P (Empty promoter)
        mp[18] = p3_P (RNAP bound)
        mp[19] = p3_P (Empty promoter)
        ---------
        Regulated:
        mp[0] = m0_P (RNAP bound)
        mp[1] = m0_E (Empty promoter)
        mp[2] = m0_R (Repressor bound)
        mp[3] = m1_P (RNAP bound)
        mp[4] = m1_E (Empty promoter)
        mp[5] = m1_R (Repressor bound)
        mp[6] = m2_P (RNAP bound)
        mp[7] = m2_E (Empty promoter)
        mp[8] = m2_R (Repressor bound)
        mp[9] = m3_P (RNAP bound)
        mp[10] = m3_E (Empty promoter)
        mp[11] = m3_R (Repressor bound)
        mp[12] = p1_P (RNAP bound)
        mp[13] = p1_E (Empty promoter)
        mp[14] = p1_R (Repressor bound)
        mp[15] = mp_P (RNAP bound)
        mp[16] = mp_E (Empty promoter)
        mp[17] = mp_R (Repressor bound)
        mp[18] = p2_P (RNAP bound)
        mp[19] = p2_E (Empty promoter)
        mp[20] = p2_R (Repressor bound)
        mp[21] = m2p_P (RNAP bound)
        mp[22] = m2p_E (Empty promoter)
        mp[23] = m2p_R (Repressor bound)
        mp[24] = mp2_P (RNAP bound)
        mp[25] = mp2_E (Empty promoter)
        mp[26] = mp2_R (Repressor bound)
        mp[27] = p3_P (RNAP bound)
        mp[28] = p3_E (Empty promoter)
        mp[29] = p3_R (Repressor bound)
    t : array-like.
        Time array
    Kmat : array-like.
        Matrix containing the transition rates between the promoter states.
    Rm : array-like.
        Matrix containing the mRNA production rate at each of the states.
    Gm : array-like.
        Matrix containing the mRNA degradation rate at each of the states.
    Rp : array-like.
        Matrix containing the protein production rate at each of the states.
    Gp : array-like.
        Matrix containing the protein degradation rate at each of the states.

    Returns
    -------
    dynamics of all mRNA and protein moments
    '''
    # Obtain the zeroth and first moment based on the size
    # of the Kmat matrix
    if Kmat.shape[0] == 2:
        m0 = mp[0:2]
        m1 = mp[2:4]
        m2 = mp[4:6]
        m3 = mp[6:8]
        p1 = mp[8:10]
        mp1 = mp[10:12]
        p2 = mp[12:14]
        m2p = mp[14:16]
        mp2 = mp[16:18]
        p3 = mp[18::]
    elif Kmat.shape[0] == 3:
        m0 = mp[0:3]
        m1 = mp[3:6]
        m2 = mp[6:9]
        m3 = mp[9:12]
        p1 = mp[12:15]
        mp1 = mp[15:18]
        p2 = mp[18:21]
        m2p = mp[21:24]
        mp2 = mp[24:27]
        p3 = mp[27::]

    # Initialize array to save all dynamics
    dmpdt = np.array([])

    # Compute the moment equations for the:
    #=== mRNA ===#
    # Zeroth moment
    dm0dt_eq = np.dot(Kmat, m0)
    dmpdt = np.append(dmpdt, dm0dt_eq)
    # <m1>
    dm1dt_eq = np.dot((Kmat - Gm), m1) + np.dot(Rm, m0)
    dmpdt = np.append(dmpdt, dm1dt_eq)
    # <m2>
    dm2dt_eq = np.dot((Kmat - 2 * Gm), m2) + np.dot((2 * Rm + Gm), m1) +\
        np.dot(Rm, m0)
    dmpdt = np.append(dmpdt, dm2dt_eq)
    # <m3>
    dm3dt_eq = np.dot((Kmat - 3 * Gm), m3) +\
        np.dot((3 * Rm + 3 * Gm), m2) +\
        np.dot((3 * Rm - Gm), m1) +\
        np.dot(Rm, m0)
    dmpdt = np.append(dmpdt, dm3dt_eq)
    #=== protein and correlations ===#
    # <p1>
    dp1dt_eq = np.dot((Kmat - Gp), p1) + np.dot(Rp, m1)
    dmpdt = np.append(dmpdt, dp1dt_eq)
    # <mp>
    dmpdt_eq = np.dot((Kmat - Gm - Gp), mp1) +\
        np.dot(Rm, p1) +\
        np.dot(Rp, m2)
    dmpdt = np.append(dmpdt, dmpdt_eq)
    # <p2>
    dp2dt_eq = np.dot((Kmat - 2 * Gp), p2) +\
        np.dot(Gp, p1) +\
        np.dot(Rp, m1) +\
        np.dot((2 * Rp), mp1)
    dmpdt = np.append(dmpdt, dp2dt_eq)
    # <m2p>
    dm2pdt_eq = np.dot((Kmat - 2 * Gm - Gp), m2p) +\
        np.dot(Rm, p1) +\
        np.dot((Rm + Gm), mp1) +\
        np.dot(Rp, m3)
    dmpdt = np.append(dmpdt, dm2pdt_eq)
    # <mp2>
    dmp2dt_eq = np.dot((Kmat - Gm - 2 * Gp), mp2) +\
        np.dot(Rm, p2) +\
        np.dot((2 * Rp), m2p) +\
        np.dot(Rp, m2) +\
        np.dot(Gp, mp1)
    dmpdt = np.append(dmpdt, dmp2dt_eq)
    # <p3>
    dp3dt_eq = np.dot((Kmat - 3 * Gp), p3) +\
        np.dot((3 * Gp), p2) -\
        np.dot(Gp, p1) +\
        np.dot((3 * Rp), mp2) +\
        np.dot((3 * Rp), mp1) +\
        np.dot(Rp, m1)
    dmpdt = np.append(dmpdt, dp3dt_eq)

    return dmpdt
```

Let's define a function that will become quite handy in future implementations of these dynamics. This function will take the output of `odeint` function and transform it into a tidy `DataFrame`.

```python
def dynamics_to_df(sol, t):
    '''
    Takes the output of the dpdt function and the vector time and returns
    a tidy pandas DataFrame with the GLOBAL moments.
    Parameters
    ----------
    sol : array-like.
        Array with 20 or 30 columns containing the dynamics of the mRNA and
        protein distribution moments.
    t : array-like.
        Time array used for integrating the differential equations
    Returns
    -------
    tidy dataframe with the GLOBAL moments
    '''
    # Define names of dataframe columns
    names = ['time', 'm1', 'm2', 'm3', 'p1', 'mp', 'p2', 'm2p', 'mp2', 'p3']

    # Initialize matrix to save global moments
    mat = np.zeros([len(t), len(names)])
    # Save time array in matrix
    mat[:, 0] = t

    # List index for columns depending on number of elements in matrix
    idx = np.arange(int(sol.shape[1]) / 10,
                    sol.shape[1], int(sol.shape[1]) / 10)

    # Loop through index and compute global moments
    for i, index in enumerate(idx):
        # Compute and save global moment
        mat[:, i+1] = np.sum(sol[:, int(index):int(index + sol.shape[1] / 10)],
                             axis=1)

    return pd.DataFrame(mat, columns=names)
```

## Two-state promoter

Having defined these functions let's first test them with the two-state unregulated promoter.

Let's define the necessary parameters.

```python
# List the parameters fit for the lacUV5 promoter
par_UV5 = dict(kp_on=5.5, kp_off=28.9, rm=87.6, gm=1)

# define protein degradation rate in units of mRNA degradation rate
gp = 0.000277 / 0.00284
par_UV5['gp'] = gp

# define rp based on the mean protein copy number per mRNA
par_UV5['rp'] = 1000 * par_UV5['gp']

kp_on = par_UV5['kp_on']
kp_off = par_UV5['kp_off']
rm = par_UV5['rm']
gm = par_UV5['gm']
gp = par_UV5['gp']
rp = par_UV5['rp']
```

Now we will define the state transition matrix $\mathbf{K}$, the mRNA production matrix $\mathbf{R}_m$, the mRNA degradation matrix $\mathbf{\Gamma}_m$, the protein production matrix $\mathbf{R}_p$, and the protein degradation matrix $\mathbf{\Gamma}_p$.

```python
# Define the rate constant matrix
Km_unreg = np.array([[-kp_off, kp_on],
                     [kp_off, -kp_on]])
# Define the mRNA production matrix
Rm_unreg = np.array([[rm, 0],
                     [0, 0]])
# Define the mRNA degradation matrix
Gm_unreg = np.array([[gm, 0],
                     [0, gm]])

# Define the protein production matrix
Rp_unreg = np.array([[rp, 0],
                     [0, rp]])
# Define the protein degradation matrix
Gp_unreg = np.array([[gp, 0],
                     [0, gp]])
```

Now we define initial conditions and time array to use for integrating the ODEs

```python
# Define time on which to perform integration
t = np.linspace(0, 60, 101)

# Define initial conditions
m0_init = [0.5, 0.5]
m1_init = [0, 0]
m2_init = [0, 0]
m3_init = [0, 0]
p1_init = [0, 0]
mp_init = [0, 0]
p2_init = [0, 0]
m2p_init = [0, 0]
mp2_init = [0, 0]
p3_init = [0, 0]

# Solve equation
mp_sol = sp.integrate.odeint(dpdt,
                             m0_init + m1_init + m2_init + m3_init +
                             p1_init + mp_init + p2_init + m2p_init +
                             mp2_init + p3_init,
                             t,
                             args=(Km_unreg, Rm_unreg, Gm_unreg,
                                   Rp_unreg, Gp_unreg))
mp_sol.shape
```

    (101, 20)

Let's test the function converting our ODE numerical integration to a `DataFrame`.

```python
df_sol = dynamics_to_df(mp_sol, t)
df_sol.head().round(1)
```

      time
      m1
      m2
      m3
      p1
      mp
      p2
      m2p
      mp2
      p3

      0
      0.0
      0.0
      0.0
      0.0
      0.0
      0.0
      0.0
      0.0
      0.0
      0.000000e+00

      1
      0.6
      6.8
      73.7
      1021.0
      234.8
      2403.3
      91597.4
      25649.9
      972671.1
      3.926675e+07

      2
      1.2
      10.1
      137.7
      2303.1
      710.9
      9025.1
      689141.1
      106746.1
      8337666.5
      6.751009e+08

      3
      1.8
      11.8
      180.3
      3277.7
      1298.2
      18041.3
      2104606.7
      217080.1
      26566097.8
      3.262837e+09

      4
      2.4
      12.8
      206.0
      3903.1
      1928.2
      27894.1
      4417669.7
      332033.6
      56570169.8
      9.351843e+09

For comparison let's compute the *steady-state* values of both the mRNA and the protein levels and save it as a pandas `Series`.

```python
with open('./two_state_moments_lambdify.dill', 'rb') as file:
    mom_unreg_dict = cloudpickle.load(file)

# Import regulated states
with open('./three_state_moments_lambdify.dill', 'rb') as file:
    mom_reg_dict = cloudpickle.load(file)
```

```python
# Define names of dataframe index
names = ['m1', 'm2', 'm3', 'p1', 'p2', 'p3']

# Initialize array to save steady state
mp_ss = np.zeros(len(names))

# Compute mRNA steady state moments
mfunc = [mom_unreg_dict[(1, 0)],
         mom_unreg_dict[(2, 0)],
         mom_unreg_dict[(3, 0)]]
for i, f in enumerate(mfunc):
    mp_ss[i] = f(kp_on, kp_off, rm, gm, rp, gp)

# Compute the protein steady state moments
pfunc = [mom_unreg_dict[(0, 1)],
         mom_unreg_dict[(0, 2)],
         mom_unreg_dict[(0, 3)]]
for i, f in enumerate(pfunc):
    mp_ss[i + 3] = f(kp_on, kp_off, rm, gm, rp, gp)

df_ss = pd.Series(mp_ss, index=names)
df_ss
```

    m1    1.400581e+01
    m2    2.392858e+02
    m3    4.756025e+03
    p1    1.400581e+04
    p2    2.000841e+08
    p3    2.914073e+12
    dtype: float64

With this we can now plot the dynamics of all moments normalized by the stead-state value in order to compare the response time of mRNA and protein moments.

```python
# Initialize figure
fig = plt.figure(figsize=(7, 5))

# Define axis to have two on the bottom and one on the top
ax = list()
ax.append(plt.subplot2grid((2, 4), (0, 1), colspan=2, rowspan=1))
ax.append(plt.subplot2grid((2, 4), (1, 0), colspan=2, rowspan=1))
ax.append(plt.subplot2grid((2, 4), (1, 2), colspan=2, rowspan=1))

# Define array with names of moments to plot
moments_m = ['m1', 'm2', 'm3']
moments_p = ['p1', 'p2', 'p3']

# Loop through groups
for i in range(len(moments_m)):
    # Plot moments moment
    ax[i].plot(df_sol['time'], df_sol[moments_m[i]] / df_ss[moments_m[i]],
              label='mRNA')
    ax[i].plot(df_sol['time'], df_sol[moments_p[i]] / df_ss[moments_p[i]],
              label='protein')

    # Label axis
    ax[i].set_xlabel('time ($\gamma_m^{-1}$ units)')
    ax[i].set_ylabel(r'$\left\langle x(t)^{:d} \right\rangle$ / '.format(i+1) +\
                     r'$\left\langle x_{{ss}}^{:d} \right\rangle$'.format(i+1))

# Include legend
ax[0].legend(fontsize=12)

plt.tight_layout()
# Save figure
plt.savefig(figdir + 'mp_unreg_dynamics.pdf', bbox_inches='tight')
```

---

## Three-state promoter

Let's now test the function for the regulated promoter. We need to define the parameters for the regulated promoter.

```python
# List the parameters fit for the lacUV5 regulated promoter
par_UV5_reg = dict(kp_on=5.5, kp_off=28.9, rm=87.6, gm=1,
                 Nns=4.6E6, ka=139, ki=0.53, epsilon=4.5)
# Define the k0 parameters in units of the mRNA degradation time
k0_norm = 2.7E-3 / 0.00284
par_UV5_reg['k0'] = k0_norm

# define protein degradation rate in units of mRNA degradation rate
gp = 0.000277 / 0.00284
par_UV5_reg['gp'] = gp

# define rp based on the mean protein copy number per mRNA
par_UV5_reg['rp'] = 1000 * par_UV5_reg['gp']

kp_on = par_UV5_reg['kp_on']
kp_off = par_UV5_reg['kp_off']
rm = par_UV5_reg['rm']
gm = par_UV5_reg['gm']
ka = par_UV5_reg['ka']
ki = par_UV5_reg['ki']
epsilon = par_UV5_reg['epsilon']
Nns = par_UV5_reg['Nns']
```

Now we will define the mRNA production matrix $\mathbf{R}_m$, the mRNA degradation matrix $\mathbf{\Gamma}_m$, the protein production matrix $\mathbf{R}_p$, and the protein degradation matrix $\mathbf{\Gamma}_p$.

```python
# Define the production matrix
Rm_reg = np.array([[rm, 0, 0],
                   [0, 0, 0],
                   [0, 0, 0]])
# Define the degradation matrix
Gm_reg = np.array([[gm, 0, 0],
                   [0, gm, 0],
                   [0, 0, gm]])

# Define the production matrix
Rp_reg = np.array([[rp, 0, 0],
                   [0, rp, 0],
                   [0, 0, rp]])

# Define the production matrix
Gp_reg = np.array([[gp, 0, 0],
                   [0, gp, 0],
                   [0, 0, gp]])
```

Let's compute the moments for the $R = 260$ strains and all operators in the absence of inducer.

```python
# Define repressors per cell
rep = 260
par_UV5_reg['rep'] = rep

# Calculate the repressor on rate including the MWC model
kr_on = k0_norm * rep * chann_cap.p_act(0, ka, ki, epsilon)

# Define energies
operators = ['Oid', 'O1', 'O2', 'O3']
energies = [-17, -15.3, -13.9, -9.7]
energy_dict = dict(zip(operators, energies))
```

Having defined the parameters that do not depend on the operator let's now loop through operators and integrate the equations. We will also compute the steady-state solution to each of the moments and save it also in a tidy `DataFrame`.

```python
# Define time on which to perform integration
t = np.linspace(0, 60, 200)

# Define initial conditions
m0_init = [0.3, 0.3, 0.4]
m1_init = [0, 0, 0]
m2_init = [0, 0, 0]
m3_init = [0, 0, 0]
p1_init = [0, 0, 0]
mp_init = [0, 0, 0]
p2_init = [0, 0, 0]
m2p_init = [0, 0, 0]
mp2_init = [0, 0, 0]
p3_init = [0, 0, 0]

# Initialize DataFrame to save the solutions
df_sol_reg = pd.DataFrame() # Dynamics
df_ss_reg = pd.DataFrame() # Steady-state

# Loop through operators integrating the equation at each point
for op in operators:
    eRA = energy_dict[op]
    # Compute the repressor off-rate based on the on-rate and
    # the binding energy
    kr_off = chann_cap.kr_off_fun(eRA, k0_norm, kp_on, kp_off, Nns)

    #=== Compute dynamics ===#
    # Define the rate constant matrix
    Km_reg = np.array([[-kp_off, kp_on, 0],
                       [kp_off, -(kp_on + kr_on), kr_off],
                       [0, kr_on, -kr_off]])
    # Solve equation
    sol = sp.integrate.odeint(dpdt,
                                m0_init + m1_init + m2_init + m3_init +\
                                p1_init + mp_init + p2_init + m2p_init +\
                                mp2_init + p3_init,
                                t,
                                args=(Km_reg, Rm_reg, Gm_reg,
                                      Rp_reg, Gp_reg))
    # Convert to pandas DataFrame and append the extra information
    df = dynamics_to_df(sol, t)
    df['operator'] = [op] * len(t)
    df['energy'] = [eRA] * len(t)
    df['repressor'] = [rep] * len(t)
    df['IPTG'] = [0] * len(t)

    # Append to global data frame
    df_sol_reg = pd.concat([df_sol_reg, df], ignore_index=True)

    #=== Compute steady-state ===#
    # Define names of dataframe index
    names = ['m1', 'm2', 'm3', 'p1', 'p2', 'p3']

    # Initialize array to save steady state
    mp_ss = np.zeros(len(names))

    # Compute mRNA steady state moments
    mfunc = [mom_reg_dict[(1, 0)],
             mom_reg_dict[(2, 0)],
             mom_reg_dict[(3, 0)]]
    for i, f in enumerate(mfunc):
        mp_ss[i] = f(kr_on, kr_off, kp_on, kp_off, rm, gm, rp, gp)

    # Compute the protein steady state moments
    pfunc = [mom_reg_dict[(0, 1)],
             mom_reg_dict[(0, 2)],
             mom_reg_dict[(0, 3)]]
    for i, f in enumerate(pfunc):
        mp_ss[i + 3] = f(kr_on, kr_off, kp_on, kp_off, rm, gm, rp, gp)

    df_ss = pd.Series(mp_ss, index=names)
    df_ss['operator'] = op
    df_ss['energy'] = eRA
    df_ss['repressor'] = rep
    df_ss['IPTG'] = 0

    # Append to global data frame
    df_ss_reg = df_ss_reg.append(df_ss, ignore_index=True)

df_sol_reg.head().round(1)
```

      time
      m1
      m2
      m3
      p1
      mp
      p2
      m2p
      mp2
      p3
      operator
      energy
      repressor
      IPTG

      0
      0.0
      0.0
      0.0
      0.0
      0.0
      0.0
      0.0
      0.0
      0.0
      0.0
      Oid
      -17.0
      260
      0

      1
      0.3
      0.7
      4.3
      39.0
      21.4
      115.5
      3206.7
      896.9
      22407.9
      588456.7
      Oid
      -17.0
      260
      0

      2
      0.6
      0.5
      2.6
      18.1
      39.1
      166.3
      11432.7
      1083.1
      67858.5
      4439009.7
      Oid
      -17.0
      260
      0

      3
      0.9
      0.4
      1.5
      8.7
      51.7
      169.0
      20515.4
      882.7
      96283.4
      11184808.5
      Oid
      -17.0
      260
      0

      4
      1.2
      0.3
      0.9
      4.4
      60.5
      151.7
      28562.6
      638.6
      104397.2
      18855645.2
      Oid
      -17.0
      260
      0

```python
df_ss_reg.round(1)
```

      IPTG
      energy
      m1
      m2
      m3
      operator
      p1
      p2
      p3
      repressor

      0
      0.0
      -17.0
      0.0
      0.1
      0.5
      Oid
      17.2
      6575.2
      4.683003e+06
      260.0

      1
      0.0
      -15.3
      0.1
      0.4
      2.7
      O1
      93.6
      42900.8
      3.402951e+07
      260.0

      2
      0.0
      -13.9
      0.4
      1.6
      11.7
      O2
      372.1
      273361.5
      2.948103e+08
      260.0

      3
      0.0
      -9.7
      9.0
      112.2
      1725.9
      O3
      9038.9
      84479685.7
      8.153365e+11
      260.0

Having this let's plot all moments

```python
# Initialize figure
fig = plt.figure(figsize=(7, 5.5))

# Define axis to have two on the bottom and one on the top
ax = list()
ax.append(plt.subplot2grid((2, 4), (0, 1), colspan=2, rowspan=1))
ax.append(plt.subplot2grid((2, 4), (1, 0), colspan=2, rowspan=1))
ax.append(plt.subplot2grid((2, 4), (1, 2), colspan=2, rowspan=1))

# Define array with names of moments to plot
moments_m = ['m1', 'm2', 'm3']
moments_p = ['p1', 'p2', 'p3']

# Define operator colors
colors = sns.color_palette('colorblind', n_colors=len(operators))
color_dict = dict(zip(operators, colors))

# Loop through moments
for i in range(len(moments_m)):

    # Loop through operators
    for op in operators:
    # Extract operator data
        data_sol = df_sol_reg[df_sol_reg['operator'] == op]
        data_ss = df_ss_reg[df_ss_reg['operator'] == op]
        # Plot moments moment
        ax[i].plot(data_sol['time'],
                   data_sol[moments_m[i]] / data_ss[moments_m[i]].values,
                  label=op, color=color_dict[op])
        ax[i].plot(data_sol['time'],
                   data_sol[moments_p[i]] / data_ss[moments_p[i]].values,
                  label='', color=color_dict[op], linestyle='--')

    # Label axis
    ax[i].set_xlabel('time ($\gamma_m^{-1}$ units)')
    ax[i].set_ylabel(r'$\left\langle x(t)^{:d} \right\rangle$ / '.format(i+1) +\
                     r'$\left\langle x_{{ss}}^{:d} \right\rangle$'.format(i+1))
    ax[i].set_ylim([0, 1.5])

# Include legend
first_legend = ax[0].legend(loc=0, ncol=2, title='operator')

# Add secondary label to distinguish protein vs mRNA.
ax2 = ax[0].add_artist(first_legend)

# "plot" a solid and a dashed line for the legend
line1, = ax[0].plot([], [], color='black', label='mRNA')
line2, = ax[0].plot([], [], color='black', linestyle='--', label='protein')
# Create another legend.
ax[0].legend(handles=[line1, line2], loc='upper right', frameon=False)

plt.tight_layout()
# Save figure
plt.savefig(figdir + 'mp_reg_dynamics.pdf', bbox_inches='tight')
```

All dynamics seem to converge to the predicted stteady state. This tells us that the numerical integration of the moment differential equations is working as expected.

## Multi-promoter dynamics.

It is known that depending on the position of a gene relative to the replication origin and the cell growth rate a gene can have multiple copies of a gene. For our specific experimental case with a doubling time of $\approx$ 60 min the *galK* locus has an average of 1.66 copies of the promoter along the cell cycle. What this means is that cells spend 2/3 of the cell cycle ($\approx$ 40 min) with one copy of the promoter and the rest ($\approx$ 20 min) with a single copy.

The way we account for that in our model is that assuming both promoters have the same mRNA production rate $r_m$, during the two-copies phase cells have a production rate of $2 r_m$, with the rest of the parameters remaining the same.

```python

```


```python
%matplotlib inline
```

Deep Learning with PyTorch
**************************

Deep Learning Building Blocks: Affine maps, non-linearities and objectives
==========================================================================

Deep learning consists of composing linearities with non-linearities in
clever ways. The introduction of non-linearities allows for powerful
models. In this section, we will play with these core components, make
up an objective function, and see how the model is trained.

Affine Maps
~~~~~~~~~~~

One of the core workhorses of deep learning is the affine map, which is
a function $f(x)$ where

\begin{align}f(x) = Ax + b\end{align}

for a matrix $A$ and vectors $x, b$. The parameters to be
learned here are $A$ and $b$. Often, $b$ is refered to
as the *bias* term.

PyTorch and most other deep learning frameworks do things a little
differently than traditional linear algebra. It maps the rows of the
input instead of the columns. That is, the $i$'th row of the
output below is the mapping of the $i$'th row of the input under
$A$, plus the bias term. Look at the example below.

```python
# Author: Robert Guthrie

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

torch.manual_seed(1)
```

```python
lin = nn.Linear(5, 3)  # maps from R^5 to R^3, parameters A, b
# data is 2x5.  A maps from 5 to 3... can we map "data" under A?
data = torch.randn(2, 5)
print(lin(data))  # yes
```

Non-Linearities
~~~~~~~~~~~~~~~

First, note the following fact, which will explain why we need
non-linearities in the first place. Suppose we have two affine maps
$f(x) = Ax + b$ and $g(x) = Cx + d$. What is
$f(g(x))$?

\begin{align}f(g(x)) = A(Cx + d) + b = ACx + (Ad + b)\end{align}

$AC$ is a matrix and $Ad + b$ is a vector, so we see that
composing affine maps gives you an affine map.

From this, you can see that if you wanted your neural network to be long
chains of affine compositions, that this adds no new power to your model
than just doing a single affine map.

If we introduce non-linearities in between the affine layers, this is no
longer the case, and we can build much more powerful models.

There are a few core non-linearities.
$\tanh(x), \sigma(x), \text{ReLU}(x)$ are the most common. You are
probably wondering: "why these functions? I can think of plenty of other
non-linearities." The reason for this is that they have gradients that
are easy to compute, and computing gradients is essential for learning.
For example

\begin{align}\frac{d\sigma}{dx} = \sigma(x)(1 - \sigma(x))\end{align}

A quick note: although you may have learned some neural networks in your
intro to AI class where $\sigma(x)$ was the default non-linearity,
typically people shy away from it in practice. This is because the
gradient *vanishes* very quickly as the absolute value of the argument
grows. Small gradients means it is hard to learn. Most people default to
tanh or ReLU.

```python
# In pytorch, most non-linearities are in torch.functional (we have it imported as F)
# Note that non-linearites typically don't have parameters like affine maps do.
# That is, they don't have weights that are updated during training.
data = torch.randn(2, 2)
print(data)
print(F.relu(data))
```

Softmax and Probabilities
~~~~~~~~~~~~~~~~~~~~~~~~~

The function $\text{Softmax}(x)$ is also just a non-linearity, but
it is special in that it usually is the last operation done in a
network. This is because it takes in a vector of real numbers and
returns a probability distribution. Its definition is as follows. Let
$x$ be a vector of real numbers (positive, negative, whatever,
there are no constraints). Then the i'th component of
$\text{Softmax}(x)$ is

\begin{align}\frac{\exp(x_i)}{\sum_j \exp(x_j)}\end{align}

It should be clear that the output is a probability distribution: each
element is non-negative and the sum over all components is 1.

You could also think of it as just applying an element-wise
exponentiation operator to the input to make everything non-negative and
then dividing by the normalization constant.

```python
# Softmax is also in torch.nn.functional
data = torch.randn(5)
print(data)
print(F.softmax(data, dim=0))
print(F.softmax(data, dim=0).sum())  # Sums to 1 because it is a distribution!
print(F.log_softmax(data, dim=0))  # theres also log_softmax
```

Objective Functions
~~~~~~~~~~~~~~~~~~~

The objective function is the function that your network is being
trained to minimize (in which case it is often called a *loss function*
or *cost function*). This proceeds by first choosing a training
instance, running it through your neural network, and then computing the
loss of the output. The parameters of the model are then updated by
taking the derivative of the loss function. Intuitively, if your model
is completely confident in its answer, and its answer is wrong, your
loss will be high. If it is very confident in its answer, and its answer
is correct, the loss will be low.

The idea behind minimizing the loss function on your training examples
is that your network will hopefully generalize well and have small loss
on unseen examples in your dev set, test set, or in production. An
example loss function is the *negative log likelihood loss*, which is a
very common objective for multi-class classification. For supervised
multi-class classification, this means training the network to minimize
the negative log probability of the correct output (or equivalently,
maximize the log probability of the correct output).

Optimization and Training
=========================

So what we can compute a loss function for an instance? What do we do
with that? We saw earlier that Tensors know how to compute gradients
with respect to the things that were used to compute it. Well,
since our loss is an Tensor, we can compute gradients with
respect to all of the parameters used to compute it! Then we can perform
standard gradient updates. Let $\theta$ be our parameters,
$L(\theta)$ the loss function, and $\eta$ a positive
learning rate. Then:

\begin{align}\theta^{(t+1)} = \theta^{(t)} - \eta \nabla_\theta L(\theta)\end{align}

There are a huge collection of algorithms and active research in
attempting to do something more than just this vanilla gradient update.
Many attempt to vary the learning rate based on what is happening at
train time. You don't need to worry about what specifically these
algorithms are doing unless you are really interested. Torch provides
many in the torch.optim package, and they are all completely
transparent. Using the simplest gradient update is the same as the more
complicated algorithms. Trying different update algorithms and different
parameters for the update algorithms (like different initial learning
rates) is important in optimizing your network's performance. Often,
just replacing vanilla SGD with an optimizer like Adam or RMSProp will
boost performance noticably.

Creating Network Components in PyTorch
======================================

Before we move on to our focus on NLP, lets do an annotated example of
building a network in PyTorch using only affine maps and
non-linearities. We will also see how to compute a loss function, using
PyTorch's built in negative log likelihood, and update parameters by
backpropagation.

All network components should inherit from nn.Module and override the
forward() method. That is about it, as far as the boilerplate is
concerned. Inheriting from nn.Module provides functionality to your
component. For example, it makes it keep track of its trainable
parameters, you can swap it between CPU and GPU with the ``.to(device)``
method, where device can be a CPU device ``torch.device("cpu")`` or CUDA
device ``torch.device("cuda:0")``.

Let's write an annotated example of a network that takes in a sparse
bag-of-words representation and outputs a probability distribution over
two labels: "English" and "Spanish". This model is just logistic
regression.

Example: Logistic Regression Bag-of-Words classifier
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

Our model will map a sparse BoW representation to log probabilities over
labels. We assign each word in the vocab an index. For example, say our
entire vocab is two words "hello" and "world", with indices 0 and 1
respectively. The BoW vector for the sentence "hello hello hello hello"
is

\begin{align}\left[ 4, 0 \right]\end{align}

For "hello world world hello", it is

\begin{align}\left[ 2, 2 \right]\end{align}

etc. In general, it is

\begin{align}\left[ \text{Count}(\text{hello}), \text{Count}(\text{world}) \right]\end{align}

Denote this BOW vector as $x$. The output of our network is:

\begin{align}\log \text{Softmax}(Ax + b)\end{align}

That is, we pass the input through an affine map and then do log
softmax.

```python
data = [("me gusta comer en la cafeteria".split(), "SPANISH"),
        ("Give it to me".split(), "ENGLISH"),
        ("No creo que sea una buena idea".split(), "SPANISH"),
        ("No it is not a good idea to get lost at sea".split(), "ENGLISH")]

test_data = [("Yo creo que si".split(), "SPANISH"),
             ("it is lost on me".split(), "ENGLISH")]

# word_to_ix maps each word in the vocab to a unique integer, which will be its
# index into the Bag of words vector
word_to_ix = {}
for sent, _ in data + test_data:
    for word in sent:
        if word not in word_to_ix:
            word_to_ix[word] = len(word_to_ix)
print(word_to_ix)

VOCAB_SIZE = len(word_to_ix)
NUM_LABELS = 2

class BoWClassifier(nn.Module):  # inheriting from nn.Module!

    def __init__(self, num_labels, vocab_size):
        # calls the init function of nn.Module.  Dont get confused by syntax,
        # just always do it in an nn.Module
        super(BoWClassifier, self).__init__()

        # Define the parameters that you will need.  In this case, we need A and b,
        # the parameters of the affine mapping.
        # Torch defines nn.Linear(), which provides the affine map.
        # Make sure you understand why the input dimension is vocab_size
        # and the output is num_labels!
        self.linear = nn.Linear(vocab_size, num_labels)

        # NOTE! The non-linearity log softmax does not have parameters! So we don't need
        # to worry about that here

    def forward(self, bow_vec):
        # Pass the input through the linear layer,
        # then pass that through log_softmax.
        # Many non-linearities and other functions are in torch.nn.functional
        return F.log_softmax(self.linear(bow_vec), dim=1)

def make_bow_vector(sentence, word_to_ix):
    vec = torch.zeros(len(word_to_ix))
    for word in sentence:
        vec[word_to_ix[word]] += 1
    return vec.view(1, -1)

def make_target(label, label_to_ix):
    return torch.LongTensor([label_to_ix[label]])

model = BoWClassifier(NUM_LABELS, VOCAB_SIZE)

# the model knows its parameters.  The first output below is A, the second is b.
# Whenever you assign a component to a class variable in the __init__ function
# of a module, which was done with the line
# self.linear = nn.Linear(...)
# Then through some Python magic from the PyTorch devs, your module
# (in this case, BoWClassifier) will store knowledge of the nn.Linear's parameters
for param in model.parameters():
    print(param)

# To run the model, pass in a BoW vector
# Here we don't need to train, so the code is wrapped in torch.no_grad()
with torch.no_grad():
    sample = data[0]
    bow_vector = make_bow_vector(sample[0], word_to_ix)
    log_probs = model(bow_vector)
    print(log_probs)
```

Which of the above values corresponds to the log probability of ENGLISH,
and which to SPANISH? We never defined it, but we need to if we want to
train the thing.

```python
label_to_ix = {"SPANISH": 0, "ENGLISH": 1}
```

So lets train! To do this, we pass instances through to get log
probabilities, compute a loss function, compute the gradient of the loss
function, and then update the parameters with a gradient step. Loss
functions are provided by Torch in the nn package. nn.NLLLoss() is the
negative log likelihood loss we want. It also defines optimization
functions in torch.optim. Here, we will just use SGD.

Note that the *input* to NLLLoss is a vector of log probabilities, and a
target label. It doesn't compute the log probabilities for us. This is
why the last layer of our network is log softmax. The loss function
nn.CrossEntropyLoss() is the same as NLLLoss(), except it does the log
softmax for you.

```python
# Run on test data before we train, just to see a before-and-after
with torch.no_grad():
    for instance, label in test_data:
        bow_vec = make_bow_vector(instance, word_to_ix)
        log_probs = model(bow_vec)
        print(log_probs)

# Print the matrix column corresponding to "creo"
print(next(model.parameters())[:, word_to_ix["creo"]])

loss_function = nn.NLLLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# Usually you want to pass over the training data several times.
# 100 is much bigger than on a real data set, but real datasets have more than
# two instances.  Usually, somewhere between 5 and 30 epochs is reasonable.
for epoch in range(100):
    for instance, label in data:
        # Step 1. Remember that PyTorch accumulates gradients.
        # We need to clear them out before each instance
        model.zero_grad()

        # Step 2. Make our BOW vector and also we must wrap the target in a
        # Tensor as an integer. For example, if the target is SPANISH, then
        # we wrap the integer 0. The loss function then knows that the 0th
        # element of the log probabilities is the log probability
        # corresponding to SPANISH
        bow_vec = make_bow_vector(instance, word_to_ix)
        target = make_target(label, label_to_ix)

        # Step 3. Run our forward pass.
        log_probs = model(bow_vec)

        # Step 4. Compute the loss, gradients, and update the parameters by
        # calling optimizer.step()
        loss = loss_function(log_probs, target)
        loss.backward()
        optimizer.step()

with torch.no_grad():
    for instance, label in test_data:
        bow_vec = make_bow_vector(instance, word_to_ix)
        log_probs = model(bow_vec)
        print(log_probs)

# Index corresponding to Spanish goes up, English goes down!
print(next(model.parameters())[:, word_to_ix["creo"]])
```

We got the right answer! You can see that the log probability for
Spanish is much higher in the first example, and the log probability for
English is much higher in the second for the test data, as it should be.

Now you see how to make a PyTorch component, pass some data through it
and do gradient updates. We are ready to dig deeper into what deep NLP
has to offer.