```python
for i in range(9):
    globals()[f'var_{i}'] = []
    for j in range(i,i+4):
        globals()[f'var_{i}'].append(f'Hello {j}')

var_8
```

    ['Hello 8', 'Hello 9', 'Hello 10', 'Hello 11']

```python
# Importing packages
import matplotlib.pyplot as plt
import pandas as pd
import ipywidgets as widgets
import numpy as np
import itertools
from ipywidgets import Layout
import seaborn as sns
import math
from IPython.display import clear_output

import time
from scipy import linalg
import scipy.optimize as optimize
import sympy as sm

from tkinter import *
import tkinter as tk

from data_gen import gen_df

from IPython.display import clear_output

```

In this project, we want to try to estimate the ratings of movies. That is, we well calculate the utility of a movie, given genres, duration, awards and so forth, and we will try to calculate the IMDB ratings for each movie, based on this utility. Finally, we will then minimize the distance between our estimated ratings and the actual ratings.

Our utility function is given by:
$$ U_i = \sum_{k=1}^{23}(\alpha_k G_{ik}) + \sum_{k=1920}^{2010} (\beta_k D_{ik}) + \gamma N_i + \delta W_i + \rho_1 L_i + \rho_2 L_i^2 $$.

Based on this utility function, we will estimate the ratings of each movies,
$$ R_i^{model} = \frac{\exp(\omega x_i)}{1 + \exp(\omega x_i)} $$

where,
$$ x_i = \big[G_1, G_2, ..., G_n, D_{1920}, D_{1930}, ..., D_{2010}, N_i, W_i, L_i, L_i^2 \big] $$
$$ \omega = \big[\alpha_1, \alpha_2, ..., \alpha_{23}, \beta_{1920}, \beta_{1930}, ..., \beta_{2010}, \gamma, \delta, \rho_1, \rho_2  \big] $$

We then use optimize methods to solve the following:

$$ \min_{\omega} \Big\{ \sum_{i=1}^{n} \left( R_i^{model} - R_i^{data} \right)^2 \Big\} $$

```python
def df_s(keep_top=None):
    """ Prepares data for optimization by creating dummmy-variables for decades, dropping movies with less than 5000 ratings,
        dropping some variables that we don't wish to use in the model, and splitting the dataset, so we have a data set
        consisting only of the variables used in the model.

    Args:
        keep_top (optional, type: int): When a number (n) is put in this option, only the top n movies a kept in the dataset. This is used to sort out movies for a better fit of the model

    Returns:
        df_X (type: Pandas dataframe): A dataframe consisting only of the varaibles that are used to calculate the rating.
        df_Y (type: Pandas dataframe): A dataframe consisting the true rating from the dataset and an index-variable, used for merging with original dataset, later in the process.

    Notice:
        The function includes global for 'df', which means that the dataframe 'df' generated, can be called outside the function. This is used so that we won't have to filter the original dataframe again, when we use it again.
    """
    global df

    # Calls the file 'imdb.csv' which is located in the repository, and contains the dataset used
    filename = 'imdb.csv'

    # A function generated earlier, which cleans the dataset [Should we include this, and do some of the filtering below in that funciton?]
    df = gen_df(filename)

    # Generate list of decades for dummies
    decade_list = [1920, 1930, 1940, 1950, 1960, 1970, 1980, 1990, 2000, 2010]

    # Decade dummies
    for i in decade_list:
        df[f'decade_{i}'] = 0
        df.loc[(df['decade'] == f'{i}s'),f'decade_{i}'] = 1

    # Filters out movies with less than 5000 ratings and drops some genres
    df = df.loc[(df['ratingCount']>=5000)]
    df = df.drop(columns=['Adult','GameShow','News','RealityTV','TalkShow'])

    # Keeps only top n movies, if this is specified when calling the function
    if keep_top != None:
        df = df.sort_values('imdbRating', ascending=False)
        df = df.iloc[:keep_top]

    # Splits the dataset into two datasets
    df_X = df.copy()
    df_Y = pd.DataFrame(df[['imdbRating', 'index']].copy())
    df_Y = df_Y.rename(columns = {'imdbRating':'rat_data'})

    # Rearrange and keep given columns
    df_X = df_X.reindex(['Action','Adventure','Animation','Biography','Comedy','Crime','Documentary',
                              'Drama','Family','Fantasy','FilmNoir','History','Horror','Music','Musical',
                              'Mystery','Romance','SciFi','Short','Sport','Thriller','War','Western',
                              'decade_1920','decade_1930','decade_1940','decade_1950','decade_1960',
                              'decade_1970','decade_1980','decade_1990','decade_2000','decade_2010',
                              'nrOfNominations','nrOfWins','duration'], axis=1)

    return df_X, df_Y
```

## Optimization
We analyze preferences for movies by minimizing the sum of the squared differences from our model prediction and the true rating, i.e.:
$$
\min_{\omega} (R_{model}-R_{data})^2
$$
By doing so, we'll find parameters that best describe movie preferences given our model. As we will discuss later, our model is not necessarily the best model to describe preferences, which might lead to these preferences not being completely true.

If the model would predict all movie ratings correctly the parameters would perfectly describe movie preferences. For example if the parameter for action-movies were negative, the prediction would be precise, if action-movies generally had a negative impact on ratings. Of course people have much different preferences, and an interesting analysis would be, to predict preferences for different people, and by this look at a distribution of these parameters. This is not included in this project, though.

To shorten the length of the optimizing process, and to sort out movies that makes it harder for the model to predict preferences, the optimization can be based on the top n movies, based on IMDb rating. When all movies are included in the optimizer the lowest ranking movies will have ratings around 2. Our rating predictor will in this case not rank any movies under 5. Thereby it is clear that the predictor model is not able to predict movies with very low rating. But it is able to predict movies with higher rating.

```python
def optimizer(keep_top=None, live_graph=True):
    """ Creates and optimizes the function which calculates rating based on variables in the dataset.

    Args:
        keep_top (optional, type: int): When a number (n) is put in this option, only the top n movies a kept in the dataset. This is used to sort out movies for a better fit of the model
        live_graph(optional, type: boolean): Produces a live graph of the optimization proces, plotting the function values against the iterations. Makes the optimizer slower, but is a good visualization of the process.

    Returns:
        result (type: scipy optimizer object): An object containing results from the optimizer among other information such as status etc.
        timer (type: int): Time of running optimizer in seconds.

    Notice:
        Uses the package scipy.optimize.minimize which optimizes functions based on a given method. Here 'Nelder-Mead' is used.
        This function defines a set of functions used in the optimizer which are not documented by docstring bu comments in code.
    """
    # "Initializes" a timer for printing time of optimization
    start = time.time()

    # Set of globals so variables can be called inside the functions defined in this function
    global fs
    global evals
    global x0
    global df
    global df_Y

    # Calculates the sum of the squared differences between rating calculated by model and ratings from the dataset.
    # This is what is to be minimized
    def sqr_diff_sum(df_X,pars):
        global df_Y
        util = df_X@pars
        df_Y['rat_model'] = 10*np.exp(util)/(1+np.exp(util))
        df_Y['sqr_diff'] = (df_Y['rat_model']-df_Y['rat_data'])**2
        return df_Y['sqr_diff'].sum()

    # List of zeroes for initial guess
    def zeros(n):
        list = [0] * n
        return list

    # Plots a live graph of optimization process, if chosen.
    if live_graph:
        # The live plot is created by taking the function value for each iteration and saving in a list
        def live_plot(evals, fs, ymax=10000, figsize=(7,5)):
            # Clears plot every time a new plot is created
            clear_output(wait=True)
            plt.figure(figsize=figsize)

            # Plots values and sets title, grid, labels, etc.
            plt.plot(evals, fs)
            plt.title('Optimizing path')
            plt.grid(True)
            plt.xlabel('Iterations')
            plt.ylabel('Function value')
            plt.xlim(0,20000)
            plt.ylim(0,ymax)

            plt.show();

    # Function to be called by optimizer each iteration. This returns function values for each iteration and plots live graph
    def collect(x):
        # Set of globals to be called in live graph-function
        global fs
        global evals
        global x0
        global fig
        global ax
        global ymax

        # Calculates function value for initial guess, for first iteration
        if evals == 0:
            fs = []
            fs.append(obj_fun(x0))

        # Calculates function values for current parameter values
        if evals != 0:
            fs.append(obj_fun(x))

        # Plots live graph if chosen
        if live_graph:
            # Calculates max y-value for axis for first iteration
            if evals == 1:
                ymax = math.ceil(obj_fun(x)/2000)*2000

            # Updates plot every 100th iteration for the first 1000 iterations
            if evals < 1000:
                if evals > 0 and evals%100 == 0:
                    live_plot(range(evals+1), fs, ymax)
            # Updates plot every 500th iteration for the rest of the process
            if evals >=1000:
                if evals%500 == 0:
                    live_plot(range(evals+1), fs, ymax)

        # Adds one to the number of iterations each iteration
        evals += 1

    # Define datasets to be used using function defined earlier
    df_X, df_Y = df_s(keep_top=keep_top)

    # Intiial guess
    x0 = zeros(36)

    # Start iteration number
    evals = 0

    # Defines function to be minimized
    obj_fun = lambda x: sqr_diff_sum(df_X,x)

    # Run optimizer
    result = optimize.minimize(obj_fun,x0,
                               method="Nelder-Mead",
                               options={"disp":True, "maxiter":50000}, # display the results
                               callback=collect
                               )

    # End timer
    end = time.time()-start

    # Returns
    return result, end
```

```python
# Run optimizer function and print results and time
result, timer = optimizer()
print(f'{" ":9s}Time: {timer:.4f} seconds')

# Create a dictionary of parameter results for each variable
variables = ['Action','Adventure','Animation','Biography','Comedy','Crime','Documentary',
            'Drama','Family','Fantasy','FilmNoir','History','Horror','Music','Musical',
            'Mystery','Romance','SciFi','Short','Sport','Thriller','War','Western',
            'decade_1920','decade_1930','decade_1940','decade_1950','decade_1960',
            'decade_1970','decade_1980','decade_1990','decade_2000','decade_2010',
            'nrOfNominations','nrOfWins','duration']

results = dict()

for j,i in enumerate(variables):
    results[i] = result.x[j]
```

## Prediction precision
Ideally our model would predict movie ratings perfectly, and thereby our parameters would perfectly describe the general preferences for movies. Whether this is possible, even with an advanced model, is not sure. People have different preferences, and two movies that would seem identical in our dataset, based on genres, duration, time of release, etc. could have very different ratings. Therefore som devitaion in our prediction from the true ratings are expected.

The left graphs above plot the mean of deviations based on genres, duration, decades, and true ratings. The right graphs plot the number of movies in the dataset based on the same groups. From these graphs we see a clear correlation between precision and number of movies in the dataset. The more movies of a given genre, duration, etc. the more precise the prediction of movie-ratings in these groups. Take drama-movies as an example. The mean of the deviation is close to zero, and movies of this genre is also quite overrepresentated in the dataset. The complete opposite case is short-movies. But we also see that sci-fi-movies are quite well predicted, even though there aren't relatively many of these movies in the dataset. The same is seen in other different genres.

This general image is seen when we group movies on decades, duration, and true rating as well.

```python
def merge_df(df_org):
    """ Merges the original dataset with the optimal solution from the optimizer, i.e. the last dataset of calculated raitngs.

    Args:
        df_org (type: Pandas dataframe): The orignal dataset of movies, containing true ratings.

    Returns:
        df_merge (type: Pandas dataframe): Original dataset combined with calculated ratings from optimal parameters.
    """

    # Merges original dataset with optimal ratings based on model
    df_merge = df_org.merge(df_Y, how='right', left_on='index', right_on='index')

    # Calculates deviations both in normal and absolute form
    df_merge['abs_diff'] = abs(df_merge['rat_model']-df_merge['rat_data'])
    df_merge['diff'] = df_merge['rat_model']-df_merge['rat_data']

    return df_merge

def _mean_genre(df,group):
    """ Calculates the mean of deviations from true ratings and ratings based on optimal parameters from optimizer.

    Args:
        df (type: Pandas dataframe): Dataframe consisting information on groups and deivations
        group (type: string): Defines which group mean of deviations are presented for. Chosen by fixed list in widget.

    Returns:
        None

    Notice:
        Only to be run through widget.interact()
    """
    # List of values used for each group
    genre_list = ['Action','Adventure','Animation','Biography','Comedy','Crime','Documentary',
                    'Drama','Family','Fantasy','FilmNoir','History','Horror','Music','Musical',
                    'Mystery','Romance','SciFi','Short','Sport','Thriller','War','Western']

    decade_list = ['1920s', '1930s', '1940s', '1950s', '1960s', '1970s', '1980s', '1990s', '2000s', '2010s']

    dur_list = [0, .5, 1, 1.5, 2, 2.5, 3, 3.5]

    rat_list = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

    mean_dict = dict()
    count_dict = dict()

    # Creates dictionaries based on chosen group
    if group == 'Genres':
        for i in genre_list:
            I = df[i] == 1
            mean_dict[i] = df.loc[I]['diff'].mean()
            count_dict[i] = df.loc[I]['index'].count()

    if group == 'Decades':
        for i in decade_list:
            I = df['decade'] == i
            mean_dict[i] = df.loc[I]['diff'].mean()
            count_dict[i] = df.loc[I]['index'].count()

    if group == 'Duration':
        for j,i in enumerate(dur_list):
            if j != len(dur_list)-1:
                # Uses values from dict, to find movies between to items in the lists.
                I = ((df['duration'] >= dur_list[j]) & (df['duration'] < dur_list[j+1]))
                mean_dict[f'{dur_list[j]:2.1f} - {dur_list[j+1]:2.1f}'] = df.loc[I]['diff'].mean()
                count_dict[f'{dur_list[j]:2.1f} - {dur_list[j+1]:2.1f}'] = df.loc[I]['index'].count()
            else:
                I = df['duration'] >= dur_list[j]
                mean_dict[f'{dur_list[j]:2.1f} {"+":5s}'] = df.loc[I]['diff'].mean()
                count_dict[f'{dur_list[j]:2.1f} {"+":5s}'] = df.loc[I]['index'].count()

    if group == 'True rating':
        for j,i in enumerate(rat_list):
            if j != len(rat_list)-1:
                # Uses values from dict, to find movies between to items in the lists.
                I = ((df['rat_data'] >= rat_list[j]) & (df['rat_data'] < rat_list[j+1]))
                mean_dict[f'{rat_list[j]:4.1f} - {rat_list[j+1]:3.1f}'] = df.loc[I]['diff'].mean()
                count_dict[f'{rat_list[j]:4.1f} - {rat_list[j+1]:3.1f}'] = df.loc[I]['index'].count()
            else:
                pass

    # Creates figure to hold two subplots
    fig1, (ax1,ax2) = plt.subplots(nrows=1,ncols=2,sharey=True,figsize=(12,5))

    # Plots means of deviations
    ax1.barh(*zip(*mean_dict.items()))
    ax1.set_ylabel(group)
    ax1.set_xlabel('Model deviation')
    ax1.grid(axis='x')

    # Plots count of movies
    ax2.barh(*zip(*count_dict.items()))
    ax2.set_xlabel('Number of observations')
    ax2.grid(axis='x')

# Calls merge-function
df_merge = merge_df(df)

# Runs interactive figure
mean_genre = widgets.interact(_mean_genre,
                             df = widgets.fixed(df_merge),
                             group = widgets.Dropdown(
                             options = ['Genres','Decades','Duration', 'True rating'],
                             description = 'Group',
                             value = 'Genres'))

```

    interactive(children=(Dropdown(description='Group', options=('Genres', 'Decades', 'Duration', 'True rating'), …


# Particles in a Box
*UW ASTR 507 Homework 1 (Notebook and framework made by Tom Wagg)*

# Contents

1. [General problem description and guidance](#prob_desc)
2. [Code framework description](#frame_desc)
3. [TODO list](#todo)
4. [Solutions](#solutions)

# General problem description and guidance

*The following cell is copied directly from the problem description, just split into bullet points*

Numerically compute the distribution of kinetic energy of a **two**-dimensional monatomic ideal gas.

- Start off with equal kinetic energy, $E$, for $N$ particles in a square box, but random positions and moving in random directions. Follow the particles until they collide (be careful how you handle the collisions).
- Conserve momentum and energy in the collisions (see, e.g., Reif, 14.1-14.2, Feynman 39) and assume that the colliding particles hard spheres (in which case they experience an elastic force directed along the line connecting the center of the spheres when they touch).
    - *Alternatively, you may instead make the approximation that when they collide the particles receive a kick in a random direction in the center-of-mass frame (make sure you conserve energy and momentum!).*
- Allow the cross section and number of particles to vary in your program. You can make the sides of the box reflecting or periodic.

# Code framework description

*This section explains the framework of the code I've set up to help with completing this problem. Note that this is just a general description, specific information is in comments in the code*

**Overview of Simulator Class**

The main part of the code is the Simulator class. If you haven't heard of [Python Classes](https://www.w3schools.com/python/python_classes.asp) before, they are essentially a way of attaching variables, functions, etc. to a specific object. In this case, the class is our box full of particles.

The class contains variables that track
- size of the box
- initial kinetic energy of each particle
- array of positions ($x$, $y$ for each particle)
- array of velocities ($v_x$, $v_y$, for each particle)
- radius of particles
- mass of each particle

This class also has a series of functions that initialise all of the particles and then handle different events. Each of these events generally happen during `run_simulation` which handles the actual evolution of the particles.

**Overview of `run_simulation()`**

For each timestep the function does the following
1. Move the particles based on their current velocities and size of timestep
2. Resolve any collisions with particles with the wall
3. Resolve any collisions between particles

Though, in reality it currently does NOT do this, you need to handle these things! See below for a list of things you need to do

# Implementation TODO list

*Before you can complete the main three problems of this problem set you'll need to implement the following features. Each are marked in the code with a comment containing "TODO"*

**Initialisation**
1. [ ] Assign random initial positions
    - Currently everything is just at the origin, they should be random across the box!
    - (You should also consider that two particles shouldn't start on top of one another)
2. [ ] Assign random initial velocities
    - Each particle has the same initial kinetic energy `E`
    - Distribute this kinetic energy randomly amongst the x and y components and in a random direction (positive or negative in each axis)
**Recurring functions**

3. [ ] Update particle positions based on current velocities in `run_simulation()`
4. [ ] Resolve any collisions with walls by implementing `resolve_wall_collisions()`
    - Don't forget to account for the fact that particles have a finite radius
    - Particles could get lucky and hit two walls in a single timestep at a corner!
5. [ ] Resolve any collisions between particles by implementing `resolve_particle_collisions()`
    - Consider how you can tell when two particles are colliding
    - You may wish to consider both where the particles currently are and *in which direction* they are moving

# Testing/visualising your simulation
We recommend that you try out your code, visualise the particles and check that everything looks right **before** you get started on answering the questions below. We've made `visualise_simulator.py` for you that should suit your purposes.

Once you've started implementing some parts, go into `visualise_simulator.py` and choose some valid values for the input. Then just run `python visualise_simulator.py` from the command line and you should see a visualisation of your particles! :D

# Solutions

## Question 1

### 1a - Code summary
Write a detailed summary of what choices you make in your code and how it works. Specifically, address the dimensions you use for velocity, mass, energy, length, and time. Discuss how you determine if a collision has occurred, what equations you solve during a collision, and how you handle the box boundaries. Email me your code, preferably written in Julia in a commented Jupyter notebook.

### 1b - Snapshot Plots

Make four plots for a case with 100 particles, each with a radius equal to 2\% of the size of the box. The first two plots should show the position of the particles (and sizes if you would like) at the initial time, and another snapshot after a long enough time that the particle velocity distribution has reached a steady-state in which the initial condition is forgotten. The second two plots should show the velocity distribution of the particles at the initial and final times.

```python
# this makes sure that you're using the latest version of your simulator
reload(simulator)

# pick a box size, initial kinetic energy and mass
E = None
size = None
mass = None

# radius is 2% of box size, 100 particles
radius = size * 0.02
N = 100

sim = simulator.Simulation(N=N, E=E, radius=radius, size=size, masses=mass, visualise=False)
sim.run_simulation(10000)
```

```python
# make some plots using `sim` results
```

## Question 2
### 2a - 2D Maxwellian Derivation

Derive the energy and velocity distribution for a Maxwellian as a function of speed and energy **in two dimensions**, $dN/dv$ and $dN/dE$, where $v = \vert \mathbf{v} \vert|$ is the speed, $\mathbf{v}$ is the velocity, and $E$ is the energy.  You can start with the result that $f(q,p) \propto e^{-(E/k_BT)}$.

\begin{align}
    f(q, p) &\propto e^{-(E / k_B T)}
\end{align}

***(edit cell above to continue derivation from the starting expression)***

### 2b - Show distribution approaches Maxwellian

Show that the energy (or speed) distribution of your simulation particles approaches a 2-D Maxwellian by plotting the speed distribution, either a histogram or cumulative distribution, averaged over several timesteps (or over several simulations) to make a smooth curve.  Overplot the 2-D Maxwellian distribution you computed in part (a), normalized to agree.

```python
# make your plot here
```

### 2c - Multiple Masses

Now make half of the masses 10 times larger. Does the velocity distribution follow a single Maxwellian? Does the energy distribution?  Explain your results.

```python
# make plots and explain them here
```

## Question 3

Initially the distribution of particle energies changes rapidly, but eventually it approaches a steady-state.

### 3a - Relaxation Criterion
Invent a numerical criterion for when your system has "relaxed", i.e. how can you tell from the distribution of particle velocities (or energies) that the system is in a statistically steady-state.

### 3b - Relaxation time derivation
Derive a formula for the relaxation time for each particle as a function of the initial velocity of each particle, $v_0$, the cross section, $\sigma$ (**which in 2-D case has units of length, not area**), and the number density of particles, $n$ (**which has units of 1/area in 2-D**).

### 3c - Comparison with Criterion
Vary the number density, particle size (i.e. cross section), and number density of particles to show that the criterion you developed in part (a) scales in the same way as the relaxation time from part (b). Demonstrate that these relations agree with a plot or two.

```python
# your code here
```

## Problem 4 (Extra Credit)
Demonstrate that $P=nkT$ holds for your simulations - try a few different values of $n$ and $kT$ (i.e. number density and mean energy of particles). Compute the mean pressure by calculating the force per length (since 2-D) on each (or just one) side of your box.  You can do this by using the definition of pressure:  force divided by length (in 2-D), and compute the forces of each particle as it reflects off of each side (if you have periodic boundary conditions, then keep track of the momentum perpendicular to the edges as the particles jump from one side of the box to the other). The average force imparted is the average change in momentum of the particle as they reflect (or twice the momentum perpendicular to the side) divided by the average time between collisions.

```python
# your code here
```


###### Content under Creative Commons Attribution license CC-BY 4.0, code under MIT license (c) 2019 Daniel Koehn, based on (c)2018 L.A. Barba, G.F. Forsyth [CFD Python](https://github.com/barbagroup/CFDPython#cfd-python), (c)2014 L.A. Barba, I. Hawke, B. Knaepen [Practical Numerical Methods with Python](https://github.com/numerical-mooc/numerical-mooc#practical-numerical-methods-with-python), also under CC-BY.

```python
# Execute this cell to load the notebook's style sheet, then ignore it
from IPython.core.display import HTML
css_file = '../style/custom.css'
HTML(open(css_file, "r").read())
```

<link href="https://fonts.googleapis.com/css?family=Merriweather:300,300i,400,400i,700,700i,900,900i" rel='stylesheet' >
<link href="https://fonts.googleapis.com/css?family=Source+Sans+Pro:300,300i,400,400i,700,700i" rel='stylesheet' >
<link href='http://fonts.googleapis.com/css?family=Source+Code+Pro:300,400' rel='stylesheet' >
<style>

@font-face {
    font-family: "Computer Modern";
    src: url('http://mirrors.ctan.org/fonts/cm-unicode/fonts/otf/cmunss.otf');
}

#notebook_panel { /* main background */
    background: rgb(245,245,245);
}

div.cell { /* set cell width */
    width: 800px;
}

div #notebook { /* centre the content */
    background: #fff; /* white background for content */
    width: 1000px;
    margin: auto;
    padding-left: 0em;
}

#notebook li { /* More space between bullet points */
margin-top:0.5em;
}

/* draw border around running cells */
div.cell.border-box-sizing.code_cell.running {
    border: 1px solid #111;
}

/* Put a solid color box around each cell and its output, visually linking them*/
div.cell.code_cell {
    background-color: rgb(256,256,256);
    border-radius: 0px;
    padding: 0.5em;
    margin-left:1em;
    margin-top: 1em;
}

div.text_cell_render{
    font-family: 'Source Sans Pro', sans-serif;
    line-height: 140%;
    font-size: 110%;
    width:680px;
    margin-left:auto;
    margin-right:auto;
}

/* Formatting for header cells */
.text_cell_render h1 {
    font-family: 'Merriweather', serif;
    font-style:regular;
    font-weight: bold;
    font-size: 250%;
    line-height: 100%;
    color: #004065;
    margin-bottom: 1em;
    margin-top: 0.5em;
    display: block;
}
.text_cell_render h2 {
    font-family: 'Merriweather', serif;
    font-weight: bold;
    font-size: 180%;
    line-height: 100%;
    color: #0096d6;
    margin-bottom: 0.5em;
    margin-top: 0.5em;
    display: block;
}

.text_cell_render h3 {
    font-family: 'Merriweather', serif;
	font-size: 150%;
    margin-top:12px;
    margin-bottom: 3px;
    font-style: regular;
    color: #008367;
}

.text_cell_render h4 {    /*Use this for captions*/
    font-family: 'Merriweather', serif;
    font-weight: 300;
    font-size: 100%;
    line-height: 120%;
    text-align: left;
    width:500px;
    margin-top: 1em;
    margin-bottom: 2em;
    margin-left: 80pt;
    font-style: regular;
}

.text_cell_render h5 {  /*Use this for small titles*/
    font-family: 'Source Sans Pro', sans-serif;
    font-weight: regular;
    font-size: 130%;
    color: #e31937;
    font-style: italic;
    margin-bottom: .5em;
    margin-top: 1em;
    display: block;
}

.text_cell_render h6 { /*use this for copyright note*/
    font-family: 'Source Code Pro', sans-serif;
    font-weight: 300;
    font-size: 9pt;
    line-height: 100%;
    color: grey;
    margin-bottom: 1px;
    margin-top: 1px;
}

    .CodeMirror{
            font-family: "Source Code Pro";
			font-size: 90%;
    }
/*    .prompt{
        display: None;
    }*/

    .warning{
        color: rgb( 240, 20, 20 )
        }
</style>

# 2D Linear Advection-Diffusion Equation

After solving the 2D linear advection problem in the [last notebook](https://nbviewer.jupyter.org/github/daniel-koehn/Differential-equations-earth-system/blob/master/08_Convection_2D/02_2D_Linear_Advection.ipynb), we have to add the diffusion equation to the problem. The resulting 2D advection-diffusion equation can be solved by separating the diffusion from the advection problem using operator splitting and applying the Marker-in-Cell method to the advection problem, while the diffusion problem is solved by finite-differences.

## 2D Linear Advection-Diffusion Problem

The linear 2D heat advection-diffusion equation reads:

\begin{equation}
\frac{\partial T}{\partial t} + v_x \frac{\partial T}{\partial x} + v_z \frac{\partial T}{\partial z}= \kappa \biggl(\frac{\partial^2 T}{\partial x^2} + \frac{\partial^2 T}{\partial z^2}\biggl), \tag{1}
\end{equation}

where $T$ denotes the 2D temperature field, $v_x, v_z$ the components of the velocity field in x- and z-direction, respectively and $\kappa$ the thermal diffusivity.

We can split the problem into an advection and diffusion part using **operator splitting** and introducing forward FD-operators for the time derivatives

\begin{equation}
\begin{split}
\frac{T^{n+1/2}_{i,j}-T^{n}_{i,j}}{\Delta t/2} &= - v_x \frac{\partial T^n}{\partial x} - v_z \frac{\partial T^n}{\partial z}\\
\frac{T^{n+1}_{i,j}-T^{n+1/2}_{i,j}}{\Delta t/2} &= \kappa \biggl(\frac{\partial^2 T^{n+1/2}}{\partial x^2} + \frac{\partial^2 T^{n+1/2}}{\partial z^2}\biggl),
\end{split}
\notag
\end{equation}

Notice, that we first update the temperature field using the advection equation by half a time step $\Delta t /2$ and use the new temperature field $T^{n+1/2}$ at time step $n+1/2$ to solve the 2D diffusion equation.

As discussed in [the previous notebook](https://nbviewer.jupyter.org/github/daniel-koehn/Differential-equations-earth-system/blob/master/08_Convection_2D/02_2D_Linear_Advection.ipynb) the 2D advection problem can be solved by the **Marker-in-Cell method**. To solve the diffusion part, we can replace the second spatial deriviatives on the RHS of the diffusion problem using the central FD scheme, leading to

\begin{equation}
\frac{T^{n+1}_{i,j} - T^{n+1/2}_{i,j}}{\Delta t/2} = \kappa \left( \frac{T^{n+1/2}_{i+1, j} - 2T^{n+1/2}_{i,j} + T^{n+1/2}_{i-1,j}}{\Delta x^2} + \frac{T^{n+1/2}_{i, j+1} - 2T^{n+1/2}_{i,j} + T^{n+1/2}_{i,j-1}}{\Delta y^2}\right)\notag
\end{equation}

Rearranging the equation to solve for the value at the next time step $T^{n+1}_{i,j}$, yields the **Forward-in-Time-Central-in-Space (FTCS)** finite difference method

\begin{equation}
T^{n+1}_{i,j}= T^{n+1/2}_{i,j} + \kappa \left( \frac{\Delta t/2}{\Delta x^2} (T^{n+1/2}_{i+1, j} - 2T^{n+1/2}_{i,j} + T^{n+1/2}_{i-1,j}) + \\\frac{\Delta t/2}{\Delta y^2} (T^{n+1/2}_{i, j+1} - 2T^{n+1/2}_{i,j} + T^{n+1/2}_{i,j-1})\right)\tag{2}
\end{equation}

which we already studied in [lecture 6](https://github.com/daniel-koehn/Differential-equations-earth-system/tree/master/06_Heat_2D#differential-equations-in-earth-sciences-2d-heat-equation). Lets combine the Marker-in-Cell method to solve the advection problem with the FTCS method to solve the diffusion problem.

As usual, we start by importing some libraries:

```python
# Import libraries
import numpy
from matplotlib import pyplot
%matplotlib inline
from scipy import interpolate
```

```python
# Set the font family and size to use for Matplotlib figures.
pyplot.rcParams['font.family'] = 'serif'
pyplot.rcParams['font.size'] = 16
```

Next, we write a function `update_diff_2D` to update the 2D diffusion equation using the FTCS method eq. (2).

```python
def update_diff_2D(T, sigma_x, sigma_z, nx, nz):

    # store old temperature field
    Tn = T.copy()

    # loop over spatial grid
    for i in range(1,nx-1):
        for j in range(1,nz-1):

            T[j, i] = (Tn[j, i] +
                        sigma_x * (Tn[j, i+1] - 2.0 * Tn[j, i] + Tn[j, i-1]) +
                        sigma_z * (Tn[j+1, i] - 2.0 * Tn[j, i] + Tn[j-1, i]))

    return T
```

In order to combine the advection with the diffusion problem, we take our Marker-in-Cell code `Adv_2D` from the [previous notebook](https://nbviewer.jupyter.org/github/daniel-koehn/Differential-equations-earth-system/blob/master/08_Convection_2D/02_2D_Linear_Advection.ipynb), rename it to `Adv_diff_2D` and include the `update_diff_2D` function to incorporate the diffusion problem:

```python
# Marker-in-cell code to solve the 2D advection equation
# ------------------------------------------------------
def Adv_diff_2D(T0, nt, dt, dx, dz, Vx, Vz, Lx1, Lx2, Lz1, Lz2, x, z, X, Z, alpha):
    """
    Computes and returns the temperature distribution
    after a given number of time steps for the 2D advection
    problem. A marker-in-cell approach with Dirichlet conditions
    on all boundaries is used in order to mitigate the effect of
    numerical diffusion.

    Parameters
    ----------
    T0 : numpy.ndarray
        The initial temperature distribution as a 2D array of floats.
    nt : integer
        Maximum number of time steps to compute.
    dt : float
        Time-step size.
    dx : float
        Grid spacing in the x direction.
    dz : float
        Grid spacing in the z direction.
    Vx : float
        x-component of the velocity field.
    Vz : float
        z-component of the velocity field.
    Lx1, Lx2 : float
        Model extension from Lx1 - Lx2.
    Lz1, Lz2 : float
        Model extension from Lz1 - Lz2.
    x, z : float
        Model coordinates as 1D arrays.
    X, Z : float
        Model coordinates as 2D arrays.

    Returns
    -------
    T : numpy.ndarray
        The temperature distribution as a 2D array of floats.
    """

    # Integrate in time.
    T = T0.copy()

    # Estimate number of grid points in x- and z-direction
    nz, nx = T.shape

    # Define number of markers and initial marker positions
    nx_mark = 4 * nx  # number of markers in x-direction
    nz_mark = 4 * nz  # number of markers in z-direction

    # Define some constants.
    sigma_x = alpha * dt / dx**2
    sigma_z = alpha * dt / dz**2

    # Time loop
    for n in range(nt):

        # initial marker positions
        x_mark = numpy.linspace(Lx1, Lx2, num=nx_mark)
        z_mark = numpy.linspace(Lz1, Lz2, num=nz_mark)
        X_mark, Z_mark = numpy.meshgrid(x_mark,z_mark)

        # Interpolate velocities from grid to marker position at timestep n
        f = interpolate.interp2d(x, z, Vx, kind='linear')
        vx_mark_n = f(x_mark, z_mark)

        f = interpolate.interp2d(x, z, Vz, kind='linear')
        vz_mark_n = f(x_mark, z_mark)

        # Interpolate temperature from grid to marker position at timestep n
        f = interpolate.interp2d(x, z, T, kind='cubic')
        T_mark = f(x_mark, z_mark)

        # Save current marker positions
        X0 = X_mark
        Z0 = Z_mark

        # Update marker position
        X_mark = X_mark + vx_mark_n * dt
        Z_mark = Z_mark + vz_mark_n * dt

        # Interpolate velocities from grid to marker position at timestep n+1
        vx_mark_n1 = interpolate.griddata((X.flatten(), Z.flatten()), Vx.flatten(), (X_mark, Z_mark), method='linear')
        vz_mark_n1 = interpolate.griddata((X.flatten(), Z.flatten()), Vz.flatten(), (X_mark, Z_mark), method='linear')

        # Replace Nan values
        mask = numpy.where(numpy.isnan(vx_mark_n1))
        vx_mark_n1[mask] = 0

        mask = numpy.where(numpy.isnan(vz_mark_n1))
        vz_mark_n1[mask] = 0

        # Update marker position with midpoint velocity
        X_mark = X0 + dt * (vx_mark_n + vx_mark_n1) / 2.
        Z_mark = Z0 + dt * (vz_mark_n + vz_mark_n1) / 2.

        # Interpolate temperature field from marker to grid positions
        T = interpolate.griddata((X_mark.flatten(), Z_mark.flatten()), T_mark.flatten(), (X, Z), method='cubic')

        # Replace Nan-values by old temperature field
        mask = numpy.where(numpy.isnan(T))
        T[mask] = T0[mask]

        # Update diffusion equation
        T = update_diff_2D(T, sigma_x, sigma_z, nx, nz)

    return T
```

Well, and here we have our 2D advection-diffusion code. That was quite easy to implement. Time to test, if the new code is working correctly.

## Temperature Advection-Diffusion in boring velocity field $V = (Z,-X)$

After implementing the code to solve the 2D linear advection-diffusion problem, let 's define some modelling parameters and initial spatial Gaussian temperature distribution

\begin{equation}
T_0 = A_0\; exp\biggl(-\frac{(x+x_{\text{shift}})^2+z^2}{\sigma^2}\biggr) \notag
\end{equation}

on the domain $x \in \{-0.5,0.5\},\; z \in \{-0.5,0.5\}$ with amplitude $A_0 = 2$, shift in x-direction $x_{\text{shift}} = 0.25$ and half-width $\sigma = 0.1$.

To advect the temperature, a velocity field is required. We start with the "boring" velocity field $V = (Z,-X)$ from the [How to sail without wind exercise](https://nbviewer.jupyter.org/github/daniel-koehn/Differential-equations-earth-system/blob/master/02_finite_difference_intro/3_fd_ODE_example_sailing_wo_wind.ipynb):

```python
# Definition of modelling parameters
# ----------------------------------
Lx = 0.5  # half-length of the model in x direction [m]
Lz = 0.5  # half-height of the model in z direction [m]
nx = 41   # number of points in the x direction
nz = 41   # number of points in the z direction
dx = 2 * Lx / (nx - 1)  # grid spacing in the x direction
dz = 2 * Lz / (nz - 1)  # grid spacing in the z direction
alpha = 1e-3 # thermal diffusivity

# Define the locations along a gridline.
x = numpy.linspace(-Lx, Lx, num=nx)
z = numpy.linspace(-Lz, Lz, num=nz)

# Define the initial temperature distribution as Gaussian function
X, Z = numpy.meshgrid(x,z) # coordinates X,Z required to define T0

sigma = 0.1
amp = 2.
xshift = 0.25

# initial temperature distribution
T0 = amp * numpy.exp(-((X + xshift)**2 + Z**2) / sigma**2.)

# Define velocity field (Vx,Vz)
Vx = Z
Vz = - X
```

Let's check if the initial temperature distribution and velocity field are correctly implemented.

```python
# Plot the initial temperature distribution and streamlines
pyplot.figure(figsize=(10.0, 7))

extent = [numpy.min(x), numpy.max(x),numpy.min(z), numpy.max(z)]
cmap = 'gist_heat'
im = pyplot.imshow(numpy.flipud(T0), extent=extent, interpolation='spline36', cmap=cmap)
stream = pyplot.streamplot(X,Z,Vx,Vz,color='w')

pyplot.xlabel('x [m]')
pyplot.ylabel('z [m]')
cbar = pyplot.colorbar(im)
cbar.set_label('Initial Temperature [°C]')
```

Looks correct, so finally we define the timestep `dt` and the number of time steps `nt` we want to advect the temperature field and run the `Adv_diff_2D` code

```python
# Set the time-step size
dt = 2.5e-2  # time-step size [s]
nt = 50      # number of time steps to compute

# Compute the temperature distribution after nt timesteps
T = Adv_diff_2D(T0, nt, dt, dx, dz, Vx, Vz, -Lx, Lx, -Lz, Lz, x, z, X, Z,alpha)
```

```python
# Plot the temperature distribution and streamlines
pyplot.figure(figsize=(10.0, 7))

extent = [numpy.min(x), numpy.max(x),numpy.min(z), numpy.max(z)]
cmap = 'gist_heat'
im = pyplot.imshow(numpy.flipud(T), extent=extent, interpolation='spline36',
                   cmap=cmap, aspect=1, vmin=numpy.min(T0), vmax=numpy.max(T0))
stream = pyplot.streamplot(X,Z,Vx,Vz,color='w')

pyplot.xlabel('x [m]')
pyplot.ylabel('z [m]')
cbar = pyplot.colorbar(im)
cbar.set_label('Temperature [°C]')
```

Notice how the temperature anomaly is advected in the velocity field, while heat is transported by conduction (diffusion), leading to a spreading of the temperature anomaly in all directions and a decrease of it 's amplitude. This becomes more obvious when visualizing the 2D advection-diffusion in the following movie, where the number of time steps is increased to `nt=400`:

```python
from IPython.display import YouTubeVideo
YouTubeVideo('yNT9-ht2HRE')
```

## What we learned:

* How to extent the 2D Marker-in-Cell advection code in order to solve the advection-diffusion problem

* Computed 2D advection-diffusion of a Gaussian temperature distribution in a simple velocity field, where the anomaly is advected in the velocity field and heat transported by heat conduction (diffusion)