# Advanced: Control your flow

Here we dive a bit deeper in advanced flow control in Noodles. Starting with a recap into for-loops, moving on to conditional evaluation of workflows and standard algorithms. This chapter will  also go a bit deeper into the teritory of functional programming. Specifically, we will see how to program sequential loops using only functions and recursion.

If you are new to the concepts of *recursion*, here is some nice material to start with:

[](https://www.youtube.com/watch?v=Qk0zUZW-U_M)

## Recap: for loops
In the *Translating Poetry* tutorial we saw how we could create parallel `for` loops in Noodles. To recap, let's reverse the words in a sentence. Assume you have the following for-loop in Python:

```python
sentence = 'the quick brown fox jumps over the lazy dog'
reverse = []

def reverse_word(word):
    return word[::-1]

for word in sentence.split():
    reverse.append(reverse_word(word))

result = ' '.join(reverse)
print(result)
```

    eht kciuq nworb xof spmuj revo eht yzal god

There is a pattern to this code that is better written as:

```python
reverse = [reverse_word(word) for word in sentence.split()]
result = ' '.join(reverse)
print(result)
```

    eht kciuq nworb xof spmuj revo eht yzal god

This last version can be translated to Noodles. Assume for some reason we want to schedule the `reverse_word` function (it takes forever to run on a single core!). Because `reverse_words` becomes a *promise*, the line with `' '.join(reverse)` also has to be captured in a scheduled function.

```python
import noodles

@noodles.schedule
def reverse_word(word):
    return word[::-1]

@noodles.schedule
def make_sentence(words):
    return ' '.join(words)

reverse_words = noodles.gather_all(
    reverse_word(word) for word in sentence.split())
workflow = make_sentence(reverse_words)
```

```python
from noodles.tutorial import display_workflows
noodles.tutorial.display_workflows(prefix='control', quick_brown_fox=workflow)
```

| quick-brown-fox |
| --- |
|  |

This example shows how we can do loops in parallel. There are cases where we will need to do loops in a serialised manner. For example, if we are handling a very large data set and all of the computation does not fit in memory when done in parallel.

There are hybrid *divide and conquer* approaches that can be implemented in Noodles. We then chunk all the work in blocks that can be executed in parallel, and stop when the first chunk gives us reason to. Divide-and-conquer can be implemented using a combination of the two looping strategies (parallel and sequential).

Sequential loops are made using recursion techniques.

## Recursion

Sequential loops can be made in Noodles using recursion. Comes the obligatory factorial function example:

```python
from noodles.tutorial import display_text

def factorial(x):
    if x == 0:
        return 1
    else:
        return factorial(x - 1) * x

display_text('100! = {}'.format(factorial(100)))
```

100! = 933262154439441526816992388562667004907159682643816214685929638 … 0000000

There is a problem with such a recursive algorithm when numbers get too high.

```python
try:
    display_text('10000! =', factorial(10000))
except RecursionError as e:
    display_text(e)
```

maximum recursion depth exceeded in comparison

Yikes! Let's head on. And translate the program to Noodles. Suppose we make `factorial` a scheduled function, we cannot multiply a *promise* with a *number* just like that (at least not in the current version of Noodles). We change the function slightly with a second argument that keeps count. This also makes the `factorial` function *tail-recursive*.

```python
@noodles.schedule
def factorial(x, acc=1):
    if x == 0:
        return acc
    else:
        return factorial(x - 1, acc * x)

result = noodles.run_single(factorial(10000))

display_text('10000! = {}'.format(result))
```

10000! = 2846259680917054518906413212119868890148051401702799230794179 … 0000000

Yeah! Noodles runs the tail-recursive function iteratively! This is actually **very important**. We'll do a little experiment. Start your system monitor (plotting a graph of your memory usage) and run the following snippets. We let every function call to factorial gobble up some memory and to be able to measure the effect of that we insert a small sleep. Fair warning: With the current setting of `gobble_size` and running 50 loops, the first version will take about **4GB** of memory. Just change the size so that a measurable fraction of your RAM is taken up by the process and you can see the result.

```python
import numpy
import time
from copy import copy

gobble_size = 10000000
```

```python
@noodles.schedule(call_by_ref=['gobble'])
def mul(x, y, gobble):
    return x*y

@noodles.schedule(call_by_ref=['gobble'])
def factorial(x, gobble):
    time.sleep(0.1)
    if x == 0:
        return 1
    else:
        return mul(factorial(x - 1, copy(gobble)), x, gobble)

gobble = numpy.zeros(gobble_size)
result = noodles.run_single(factorial(50, gobble))
```

We passed the `gobble` argument by reference. This prevents Noodles from copying the array when creating the workflow. If you have functions that take large arrays as input *and you don't change the value of the array in between calls* this is a sensible thing to do.
On my machine, running only 10 loops, this gives the following result:

Try to understand why this happens. We have reserved a NumPy array with `gobble_size` ($10^7$) floating points of 8 bytes each. The total size in bytes of this array is $8 \times 10^7\ MB$. In each recursive call to `factorial` the array is copied, so in total this will use $10 \cdot 8 \times 10^7\ MB = 800\ MB$ of memory!

The next version is **tail-recursive**. This should barely make a dent in your memory usage!

```python
@noodles.schedule(call_by_ref=['gobble'])
def factorial_tr(x, acc=1, gobble=None):
    time.sleep(0.1)
    if x == 0:
        return acc
    else:
        return factorial_tr(x - 1, mul(acc, x, gobble), copy(gobble))

gobble = numpy.zeros(gobble_size)
result = noodles.run_single(factorial_tr(50, gobble=gobble))
```

Now, the `factorial` function is still recursive. However, since returning a call to the `factorial` function is last thing we do, the intermediate results can be safely thrown away. We'll have in memory the original reference to `gobble` and one version in the Noodles run-time for the last time `factorial` returned a workflow where `gobble.copy()` was one of the arguments.
In total this gives a memory consumption of $160\ MB$ (plus a little extra for the Python run-time itself). We see peeks that reach over $250\ MB$ in the graph: this is where `gobble` is being copied, after which the garbage collector deletes the old array.

Try to understand why this happens. In the first case the function returns a new workflow to be evaluated. This workflow has two nodes:

```python
display_workflows(
    prefix='control',
    factorial_one=noodles.unwrap(factorial)(10, '<memory gobble>'))
```

| factorial-one |
| --- |
|  |

To evaluate this workflow, Noodles first runs the top node `factorial(9, '<memory gobble>')`. When the answer for this function is obtained it is inserted into the slot for `mul(-, 10)`. Until the entire workflow is evaluated, the `<memory gobble>` remains in memory. Before this happens the factorial function is called which copies the gobble and creates a new workflow! We can write this out by expanding our algorithm symbolically $f(x) = x \cdot f(x-1)$:

$$\begin{align}
f(10) &= 10 \cdot f(9)\\
&= 10 \cdot (9 \cdot f(8))\\
&= 10 \cdot (9 \cdot (8 \cdot f(7)))\\
&\dots\\
&= 10 \cdot (9 \cdot (8 \cdot (7 \cdot (6 \cdot (5 \cdot (4 \cdot (3 \cdot (2 \cdot 1))))))))\\
&= 10 \cdot (9 \cdot (8 \cdot (7 \cdot (6 \cdot (5 \cdot (4 \cdot (3 \cdot 2)))))))\\
&= 10 \cdot (9 \cdot (8 \cdot (7 \cdot (6 \cdot (5 \cdot (4 \cdot 6))))))\\
&\dots
\end{align}$$

Now for the tail-recursive version, the workflow looks a bit different:

```python
display_workflows(
    prefix='control',
    tail_recursive_factorial=noodles.unwrap(factorial_tr)(10, gobble='<memory gobble>'))
```

| tail-recursive-factorial |
| --- |
|  |

First the `mul(1, 10, '<memory gobble>')` is evaluated. Its result is inserted into the empty slot in the call to `factorial_tr`. This call returns a new workflow with a new copy of `<memory gobble>`. This time however, the old workflow can be safely deleted. Again, it helps to look at the algorithm symbolically, given
$f(x, a) = f(x-1, x \cdot a)$:

$$\begin{align}
f(10, 1) &= f(9, (10 \cdot 1))\\
&= f(9, 10)\\
&= f(8, (9 \cdot 10))\\
&= f(8, 90)\\
&= f(7, (8 \cdot 90))\\
&\dots
\end{align}$$

## Conditional evaluation

But Python has more statements for flow control! The conditional execution of code is regulated through the `if` statement. You may want to make the exection of parts of your workflow conditional based on intermediate results. One such instance may look like this:

```python
@noodles.schedule
def method_one(x):
    pass

@noodles.schedule
def method_two(x):
    pass

@noodles.schedule
def what_to_do(x):
    if condition(x):
        return method_one(x)
    else:
        return method_two(x)
```

We've put the `if`-statement inside the scheduled function `what_to_do`. This returns a new workflow depending on the value of `x`. We can no longer get a nice single graph picture of the workflow, because the workflow doesn't exist! (there is no spoon ...) We can work through a small example from the Python tutorial: computing prime numbers.

```python
for n in range(2, 10):
    for x in range(2, n):
        if n % x == 0:
            print(n, 'equals', x, '*', n//x)
            break
    else:
        # loop fell through without finding a factor
        print(n, 'is a prime number')
```

    2 is a prime number
    3 is a prime number
    4 equals 2 * 2
    5 is a prime number
    6 equals 2 * 3
    7 is a prime number
    8 equals 2 * 4
    9 equals 3 * 3

The core computation in this example is the `n % x == 0` bit. So we start by creating a scheduled function that does that.

```python
@noodles.schedule
def divides(n, x):
    return n % x == 0
```

Noodles can parallelize the inner loop, but this gives a problem: how do we know when to stop? There is no way to get it both ways.

First, we'll see how to do the parallel solution. We'll compute the `divides(n, x)` function for the values of `n` and `x` and then filter out those where `divides` gave `False`. This last step is done using the `compress` function.

```python
@noodles.schedule
def compress(lst):
    """Takes a list of pairs, returns a list of
    first elements of those pairs for which the
    second element is thruthy."""
    return [a for a, b in lst if b]
```

Using the `compress` function we can write the Noodlified parallel version of the `filter` function. We'll call it `p_filter` for *parallel filter*.

```python
?filter
```

    Init signature: filter(self, /, *args, **kwargs)
    Docstring:
    filter(function or None, iterable) --> filter object

    Return an iterator yielding those items of iterable for which function(item)
    is true. If function is None, return the items that are true.
    Type:           type
    Subclasses:

Using the generic `p_filter` function we then write the function `find_factors` that finds all integer factors of a number in parallel. Both `p_filter` and `find_factors` won't be scheduled functions. Rather, together they build the workflow that solves our problem.

```python
def p_filter(f, lst):
    return compress(noodles.gather_all(
        noodles.gather(x, f(x)) for x in lst))

def find_factors(n):
    return p_filter(lambda x: divides(n, x), range(2, n))
```

```python
display_workflows(prefix='control', factors=find_factors(5))
```

| factors |
| --- |
|  |

No we can run this workflow for all the numbers we like.

```python
result = noodles.run_parallel(
    noodles.gather_all(noodles.gather(n, find_factors(n))
                       for n in range(2, 10)),
    n_threads=4)

for n, factors in result:
    if factors:
        print(n, 'equals', ', '.join(
            '{}*{}'.format(x, n//x) for x in factors))
    else:
        print(n, 'is prime')
```

    2 is prime
    3 is prime
    4 equals 2*2
    5 is prime
    6 equals 2*3, 3*2
    7 is prime
    8 equals 2*4, 4*2
    9 equals 3*3

Few! We managed, but if all we wanted to do is find primes, we did way too much work; we also found all factors of the numbers. We had to write some boiler plate code. Argh, this tutorial was supposed to be on flow control! We move on to the sequential version. Wait, I hear you think, we were using Noodles to do things in parallel!?? Why make an effort to do sequential work? Well, we'll need it to implement the divide-and-conquer strategy, among other things. Noodles is not only a framework for parallel programming, but it also works concurrent. In the context of a larger workflow we may still want to make decision steps on a sequential basis, while another component of the workflow is happily churning out numbers.

## Find-first

Previously we saw the definition of a Noodlified `filter` function. How can we write a `find_first` that stops after finding a first match? If we look at the workflow that `p_filter` produces, we see that all predicates are already present in the workflow and will be computed concurrently. We now write a sequential version. We may achieve sequential looping through recursion like this:

```python
def find_first(f, lst):
    if not lst:
        return None
    elif f(lst[0]):
        return lst[0]
    else:
        return find_first(f, lst[1:])
```

However, if `f` is a scheduled function `f(lst[0])` will give a promise, and this routine will fail.

```python
@noodles.schedule
def find_first_helper(f, lst, first):
    if first:
        return lst[0]
    elif len(lst) == 1:
        return None
    else:
        return find_first_helper(f, lst[1:], f(lst[1]))

def find_first(f, lst):
    return find_first_helper(f, lst, f(lst[0]))
```

```python
noodles.run_single(find_first(lambda x: divides(77, x), range(2, 63)))
```

    7

That works. Now suppose the input list is somewhat harder to compute; every element is the result of a workflow.

## Appendix: creating memory profile plots

```python
%%writefile test-tail-recursion.py
import numpy
import noodles
import time
from copy import copy

@noodles.schedule(call_by_ref=['gobble'])
def factorial_tr(x, acc=1, gobble=None):
    time.sleep(0.1)
    if x == 0:
        return acc
    else:
        return factorial_tr(x - 1, acc * x, copy(gobble))

gobble_size = 10000000
gobble = numpy.zeros(gobble_size)
result = noodles.run_single(factorial_tr(10, gobble=gobble))
```

    Writing test-tail-recursion.py

```python
%%writefile test-recursion.py
import numpy
import noodles
import time
from copy import copy

@noodles.schedule(call_by_ref=['gobble'])
def mul(x, y, gobble):
    return x*y

@noodles.schedule(call_by_ref=['gobble'])
def factorial(x, gobble):
    time.sleep(0.1)
    if numpy.all(x == 0):
        return numpy.ones_like(x)
    else:
        return mul(factorial(x - 1, copy(gobble)), x, gobble)

gobble_size = 10000000
gobble = numpy.zeros(gobble_size)
result = noodles.run_single(factorial(10, gobble))
```

    Writing test-recursion.py

```python
!pip install matplotlib
!pip install memory_profiler
```

    Requirement already satisfied: matplotlib in /home/johannes/.local/share/workon/threepop/lib/python3.7/site-packages (3.1.1)
    Requirement already satisfied: python-dateutil>=2.1 in /home/johannes/.local/share/workon/threepop/lib/python3.7/site-packages (from matplotlib) (2.8.0)
    Requirement already satisfied: cycler>=0.10 in /home/johannes/.local/share/workon/threepop/lib/python3.7/site-packages (from matplotlib) (0.10.0)
    Requirement already satisfied: kiwisolver>=1.0.1 in /home/johannes/.local/share/workon/threepop/lib/python3.7/site-packages (from matplotlib) (1.1.0)
    Requirement already satisfied: numpy>=1.11 in /home/johannes/.local/share/workon/threepop/lib/python3.7/site-packages (from matplotlib) (1.17.0)
    Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /home/johannes/.local/share/workon/threepop/lib/python3.7/site-packages (from matplotlib) (2.4.2)
    Requirement already satisfied: six>=1.5 in /home/johannes/.local/share/workon/threepop/lib/python3.7/site-packages (from python-dateutil>=2.1->matplotlib) (1.12.0)
    Requirement already satisfied: setuptools in /home/johannes/.local/share/workon/threepop/lib/python3.7/site-packages (from kiwisolver>=1.0.1->matplotlib) (41.2.0)
    Collecting memory_profiler
      Using cached https://files.pythonhosted.org/packages/9f/fe/1fca7273dd111108f204a686b12a12b6422d405fe4614087aa7d5a66ea87/memory_profiler-0.55.0.tar.gz
    Collecting psutil (from memory_profiler)
      Using cached https://files.pythonhosted.org/packages/1c/ca/5b8c1fe032a458c2c4bcbe509d1401dca9dda35c7fc46b36bb81c2834740/psutil-5.6.3.tar.gz
    Installing collected packages: psutil, memory-profiler
      Running setup.py install for psutil ... done
      Running setup.py install for memory-profiler ... done
    Successfully installed memory-profiler-0.55.0 psutil-5.6.3

```bash
%%bash
rm -f mprofile_*.dat
mprof run -T 0.001 python ./test-tail-recursion.py
mprof run -T 0.001 python ./test-recursion.py
```

    mprof: Sampling memory every 0.001s
    mprof: Sampling memory every 0.001s

```python
from pathlib import Path
from matplotlib import pyplot as plt

plt.rcParams['font.family'] = 'serif'

def read_mprof(filename):
    lines = list(open(filename, 'r'))
    cmd = filter(lambda l: l[:3] == 'CMD', lines)
    mem = filter(lambda l: l[:3] == 'MEM', lines)
    data = numpy.array([list(map(float, l.split()[1:])) for l in mem])
    data[:,1] -= data[0,1]
    data[:,0] *= 1024**2
    return cmd, data

def plot_mprof(filename):
    cmd, data = read_mprof(filename)
    if 'tail' in next(cmd):
        figname = 'tail-recursion'
    else:
        figname = 'recursion'

    plt.plot(data[:,1], data[:,0] / 1e6)
    plt.xlabel('time (s)')
    plt.ylabel('memory usage (MB)')
    plt.title(figname)
    plt.savefig('control-' + figname + '-raw.svg', bbox_inches='tight')
    plt.show()

files = list(Path('.').glob('mprofile_*.dat'))
for f in files:
    plot_mprof(f)
    plt.close()
```


# Regression model—biophony

This notebook explores and models the data collected from recordings of the natural acoustic environment over the urban-rural gradient near Innsbruck, Austria. The models are implemented as Bayesian models with the PyMC3 probabilistic programming library.

References:

https://github.com/fonnesbeck/multilevel_modeling

Gelman, A., & Hill, J. (2006). Data Analysis Using Regression and Multilevel/Hierarchical Models (1st ed.). Cambridge University Press.

#### Import statements

```python
import warnings
warnings.filterwarnings('ignore')
```

```python
import pandas
import numpy
from os import path
```

```python
%matplotlib inline
from matplotlib import pyplot
from matplotlib.patches import Rectangle
import seaborn
import mpld3
from mpld3 import plugins
```

```python
from pymc3 import glm, Model, NUTS, sample, stats, \
                  forestplot, traceplot, plot_posterior, summary, \
                  Normal, Uniform, Deterministic, StudentT
from pymc3.backends import SQLite
```

#### Plot settings

```python
from matplotlib import rcParams
```

```python
rcParams['font.sans-serif']
```

    ['DejaVu Sans',
     'Bitstream Vera Sans',
     'Computer Modern Sans Serif',
     'Lucida Grande',
     'Verdana',
     'Geneva',
     'Lucid',
     'Arial',
     'Helvetica',
     'Avant Garde',
     'sans-serif']

```python
rcParams['font.sans-serif'] = ['Helvetica',
                               'Arial',
                               'Bitstream Vera Sans',
                               'DejaVu Sans',
                               'Lucida Grande',
                               'Verdana',
                               'Geneva',
                               'Lucid',
                               'Avant Garde',
                               'sans-serif']
```

#### Variable definitions

```python
data_filepath = "/Users/Jake/OneDrive/Documents/alpine soundscapes/data/dataset.csv"
```

```python
trace_output_path = "/Users/Jake/OneDrive/Documents/alpine soundscapes/data/model traces/biophony"
```

```python
seaborn_blue = seaborn.color_palette()[0]
```

## Load data

```python
data = pandas.read_csv(data_filepath)
data = data.loc[data.site<=30]
```

sort data by site and then by visit

```python
data_sorted = data.sort_values(by=['site', 'sound']).reset_index(drop=True)
```

transform variables (mean center)

```python
column_list = ['sel', 'sel_anthrophony', 'sel_biophony', 'biophony', 'week',
    'building_50m', 'pavement_50m', 'forest_50m', 'field_50m',
    'building_100m', 'pavement_100m', 'forest_100m', 'field_100m',
    'building_200m', 'pavement_200m', 'forest_200m', 'field_200m',
    'building_500m', 'pavement_500m', 'forest_500m', 'field_500m',
    'd2n_50m', 'd2n_100m', 'd2n_200m', 'd2n_500m',
    'temperature', 'wind_speed', 'pressure', 'bus_stop',
    'construction', 'crossing', 'cycleway', 'elevator', 'escape', 'footway',
    'living_street', 'motorway', 'motorway_link', 'path', 'pedestrian',
    'platform', 'primary_road', 'primary_link', 'proposed', 'residential',
    'rest_area', 'secondary', 'secondary_link', 'service', 'services',
    'steps', 'tertiary', 'tertiary_link', 'track', 'unclassified', 'combo']

data_centered = data_sorted.copy()
for column in column_list:
    data_centered[column] = data_sorted[column] - data_sorted[column].mean()

# copy land_composite values
#land_composites = ['land_composite_50m', 'land_composite_100m', 'land_composite_200m', 'land_composite_500m']
#for column in land_composites:
#    data_centered[column] = data_sorted[column]
```

create sites variable for PyMC3 models

```python
sites = numpy.copy(data_sorted.site.values) - 1
```

## Model 0 - emtpy model

$$
\begin{align}
y_{ts} \sim \mathcal{N}(\alpha_s + \epsilon_t, \sigma_y^2) \\
\alpha_s \sim \mathcal{N}(M + \epsilon_s, \sigma_\alpha^2) \\
\end{align}
$$

```python
with Model() as model0:

    # Priors
    mu_grand = Normal('mu_grand', mu=0., tau=0.0001)
    sigma_a = Uniform('sigma_a', lower=0, upper=100)
    tau_a = sigma_a**-2

    # Random intercepts
    a = Normal('a', mu=mu_grand, tau=tau_a, shape=len(set(sites)))

    # Model error
    sigma_y = Uniform('sigma_y', lower=0, upper=100)
    tau_y = sigma_y**-2

    # Expected value
    y_hat = a[sites]

    # Data likelihood
    y_like = Normal('y_like', mu=y_hat, tau=tau_y, observed=data_centered.biophony)

    # sample model
    backend = SQLite(path.join(trace_output_path, "model0.sqlite"))
    model0_samples = sample(draws=10000, step=NUTS(), random_seed=1, trace=backend)
```

    100%|██████████| 10500/10500 [00:19<00:00, 547.60it/s]

```python
fig, ax = pyplot.subplots()

# organize results
model0_data = pandas.DataFrame({'site': data_sorted.site.unique(),
                                 'site_name': data_sorted.site_name.unique()})
model0_data['d2n_200m'] = data_sorted.d2n_200m.unique()
model0_data['quantiles'] = [stats.quantiles(model0_samples.a[:5000, i]) for i in range(len(set(sites)))]

# plot quantiles
for i, row in model0_data.sort_values(by='d2n_200m').iterrows():
    x = row['d2n_200m']
    ax.plot([x, x], [row['quantiles'][2.5], row['quantiles'][97.5]], color='black', linewidth=0.5)
    ax.plot([x, x], [row['quantiles'][25], row['quantiles'][75]], color='black', linewidth=1)
    ax.scatter([x], [row['quantiles'][50]], color='black', marker='o')

# format plot
l1 = ax.set_xlim([-0.1, 0.75])
xl = ax.set_xlabel("d2n within 200 meters (percent area)")
yl = ax.set_ylabel("percent biophony (difference from grand mean)")
```

## Model 1—time and site predictors

$$
\begin{align}
\text{level 1} \\
y_{ts} \sim \mathcal{N}(\alpha_s + \beta_s T_t, \sigma_y^2) \\
\text{level 2} \\
\alpha_s \sim \mathcal{N}(\gamma_\alpha + \gamma_{\alpha s} L_s, \sigma_\alpha^2) \\
\beta_s \sim \mathcal{N}(\gamma_\beta + \gamma_{\beta s} L_s, \sigma_\beta^2) \\
\end{align}
$$

```python
site_predictors = [
#    'building_50m', 'pavement_50m', 'forest_50m', 'field_50m',
#    'building_100m', 'pavement_100m', 'forest_100m', 'field_100m',
#    'building_200m', 'pavement_200m', 'forest_200m', 'field_200m',
#    'building_500m', 'pavement_500m', 'forest_500m', 'field_500m',
    'd2n_50m',
    'd2n_100m',
    'd2n_200m', 'd2n_500m',
]
for predictor in site_predictors:
    with Model() as model_1:

        # intercept
        g_a = Normal('g_a', mu=0, tau=0.001)
        g_as = Normal('g_as', mu=0, tau=0.001)
        sigma_a = Uniform('sigma_a', lower=0, upper=100)
        tau_a = sigma_a**-2
        mu_a = g_a + (g_as * data_centered.groupby('site')[predictor].mean())
        a = Normal('a', mu=mu_a, tau=tau_a, shape=len(set(sites)))

        # slope
        g_b = Normal('g_b', mu=0, tau=0.001)
        g_bs = Normal('g_bs', mu=0, tau=0.001)
        sigma_b = Uniform('sigma_b', lower=0, upper=100)
        tau_b = sigma_b**-2
        mu_b = g_b + (g_bs * data_centered.groupby('site')[predictor].mean())
        b = Normal('b', mu=mu_b, tau=tau_b, shape=len(set(sites)))

        # model error (data-level)
        sigma_y = Uniform('sigma_y', lower=0, upper=100)
        tau_y = sigma_y**-2

        # expected values
        y_hat = a[sites] + (b[sites] * data_centered.week)

        # likelihood
        y_like = Normal('y_like', mu=y_hat, tau=tau_y, observed=data_centered.biophony)

        # simulated
        #y_sim = Normal('y_sim', mu=y_hat, tau=tau_y, shape=y_hat.tag.test_value.shape)

        # sample model
        backend = SQLite(path.join(trace_output_path, "model1_{}.sqlite".format(predictor)))
        model_1_samples = sample(draws=10000, step=NUTS(), random_seed=1, trace=backend)
```

    100%|██████████| 10500/10500 [01:06<00:00, 157.64it/s]
    100%|██████████| 10500/10500 [01:07<00:00, 156.01it/s]
    100%|██████████| 10500/10500 [07:31<00:00, 23.28it/s]
    100%|██████████| 10500/10500 [01:17<00:00, 134.66it/s]

```python
len(model_1_samples['a'][:5000][:,1])
```

    5000

```python
fig = pyplot.figure()

fig.set_figwidth(6.85)
fig.set_figheight(6.85/2)

ax_a = pyplot.subplot2grid((1, 2), (0, 0), rowspan=1, colspan=1)
ax_b = pyplot.subplot2grid((1, 2), (0, 1), rowspan=1, colspan=1, sharex=ax_a)

fig.subplots_adjust(left=0, bottom=0, right=1, top=1)

# organize results
model_1_data = pandas.DataFrame({'site': data_sorted.site.unique(),
                                 'site_name': data_sorted.site_name.unique()})
model_1_data['d2n_200m'] = data_sorted.d2n_200m.unique()
model_1_data['quantiles_a'] = [stats.quantiles(model_1_samples['a'][:5000][:, i]) for i in range(len(set(sites)))]
model_1_data['quantiles_b'] = [stats.quantiles(model_1_samples['b'][:5000][:, i]) for i in range(len(set(sites)))]

# plot quantiles
for i, row in model_1_data.sort_values(by='d2n_200m').iterrows():
    x = row['d2n_200m']
    #ax_a.plot([x, x], [row['quantiles_a'][2.5], row['quantiles_a'][97.5]], color='black', linewidth=0.5)
    ax_a.plot([x, x], [row['quantiles_a'][25], row['quantiles_a'][75]], color='black', linewidth=1)
    ax_a.scatter([x], [row['quantiles_a'][50]], color='black', marker='o')

# format plot
l1 = ax_a.set_xlim([0, 1])
xl = ax_a.set_xlabel("d2n within 200 meters (percent area)")
yl = ax_a.set_ylabel("percent biophony (difference from grand mean)")

# plot quantiles
for i, row in model_1_data.sort_values(by='d2n_200m').iterrows():
    x = row['d2n_200m']
    #ax_b.plot([x, x], [row['quantiles_b'][2.5], row['quantiles_b'][97.5]], color='black', linewidth=0.5)
    ax_b.plot([x, x], [row['quantiles_b'][25], row['quantiles_b'][75]], color='black', linewidth=1)
    ax_b.scatter([x], [row['quantiles_b'][50]], color='black', marker='o')

# format plot
l1 = ax_b.set_xlim([0, 1])
xl = ax_b.set_xlabel("d2n within 200 meters (percent area)")
yl = ax_b.set_ylabel("rate of change of biophony (percent/week)")
```

## Model 2—environmental predictors

$$
\begin{align}
\text{level 1} \\
y_{ts} \sim \mathcal{N}(\alpha_s + \beta_s T_t, \sigma_y^2) \\
\text{level 2} \\
\alpha_s \sim \mathcal{N}(\gamma_\alpha + \gamma_{\alpha s} L_s, \sigma_\alpha^2) \\
\beta_s \sim \mathcal{N}(\gamma_\beta + \gamma_{\beta s} L_s, \sigma_\beta^2) \\
\end{align}
$$

```python
measurement_predictors = [
    'temperature', 'wind_speed', 'precipitation', 'pressure',
]
for predictor in measurement_predictors:
    with Model() as model2a:

        # intercept
        g_a = Normal('g_a', mu=0, tau=0.001)
        g_as = Normal('g_as', mu=0, tau=0.001)
        sigma_a = Uniform('sigma_a', lower=0, upper=100)
        tau_a = sigma_a**-2
        mu_a = g_a + (g_as * data_centered.groupby('site')['forest_200m'].mean())
        a = Normal('a', mu=mu_a, tau=tau_a, shape=len(set(sites)))

        # time slope
        g_b = Normal('g_b', mu=0, tau=0.001)
        g_bs = Normal('g_bs', mu=0, tau=0.001)
        sigma_b = Uniform('sigma_b', lower=0, upper=100)
        tau_b = sigma_b**-2
        mu_b = g_b + (g_bs * data_centered.groupby('site')['forest_200m'].mean())
        b = Normal('b', mu=mu_b, tau=tau_b, shape=len(set(sites)))

        # temp slope
        #g_c = Normal('g_c', mu=0, tau=0.001)
        #g_cs = Normal('g_cs', mu=0, tau=0.001)
        #sigma_c = Uniform('sigma_c', lower=0, upper=100)
        #tau_c = sigma_c**-2
        #mu_c = g_c + (g_cs * data_centered.groupby('site')['forest_200m'].mean())
        #c = Normal('c', mu=mu_c, tau=tau_c, shape=len(set(sites)))
        c = Uniform('c', lower=-100, upper=100)

        # model error (data-level)
        sigma_y = Uniform('sigma_y', lower=0, upper=100)
        tau_y = sigma_y**-2

        # expected values
        y_hat = a[sites] + (b[sites] * data_centered.week) + (c * data_centered[predictor])

        # likelihood
        y_like = Normal('y_like', mu=y_hat, tau=tau_y, observed=data_centered.biophony)

        # simulated
        #y_sim = Normal('y_sim', mu=y_hat, tau=tau_y, shape=y_hat.tag.test_value.shape)

        # sample model
        backend = SQLite(path.join(trace_output_path, "model2a_{0}.sqlite".format(predictor)))
        model_2_samples = sample(draws=10000, step=NUTS(), random_seed=1, trace=backend)
```

```python
measurement_predictors = [
    'temperature', 'wind_speed', 'precipitation', 'pressure',
]
for predictor in measurement_predictors:
    with Model() as model2b:

        # intercept
        g_a = Normal('g_a', mu=0, tau=0.001)
        g_as = Normal('g_as', mu=0, tau=0.001)
        sigma_a = Uniform('sigma_a', lower=0, upper=100)
        tau_a = sigma_a**-2
        mu_a = g_a + (g_as * data_centered.groupby('site')['forest_200m'].mean())
        a = Normal('a', mu=mu_a, tau=tau_a, shape=len(set(sites)))

        # time slope
        g_b = Normal('g_b', mu=0, tau=0.001)
        g_bs = Normal('g_bs', mu=0, tau=0.001)
        sigma_b = Uniform('sigma_b', lower=0, upper=100)
        tau_b = sigma_b**-2
        mu_b = g_b + (g_bs * data_centered.groupby('site')['forest_200m'].mean())
        b = Normal('b', mu=mu_b, tau=tau_b, shape=len(set(sites)))

        # predictor slope
        g_c = Normal('g_c', mu=0, tau=0.001)
        g_cs = Normal('g_cs', mu=0, tau=0.001)
        sigma_c = Uniform('sigma_c', lower=0, upper=100)
        tau_c = sigma_c**-2
        mu_c = g_c + (g_cs * data_centered.groupby('site')['forest_200m'].mean())
        c = Normal('c', mu=mu_c, tau=tau_c, shape=len(set(sites)))

        # model error (data-level)
        sigma_y = Uniform('sigma_y', lower=0, upper=100)
        tau_y = sigma_y**-2

        # expected values
        y_hat = a[sites] + (b[sites] * data_centered.week) + (c[sites] * data_centered[predictor])

        # likelihood
        y_like = Normal('y_like', mu=y_hat, tau=tau_y, observed=data_centered.biophony)

        # simulated
        #y_sim = Normal('y_sim', mu=y_hat, tau=tau_y, shape=y_hat.tag.test_value.shape)

        # sample model
        backend = SQLite(path.join(trace_output_path, "model2b_{0}.sqlite".format(predictor)))
        model_2_samples = sample(draws=10000, step=NUTS(), random_seed=1, trace=backend)
```

```python
fig, ax = pyplot.subplots()

# organize results
model_2_data = pandas.DataFrame({'site': data_sorted.site.unique(),
                                 'site_name': data_sorted.site_name.unique()})
model_2_data['forest_200m'] = data_sorted.forest_200m.unique()
model_2_data['quantiles'] = [stats.quantiles(model_2_samples['c'][:1000][:, i]) for i in range(len(set(sites)))]

# plot quantiles
for i, row in model_2_data.sort_values(by='forest_200m').iterrows():
    x = row['forest_200m']
    ax.plot([x, x], [row['quantiles'][2.5], row['quantiles'][97.5]], color='black', linewidth=0.5)
    ax.plot([x, x], [row['quantiles'][25], row['quantiles'][75]], color='black', linewidth=1)
    ax.scatter([x], [row['quantiles'][50]], color='black', marker='o')

# format plot
l1 = ax.set_xlim([0, 100])
xl = ax.set_xlabel("forest land cover within 200 meters (percent area)")
yl = ax.set_ylabel("percent biophony (difference from grand mean)")
```


# Grey radiation modeling with climlab

This notebook is part of [The Climate Laboratory](https://brian-rose.github.io/ClimateLaboratoryBook) by [Brian E. J. Rose](http://www.atmos.albany.edu/facstaff/brose/index.html), University at Albany.

____________

## 1. Introducing `climlab`
____________

``climlab`` is a flexible engine for process-oriented climate modeling.
It is based on a very general concept of a model as a collection of individual,
interacting processes. ``climlab`` defines a base class called ``Process``, which
can contain an arbitrarily complex tree of sub-processes (each also some
sub-class of ``Process``). Every climate process (radiative, dynamical,
physical, turbulent, convective, chemical, etc.) can be simulated as a stand-alone
process model given appropriate input, or as a sub-process of a more complex model.
New classes of model can easily be defined and run interactively by putting together an
appropriate collection of sub-processes.

``climlab`` is a work-in-progress, and the code base will evolve substantially over the course of this semester.
The latest code can always be found on ``github``:

https://github.com/brian-rose/climlab

You are strongly encouraged to clone the ``climlab`` repository and use ``git`` to keep your local copy up-to-date.

Running this notebook requires that ``climlab`` is already installed on your system.

```python
%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import xarray as xr
from numpy import cos, deg2rad, log
import climlab
```

____________

## 2. Using `climlab` to implement the two-layer leaky greenhouse model
____________

One of the things that ``climlab`` is set up to do is the grey-radiation modeling we have already been discussing.

Since we already derived a [complete analytical solution to the two-layer leaky greenhouse model](Lecture06 -- Elementary greenhouse models.ipynb), we will use this to validate the `climlab` code.

### Validation

We want to verify that the model reproduces the observed OLR given observed temperatures, and the absorptivity that we tuned in the analytical model. The target numbers are:

\begin{align}
T_s &= 288 \text{ K} \\
T_0 &= 275 \text{ K} \\
T_1 &= 230 \text{ K} \\
\end{align}

$$ \epsilon = 0.586  $$

$$ OLR = 238.5 \text{ W m}^{-2} $$

### Initialize a model in `climlab`
The first thing we do is create a new model.

The following example code is sparsely commented but will hopefully orient you on the basics of defining and working with a `climlab Process` object.

```python
#  Test in a 2-layer atmosphere
col = climlab.GreyRadiationModel(num_lev=2)
print( col)
```

    climlab Process of type <class 'climlab.model.column.GreyRadiationModel'>.
    State variables and domain shapes:
      Ts: (1,)
      Tatm: (2,)
    The subprocess tree:
    Untitled: <class 'climlab.model.column.GreyRadiationModel'>
       LW: <class 'climlab.radiation.greygas.GreyGas'>
       SW: <class 'climlab.radiation.greygas.GreyGasSW'>
       insolation: <class 'climlab.radiation.insolation.FixedInsolation'>

```python
col.subprocess
```

    AttrDict({'LW': <climlab.radiation.greygas.GreyGas object at 0x104c67fa0>, 'SW': <climlab.radiation.greygas.GreyGasSW object at 0x11df6e550>, 'insolation': <climlab.radiation.insolation.FixedInsolation object at 0x11df6e940>})

Every item in the above dictionary is itself an instance of the `climlab.Process` object:

```python
print( col.subprocess['LW'])
```

    climlab Process of type <class 'climlab.radiation.greygas.GreyGas'>.
    State variables and domain shapes:
      Ts: (1,)
      Tatm: (2,)
    The subprocess tree:
    Untitled: <class 'climlab.radiation.greygas.GreyGas'>

The `state` dictionary holds the state variables of the model. In this case, temperatures:

```python
climlab.to_xarray(col.state)
```

<xarray.Dataset>
Dimensions:       (depth: 1, depth_bounds: 2, lev: 2, lev_bounds: 3)
Coordinates:
  * depth         (depth) float64 0.5
  * depth_bounds  (depth_bounds) float64 0.0 1.0
  * lev           (lev) float64 250.0 750.0
  * lev_bounds    (lev_bounds) float64 0.0 500.0 1e+03
Data variables:
    Ts            (depth) float64 288.0
    Tatm          (lev) float64 200.0 278.0

Access these either through dictionary methods or as attributes of the model object:

```python
print( col.state['Ts'])
print( col.Ts)
col.Ts is col.state['Ts']
```

    [288.]
    [288.]

    True

Now we are assigning the "observed" temperatures to our model state:

```python
col.Ts[:] = 288.
col.Tatm[:] = np.array([230., 275.])
climlab.to_xarray(col.state)
```

<xarray.Dataset>
Dimensions:       (depth: 1, depth_bounds: 2, lev: 2, lev_bounds: 3)
Coordinates:
  * depth         (depth) float64 0.5
  * depth_bounds  (depth_bounds) float64 0.0 1.0
  * lev           (lev) float64 250.0 750.0
  * lev_bounds    (lev_bounds) float64 0.0 500.0 1e+03
Data variables:
    Ts            (depth) float64 288.0
    Tatm          (lev) float64 230.0 275.0

```python
LW = col.subprocess['LW']
print(LW)
```

    climlab Process of type <class 'climlab.radiation.greygas.GreyGas'>.
    State variables and domain shapes:
      Ts: (1,)
      Tatm: (2,)
    The subprocess tree:
    Untitled: <class 'climlab.radiation.greygas.GreyGas'>

```python
LW.absorptivity
```

    Field([0.47737425, 0.47737425])

```python
#  copying the tuned value of epsilon from Lecture 6 notes
LW.absorptivity = 0.586
LW.absorptivity
```

    Field([0.586, 0.586])

```python
#  This does all the calculations that would be performed at each time step,
#  but doesn't actually update the temperatures
col.compute_diagnostics()
#  Print out the dictionary
col.diagnostics
```

    {'flux_from_sfc': Field([102.0487]),
     'flux_to_sfc': array([341.3]),
     'flux_to_space': array([102.0487]),
     'absorbed': array([0., 0.]),
     'absorbed_total': 0.0,
     'emission': Field([0., 0.]),
     'emission_sfc': Field([0.]),
     'flux_reflected_up': array([  0.    ,   0.    , 102.0487]),
     'insolation': Field([341.3]),
     'coszen': Field([1.]),
     'OLR': array([238.5247311]),
     'LW_down_sfc': array([228.53426769]),
     'LW_up_sfc': Field([390.10502995]),
     'LW_absorbed_sfc': Field([-161.57076227]),
     'LW_absorbed_atm': array([ 20.02990881, -96.98387764]),
     'LW_emission': Field([ 92.98664086, 190.03779837]),
     'ASR': Field([239.2513]),
     'SW_absorbed_sfc': Field([239.2513]),
     'SW_absorbed_atm': array([0., 0.]),
     'SW_up_sfc': Field([102.0487]),
     'SW_up_TOA': array([102.0487]),
     'SW_down_TOA': Field([341.3]),
     'planetary_albedo': Field([0.299])}

```python
#  Check OLR against our analytical solution
col.OLR
```

    array([238.5247311])

```python
# Like the state variables, the diagnostics can also be accessed in two different ways
col.diagnostics['OLR']
```

    array([238.5247311])

```python
col.state
```

    AttrDict({'Ts': Field([288.]), 'Tatm': Field([230., 275.])})

```python
# perform a single time step
col.step_forward()
```

```python
col.state
```

    AttrDict({'Ts': Field([289.60514636]), 'Tatm': Field([230.33784312, 273.3641795 ])})

We just stepped forward one discreet unit in time. Because we didn't specify a timestep when we created the model, it is set to a default value:

```python
col.timestep
```

    86400.0

which is 1 day (expressed in seconds).

Now we will integrate the model out to equilibrium.

We could easily write a loop to call the `step_forward()` method many times.

Or use a handy shortcut that allows us to specify the integration length in physical time units:

```python
# integrate out to radiative equilibrium
col.integrate_years(2.)
```

    Integrating for 730 steps, 730.4844 days, or 2.0 years.
    Total elapsed time is 2.0014116660123062 years.

```python
# Check for equilibrium
col.ASR - col.OLR
```

    Field([-2.34514516e-07])

```python
#  The temperatures at radiative equilibrium
col.state
```

    AttrDict({'Ts': Field([296.38447748]), 'Tatm': Field([233.72131685, 262.28540231])})

Compare these to the analytical solutions for radiative equilibrium with $\epsilon = 0.58$:

\begin{align}
T_s &= 296.4 \text{ K} \\
T_0 &= 262.3 \text{ K} \\
T_1 &= 233.8 \text{ K} \\
\end{align}

So it looks like `climlab` agrees with our analytical results to within 0.1 K. That's good.

____________

## 3. The observed annual, global mean temperature profile
____________

We want to model the OLR in a column whose temperatures match observations. As we've done before, we'll calculate the global, annual mean air temperature from the NCEP Reanalysis data.

```python
## The NOAA ESRL server is shutdown! January 2019
## This will try to read the data over the internet.
ncep_filename = 'air.mon.1981-2010.ltm.nc'
##  to read over internet
ncep_url = "http://www.esrl.noaa.gov/psd/thredds/dodsC/Datasets/ncep.reanalysis.derived/pressure/"
path = ncep_url
##  Open handle to data
ncep_air = xr.open_dataset( path + ncep_filename, decode_times=False )

#url = 'http://apdrc.soest.hawaii.edu:80/dods/public_data/Reanalysis_Data/NCEP/NCEP/clima/pressure/air'
#air = xr.open_dataset(url)
# The name of the vertical axis is different than the NOAA ESRL version..
#ncep_air = air.rename({'lev': 'level'})
print( ncep_air)
```

    <xarray.Dataset>
    Dimensions:             (lat: 73, level: 17, lon: 144, nbnds: 2, time: 12)
    Coordinates:
      * level               (level) float32 1000.0 925.0 850.0 ... 30.0 20.0 10.0
      * lon                 (lon) float32 0.0 2.5 5.0 7.5 ... 352.5 355.0 357.5
      * time                (time) float64 -6.571e+05 -6.57e+05 ... -6.567e+05
      * lat                 (lat) float32 90.0 87.5 85.0 82.5 ... -85.0 -87.5 -90.0
    Dimensions without coordinates: nbnds
    Data variables:
        climatology_bounds  (time, nbnds) float64 ...
        air                 (time, level, lat, lon) float32 ...
        valid_yr_count      (time, level, lat, lon) float32 ...
    Attributes:
        description:                     Data from NCEP initialized reanalysis (4...
        platform:                       Model
        Conventions:                    COARDS
        not_missing_threshold_percent:  minimum 3% values input to have non-missi...
        history:                        Created 2011/07/12 by doMonthLTM\nConvert...
        title:                          monthly ltm air from the NCEP Reanalysis
        dataset_title:                  NCEP-NCAR Reanalysis 1
        References:                     http://www.psl.noaa.gov/data/gridded/data...

```python
#  Take global, annual average and convert to Kelvin
weight = cos(deg2rad(ncep_air.lat)) / cos(deg2rad(ncep_air.lat)).mean(dim='lat')
Tglobal = (ncep_air.air * weight).mean(dim=('lat','lon','time'))
print( Tglobal)
```

    <xarray.DataArray (level: 17)>
    array([ 15.179084  ,  11.207003  ,   7.8383274 ,   0.21994135,
            -6.4483433 , -14.888848  , -25.570469  , -39.36969   ,
           -46.797905  , -53.652245  , -60.56356   , -67.006065  ,
           -65.53293   , -61.48664   , -55.853584  , -51.593952  ,
           -43.21999   ], dtype=float32)
    Coordinates:
      * level    (level) float32 1000.0 925.0 850.0 700.0 ... 50.0 30.0 20.0 10.0

We're going to convert this to degrees Kelvin, using a handy list of pre-defined constants in `climlab.constants`

```python
climlab.constants.tempCtoK
```

    273.15

```python
Tglobal += climlab.constants.tempCtoK
print(Tglobal)
```

    <xarray.DataArray (level: 17)>
    array([288.32907, 284.357  , 280.9883 , 273.36993, 266.70166, 258.26114,
           247.57953, 233.7803 , 226.35208, 219.49774, 212.58643, 206.14392,
           207.61707, 211.66336, 217.29642, 221.55605, 229.93001],
          dtype=float32)
    Coordinates:
      * level    (level) float32 1000.0 925.0 850.0 700.0 ... 50.0 30.0 20.0 10.0

```python
#  A handy re-usable routine for making a plot of the temperature profiles
#  We will plot temperatures with respect to log(pressure) to get a height-like coordinate

def zstar(lev):
    return -np.log(lev / climlab.constants.ps)

def plot_soundings(result_list, name_list, plot_obs=True, fixed_range=True):
    color_cycle=['r', 'g', 'b', 'y']
    # col is either a column model object or a list of column model objects
    #if isinstance(state_list, climlab.Process):
    #    # make a list with a single item
    #    collist = [collist]
    fig, ax = plt.subplots(figsize=(9,9))
    if plot_obs:
        ax.plot(Tglobal, zstar(Tglobal.level), color='k', label='Observed')
    for i, state in enumerate(result_list):
        Tatm = state['Tatm']
        lev = Tatm.domain.axes['lev'].points
        Ts = state['Ts']
        ax.plot(Tatm, zstar(lev), color=color_cycle[i], label=name_list[i])
        ax.plot(Ts, 0, 'o', markersize=12, color=color_cycle[i])
    #ax.invert_yaxis()
    yticks = np.array([1000., 750., 500., 250., 100., 50., 20., 10., 5.])
    ax.set_yticks(-np.log(yticks/1000.))
    ax.set_yticklabels(yticks)
    ax.set_xlabel('Temperature (K)', fontsize=14)
    ax.set_ylabel('Pressure (hPa)', fontsize=14)
    ax.grid()
    ax.legend()
    if fixed_range:
        ax.set_xlim([200, 300])
        ax.set_ylim(zstar(np.array([1000., 5.])))
    #ax2 = ax.twinx()

    return ax
```

```python
plot_soundings([],[] );
```

____________

## 4. A 30-layer model using the observed temperatures
____________

```python
#  initialize a grey radiation model with 30 levels
col = climlab.GreyRadiationModel()
print(col)
```

    climlab Process of type <class 'climlab.model.column.GreyRadiationModel'>.
    State variables and domain shapes:
      Ts: (1,)
      Tatm: (30,)
    The subprocess tree:
    Untitled: <class 'climlab.model.column.GreyRadiationModel'>
       LW: <class 'climlab.radiation.greygas.GreyGas'>
       SW: <class 'climlab.radiation.greygas.GreyGasSW'>
       insolation: <class 'climlab.radiation.insolation.FixedInsolation'>

```python
col.lev
```

    array([ 16.66666667,  50.        ,  83.33333333, 116.66666667,
           150.        , 183.33333333, 216.66666667, 250.        ,
           283.33333333, 316.66666667, 350.        , 383.33333333,
           416.66666667, 450.        , 483.33333333, 516.66666667,
           550.        , 583.33333333, 616.66666667, 650.        ,
           683.33333333, 716.66666667, 750.        , 783.33333333,
           816.66666667, 850.        , 883.33333333, 916.66666667,
           950.        , 983.33333333])

```python
col.lev_bounds
```

    array([   0.        ,   33.33333333,   66.66666667,  100.        ,
            133.33333333,  166.66666667,  200.        ,  233.33333333,
            266.66666667,  300.        ,  333.33333333,  366.66666667,
            400.        ,  433.33333333,  466.66666667,  500.        ,
            533.33333333,  566.66666667,  600.        ,  633.33333333,
            666.66666667,  700.        ,  733.33333333,  766.66666667,
            800.        ,  833.33333333,  866.66666667,  900.        ,
            933.33333333,  966.66666667, 1000.        ])

```python
# interpolate to 30 evenly spaced pressure levels
lev = col.lev
Tinterp = np.interp(lev, np.flipud(Tglobal.level), np.flipud(Tglobal))
Tinterp
#  Need to 'flipud' because the interpolation routine
#  needs the pressure data to be in increasing order
```

    array([224.34736633, 211.6633606 , 206.96233453, 208.29142253,
           212.58642578, 217.19396973, 221.78252157, 226.3520813 ,
           231.30422974, 236.08017476, 240.67991638, 245.279658  ,
           249.35979716, 252.92033386, 256.48087056, 259.66789246,
           262.48139954, 265.29490662, 267.81303914, 270.03579712,
           272.25855509, 274.21642049, 275.90939331, 277.60236613,
           279.29533895, 280.98831177, 282.48550415, 283.98269653,
           285.6810201 , 287.4463874 ])

```python
# Initialize model with observed temperatures
col.Ts[:] = Tglobal[0]
col.Tatm[:] = Tinterp
```

```python
# This should look just like the observations
result_list = [col.state]
name_list = ['Observed, interpolated']
plot_soundings(result_list, name_list);
```

### Tune absorptivity to get observed OLR

```python
col.compute_diagnostics()
col.OLR
```

    array([263.15000104])

```python
# Need to tune absorptivity to get OLR = 238.5
epsarray = np.linspace(0.01, 0.1, 100)
OLRarray = np.zeros_like(epsarray)
```

```python
for i in range(epsarray.size):
    col.subprocess['LW'].absorptivity = epsarray[i]
    col.compute_diagnostics()
    OLRarray[i] = col.OLR

plt.plot(epsarray, OLRarray)
plt.grid()
plt.xlabel('epsilon')
plt.ylabel('OLR')
```

The necessary value seems to lie near 0.055 or so.

We can be more precise with a numerical root-finder.

```python
def OLRanom(eps):
    col.subprocess['LW'].absorptivity = eps
    col.compute_diagnostics()
    return col.OLR - 238.5
```

```python
# Use numerical root-finding to get the equilibria
from scipy.optimize import brentq
# brentq is a root-finding function
#  Need to give it a function and two end-points
#  It will look for a zero of the function between those end-points
eps = brentq(OLRanom, 0.01, 0.1)
print( eps)
```

    0.05369073166506245

```python
col.subprocess.LW.absorptivity = eps
col.subprocess.LW.absorptivity
```

    Field([0.05369073, 0.05369073, 0.05369073, 0.05369073, 0.05369073,
           0.05369073, 0.05369073, 0.05369073, 0.05369073, 0.05369073,
           0.05369073, 0.05369073, 0.05369073, 0.05369073, 0.05369073,
           0.05369073, 0.05369073, 0.05369073, 0.05369073, 0.05369073,
           0.05369073, 0.05369073, 0.05369073, 0.05369073, 0.05369073,
           0.05369073, 0.05369073, 0.05369073, 0.05369073, 0.05369073])

```python
col.compute_diagnostics()
col.OLR
```

    array([238.5])

____________

## 5. Radiative forcing in the 30-layer model
____________

Let's compute radiative forcing for a **2% increase in absorptivity**.

```python
#  clone our model using a built-in climlab function
col2 = climlab.process_like(col)
print(col2)
```

    climlab Process of type <class 'climlab.model.column.GreyRadiationModel'>.
    State variables and domain shapes:
      Ts: (1,)
      Tatm: (30,)
    The subprocess tree:
    Untitled: <class 'climlab.model.column.GreyRadiationModel'>
       LW: <class 'climlab.radiation.greygas.GreyGas'>
       SW: <class 'climlab.radiation.greygas.GreyGasSW'>
       insolation: <class 'climlab.radiation.insolation.FixedInsolation'>

```python
col2.subprocess['LW'].absorptivity *= 1.02
col2.subprocess['LW'].absorptivity
```

    Field([0.05476455, 0.05476455, 0.05476455, 0.05476455, 0.05476455,
           0.05476455, 0.05476455, 0.05476455, 0.05476455, 0.05476455,
           0.05476455, 0.05476455, 0.05476455, 0.05476455, 0.05476455,
           0.05476455, 0.05476455, 0.05476455, 0.05476455, 0.05476455,
           0.05476455, 0.05476455, 0.05476455, 0.05476455, 0.05476455,
           0.05476455, 0.05476455, 0.05476455, 0.05476455, 0.05476455])

```python
#  Radiative forcing by definition is the change in TOA radiative flux,
# HOLDING THE TEMPERATURES FIXED.
col2.Ts - col.Ts
```

    Field([0.])

```python
col2.Tatm - col.Tatm
```

    Field([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
           0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])

```python
col2.compute_diagnostics()
col2.OLR
```

    array([236.65384137])

The OLR decreased after we added the extra absorbers, as we expect. Now we can calculate the Radiative Forcing:

```python
RF = -(col2.OLR - col.OLR)
print( 'The radiative forcing is %.2f W/m2.' %RF)
```

    The radiative forcing is 1.85 W/m2.

____________

## 6. Radiative equilibrium in the 30-layer model
____________

```python
re = climlab.process_like(col)
```

```python
#  To get to equilibrium, we just time-step the model forward long enough
re.integrate_years(1.)
```

    Integrating for 365 steps, 365.2422 days, or 1.0 years.
    Total elapsed time is 0.9993368783782377 years.

```python
#  Check for energy balance
print( 'The net downward radiative flux at TOA is %.4f W/m2.' %(re.ASR - re.OLR))
```

    The net downward radiative flux at TOA is -0.0015 W/m2.

```python
result_list.append(re.state)
name_list.append('Radiative equilibrium (grey gas)')
plot_soundings(result_list, name_list)
```

Some properties of the **radiative equilibrium** temperature profile:

- The surface is warmer than observed.
- The lower troposphere is colder than observed.
- Very cold air is sitting immediately above the warm surface.
- There is no tropopause, no stratosphere.

____________

## 7. Radiative-Convective Equilibrium in the 30-layer model
____________

We recognize that the large drop in temperature just above the surface is unphysical. Parcels of air in direct contact with the ground will be warmed by mechansisms other than radiative transfer.

These warm air parcels will then become buoyant, and will convect upward, mixing their heat content with the environment.

We **parameterize** the statistical effects of this mixing through a **convective adjustment**.

At each timestep, our model checks for any locations at which the **lapse rate** exceeds some threshold. Unstable layers are removed through an energy-conserving mixing formula.

This process is assumed to be fast relative to radiative heating. In the model, it is instantaneous.

### Add the convective adjustment as an additional subprocess

```python
#  Here is the existing model
print(re)
```

    climlab Process of type <class 'climlab.model.column.GreyRadiationModel'>.
    State variables and domain shapes:
      Ts: (1,)
      Tatm: (30,)
    The subprocess tree:
    Untitled: <class 'climlab.model.column.GreyRadiationModel'>
       LW: <class 'climlab.radiation.greygas.GreyGas'>
       SW: <class 'climlab.radiation.greygas.GreyGasSW'>
       insolation: <class 'climlab.radiation.insolation.FixedInsolation'>

```python
#  First we make a new clone
rce = climlab.process_like(re)
#  Then create a new ConvectiveAdjustment process
conv = climlab.convection.ConvectiveAdjustment(state=rce.state,
                                               adj_lapse_rate=6.)
#  And add it to our model
rce.add_subprocess('Convective Adjustment', conv)
print( rce)
```

    climlab Process of type <class 'climlab.model.column.GreyRadiationModel'>.
    State variables and domain shapes:
      Ts: (1,)
      Tatm: (30,)
    The subprocess tree:
    Untitled: <class 'climlab.model.column.GreyRadiationModel'>
       LW: <class 'climlab.radiation.greygas.GreyGas'>
       SW: <class 'climlab.radiation.greygas.GreyGasSW'>
       insolation: <class 'climlab.radiation.insolation.FixedInsolation'>
       Convective Adjustment: <class 'climlab.convection.convadj.ConvectiveAdjustment'>

This model is exactly like our previous models, except for one additional subprocess called ``Convective Adjustment``.

We passed a parameter ``adj_lapse_rate`` (in K / km) that sets the neutrally stable lapse rate -- in this case, 6 K / km.

This number is chosed to very loosely represent the net effect of **moist convection**.

```python
#  Run out to equilibrium
rce.integrate_years(1.)
```

    Integrating for 365 steps, 365.2422 days, or 1.0 years.
    Total elapsed time is 1.9986737567564754 years.

```python
#  Check for energy balance
rce.ASR - rce.OLR
```

    Field([0.0007796])

```python
result_list.append(rce.state)
name_list.append('Radiatve-Convective equilibrium (grey gas)')
```

```python
plot_soundings(result_list, name_list)
```

Introducing convective adjustment into the model cools the surface quite a bit (compared to Radiative Equilibrium, in green here) -- and warms the lower troposphere. It gives us a MUCH better fit to observations.

But of course we still have no stratosphere.

____________

## 8. Putting stratospheric ozone in the grey-gas model
____________

Our model has no equivalent of the stratosphere, where temperature increases with height. That's because our model has been completely transparent to shortwave radiation up until now.

We can load the observed ozone climatology from the input files for the CESM model:

```python
datapath = "http://thredds.atmos.albany.edu:8080/thredds/dodsC/CESMA/"
ozone = xr.open_dataset( datapath + "som_input/ozone_1.9x2.5_L26_2000clim_c091112.nc")
```

```python
print(ozone)
```

    <xarray.Dataset>
    Dimensions:  (ilev: 27, lat: 96, lev: 26, lon: 144, slat: 95, slon: 144, time: 12)
    Coordinates:
      * ilev     (ilev) float64 2.194 4.895 9.882 18.05 ... 903.3 956.0 985.1 1e+03
      * lat      (lat) float64 -90.0 -88.11 -86.21 -84.32 ... 84.32 86.21 88.11 90.0
      * lev      (lev) float64 3.545 7.389 13.97 23.94 ... 867.2 929.6 970.6 992.6
      * lon      (lon) float64 0.0 2.5 5.0 7.5 10.0 ... 350.0 352.5 355.0 357.5
      * slat     (slat) float64 -89.05 -87.16 -85.26 -83.37 ... 85.26 87.16 89.05
      * slon     (slon) float64 -1.25 1.25 3.75 6.25 ... 348.8 351.2 353.8 356.2
      * time     (time) object 2000-01-15 00:00:00 ... 2000-12-15 00:00:00
    Data variables:
        P0       float64 ...
        date     (time) int32 ...
        datesec  (time) int32 ...
        gw       (lat) float64 ...
        hyai     (ilev) float64 ...
        hyam     (lev) float64 ...
        hybi     (ilev) float64 ...
        hybm     (lev) float64 ...
        w_stag   (slat) float64 ...
        O3       (time, lev, lat, lon) float32 ...
        PS       (time, lat, lon) float32 ...
    Attributes:
        Conventions:                     CF-1.0
        source:                          CAM
        case:                            ar5_cam_1850-2000_03
        title:
        logname:                         lamar
        host:                            be0207en.ucar.ed
        Version:                         $Name$
        revision_Id:                     $Id$
        initial_file:                    /ptmp/lamar/cam-inputs/ar5_cam_1850_06.c...
        topography_file:                 /fs/cgd/csm/inputdata/atm/cam/topo/USGS-...
        history:                         Wed Jan  5 14:56:28 2011: ncks -d time,1...
        nco_openmp_thread_number:        1
        NCO:                             4.0.5
        DODS_EXTRA.Unlimited_Dimension:  time

The pressure levels in this dataset are:

```python
print(ozone.lev)
```

    <xarray.DataArray 'lev' (lev: 26)>
    array([  3.544638,   7.388814,  13.967214,  23.944625,  37.23029 ,  53.114605,
            70.05915 ,  85.439115, 100.514695, 118.250335, 139.115395, 163.66207 ,
           192.539935, 226.513265, 266.481155, 313.501265, 368.81798 , 433.895225,
           510.455255, 600.5242  , 696.79629 , 787.70206 , 867.16076 , 929.648875,
           970.55483 , 992.5561  ])
    Coordinates:
      * lev      (lev) float64 3.545 7.389 13.97 23.94 ... 867.2 929.6 970.6 992.6
    Attributes:
        long_name:      hybrid level at midpoints (1000*(A+B))
        units:          level
        positive:       down
        standard_name:  atmosphere_hybrid_sigma_pressure_coordinate
        formula_terms:  a: hyam b: hybm p0: P0 ps: PS

### Take the global average of the ozone climatology, and plot it as a function of pressure (or height)

```python
#  Take global, annual average and convert to Kelvin
weight_ozone = cos(deg2rad(ozone.lat)) / cos(deg2rad(ozone.lat)).mean(dim='lat')
O3_global = (ozone.O3 * weight_ozone).mean(dim=('lat','lon','time'))
print(O3_global)
```

    <xarray.DataArray (lev: 26)>
    array([7.82792878e-06, 8.64150529e-06, 7.58940028e-06, 5.24567145e-06,
           3.17761574e-06, 1.82320006e-06, 9.80756960e-07, 6.22870516e-07,
           4.47620550e-07, 3.34481169e-07, 2.62570302e-07, 2.07898125e-07,
           1.57074555e-07, 1.12425545e-07, 8.06004999e-08, 6.27826498e-08,
           5.42990561e-08, 4.99506089e-08, 4.60075681e-08, 4.22977789e-08,
           3.80559071e-08, 3.38768568e-08, 3.12171619e-08, 2.97807119e-08,
           2.87980968e-08, 2.75429934e-08])
    Coordinates:
      * lev      (lev) float64 3.545 7.389 13.97 23.94 ... 867.2 929.6 970.6 992.6

```python
ax = plt.figure(figsize=(10,8)).add_subplot(111)
ax.plot( O3_global * 1.E6, -np.log(ozone.lev/climlab.constants.ps) )
ax.set_xlabel('Ozone (ppm)', fontsize=16)
ax.set_ylabel('Pressure (hPa)', fontsize=16 )
yticks = np.array([1000., 750., 500., 250., 100., 50., 20., 10., 5.])
ax.set_yticks(-np.log(yticks/1000.))
ax.set_yticklabels(yticks)
ax.grid()
ax.set_title('Global, annual mean ozone concentration', fontsize = 24);
```

This shows that most of the ozone is indeed in the stratosphere, and peaks near the top of the stratosphere.

Now create a new column model object **on the same pressure levels as the ozone data**.  We are also going set an adjusted lapse rate of 6 K / km.

```python
# the RadiativeConvectiveModel is pre-defined in climlab
#  It contains the same components are our previous model
#   But here we are specifying a different set of vertical levels.
oz_col = climlab.RadiativeConvectiveModel(lev = ozone.lev, adj_lapse_rate=6)
print(oz_col)
```

    climlab Process of type <class 'climlab.model.column.RadiativeConvectiveModel'>.
    State variables and domain shapes:
      Ts: (1,)
      Tatm: (26,)
    The subprocess tree:
    Untitled: <class 'climlab.model.column.RadiativeConvectiveModel'>
       LW: <class 'climlab.radiation.greygas.GreyGas'>
       SW: <class 'climlab.radiation.greygas.GreyGasSW'>
       insolation: <class 'climlab.radiation.insolation.FixedInsolation'>
       convective adjustment: <class 'climlab.convection.convadj.ConvectiveAdjustment'>

Now we will do something new: let the column absorb some shortwave radiation. We will assume that the shortwave absorptivity is proportional to the ozone concentration we plotted above.

Now we need to weight the absorptivity by the pressure (mass) of each layer.

```python
# This number is an arbitrary parameter that scales how absorptive we are making the ozone
# in our grey gas model
ozonefactor = 75
dp = oz_col.Tatm.domain.lev.delta
epsSW = O3_global.values * dp * ozonefactor
```

We want to use the field `epsSW` as the absorptivity for our SW radiation model.

Let's see what the absorptivity is current set to:

```python
print(oz_col.subprocess['SW'].absorptivity)
```

    [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
     0. 0.]

It defaults to zero.

Before changing this (putting in the ozone), let's take a look at the shortwave absorption in the column:

```python
oz_col.compute_diagnostics()
```

```python
oz_col.diagnostics['SW_absorbed_atm']
```

    array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
           0., 0., 0., 0., 0., 0., 0., 0., 0.])

Let's now put in the ozone:

```python
oz_col.subprocess['SW'].absorptivity = epsSW
print(oz_col.subprocess['SW'].absorptivity)
```

    [3.20948549e-03 3.37750296e-03 4.71182551e-03 4.57614201e-03
     3.47591203e-03 2.24450924e-03 1.18884331e-03 7.11369788e-04
     5.50761612e-04 4.84170273e-04 4.47141487e-04 4.16507314e-04
     3.70212131e-04 3.11733073e-04 2.62922861e-04 2.40936639e-04
     2.45147939e-04 2.65307555e-04 2.87482272e-04 2.95567946e-04
     2.67120872e-04 2.16427978e-04 1.66169127e-04 1.15468088e-04
     6.79353134e-05 3.81013280e-05]

Let's check how this changes the SW absorption:

```python
oz_col.compute_diagnostics()
oz_col.SW_absorbed_atm
```

    array([1.40571621, 1.47671285, 2.05685574, 1.99236565, 1.50916587,
           0.97239284, 0.51429089, 0.30750441, 0.23797743, 0.20913838,
           0.19309134, 0.17981756, 0.1597929 , 0.13452308, 0.11343944,
           0.10393806, 0.10574109, 0.11442201, 0.12396838, 0.12743553,
           0.11515129, 0.09328389, 0.07161233, 0.04975714, 0.02927232,
           0.01641659])

It is now non-zero, and largest near the top of the column (also top of the array) where the ozone concentration is highest.

Now it's time to run the model out to radiative-convective equilibrium

```python
oz_col.integrate_years(1.)
```

    Integrating for 365 steps, 365.2422 days, or 1.0 years.
    Total elapsed time is 0.9993368783782377 years.

```python
print(oz_col.ASR - oz_col.OLR)
```

    [-0.00396053]

And let's now see what we got!

```python
result_list.append(oz_col.state)
name_list.append('Radiative-Convective equilibrium with O3')
```

```python
#  Make a plot to compare observations, Radiative Equilibrium, Radiative-Convective Equilibrium, and RCE with ozone!
plot_soundings(result_list, name_list)
```

And we finally have something that looks looks like the tropopause, with temperature increasing above at approximately the correct rate.

There are still plenty of discrepancies between this model solution and the observations, including:

- Tropopause temperature is too warm, by about 15 degrees.
- Surface temperature is too cold

There are a number of parameters we might adjust if we wanted to improve the fit, including:

- Longwave absorptivity
- Surface albedo

Feel free to experiment! (That's what models are for, after all).

### The take home message

The dominant effect of stratospheric ozone is to vastly increase the radiative equilibrium temperature in the ozone layer. The temperature needs to be higher so that the longwave emission can balance the shortwave absorption.

Without ozone to absorb incoming solar radiation, the **temperature does not increase with height**.

This simple grey-gas model illustrates this principle very clearly.

____________

## Credits

This notebook is part of [The Climate Laboratory](https://brian-rose.github.io/ClimateLaboratoryBook), an open-source textbook developed and maintained by [Brian E. J. Rose](http://www.atmos.albany.edu/facstaff/brose/index.html), University at Albany.

It is licensed for free and open consumption under the
[Creative Commons Attribution 4.0 International (CC BY 4.0)](https://creativecommons.org/licenses/by/4.0/) license.

Development of these notes and the [climlab software](https://github.com/brian-rose/climlab) is partially supported by the National Science Foundation under award AGS-1455071 to Brian Rose. Any opinions, findings, conclusions or recommendations expressed here are mine and do not necessarily reflect the views of the National Science Foundation.
____________

```python

```


# Лабораторная работа №2

Численная реализация оптического преобразования Фурье на основе быстрого преобразования Фурье.

> Вариант 7

```python
import matplotlib.pyplot as plt
import numpy as np
import scipy.integrate
import sympy

%matplotlib inline
```

```python
def versions():
    from datetime import datetime
    print("Notebook date:", datetime.now())

    import numpy as np
    print("NumPy:", np.__version__)

    import scipy as sc
    print("SciPy:", sc.__version__)

    import matplotlib
    print("Matplotlib:", matplotlib.__version__)

    import sys
    print("Python:", sys.version)

versions()
```

    Notebook date: 2020-04-30 23:37:14.175021
    NumPy: 1.18.3
    SciPy: 1.4.1
    Matplotlib: 3.2.1
    Python: 3.7.3 (default, Apr 09 2019, 05:18:21) [GCC]

```python
a = 5
N = 128
M = 2048
```

### 1. Реализация финитного преобразования с помощью БПФ

```python
def finite_fft(vector_f, a, b, N, M):
    # 1. Append zeroes vectors to M size.
    zeros = np.zeros(int((M - N) / 2))

    vector_f = np.concatenate((zeros, vector_f, zeros), axis=None)

    # 2. Swap parts of vector f.
    center = int(len(vector_f) / 2)

    vector_f = np.concatenate((vector_f[center:], vector_f[:center]))

    # 3. Apply FFT.
    vector_F = np.fft.fft(vector_f, M) * (b - a) / (N - 1)

    # 4. Swap parts of vector F.
    center = int(len(vector_F) / 2)

    vector_F = np.concatenate((vector_F[center:], vector_F[:center]))

    # 5. Cut center.
    vector_F = vector_F[int((M - N) / 2) : int((M - N) / 2 + N)]

    # 6. Recount domain of a function.
    _new_domain = N ** 2 / (4 * a * M)
    func_domain = sorted([-_new_domain, _new_domain])

    return vector_F, func_domain
```

### 2. Гауссов пучок

```python
def gen_gaussian_beam(x):
    return np.exp(-(x ** 2))

lspace = np.linspace(-a, a, num=N, endpoint=False)
gaussian_beam = gen_gaussian_beam(lspace)

_, arr = plt.subplots(1, 2, figsize=(15, 5))

arr[0].plot(lspace, np.absolute(gaussian_beam), color="r", linewidth=3.0, label="Амплитуда")
arr[0].set_title(r"$\exp(-x^2)$")
arr[0].legend()
arr[0].grid()

arr[1].plot(lspace, np.angle(gaussian_beam), color="b", linewidth=3.0, label="Фаза")
arr[1].set_title(r"$\exp(-x^2)$")
arr[1].legend()
arr[1].grid()

plt.show()
```

### 3. Правильность реализации финитного преобразования с помощью БПФ

```python
fft_gaussian_beam, domain = finite_fft(
    gen_gaussian_beam(np.linspace(-a, a, num=N, endpoint=False)), a=-a, b=a, N=N, M=M
)

lspace = np.linspace(domain[0], domain[1], num=N, endpoint=False)

_, arr = plt.subplots(1, 2, figsize=(15, 5))

arr[0].plot(
    lspace, np.absolute(fft_gaussian_beam), color="r", linewidth=3.0, label="Амплитуда",
)
arr[0].set_title(r"$\exp(-x^2)$")
arr[0].legend()
arr[0].grid()

arr[1].plot(lspace, np.angle(fft_gaussian_beam), color="b", linewidth=3.0, label="Фаза")
arr[1].set_title(r"$\exp(-x^2)$")
arr[1].legend()
arr[1].grid()

plt.show()
```

### 4. Реализация финитного преобразования с помощью численного метода

```python
def quad(func, a, b, N):
    h = (b - a) / N

    return ((func(a) + func(b)) / 2 + sum([func(a + h * i) for i in range(N)])) * h

def numerical(func, a, b, N, M):
    # Create new function domain.
    _new_domain = N ** 2 / (4 * a * M)
    func_domain = sorted([-_new_domain, _new_domain])

    # Step for function.
    func_step = (func_domain[1] - func_domain[0]) / (N - 1)

    F = np.zeros((N,), dtype=np.complex)

    for i in range(len(F)):
        u = func_domain[0] + i * func_step

        F[i] = quad(lambda x: func(x) * np.exp(-2 * np.pi * u * x * 1j), a=a, b=b, N=N)

    return F, func_domain
```

```python
numerical_gaussian_beam, domain = numerical(gen_gaussian_beam, a=-a, b=a, N=N, M=M)

lspace = np.linspace(domain[0], domain[1], num=N, endpoint=False)

_, arr = plt.subplots(1, 2, figsize=(15, 5))

arr[0].plot(lspace, np.absolute(numerical_gaussian_beam), color="r", linewidth=3.0, label="Амплитуда")
arr[0].set_title(r"$\exp(-x^2)$")
arr[0].legend()
arr[0].grid()

arr[1].plot(lspace, np.angle(numerical_gaussian_beam), color="b", linewidth=3.0, label="Фаза")
arr[1].set_title(r"$\exp(-x^2)$")
arr[1].legend()
arr[1].grid()

plt.show()
```

### 5. Результаты двух разных реализаций преобразования

```python
_, arr = plt.subplots(1, 2, figsize=(15, 5))

arr[0].plot(lspace, np.absolute(fft_gaussian_beam), color="r", linewidth=3.0, label="Finite FFT, Амплитуда")
arr[0].plot(lspace, np.absolute(numerical_gaussian_beam), color="b", linewidth=3.0, label="Finite Numerical, Амплитуда")
arr[0].grid()
arr[0].legend()

arr[1].plot(lspace, np.angle(fft_gaussian_beam), color="r", linewidth=3.0, label="Finite FFT, Фаза")
arr[1].plot(lspace, np.angle(numerical_gaussian_beam), color="b", linewidth=3.0, label="Finite Numerical, Фаза")
arr[1].grid()
arr[1].legend()

plt.show()
```

### 6. Световое поле

#### Исходное поле

Вариант 7: $x^2$

```python
def gen_light_field(x):
    return x ** 2
```

```python
lspace = np.linspace(-a, a, num=N, endpoint=False)
light_field = gen_light_field(lspace)

_, arr = plt.subplots(1, 2, figsize=(15, 5))

arr[0].plot(lspace, np.absolute(light_field), color="r", linewidth=3.0, label="Амплитуда")
arr[0].set_title(r"$x^2$")
arr[0].legend()
arr[0].grid()

arr[1].plot(lspace, np.angle(light_field), color="b", linewidth=3.0, label="Фаза")
arr[1].set_title(r"$x^2$")
arr[1].legend()
arr[1].grid()

plt.show()
```

#### Преобразование

```python
fft_light_field, domain = finite_fft(gen_light_field(np.linspace(-a, a, num=N, endpoint=False)), a=-a, b=a, N=N, M=M)

lspace = np.linspace(domain[0], domain[1], num=N, endpoint=False)

_, arr = plt.subplots(1, 2, figsize=(15, 5))

arr[0].plot(lspace, np.absolute(fft_light_field), color="r", linewidth=3.0, label="Амплитуда")
arr[0].set_title(r"$x^2$")
arr[0].legend()
arr[0].grid()

arr[1].plot(lspace, np.angle(fft_light_field), color="b", linewidth=3.0, label="Фаза")
arr[1].set_title(r"$x^2$")
arr[1].legend()
arr[1].grid()

plt.show()
```

### 7. Аналитический результат преобразования

```python
_x = sympy.Symbol("x")
_u = sympy.Symbol("u")
_a = sympy.Symbol("a")

integrate = sympy.integrate(
    gen_light_field(_x) * sympy.exp(-2 * sympy.pi * 1j * _u * _x),
    (_x, -_a, _a),
).simplify()
```

#### Аналитическое решение в общем виде

```python
integrate
```

$\displaystyle \begin{cases} \frac{\left(0.5 i \pi^{2} a^{2} u^{2} + 0.5 \pi a u + \left(- 0.5 i \pi^{2} a^{2} u^{2} + 0.5 \pi a u + 0.25 i\right) e^{4.0 i \pi a u} - 0.25 i\right) e^{- 2.0 i \pi a u}}{\pi^{3} u^{3}} & \text{for}\: u > -\infty \wedge u < \infty \wedge u \neq 0 \\\frac{2 a^{3}}{3} & \text{otherwise} \end{cases}$

#### Аналитическое решение после подстановки $a$

```python
integrate.subs(_a, a)
```

$\displaystyle \begin{cases} \frac{\left(12.5 i \pi^{2} u^{2} + 2.5 \pi u + \left(- 12.5 i \pi^{2} u^{2} + 2.5 \pi u + 0.25 i\right) e^{20.0 i \pi u} - 0.25 i\right) e^{- 10.0 i \pi u}}{\pi^{3} u^{3}} & \text{for}\: u > -\infty \wedge u < \infty \wedge u \neq 0 \\\frac{250}{3} & \text{otherwise} \end{cases}$

```python
analytical_function = sympy.lambdify(
    _u, integrate.subs(_a, a).args[0][0], "numpy"
)
```

#### Сравнение графиков

```python
fft_light_field, domain = finite_fft(gen_light_field(np.linspace(-a, a, num=N, endpoint=False)), a=-a, b=a, N=N, M=M)

lspace = np.linspace(domain[0], domain[1], num=N, endpoint=False)

analytical_result = analytical_function(lspace)

_, arr = plt.subplots(1, 2, figsize=(15, 5))

arr[0].plot(lspace, np.absolute(fft_light_field), color="r", linewidth=3.0, label="Finite FFT, Амплитуда")
arr[0].plot(lspace, np.absolute(analytical_result), color="b", linewidth=3.0, label="Analytical, Амплитуда")
arr[0].grid()
arr[0].legend()

arr[1].plot(lspace, np.angle(fft_light_field), color="r", linewidth=3.0, label="Finite FFT, Фаза")
arr[1].plot(lspace, np.angle(analytical_result), color="b", linewidth=3.0, label="Analytical, Фаза")
arr[1].grid()
arr[1].legend()

plt.show()
```

# Двумерный случай

## 1. Гауссов пучок

```python
def gen_gaussian_beam_2d(x, y):
    return np.exp(-x ** 2 - y ** 2)

lspace = np.linspace(-a, a, num=N, endpoint=False)
X, Y = np.meshgrid(lspace, lspace)

gaussian_beam_2d = gen_gaussian_beam_2d(X, Y).astype(np.complex)

# - - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -

fig = plt.figure(figsize=(22, 8))

ax = fig.add_subplot(121)
plt.imshow(np.absolute(gaussian_beam_2d))
plt.title(r"Амплитуда $\exp(-x^2 -y^2)$")
plt.colorbar()

ax = fig.add_subplot(122)
plt.imshow(np.angle(gaussian_beam_2d))
plt.title(r"Фаза $\exp(-x^2 -y^2)$")
plt.colorbar()

plt.show()
```

## 2. Правильность реализации финитного преобразования с помощью БПФ

```python
def finite_fft_2d(matrix, a, b, N, M):
    fft_gaussian_beam_2d = matrix.copy()

    # Iterate over rows.
    for k, row in enumerate(fft_gaussian_beam_2d):
        fft_gaussian_beam_2d[k], domain = finite_fft(row, a=-a, b=a, N=N, M=M)

    # Iterate over cols.
    fft_gaussian_beam_2d = np.rot90(fft_gaussian_beam_2d, k = -1)

    for k, row in enumerate(fft_gaussian_beam_2d):
        fft_gaussian_beam_2d[k], domain = finite_fft(row, a=-a, b=a, N=N, M=M)

    fft_gaussian_beam_2d = np.rot90(fft_gaussian_beam_2d, k = 1)

    return fft_gaussian_beam_2d, domain
```

```python
lspace = np.linspace(-a, a, num=N, endpoint=False)

X, Y = np.meshgrid(lspace, lspace)

gaussian_beam_2d = gen_gaussian_beam_2d(X, Y).astype(np.complex)

fft_gaussian_beam_2d, domain = finite_fft_2d(gaussian_beam_2d, a=-a, b=a, N=N, M=M)

# - - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -

fig = plt.figure(figsize=(22, 8))

ax = fig.add_subplot(121)
plt.imshow(np.absolute(fft_gaussian_beam_2d))
plt.title(r"Амплитуда $\exp(-x^2 -y^2)$")
plt.colorbar()

ax = fig.add_subplot(122)
plt.imshow(np.angle(fft_gaussian_beam_2d))
plt.title(r"Фаза $\exp(-x^2 -y^2)$")
plt.colorbar()

plt.show()
```

## 3. Световое поле

```python
def gen_light_field_2d(x, y):
    return x ** 2 * y ** 2
```

#### Исходное поле

```python
lspace = np.linspace(-a, a, num=N, endpoint=False)

X, Y = np.meshgrid(lspace, lspace)

light_field_2d = gen_light_field_2d(X, Y).astype(np.complex)

# - - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -

fig = plt.figure(figsize=(22, 8))

ax = fig.add_subplot(121)
plt.imshow(np.absolute(light_field_2d))
plt.title(r"Амплитуда $x^2 \cdot y ^ 2$")
plt.colorbar()

ax = fig.add_subplot(122)
plt.imshow(np.angle(light_field_2d))
plt.title(r"Фаза $x^2 \cdot y ^ 2$")
plt.colorbar()

plt.show()
```

#### Преобразование

```python
fft_gaussian_beam_2d, domain = finite_fft_2d(light_field_2d, a=-a, b=a, N=N, M=M)

fig = plt.figure(figsize=(22, 8))

ax = fig.add_subplot(121)
plt.imshow(np.absolute(fft_gaussian_beam_2d))
plt.title(r"Амплитуда $x^2 \cdot y ^ 2$")
plt.colorbar()

ax = fig.add_subplot(122)
plt.imshow(np.angle(fft_gaussian_beam_2d))
plt.title(r"Фаза $x^2 \cdot y ^ 2$")
plt.colorbar()

plt.show()
```

## 4. Аналитический результат преобразования

```python
_x = sympy.Symbol("x")
_y = sympy.Symbol("y")
_u = sympy.Symbol("u")
_v = sympy.Symbol("v")
_a = sympy.Symbol("a")
```

#### Вычисление повторного интеграла: интегрирование по dx

```python
integrate_dx = sympy.integrate(gen_light_field_2d(_x, _y) * sympy.exp(-2 * sympy.pi * 1j * (_u * _x + _v * _y)), (_x, -_a, _a)).subs(_a, a)
```

#### Интегрирование по dy

```python
integrate_dy = sympy.integrate(integrate_dx.args[0][0], (_y, -_a, _a)).simplify()
```

#### Аналитическое решение в общем виде

```python
integrate_dy
```

$\displaystyle \begin{cases} \frac{\left(\left(- 6.25 \pi^{4} a^{2} u^{2} v^{2} e^{20.0 i \pi u} + 6.25 \pi^{4} a^{2} u^{2} v^{2} e^{40.0 i \pi u} + 1.25 i \pi^{3} a^{2} u v^{2} e^{20.0 i \pi u} + 1.25 i \pi^{3} a^{2} u v^{2} e^{40.0 i \pi u} + 0.125 \pi^{2} a^{2} v^{2} e^{20.0 i \pi u} - 0.125 \pi^{2} a^{2} v^{2} e^{40.0 i \pi u} + 6.25 i \pi^{3} a u^{2} v e^{20.0 i \pi u} - 6.25 i \pi^{3} a u^{2} v e^{40.0 i \pi u} + 1.25 \pi^{2} a u v e^{20.0 i \pi u} + 1.25 \pi^{2} a u v e^{40.0 i \pi u} - 0.125 i \pi a v e^{20.0 i \pi u} + 0.125 i \pi a v e^{40.0 i \pi u} + 3.125 \pi^{2} u^{2} e^{20.0 i \pi u} - 3.125 \pi^{2} u^{2} e^{40.0 i \pi u} - 0.625 i \pi u e^{20.0 i \pi u} - 0.625 i \pi u e^{40.0 i \pi u} - 0.0625 e^{20.0 i \pi u} + 0.0625 e^{40.0 i \pi u}\right) e^{30.0 i \pi u} + \left(6.25 \pi^{4} a^{2} u^{2} v^{2} e^{20.0 i \pi u} - 6.25 \pi^{4} a^{2} u^{2} v^{2} e^{40.0 i \pi u} - 1.25 i \pi^{3} a^{2} u v^{2} e^{20.0 i \pi u} - 1.25 i \pi^{3} a^{2} u v^{2} e^{40.0 i \pi u} - 0.125 \pi^{2} a^{2} v^{2} e^{20.0 i \pi u} + 0.125 \pi^{2} a^{2} v^{2} e^{40.0 i \pi u} + 6.25 i \pi^{3} a u^{2} v e^{20.0 i \pi u} - 6.25 i \pi^{3} a u^{2} v e^{40.0 i \pi u} + 1.25 \pi^{2} a u v e^{20.0 i \pi u} + 1.25 \pi^{2} a u v e^{40.0 i \pi u} - 0.125 i \pi a v e^{20.0 i \pi u} + 0.125 i \pi a v e^{40.0 i \pi u} - 3.125 \pi^{2} u^{2} e^{20.0 i \pi u} + 3.125 \pi^{2} u^{2} e^{40.0 i \pi u} + 0.625 i \pi u e^{20.0 i \pi u} + 0.625 i \pi u e^{40.0 i \pi u} + 0.0625 e^{20.0 i \pi u} - 0.0625 e^{40.0 i \pi u}\right) e^{i \pi \left(4.0 a v + 30.0 u\right)}\right) e^{- i \pi \left(2.0 a v + 60.0 u\right)}}{\pi^{6} u^{3} v^{3}} & \text{for}\: \pi^{15} u^{9} v^{6} e^{30.0 i \pi u} \neq 0 \\\frac{a^{3} \left(8.33333333333333 i \pi^{2} u^{2} e^{20.0 i \pi u} - 8.33333333333333 i \pi^{2} u^{2} e^{40.0 i \pi u} + 1.66666666666667 \pi u e^{20.0 i \pi u} + 1.66666666666667 \pi u e^{40.0 i \pi u} - 0.166666666666667 i e^{20.0 i \pi u} + 0.166666666666667 i e^{40.0 i \pi u}\right) e^{- 30.0 i \pi u}}{\pi^{3} u^{3}} & \text{otherwise} \end{cases}$

#### Аналитическое решение после подстановки $a$

```python
integrate_dy.subs(_a, a)
```

$\displaystyle \begin{cases} \frac{\left(\left(- 156.25 \pi^{4} u^{2} v^{2} e^{20.0 i \pi u} + 156.25 \pi^{4} u^{2} v^{2} e^{40.0 i \pi u} + 31.25 i \pi^{3} u^{2} v e^{20.0 i \pi u} - 31.25 i \pi^{3} u^{2} v e^{40.0 i \pi u} + 3.125 \pi^{2} u^{2} e^{20.0 i \pi u} - 3.125 \pi^{2} u^{2} e^{40.0 i \pi u} + 31.25 i \pi^{3} u v^{2} e^{20.0 i \pi u} + 31.25 i \pi^{3} u v^{2} e^{40.0 i \pi u} + 6.25 \pi^{2} u v e^{20.0 i \pi u} + 6.25 \pi^{2} u v e^{40.0 i \pi u} - 0.625 i \pi u e^{20.0 i \pi u} - 0.625 i \pi u e^{40.0 i \pi u} + 3.125 \pi^{2} v^{2} e^{20.0 i \pi u} - 3.125 \pi^{2} v^{2} e^{40.0 i \pi u} - 0.625 i \pi v e^{20.0 i \pi u} + 0.625 i \pi v e^{40.0 i \pi u} - 0.0625 e^{20.0 i \pi u} + 0.0625 e^{40.0 i \pi u}\right) e^{30.0 i \pi u} + \left(156.25 \pi^{4} u^{2} v^{2} e^{20.0 i \pi u} - 156.25 \pi^{4} u^{2} v^{2} e^{40.0 i \pi u} + 31.25 i \pi^{3} u^{2} v e^{20.0 i \pi u} - 31.25 i \pi^{3} u^{2} v e^{40.0 i \pi u} - 3.125 \pi^{2} u^{2} e^{20.0 i \pi u} + 3.125 \pi^{2} u^{2} e^{40.0 i \pi u} - 31.25 i \pi^{3} u v^{2} e^{20.0 i \pi u} - 31.25 i \pi^{3} u v^{2} e^{40.0 i \pi u} + 6.25 \pi^{2} u v e^{20.0 i \pi u} + 6.25 \pi^{2} u v e^{40.0 i \pi u} + 0.625 i \pi u e^{20.0 i \pi u} + 0.625 i \pi u e^{40.0 i \pi u} - 3.125 \pi^{2} v^{2} e^{20.0 i \pi u} + 3.125 \pi^{2} v^{2} e^{40.0 i \pi u} - 0.625 i \pi v e^{20.0 i \pi u} + 0.625 i \pi v e^{40.0 i \pi u} + 0.0625 e^{20.0 i \pi u} - 0.0625 e^{40.0 i \pi u}\right) e^{i \pi \left(30.0 u + 20.0 v\right)}\right) e^{- i \pi \left(60.0 u + 10.0 v\right)}}{\pi^{6} u^{3} v^{3}} & \text{for}\: \pi^{15} u^{9} v^{6} e^{30.0 i \pi u} \neq 0 \\\frac{125 \left(8.33333333333333 i \pi^{2} u^{2} e^{20.0 i \pi u} - 8.33333333333333 i \pi^{2} u^{2} e^{40.0 i \pi u} + 1.66666666666667 \pi u e^{20.0 i \pi u} + 1.66666666666667 \pi u e^{40.0 i \pi u} - 0.166666666666667 i e^{20.0 i \pi u} + 0.166666666666667 i e^{40.0 i \pi u}\right) e^{- 30.0 i \pi u}}{\pi^{3} u^{3}} & \text{otherwise} \end{cases}$

```python
analytical_function_2d = sympy.lambdify([_u, _v], integrate_dy.subs(_a, a).args[0][0], "numpy")
```

```python
fft_gaussian_beam_2d, domain = finite_fft_2d(light_field_2d, a=-a, b=a, N=N, M=M)

# - - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -- - -

lspace = np.linspace(domain[0], domain[1], num=N)

X, Y = np.meshgrid(lspace, lspace)

# Unused analytical function.
# analytical_result_2d = analytical_function_2d(X, Y)

analytical_result_2d = analytical_function(X) * analytical_function(Y)

fig = plt.figure(figsize=(22, 8))

ax = fig.add_subplot(121)
plt.imshow(np.absolute(analytical_result_2d))
plt.title(r"Амплитуда $x^2 \cdot y ^ 2$")
plt.colorbar()

ax = fig.add_subplot(122)
plt.imshow(np.angle(analytical_result_2d))
plt.title(r"Фаза $x^2 \cdot y ^ 2$")
plt.colorbar()

plt.show()
```