## Machine Translation
A sequence-to-sequence model is a model that takes a sequence of items and outputs another sequence of items using two networks that are trained end-to-end. This is perfect for machine translation since input sequences are directly related to output sequences. We will looking at preparing a dataset for Machine Translation task and implementing a seq2seq model. We will be using the parallel corpus available from [here](ftp://ftp.monash.edu/pub/nihongo/examples.utf.gz)

### Concepts covered here:
- Encoder-Decoder architecture: We will be using a encoder to encode input sequences from one language in a latent space and use the encoding to generate words in the target language one token at a time. The encoder and the decoder are an RNNs.

- Attention: We will be looking at the attention mechanism described in [Bahdanau et al., 2015](https://arxiv.org/pdf/1409.0473.pdf). We will explore the theory, the intuition behind it and how to move from theory and intuition to implementation.

```
%matplotlib inline
from pathlib import Path
import re,string
import numpy as np
from google.colab import files
```

```
!mkdir data
```

```
path = Path('data')
```

We will need to download the translation corpus and install a tokenizer for Japanese text.

```
#get corpus
#!wget ftp://ftp.monash.edu/pub/nihongo/examples.utf.gz
#decompress it and move it to data folder
#!gunzip examples.utf.gz
#!mv examples.utf data/
#install dependencies for mecab tokenizer
#!sudo apt install swig
#!sudo apt install mecab
#!sudo apt install libmecab-dev
#!sudo apt install mecab-ipadic-utf8
#!sudo pip3 install mecab-python3
```

```
corpus = (path/'examples.utf').open().readlines()
corpus[:5]
```

    ['A: ムーリエルは２０歳になりました。\tMuiriel is 20 now.#ID=1282_4707\n',
     'B: は 二十歳(はたち){２０歳} になる[01]{になりました}\n',
     'A: すぐに戻ります。\tI will be back soon.#ID=1284_4709\n',
     'B: 直ぐに{すぐに} 戻る{戻ります}\n',
     'A: すぐに諦めて昼寝をするかも知れない。\tI may give up soon and just nap instead.#ID=1300_4727\n']

```
def make_corpus(corpus_path):
    corpus = corpus_path.open().readlines()
    en,ja = [],[]
    pat = r'#ID.+\n'
    for c in corpus:
        if 'A: ' in c:
            clean_c = c.replace('A: ','')
            res = re.search(pat,clean_c)
            clean_c = clean_c.replace(res.group(0),'').split('\t')
            ja.append(clean_c[0])
            en.append(clean_c[1])
    return en,ja
```

```
en, ja = make_corpus(path/'examples.utf')
en[:2],ja[:2]
```

    (['Muiriel is 20 now.', 'I will be back soon.'],
     ['ムーリエルは２０歳になりました。', 'すぐに戻ります。'])

```
import MeCab
```

```
tagger = MeCab.Tagger('-Owakati')

def ja_tokenizer(text):
    result = tagger.parse(text)
    words = result.split()
    if len(words) ==0: return []
    if words[-1] == '\n':return words[:-1]
    return words
```

```
ja_tokenizer(ja[0])
```

    ['ムーリエル', 'は', '２', '０', '歳', 'に', 'なり', 'まし', 'た', '。']

```
import spacy
from spacy.symbols import ORTH
```

```
en_tok = spacy.load('en')
```

```
def en_tokenizer(text):
    text = text.lower()
    return [t.text for t in en_tok.tokenizer(text)]
```

```
en_tokenizer(en[0])
```

    ['muiriel', 'is', '20', 'now', '.']

```
en_toks = [en_tokenizer(text) for text in en]
ja_toks = [ja_tokenizer(text) for text in ja]
en_toks[:2], ja_toks[:2]
```

    ([['muiriel', 'is', '20', 'now', '.'],
      ['i', 'will', 'be', 'back', 'soon', '.']],
     [['ムーリエル', 'は', '２', '０', '歳', 'に', 'なり', 'まし', 'た', '。'],
      ['すぐ', 'に', '戻り', 'ます', '。']])

```
len(en_toks), len(ja_toks)
```

    (149785, 149785)

```
from collections import Counter,defaultdict
```

```
def numericalize_tok(tokens, max_vocab=50000, min_freq=0, unk_tok="xxunk", pad_tok="xxpad", bos_tok="xxbos", eos_tok="xxeos"):
    if isinstance(tokens, str):
        raise ValueError("Expected to receive a list of tokens. Received a string instead")
    if isinstance(tokens[0], list):
        tokens = [p for o in tokens for p in o]
    freq = Counter(tokens)
    int2tok = [o for o,c in freq.most_common(max_vocab) if c>min_freq]
    unk_id = 3
    int2tok.insert(0, bos_tok)
    int2tok.insert(1, pad_tok)
    int2tok.insert(2, eos_tok)
    int2tok.insert(unk_id, unk_tok)
    tok2int = defaultdict(lambda:unk_id, {v:k for k,v in enumerate(int2tok)})
    return int2tok, tok2int
```

```
int2j,j2int = numericalize_tok(ja_toks)
int2en,en2int = numericalize_tok(en_toks)
```

```
len(int2j), len(int2en)
```

    (31813, 21393)

```
import pickle
```

```
pickle.dump(int2j,(path/'int2j.pkl').open('wb'))
pickle.dump(int2en,(path/'int2en.pkl').open('wb'))
```

```
int2j = pickle.load((path/'int2j.pkl').open('rb'))
int2en = pickle.load((path/'int2en.pkl').open('rb'))
j2int = defaultdict(lambda:3, {v:k for k,v in enumerate(int2j)})
en2int = defaultdict(lambda:3, {v:k for k,v in enumerate(int2en)})
```

```
len(int2j), len(int2en)
```

    (12677, 9290)

```
j_ids = np.array([[0]+[j2int[o] for o in sent]+[2] for sent in ja_toks])
en_ids = np.array([[0]+[en2int[o] for o in sent]+[2] for sent in en_toks])
len(j_ids),len(en_ids), j_ids[10],en_ids[10]
```

    (149785,
     149785,
     [0,
      48,
      6,
      4891,
      5,
      109,
      11,
      143,
      10,
      83,
      8,
      57,
      86,
      1798,
      7,
      2146,
      232,
      255,
      47,
      36,
      4,
      2],
     [0, 114, 2251, 107, 38, 97, 85, 2649, 77, 28, 356, 4, 2])

```
np.random.seed(42)
```

```
trn_keep = np.random.rand(len(en_ids))>0.1
en_trn,j_trn = en_ids[trn_keep],j_ids[trn_keep]
en_val,j_val = en_ids[~trn_keep],j_ids[~trn_keep]
len(en_trn),len(en_val)
```

    (134774, 15011)

```
import torch
import torch.nn as nn
import torch.nn.functional as F

from torch.autograd.variable import Variable
from torch.utils.data import Dataset,DataLoader
```

```
from numpy import array as A
```

```
class Seq2SeqDataset(Dataset):
    def __init__(self, x, y): self.x,self.y = x,y
    def __getitem__(self, idx): return A(self.x[idx]), A(self.y[idx])
    def __len__(self): return len(self.x)
```

```
trn_ds = Seq2SeqDataset(en_trn,j_trn)
val_ds = Seq2SeqDataset(en_val,j_val)

bs = 120

trn_dl = DataLoader(trn_ds,batch_size=bs,shuffle=True)
val_dl = DataLoader(val_ds,batch_size=bs)
```

```
x, y = next(iter(val_dl))
x.size(), y.size()
```

    (torch.Size([120, 25]), torch.Size([120, 25]))

```
from keras.preprocessing.sequence import pad_sequences
```

```
enlen_90 = int(np.percentile([len(o) for o in en_ids], 99))
jalen_90 = int(np.percentile([len(o) for o in j_ids], 99))
enlen_90,jalen_90
```

    (25, 29)

```
j_ids = pad_sequences(j_ids, maxlen=29, dtype='int32', padding='post', truncating='post', value=1)
en_ids = pad_sequences(en_ids, maxlen=25, dtype='int32', padding='post', truncating='post', value=1)
```

```
trn_keep = np.random.rand(len(en_ids))>0.1
en_trn,j_trn = en_ids[trn_keep],j_ids[trn_keep]
en_val,j_val = en_ids[~trn_keep],j_ids[~trn_keep]
len(en_trn),len(en_val)
```

    (134746, 15039)

```
trn_ds = Seq2SeqDataset(en_trn,j_trn)
val_ds = Seq2SeqDataset(en_val,j_val)

bs = 120

trn_dl = DataLoader(trn_ds,batch_size=bs,shuffle=True)
val_dl = DataLoader(val_ds,batch_size=bs)
```

```
x, y = next(iter(val_dl))
x.size(), y.size()
```

    (torch.Size([120, 25]), torch.Size([120, 29]))

```
## load fasttext vectors
#code from here:https://github.com/facebookresearch/fastText/blob/master/docs/crawl-vectors.md
import io
def load_vectors(fname):
    fin = io.open(fname, 'r', encoding='utf-8', newline='\n', errors='ignore')
    header = fin.readline().split()
    n, d = int(header[0]), int(header[1])
    data = {}
    for line in fin:
        tokens = line.rstrip().split(' ')
        data[tokens[0]] = np.array(tokens[1:], dtype=float)
    return data, int(n), int(d)
```

```
# get word vectors
!wget https://dl.fbaipublicfiles.com/fasttext/vectors-english/wiki-news-300d-1M.vec.zip
!wget https://dl.fbaipublicfiles.com/fasttext/vectors-crawl/cc.ja.300.vec.gz
!unzip wiki-news-300d-1M.vec.zip
!gunzip cc.ja.300.vec.gz
!mv wiki-news-300d-1M.vec data/
!mv cc.ja.300.vec data/
```

    --2019-03-15 09:27:27--  https://dl.fbaipublicfiles.com/fasttext/vectors-english/wiki-news-300d-1M.vec.zip
    Resolving dl.fbaipublicfiles.com (dl.fbaipublicfiles.com)... 104.20.22.166, 104.20.6.166, 2606:4700:10::6814:6a6, ...
    Connecting to dl.fbaipublicfiles.com (dl.fbaipublicfiles.com)|104.20.22.166|:443... connected.
    HTTP request sent, awaiting response... 200 OK
    Length: 681808098 (650M) [application/zip]
    Saving to: ‘wiki-news-300d-1M.vec.zip’

    wiki-news-300d-1M.v 100%[===================>] 650.22M  30.1MB/s    in 22s

    2019-03-15 09:27:49 (29.4 MB/s) - ‘wiki-news-300d-1M.vec.zip’ saved [681808098/681808098]

    --2019-03-15 09:27:51--  https://dl.fbaipublicfiles.com/fasttext/vectors-crawl/cc.ja.300.vec.gz
    Resolving dl.fbaipublicfiles.com (dl.fbaipublicfiles.com)... 104.20.22.166, 104.20.6.166, 2606:4700:10::6814:6a6, ...
    Connecting to dl.fbaipublicfiles.com (dl.fbaipublicfiles.com)|104.20.22.166|:443... connected.
    HTTP request sent, awaiting response... 200 OK
    Length: 1279641604 (1.2G) [binary/octet-stream]
    Saving to: ‘cc.ja.300.vec.gz’

    cc.ja.300.vec.gz    100%[===================>]   1.19G  30.1MB/s    in 41s

    2019-03-15 09:28:32 (29.9 MB/s) - ‘cc.ja.300.vec.gz’ saved [1279641604/1279641604]

    Archive:  wiki-news-300d-1M.vec.zip
      inflating: wiki-news-300d-1M.vec

```
en_vecs,_,dim_en_vec = load_vectors('data/wiki-news-300d-1M.vec')
j_vecs,_,dim_j_vec = load_vectors('data/cc.ja.300.vec')
```

```
def create_emb(vecs, itos, em_sz):
    emb = nn.Embedding(len(itos), em_sz, padding_idx=1)
    if vecs is None: return emb
    wgts = emb.weight.data
    miss = []
    for i,w in enumerate(itos):
        try: wgts[i] = torch.from_numpy(vecs[w])
        except: miss.append(w)
    print('Number of unknowns in data: {}'.format(len(miss)))
    return emb

```

```
def V(tensor,req_grad=True):
    if torch.cuda.is_available():return Variable(tensor.cuda())
    else: return Variable(tensor)
```

### RNN Visualization

[RNN Visualization](http://jalammar.github.io/images/RNN_1.mp4)

### Seq2Seq Architecture
> ...a multilayered Long Short-Term Memory (LSTM) to map the input sequence to a vector of a fixed dimensionality, and then another deep LSTM to decode the target sequence from the vector.

[Sutskever et al., 2014](https://papers.nips.cc/paper/5346-sequence-to-sequence-learning-with-neural-networks.pdf)

Image:[from here](https://lilianweng.github.io/lil-log/2018/06/24/attention-attention.html)

```
class Seq2Seq(nn.Module):
    def __init__(self,int2en,int2j,em_sz,j_vecs=None,en_vecs=None,nh=128,out_sl=25,dropf=1,nl=2):
        super().__init__()
        #encoder
        self.nl,self.nh,self.em_sz,self.out_sl = nl,nh,em_sz,out_sl
        self.emb_enc = create_emb(en_vecs,int2en,em_sz)
        self.emb_drop = nn.Dropout(0.15*dropf)
        self.encoder = nn.GRU(em_sz,nh,num_layers=nl,dropout=0.25*dropf, bidirectional=True)
        #decoder
        self.emb_dec = create_emb(j_vecs,int2j,em_sz)
        self.decoder = nn.GRU(em_sz,nh*2,num_layers=nl,dropout=0.25*dropf)
        self.out_drop = nn.Dropout(0.35*dropf)
        self.out = nn.Linear(nh*2,len(int2j))

    def forward(self,inp,y=None):
        sl, bs = inp.size()
        emb_in = self.emb_drop(self.emb_enc(inp))
        h_n = self.initHidden(bs)
        enc_out, h_n = self.encoder(emb_in,h_n)
        h_n = h_n.view(2,2,bs,-1).permute(0,2,1,3).contiguous().view(self.nl,bs,-1)

        dec_inp = V(torch.zeros(bs).long())
        res = []
        for i in range(self.out_sl):
            dec_emb = self.emb_dec(dec_inp)
            outp,h_n = self.decoder(dec_emb.unsqueeze(0),h_n)
            outp = F.log_softmax(self.out(self.out_drop(outp[0])),dim=-1)
            res.append(outp)
            dec_inp = outp.data.max(1)[1]
            if (dec_inp==1).all(): break
        return torch.stack(res)

    def initHidden(self,bs):
        return V(torch.zeros([self.nl*2,bs,self.nh]))
```

### Why Bidirectional Encoder
> Finally, we found that reversing the order of the words in all source sentences (butnot target sentences) improved the LSTM’s performance markedly, because doing so introduced many short term dependencies between the source and the targetsentence which made the optimization problem easier.

[Sutskever et al., 2014](https://papers.nips.cc/paper/5346-sequence-to-sequence-learning-with-neural-networks.pdf)

Input and output sequences may not directly map to each other so preserving information from both passes of the input sequence will help learn how tokens relate to each other. For example in a translation task, subject and object can be in opposite positions depending on the language structure.

```
seq2seq = Seq2Seq(int2en,int2j,300,en_vecs=None,j_vecs=None)
seq2seq.cuda()
```

    Seq2Seq(
      (emb_enc): Embedding(21393, 300, padding_idx=1)
      (emb_drop): Dropout(p=0.15)
      (encoder): GRU(300, 128, num_layers=2, dropout=0.25, bidirectional=True)
      (emb_dec): Embedding(31813, 300, padding_idx=1)
      (decoder): GRU(300, 256, num_layers=2, dropout=0.25)
      (out_drop): Dropout(p=0.35)
      (out): Linear(in_features=256, out_features=31813, bias=True)
    )

```
out = seq2seq(V(x.transpose(1,0).long()))
out.size()
```

    torch.Size([25, 120, 31813])

Loss Function: We will use Cross Entropy Loss as we are trying to classify ot the correct words. Cross entropy loss can be simplified to:

`cross_entropy = sum(-log(y_pred) for y_pred in y_preds)`

where `y_pred` is the likelihood of the target class predicted by the model. This is a good loss function for classification because if the likelihood of the correct class is low, the loss value goes up and if it is high, the loss value goes down.

```
def seq2seq_loss(input, target):
    sl,bs = target.size()
    sl_in,bs_in,nc = input.size()
    if sl>sl_in: input = F.pad(input, (0,0,0,0,0,sl-sl_in))
    input = input[:sl]
    return F.cross_entropy(input.view(-1,nc), target.view(-1))
```

```
seq2seq_loss(out,V(y.transpose(1,0).long()))
```

    tensor(10.3560, device='cuda:0', grad_fn=<NllLossBackward>)

```
def step(x, y, epoch, m, crit, opt, clip=None):
    output = m(x, y)
    if isinstance(output,tuple): output = output[0]
    opt.zero_grad()
    loss = crit(output, y)
    loss.backward()
    if clip:
        nn.utils.clip_grad_norm_(m.parameters(), clip)
    opt.step()
    return loss.data.item()
```

```
from tqdm import tqdm
```

```
def train(trn_dl,val_dl,model,crit,opt,epochs=10,clip=None):
    for epoch in range(epochs):
        loss_val = loss_trn = 0
        with tqdm(total=len(trn_dl)) as pbar:
            model.train()
            for i, ds in enumerate(trn_dl):
                x, y = ds
                #if isinstance(x,tuple): x = x[0]
                x, y = x.transpose(1,0), y.transpose(1,0)
                loss = step(V(x.long()),V(y.long()),epoch,model,crit,opt)
                loss_trn += loss
                pbar.update()
        model.eval()
        for i, ds in enumerate(val_dl):
            with torch.no_grad():
                x, y = ds
                #if isinstance(x,tuple): x = x[0]
                x, y = x.transpose(1,0), y.transpose(1,0)
                out = model(V(x.long()))
                if isinstance(out,tuple): out = out[0]
                loss_val+= crit(out, V(y.long()))
                #loss_val +=loss
        print(f'Epoch: {epoch} trn loss: {loss_trn/len(trn_dl)} val loss: {loss_val/len(val_dl)}')
```

```
from torch import optim
```

```
opt = optim.Adam(seq2seq.parameters(),lr=3e-3,betas=(0.7,0.8))
```

```
train(trn_dl,val_dl,seq2seq,seq2seq_loss,opt,epochs=10)
```

    100%|██████████| 1124/1124 [02:40<00:00,  7.02it/s]
      0%|          | 1/1124 [00:00<03:05,  6.06it/s]

    Epoch: 0 trn loss: 4.750732962983359 val loss: 4.082951068878174

    100%|██████████| 1124/1124 [02:52<00:00,  6.51it/s]
      0%|          | 1/1124 [00:00<03:11,  5.86it/s]

    Epoch: 1 trn loss: 3.5997307712073003 val loss: 3.3246185779571533

    100%|██████████| 1124/1124 [02:55<00:00,  6.39it/s]
      0%|          | 1/1124 [00:00<03:04,  6.08it/s]

    Epoch: 2 trn loss: 3.1077632744965604 val loss: 3.4751904010772705

    100%|██████████| 1124/1124 [02:57<00:00,  6.34it/s]
      0%|          | 1/1124 [00:00<02:50,  6.60it/s]

    Epoch: 3 trn loss: 2.8518904808153036 val loss: 3.94274640083313

    100%|██████████| 1124/1124 [02:58<00:00,  6.30it/s]
      0%|          | 1/1124 [00:00<03:06,  6.01it/s]

    Epoch: 4 trn loss: 2.7045018161743135 val loss: 2.766321897506714

    100%|██████████| 1124/1124 [02:58<00:00,  6.29it/s]
      0%|          | 1/1124 [00:00<03:06,  6.01it/s]

    Epoch: 5 trn loss: 2.5743112220458713 val loss: 2.752124071121216

    100%|██████████| 1124/1124 [02:58<00:00,  6.28it/s]
      0%|          | 1/1124 [00:00<03:04,  6.08it/s]

    Epoch: 6 trn loss: 2.5332621384769998 val loss: 2.935807228088379

    100%|██████████| 1124/1124 [02:59<00:00,  6.27it/s]
      0%|          | 1/1124 [00:00<03:03,  6.11it/s]

    Epoch: 7 trn loss: 2.500852593323514 val loss: 3.133931875228882

    100%|██████████| 1124/1124 [02:59<00:00,  6.99it/s]
      0%|          | 1/1124 [00:00<03:04,  6.10it/s]

    Epoch: 8 trn loss: 2.458315055141245 val loss: 2.8610451221466064

    100%|██████████| 1124/1124 [02:59<00:00,  6.26it/s]

    Epoch: 9 trn loss: 2.456163497369909 val loss: 3.016350269317627

```
def produce_out(val_dl, model,int2en,int2j,interval=(20,30)):
    model.eval()
    x,y = next(iter(val_dl))
    x, y = x.transpose(1,0), y.transpose(1,0)
    probs = seq2seq(V(x.long()))
    if isinstance(probs,tuple): probs = probs[0]
    preds = A(probs.max(2)[1].cpu())
    for i in range(interval[0],interval[1]):
        print(' '.join([int2en[o] for o in x[:,i] if o not in [0,1,2]]))
        print(''.join([int2j[o] for o in y[:,i] if o not in [0,1,2]]))
        print(''.join([int2j[o] for o in preds[:,i] if o not in [0,1,2]]))
        print()
```

```
produce_out(trn_dl,seq2seq,int2en,int2j)
```

    the flower garden needs watering .
    その花壇は水をやる必要がある。
    その花は庭がががが。

    country life is very peaceful in comparison with city life .
    田舎での生活は、都会生活と比較してとても穏やかだ。
    国ののはははのののににに。。

    you are beautifully dressed .
    あなたはとても美しいドレスを着ていらっしゃいますね。
    あなたははなををててている。

    it 's quite all right .
    全くかまいません。
    それはですだ。

    the president put off visiting japan .
    大統領は訪日を延期しました。
    大統領は日本にををしたた

    he has two brothers , one lives in osaka and the other in kobe .
    彼には兄弟が二人いて、一人は大阪で、もう一人は神戸で暮らしている。
    彼はは２人２人ががが、人人でで。。

    a famous architect built this house .
    有名な建築家がこの家を建てた。
    そのな家家家家家家をたた。

    i could not but think that he had died .
    彼は死んでしまったと考えざるを得なかった。
    私は彼たたたはははたた。

    i enjoyed watching the easter parade .
    私は復活祭のパレードを見て楽しんだ。
    私はそのをを見を見を見。。

    i will not allow you to be ill - treated .
    君が虐待されているのを放ってはいられない。
    私は病気にににににない。。。

```
torch.save(seq2seq.state_dict(),open(path/'translate_seq2seq.pth','wb'))
```

```
seq2seq.load_state_dict(torch.load('translate_seq2seq.pth', map_location=lambda storage, loc: storage))
```

### Seq2Seq w/ Attention
A critical and apparent disadvantage of this fixed-length context vector design is incapability of remembering long sentences. Often it has forgotten the first part once it completes processing the whole input. The attention mechanism was born [Bahdanau et al., 2015](https://arxiv.org/pdf/1409.0473.pdf) to resolve this problem.

Given the following vectors:
\begin{align}
\boldsymbol{x} = \{x_1,x_2,x_3,\ldots,x_n\} \\
\boldsymbol{y} = \{y_1,y_2,y_3,\ldots,y_m\} \\
\end{align}
The hidden state from the Bidir encoder is given by:
\begin{align}
\boldsymbol{h}_i = [\overrightarrow{\boldsymbol{h}}_i^\top; \overleftarrow{\boldsymbol{h}}_i^\top]^\top, i=1,\dots,n \\
\end{align}

The hidden state from the decoder at time $t$ is given by:  $s_t$

The score for at time $t$:

\begin{aligned}
\text{score}(\boldsymbol{s}_t, \boldsymbol{h}_i) = \mathbf{v}_a^\top \tanh(\mathbf{W}_a[\boldsymbol{s}_t; \boldsymbol{h}_i])
\end{aligned}

where both $v_a$ and $W_a$ are weight matrices to be learned in the alignment model.

\begin{aligned}
\mathbf{c}_t &= \sum_{i=1}^n \alpha_{t,i} \boldsymbol{h}_i & \small{\text{; Context vector for output }y_t}\\
\alpha_{t,i} &= \text{align}(y_t, x_i) & \small{\text{; How well two words }y_t\text{ and }x_i\text{ are aligned.}}\\
&= \frac{\exp(\text{score}(\boldsymbol{s}_{t-1}, \boldsymbol{h}_i))}{\sum_{i'=1}^n \exp(\text{score}(\boldsymbol{s}_{t-1}, \boldsymbol{h}_{i'}))} & \small{\text{; Softmax of some predefined alignment score.}}.
\end{aligned}

Equations borrowed from [here](https://lilianweng.github.io/lil-log/2018/06/24/attention-attention.html)

[Attention Visualization](http://jalammar.github.io/images/attention_process.mp4)

```
import math,random

def rand_t(*sz): return torch.randn(sz)/math.sqrt(sz[0])
def rand_p(*sz): return nn.Parameter(rand_t(*sz))
```

```
class Seq2SeqAttention(nn.Module):
    def __init__(self,int2en,int2j,em_sz,j_vecs=None,en_vecs=None,nh=128,out_sl=25,dropf=1,nl=2):
        super().__init__()
        #encoder
        self.nl,self.nh,self.em_sz,self.out_sl = nl,nh,em_sz,out_sl
        self.emb_enc = create_emb(en_vecs,int2en,em_sz)
        self.emb_drop = nn.Dropout(0.15*dropf)
        self.encoder = nn.GRU(em_sz,nh,num_layers=nl,dropout=0.25*dropf, bidirectional=True)
        #decoder
        self.emb_dec = create_emb(j_vecs,int2j,em_sz)
        self.decoder = nn.GRU(em_sz,nh*2,num_layers=nl,dropout=0.25*dropf)
        self.out_drop = nn.Dropout(0.35*dropf)
        self.out = nn.Linear(nh*2,len(int2j))
        #attention layer
        self.W1 = rand_p(nh*2, nh*2) #parameter
        self.l2 = nn.Linear(nh*2, nh*2)
        self.l3 = nn.Linear(em_sz+nh*2, em_sz)
        self.V = rand_p(nh*2) #parameter

    def forward(self,inp,y=None):
        sl, bs = inp.size()
        emb_in = self.emb_drop(self.emb_enc(inp))
        h_n = self.initHidden(bs)
        enc_out, h_n = self.encoder(emb_in,h_n)
        h_n = h_n.view(2,2,bs,-1).permute(0,2,1,3).contiguous().view(self.nl,bs,-1)

        dec_inp = V(torch.zeros(bs).long())
        res,attns = [], []
        #multiply by parameter
        w1e = enc_out @ self.W1
        for i in range(self.out_sl):
            #linear layer
            w2h = self.l2(h_n[-1])
            #non-linear activation to calculate score
            u = torch.tanh(w1e + w2h)
            #softmax to make them into probs
            a = F.softmax(u @ self.V, 0)
            attns.append(a)
            #multiply each vector by scores and then add them up
            Xa = (a.unsqueeze(2) * enc_out).sum(0)
            dec_emb = self.emb_dec(dec_inp)
            #linear layer to reduce dimensions
            wgt_enc = self.l3(torch.cat([dec_emb, Xa], 1))
            outp,h_n = self.decoder(wgt_enc.unsqueeze(0),h_n)
            outp = F.log_softmax(self.out(self.out_drop(outp[0])),dim=-1)
            res.append(outp)
            dec_inp = outp.data.max(1)[1]
            if (random.random() > 0.5) and y is not None: dec_inp=y[i]
            if (dec_inp==1).all(): break
        return torch.stack(res),attns

    def initHidden(self,bs):
        return V(torch.zeros([self.nl*2,bs,self.nh]))
```

```
seq2seq = Seq2SeqAttention(int2en,int2j,300,en_vecs=None,j_vecs=None)
seq2seq.cuda()
```

    Seq2SeqAttention(
      (emb_enc): Embedding(21393, 300, padding_idx=1)
      (emb_drop): Dropout(p=0.15)
      (encoder): GRU(300, 128, num_layers=2, dropout=0.25, bidirectional=True)
      (emb_dec): Embedding(31813, 300, padding_idx=1)
      (decoder): GRU(300, 256, num_layers=2, dropout=0.25)
      (out_drop): Dropout(p=0.35)
      (out): Linear(in_features=256, out_features=31813, bias=True)
      (l2): Linear(in_features=256, out_features=256, bias=True)
      (l3): Linear(in_features=556, out_features=300, bias=True)
    )

```
opt = optim.Adam(seq2seq.parameters(),lr=3e-3,betas=(0.7,0.8))
```

```
train(trn_dl,val_dl,seq2seq,seq2seq_loss,opt,epochs=10)
```

    100%|██████████| 1124/1124 [03:21<00:00,  6.28it/s]
      0%|          | 1/1124 [00:00<03:27,  5.42it/s]

    Epoch: 0 trn loss: 2.797933133683595 val loss: 4.138897895812988

    100%|██████████| 1124/1124 [03:26<00:00,  6.24it/s]
      0%|          | 1/1124 [00:00<03:32,  5.28it/s]

    Epoch: 1 trn loss: 2.16145394790215 val loss: 2.8873095512390137

    100%|██████████| 1124/1124 [03:26<00:00,  6.17it/s]
      0%|          | 1/1124 [00:00<03:29,  5.35it/s]

    Epoch: 2 trn loss: 2.0703822573733075 val loss: 3.0309219360351562

    100%|██████████| 1124/1124 [03:26<00:00,  6.07it/s]
      0%|          | 1/1124 [00:00<03:28,  5.38it/s]

    Epoch: 3 trn loss: 2.07098806148322 val loss: 2.844566822052002

    100%|██████████| 1124/1124 [03:26<00:00,  6.07it/s]
      0%|          | 1/1124 [00:00<03:28,  5.38it/s]

    Epoch: 4 trn loss: 2.088127380906475 val loss: 2.8696582317352295

    100%|██████████| 1124/1124 [03:26<00:00,  6.26it/s]
      0%|          | 1/1124 [00:00<03:29,  5.36it/s]

    Epoch: 5 trn loss: 2.093448414603162 val loss: 2.7895729541778564

    100%|██████████| 1124/1124 [03:26<00:00,  6.12it/s]
      0%|          | 1/1124 [00:00<03:29,  5.37it/s]

    Epoch: 6 trn loss: 2.1055155513125383 val loss: 2.9072670936584473

    100%|██████████| 1124/1124 [03:26<00:00,  5.45it/s]
      0%|          | 1/1124 [00:00<03:30,  5.34it/s]

    Epoch: 7 trn loss: 2.11192792569191 val loss: 2.7808303833007812

    100%|██████████| 1124/1124 [03:26<00:00,  6.23it/s]
      0%|          | 1/1124 [00:00<03:29,  5.37it/s]

    Epoch: 8 trn loss: 2.116793169139543 val loss: 2.8389084339141846

    100%|██████████| 1124/1124 [03:26<00:00,  6.13it/s]

    Epoch: 9 trn loss: 2.1336968614325404 val loss: 2.8806753158569336

```
produce_out(trn_dl,seq2seq,int2en,int2j)
```

    mingle your joys sometimes with your earnest occupation .
    ときに君の喜びと君の真剣な職業とを交流せしめよ。
    あなたの仕事をはにはををする。

    a tunnel has been bored through the mountain .
    山を掘り抜いてトンネルが造られた。
    そのはは山ののをた。

    please refrain from smoking .
    どうぞタバコを控えてください。
    タバコをタバコをください。

    i am poor at tennis .
    私はテニスが下手だ。
    私はテニスがテニスが。

    i still owe my brother the ten dollars that he lent me last week .
    先週弟が貸してくれた１０ドル、借りたままだ。
    私は私ののののののののを１週間、ていた。

    there is a rumor that she got married .
    彼女が結婚したといううわさがある。
    彼女は結婚結婚した。

    he tried to appeal .
    彼は訴えようとした。
    彼はうとした。

    day is breaking .
    夜が明けかけてきた。
    日は日日。

    though she looks like his older sister , the fact is that she is his mother .
    彼女は彼の姉のように見えるが、実は母親なのだ。
    彼女は彼の妹に、て、、は、彼の。

    his failure in business left him penniless .
    彼は事業に失敗して一文なしになった。
    彼の失敗は失敗したのは失敗した。

```
torch.save(seq2seq.state_dict(),open('translate_seq2seq_attention.pth','wb'))
```

```
seq2seq.load_state_dict(torch.load('translate_seq2seq_attention.pth', map_location=lambda storage, loc: storage))
```

```
out, atts = seq2seq(x.long().cuda())
```

```
import matplotlib.pyplot as plt
plt.switch_backend('agg')
import matplotlib.ticker as ticker
```

```
torch.stack(atts).size()
```

    torch.Size([25, 120, 25])

```
plt.matshow(torch.stack(atts)[:,5,:].detach().cpu().numpy())
```

```

```


```python
import numpy as np
import scipy.special as sci
import matplotlib.pyplot as plt
from scipy import stats # linregress
import pandas as pd
from IPython.display import Latex
```

## Lecture 10:  Reactive Mass Transport

_(The contents presented in this section were re-developed principally by Dr. P. K. Yadav. The original contents are from Prof. Rudolf Liedl)_

---

The last lecture dealt with the  conservative transport processes and  quantified the mass flow and flux emanating from those processes. The effects of these processes were evaluated as an isolated processes and as joint transport process.

The last lecture dealt with the  conservative transport processes and  quantified the mass flow and flux emanating from those processes. The effects of these processes were evaluated as an isolated processes and as joint transport process.

```{admonition} Last lecture important conclusions
>  $J_{adv}>J_{dis}>>J_{diff}$ is normal aquifers.
> $J_{diff}$ may only be useful as an individual processes in special aquifers, e.g., clayey aquifers.
> In general aquifers, hydrodynamic dispersion $J_{hyd} = J_{dis} + J_{diff}$ is used in the analysis of solute transport process.
```

Finally, the last chapter introduced _Concentration Profile_ $(C-t)$ and _Breakthrough Curve_ $(C-x)$ to visually evaluate solute transport in aquifers using _Concentration_ $(C)$, a process output, as a function of _time_ $(t)$ and _space_ $(x)$.

Finally, the last chapter introduced _Concentration Profile_ $(C-t)$ and _Breakthrough Curve_ $(C-x)$ to visually evaluate solute transport in aquifers using _Concentration_ $(C)$, a process output, as a function of _time_ $(t)$ and _space_ $(x)$.

This section focuses on the _reactive transport processes,_ which as already discussed involves the transport of solute with _reaction_ processes. This course being an introductory groundwater course, _sorption_ and _degradation_ are the only two reaction types introduced and combined with the conservative transport processes- _advection_ and _dispersion_. Eventully, the section evaluates the joint action of conservative transport and reactive processes limiting to 1-D scenario.

The section will, however, first deal with 3-D effects of dispersive process, which is more important to quantify the reactive processes.

## Dispersive Mass Flow in 3-D

In the last section we saw that concentration gradient $\frac{\Delta C}{\Delta L}$ drives dispersive and diffusive solute transport process.

However, in the natural aquifers $\frac{\Delta C}{\Delta L}$ is normally varying with space $(x,y,z)$ and time $(t)$. Therefore, a differential operator $\big(\frac{\mathrm{d}}{\mathrm{d} x}\big)$ is more suitable representation of gradient than the difference operator $\frac{\Delta C}{\Delta x}$. The differential operator also generalizes the gradient case.

Considering the differential operator, the diffusive mass flow and diffusive mass flux (= mass flow per unit area) in 1D is then expressed as:

$$
J_{diff} = - n_e \cdot A \cdot D_p \cdot \frac{\mathrm{d} C}{\mathrm{d}x}
$$

and
$$
j_{diff} = - n_e \cdot D_p \cdot \frac{\mathrm{d} C}{\mathrm{d}x}
$$

Likewise, the dispersive mass flow and dispersive mass flux (1-D) is:

$$
J_{disp, h} = - n_e \cdot A \cdot D_{disp} \cdot \frac{\mathrm{d} C}{\mathrm{d}x}
$$

$$
j_{disp, h} = - n_e \cdot D_{disp} \cdot \frac{\mathrm{d} C}{\mathrm{d}x}
$$

The examples of these relations are presented in last section (@Alex pls. link)

### 3-D concentration Gradient

The concentration gradient  $\frac{\mathrm{d}C}{\mathrm{d}x}$  for 1-D solute transport problems is uni-directional, i.e, direction is fixed, and thus only the magnitude of the gradient is the important factor. However in higher dimensions, 2-D or 3-D, solute transport problems, the _direction_ of gradient along with it's _magnitude_ in that direction has to be specified. Thus, for higher dimension solute transport problems, the _concentration gradient_ becomes **concentration vector**, i.e., a quantity providing both magnitude and direction.

Thus, the representation of concentration gradient in Cartesian coordinate in 2-D and 3-D is:

$$
\mathrm{grad}C = \nabla C=  \begin{pmatrix}
\frac{\partial C}{\partial x}\\
\frac{\partial C}{\partial y}
\end{pmatrix}
$$

and

$$
\mathrm{grad} C = \nabla C=  \begin{pmatrix}
\frac{\partial C}{\partial x}\\
\frac{\partial C}{\partial y}\\
\frac{\partial C}{\partial z}
\end{pmatrix}
$$

The $\nabla$, the inverted Delta symbol, is called the **del** or **nabla** operator. The vector **grad$C$** in the above relations points in the direction of the _steepest increase_ of $C$. However, for the **Hydrogeologists**, the concentration gradients as well the grad$C$ points to the _steepest decrease_ of $C$.

**@Anne/Sophie - can we try to find a simple example finding grad$C$ in 2-D/3-D**

### Isotropic and Anisotropic  Dispersion

Corresponding the expression for the concentration gradient at higher dimensions, the expression for mass flow and flux becomes:

$$
J_{disp,\, h} = \begin{pmatrix}
J_{disp,\, hx}\\
J_{disp,\, hy}\\
J_{disp,\, hz}
\end{pmatrix}
$$
and the 3-D mass flux is:

$$
j_{disp,\, h} = \begin{pmatrix}
j_{disp,\, hx}\\
j_{disp,\, hy}\\
j_{disp,\, hz}
\end{pmatrix}
$$

The subscript in ${{disp, h}}$ refers to _hydrodynamic dispersion_ which is sum of _mechanical dispersion_ and _diffusion_. Likewise, the subscript ${{disp,\, hx}}$, ${{disp,\, hy}}$ and ${{disp,\, hz}}$ refers to dispersion components along the Cartesian coordinates. The corresponding mass flow and mass flux in the higher dimension is then:

$$
J_{disp,\, h} = - n_e \cdot A \cdot D_{hyd} \cdot \text{grad}C
$$

and

###  Isotropic and Anisotropic  Dispersion

Corresponding the expression for the concentration gradient at higher dimensions, the expression for mass flow and flux becomes:

$$
J_{disp,\, h} = \begin{pmatrix}
J_{disp,\, hx}\\
J_{disp,\, hy}\\
J_{disp,\, hz}
\end{pmatrix}
$$
and the 3-D mass flux is:

$$
j_{disp,\, h} = \begin{pmatrix}
j_{disp,\, hx}\\
j_{disp,\, hy}\\
j_{disp,\, hz}
\end{pmatrix}
$$

The subscript in ${{disp,\, h}}$ refers to _hydrodynamic dispersion_ which is sum of _mechanical dispersion_ and _diffusion_. Likewise, the subscript ${{disp,\, hx}}$, ${{disp,\, hy}}$ and ${{disp,\, hz}}$ refers to dispersion components along the Cartesian coordinates. The corresponding mass flow and mass flux in the higher dimension is then:

$$
J_{disp,\, h} = - n_e \cdot A \cdot D_{hyd} \cdot \text{grad}C
$$

and

$$
j_{disp,\, h} = - n_e \cdot A \cdot D_{hyd} \cdot \text{grad}C
$$

The **isotropic dispersion**, rather an _exceptional,_ the $D_{hyd}$ in this case is:

$$
D_{hyd} = \alpha \cdot |v| + n_e \cdot D
$$
where, $D$ is direction independent dispersion coefficient and $\alpha$ $[L]$ is dispersivity, which in:

**heterogeneous aquifer**: $\alpha = \alpha(x,y,z)$ and in

**homogeneous aquifer**: $\alpha = \text{constant}$

For more practical cases and in normal aquifers, the 2-D and 3-D the dispersion for solute transport is _direction dependent,_ i.e. **anisotropic**. Hence the $D_{hyd}$ is not an scalar quantity but a _matrix (tensor),_ which relates the concentration gradient (vector) to the dispersive mass flow (vector). However, if the princopal axes of the dispersion tensor $D_{hyd}$ is made to coincide with the axes of a Cartesian coordinate system _and_ the groundwater flow is considered uniform along the $x-$axis, the dispersive mass flux can be obtained from

$$
\begin{pmatrix} J_x \\ j_y \\ j_z \end{pmatrix} =
\begin{pmatrix} \alpha_L \cdot v_x + n_e \cdot D & 0 & 0 \\
0 & \alpha_{Th} \cdot v_x + n_e \cdot D & 0\\
0 & 0 & \alpha_{Tv} \cdot v_x + n_e \cdot D
\end{pmatrix}
\cdot
\begin{pmatrix} \frac{\partial C}{\partial x} \\ \frac{\partial C}{\partial y} \\ \frac{\partial C}{\partial z} \end{pmatrix}
$$

with $\alpha_L$, $\alpha_{Th}$ and $\alpha_{Tv}$ are longitudinal dispersivity, horizontal transverse dispversity and vertical transverse dispersivity, respectively. The statistical analysis of dispersivity data shows that $\alpha_{L}>\alpha_{Th}>\alpha_{Tv}$ and the values differ by roughly an order of magnitude. This, however, is just a rule of thumb.

@Anne/@Sophie can you find a very simple numerical example for caculating dispersivity or something that explains the above relation.

```python
# Analytical solution from Bear (1976) - Line source, 1st-type input and infinte plane

# Input (values can be changed)
Co = 1 # mg/L, input concentration
Dx = 3 # m, Dispersion in x direction
Dy = Dx/10 # m
v = 0.05 # m/d
Q = 10 # m^3/d

## domain dimension and descritization (values can be changed)
xmin = -100; xmax= 101
ymin = 0.1; ymax = 11
[x, y] = np.meshgrid(np.linspace(xmin, xmax, 1000), np.linspace(ymin, ymax, 100)) # mesh

# Bear (1976) solution Implementation
#"k0: Modified Bessel function of second type and zero Order"

term1 = (Co*Q)/(2*np.pi* np.sqrt(Dx*Dy))
term2 = (x*v)/(2*Dx)
args = (v**2*x**2)/(4*Dx**2) + (v**2*y**2)/(4*Dx*Dy)
sol = term1*np.exp(term2)*sci.k0(args)

# plots
fig, ax = plt.subplots()
CS = ax.contour(x,y,sol, cmap='flag')
ax.clabel(CS, inline=1, fontsize= 10)
CB = fig.colorbar(CS, shrink=0.8, extend='both')
```

## Equilibrium Sorption ##

---

A reactive transport system can include a single reactive process, e.g., degradation, or combination of several multiple reactive processes, e.g. degradation and sorption. The inclusion of the reactive process(es) in the transport studies are site specific. The important to note is that an inclusion of a reactive process increases the complexity of transport problem.

In this course we limit ourselves with the following two types of reaction processes:

**1. Sorption**

**2. Degradation**

Acid-base reaction, precipitation-dissolution reaction, organic combustion etc. are among the reactions type that can be part of the reactive process individually or in any combination.

Also an important distinction is the rate or speed of the reaction. One distinguishes between time-dependent reaction (kinetics) or time-independent reactions (steady-state or equilibrium). Special reaction rates such as instantaneous reaction (extremely fast reaction) can also be part of the reaction process in the transport system.

### Sorption Basics

**Sorption** is a rather a general term used to indicate both **adsorption** and **absorption**. But in this course the _sorption_ refers to only _adsorption._

**Adsorption** can be more formally defined as the process of accumulation of dissolved chemicals on the surface of a solid, e.g., accumulation of a chemicals dissolved in groundwater on the surface of the aquifer material.

The figure below clarifies the _adsorption_ process.

In the figure _chemical in solution_ (the circular objects) more often called **solute** in the water is found to attach is the solid surface. The figure presents the following two important terms part of the adsorption process:

> **Adsorbent**: The solid onto which the chemicals are attached. More formally, _adsorbents_ provide adsorption sites for solutes.

> **Adsorbate**: These are solutes that are attached on the _adsorbent._

Based on the figure, _adsorption_ can be considered as a partition process that divides the chemical originally present in water between adsorbent and water.

Quite often adsorption is a reversible process, i.e., adsorbed chemicals can get back to water phase. This process is called **desorption**.

Speaking about _equilibrium,_ this is reached when

> _adsorption rate_  $\rightleftharpoons $ _desorption rate_

Adsorption in groundwater is often a rapid process. Although sorption kinetics can be important, the description in this introductory level course is limited to equilibrium sorption. Thus, we learn next to quantify equilibrium sorption.

### Adsorption Isotherms

The adsorption process that has reached equilibrium can be relatively easily quantified with the use of empirical models called **isotherms**. These models are often simple algebraic equation that relates solute concentrations partitioned between the adsorbate and adsorbent at constant temperature. More than 15 different _isotherm_ models can be found in the literature. However, in groundwater reactive transport studies the following three are the two most commonly used isotherms:

1. **Henry or Linear isotherm**

2. **Freundlich isotherm**

For quantification, laboratory based experiments are performed using solids from subsurface and chemicals of interest. The laboratory observations are then graphically fitted with empirical isotherm models to quantify adsorption properties. Figure below shows isotherms that are particularly observed in groundwater transport studies. As can be observed in the figure sorption coefficient ($K$) is the common quantities obtained from isotherm models.

### Henry Isotherm

The **Henry isotherm** (Henry, 1803) is based on the idea of a _linear_ relationship between the solute concentration **$C$** and the _adsorbate:adsorbent_ mass ratio $C_a$. Henry isotherm is quite often also called _linear isotherm_ or the $K_d$ model. Mathematically, the Henry isotherm is:

$$
C_a = K_d \cdot C
$$

with

$C$ = solute concentration [ML$^{-3}$]

$C_a$ = mass ratio adsorbate:adsorbent [M:M]

$K_d$ = distribution or partitioning coefficient [L$^3$M$^{-1}$].

Often symbols $C_s$ or $s$ are used instead of $C_a$.

The Henry model has been most widely used in groundwater transport studies. This is largely because of the simplicity (see equation) of the model and it's applicability in representing adsorption process more generally observed in groundwater studies. $K_d$, the partitioning coefficient, is particularly used in groundwater transport studies. It is equal to the slope of the Henry isotherm.

```python
# Example of Henry isotherm (Source: Fetter et al. 2018)

# Following sorption data are available:

C = np.array([7, 15, 174, 249, 362]) # ug/L, Eq. concentration
Ca = np.array([2, 4, 33, 50, 70]) # ug/g, Eq. sorbed mass

# linear- fit y = m*x+c

slope, intercept, r_value, p_value, std_err = stats.linregress(C, Ca)
print("slope: %f    intercept: %f  R-squared: %f" % (slope, intercept, r_value**2))
fit_line = slope*C + intercept

#plot

plt.scatter(C, Ca, label= "Original data") # data plot
plt.plot(C, fit_line, color = "red", label = "fit-line")
plt.legend(); plt.xlabel(r"Equilibrium Aqueous Concentration, $C$ ($\mu$g/L) ")
plt.ylabel(r"Mass sorbed per unit absorbent weight, $C_a$ ($\mu$g/g) ");
plt.text(0, 50, '$C_a=%0.5s C + %0.5s$'%(slope, intercept), fontsize=10)
plt.text(0, 40, '$R^2=%0.5s $'%(r_value**2), fontsize=10)

# Output

Latex("The required partition coefficient = slope, $K_{d}$= %0.5s L/g " % slope)
```

### Freundlich Isotherm

**Freundlich isotherm** (Freundlich, 1907) is a more general isotherm. It is based on the idea of a power law, i.e., includes also the non-linear behaviour, relating the solute concentration $C$ to the adsorbate:adsorbent mass ration $C_a$. The isotherm is mathematically given as

$$
C_a = K_{Fr} \cdot C^N
$$

with

$C$ = solute concentration [ML$^{-3}$]

$C_a$ = mass ratio adsorbate:adsorbent [M:M]

$n$ = Freundlich exponent [-]

$K_{Fr}$ = Freundlich partitioning coefficient [(M:M)/(M/L$^3)^n$].

The Freundlich isotherm equation can be easily linearized by applying logarithmic transformation of the equation, which gives

\begin{eqnarray}
\log C_a = \log K_{Fr} + n\cdot \log C
\end{eqnarray}

The above equation resembles the straight line equation $y = b + a \cdot x b$, in which $b\equiv \log K_{Fr}$ is the intercept and $a\equiv n$ the slope. Thus, from fitting the adsoprtion experimental results with the above equation, both $n$ and $K_{Fr}$ can be obtained.

```python
# Example of Freundlich isotherm

# Following sorption data are available:

Cf= np.array([23.6, 6.67, 3.26, 0.322, 0.169, 0.114]) # mg/L, Eq. concentration
Caf = np.array([737, 450, 318, 121, 85.2, 75.8]) # mg/g, Eq. sorbed mass

logCf = np.log10(Cf) # log10 transformation of data
logCaf = np.log10(Caf)

# fitting: y = mx +c
slope, intercept, r_value, p_value, std_err = stats.linregress(logCf, logCaf)
print("slope: %f    intercept: %f  R-squared: %f" % (slope, intercept, r_value**2))
fit_line = slope*logCf + intercept

# plots
plt.figure(figsize=(10,4))

plt.subplot(121)
plt.plot(Cf, Caf, "*--", label= "Original data")
plt.legend(); plt.xlabel(r"Eq. Aq. Conc., $C$ ($mg$g/L) ");
plt.ylabel(r"Mass sorbed/absorbent weight, $C_a$ ($mg$g/g) ");

plt.subplot(122)
plt.scatter(logCf, logCaf, label="Log transformed data")
plt.plot(logCf, fit_line, color="red", label= "linear fit line")
plt.legend(); plt.xlabel(r"Eq. Aq. Conc., $\log C$ ($mg$g/L) ");
plt.ylabel(r"Mass sorbed/absorbent weight, $\log C_a$ ($mg$/g) ");
plt.text(-1, 2.6, '$C_a=%0.5s C + %0.5s$'%(slope, intercept), fontsize=10)
plt.text(-1, 2.5, '$R^2=%0.5s $'%(r_value**2), fontsize=10)
plt.subplots_adjust(wspace=0.35)

Latex("$K_{Fr}$ = %0.5s (mg/g)$^{1/n}$(mg/L) and   $n$ = %0.4s"  % (10**intercept, slope))
```

### Retardation Factor (for Henry Isotherm)

The net effect of adsorption is the retarded movement of solute in comparison to the average flow of the groundwater. The term **Retardation Factor** $(R)$ is defined that quantifies the retarded movement of solute. The formulation of $R$ is based on the type of isotherm. For Henry isotherm $R$ can be straightforwardly calculated with the help of a mass budget.

For this purpose, an aquifer volume $V$ with the effective porosity $n_e$ is considered (see fig. below)

The steps involved are:

- Total volume: $V$
- Water volume: $n_e \cdot V$
- Mass of dissolved chemical: $n_e \cdot V \cdot C$
- Volume of solid: $(1-n_e)\cdot V$
- Density of solid material: $\rho$
- Mass of solid: $\rho \cdot(1-n_e)\cdot V$
- Mass of adsorbate: $\rho \cdot(1-n_e)\cdot V\cdot C_a$ = $(1-n_e)\cdot\rho \cdot V \cdot K_d \cdot C$
- Total mass: $n_e\cdot V \cdot C + (1-n_e)\cdot\rho \cdot V\cdot K_d \cdot C = n_e \cdot R \cdot V \cdot C$

with _Retardation factor_ $$R = 1 + \frac{1-n_e}{n_e}\cdot \rho \cdot K_d$$

The expression for $R$ can be further modified  by using bulk density $\rho_b$ $= (1-n_e)\cdot \rho$ = mass of solid/total volume. This leads to

$$
R = 1+\frac{\rho_b}{n_e} \cdot K_d
$$

As can be observed from the equation, $R = 1$ when there is no adsorption, i.e., when $K_d= 0$.

@ Anne @ Sophie pls. provide a very short numerical example on R

## Degradation

**Degradation** leads to alteration or transformation of chemical structure of chemicals. This contrasts to adsorption in which chemical structure is not altered. In adsorption (or desorption) the original chemical is partitioned between the solid particles and water. It is _degradation_ that eventually lead to removal of the _original_ chemical from the groundwater. The transformation of original chemical, due to degradation, results to so-called _daughter products (metabolites)._ The new chemical(s) can make groundwater more suitable (decrease contamination) or further contaminate it.

In groundwater studies, degradation can appear as:

- **Radioactive decay**
- **Microbial degradation (bio-degradation)**
- **Chemical degradation**

There are several approaches to quantify degradation process.
A common aspect to most of them is the assumption of _time-dependency_ (or _Kinetics_ ).

### $n^{th}$ - Order Degradation Kinetics

The general equation for the degradation kinetics is:

$$
\frac{\text{d}C}{\text{d} t} = - \lambda \cdot C^n
$$

with $t$ = time [t]

$C$ = solute concentration [ML$^{-3}$]

$n$ = order of the degradation kinetics [ - ] ($n\geq 0)$

$\lambda$ = degradation rate constant [(ML$^{-3})^{(1-n)}$T$^{-1}$].

Considering the initial concentration (or input concentration) $C_0$, the solutions of the kinetics equation are:

$$
C(t) = C_0\cdot e^{-\lambda \cdot t} \: \: \: \text{if }\: n = 1
$$

and

$$
C(t) = [C_0^{1-n} - (1-n)\cdot \lambda t]^{\frac{1}{1-n}} \:\:\: \text{if }\: n\neq 1
$$

The **half life** $(T_{1/2})$, which is the time span elapsing until the initial concentration $C_0$ is reduced by half, is an important time-scale in the degradation analysis. $T_{1/2}$ is $C_0$ dependent in nearly all cases with an exception for 1$^\text{st}$- order degradation kinetics. 0$^{th}$-order and the 1$^\text{st}$- order degradation kinetics are most commonly observed in groundwater studies. The $(T_{1/2})$ of these orders are:

$$
T_{1/2} = \frac{C_0}{2\cdot \lambda} \:\:\: \text{for } \:0^{\text{th}}\text{-order}
$$

$$
T_{1/2} = \frac{\ln 2}{\lambda} \:\:\: \text{for } \:1^{\text{st}}\text{-order}
$$

As can be observed above $T_{1/2}$ is independent of concentration for the 1$^{\text{st}}$-order degradation kinetics.

Another important properties of the degradation kinetics is that for $n\geq 1$ the solute concentration _asymptotically_ approaches zero, whereas for $n<1$, the solute concentration actually reaches zero

```python
# behaviour of degradation kinetics

#input - you may change the values
Co = 1 # mg/L, initial concentration
la = 0.003 # unit is order dependent. For n=1, 1/t

# main equation
Z_order = lambda t: Co* np.exp(-la*t) # for n = 0
F_order = lambda t: Co-la*t # for n = 1

# simulation for t
t = np.linspace(1,1000, 1000) # 1000 time units
Z_results = Z_order(t)
F_results = F_order(t)

# plots
plt.figure(figsize=(10,4))

# n = 1
plt.subplot(121)
plt.plot(t, Z_results)
plt.ylim(0, Co); plt.xlim(0)
plt.text(400, Co*0.8, r"$C(t) = C_0 \cdot e^{-\lambda \cdot t} $", fontsize = 12)
plt.text (400, Co*0.9, r"1$^{st}$-order kinetics", color = "red", fontsize = 12)
plt.text(0, Co/2, r"$T_{1/2}= \frac{\ln 2}{\lambda}$", color= "red", fontsize=14)
plt.xlabel("Time, t  (days)"); plt.ylabel(r"Concentration, $C(t)$ (mg/L)")

# n = 0
plt.subplot(122)
plt.plot(t, F_results)
plt.ylim(0, Co); plt.xlim(0)
plt.text(400, Co*0.8, r"$C(t) = C_0 \cdot-\lambda \cdot t} $", fontsize = 12)
plt.text (400, Co*0.9, r"0$^{th}$-order kinetics", color = "red", fontsize = 12)
plt.text(0, Co/2, r"$T_{1/2}= \frac{C_0}{2\cdot \lambda}$", color= "red", fontsize=14)
plt.xlabel("Time, t  (days)"); plt.ylabel(r"Concentration, $C(t)$ (mg/L)")

plt.subplots_adjust(wspace=0.35)
```

### Radioactive decay

Radioactive decay is degradation of a chemical due to radiation. The radioactive decay is limited to radioactive chemicals such as Cobalt, Cesium, Iodine. This decay obeys the 1$^\text{st}$- order degradation kinetics and therefore the half-life is $T_{1/2} = \frac{\ln 2}{\lambda}$. $T_{1/2}$ is characteristic property of radioactive chemicals and it can be used to compute degradation rate ($\lambda$).

```python
# Example of Radioactive decaly

#experimental results

t = [0, 1, 2, 5, 10, 20, 28 ] # yr, time
Co_60 = [10, 8.76, 7.68, 5.17, 2.68, 0.72, 0.25] # mg/L, Cobalt 60 conc.
So_90 = [10, 9.76, 9.52, 8.84, 7.81, 6.10, 5] # mg/L, Strontium 90 Conc.

z_list = list(zip(t, Co_60, So_90))

Cols= ["time (a)", "Cobalt 60 (mg/L)", "Strontium 90 (mg/L)"]
df = pd.DataFrame(z_list, columns=Cols)
print(df)

# computing
TH_Co60 = 28 # yr, Half life of Cobalt 60
TH_St90 = 5.26 # yr, Half life of Strontium 90
la_Co60 = np.log(2)/TH_Co60 # 1/yr, degradation rate of Cobalt 60
la_St90 = np.log(2)/TH_St90 # 1/yr, degradation rate of Strontium 90

# visualize
plt.plot(t, Co_60, "o--", label = "Cobalt 60")
plt.plot(t, So_90, "v--", label= "Strontium 90")
plt.xlabel("Time (years)"); plt.ylabel("Concentration (mg/L)")
plt.legend();

Latex("The degradation rate ($\lambda$) for Cobalt 60 = %0.5s 1/y and for Strontium 90 = %0.5s 1/y"  % (la_Co60, la_St90))
```

## Joint Action of Conservative and Reactive Transport (1D)

### Concentration Profile

Figure below presents the joint action of conservative transport with equilibrium sorption (linear isotherm) and degradation. The figure shows the solute concentration $C$ (in water) at the same time-levels for various combinations of acting processes.

The figure can be explained in the following way:

(A): The solute is initially present at constant concentration in a limited area.

(B): Solute spreads only due to advection. Due to absence of dispersion there is no (1D) spreading effect.

(C): Inclusion of dispersion process causes spread of concentration. As retardation is absence the front centreline remains unchanged

(D): The inclusion of retardation ($R$) with advection and dispersion leads to removal of chemicals from water and as well the retarded movement of the chemical front.

(E): The inclusion of retardation along with degradation and conservative transport process leads to high removal of chemical from water.

### Breakthrough Curve

Breakthrough curves provide a _time-dependent_ spread of chemicals in the groundwater. The inclusion of multiple processes are normally solved using numerical models. Analytical models are available for limited processes and simplified problems. A 1-D analytical solution by Kinzelbach (1987) provide a transient (time-dependent) solution of reactive transport problem with inclusion of equilibrium linear sorption represented by retardation $(R)$, first-order degradation rate $(\lambda)$ and the conservative transport quantities - dispersion $(D)$ and advection. The solution is given as:

$$
C(x,t) = C_0 \cdot \exp(-\lambda\cdot t)\bigg(1- \frac{1}{2}\text{erfc}\bigg(\frac{R\cdot x - v\cdot t}{2\cdot\sqrt{D\cdot R \cdot t}}\bigg) - \frac{1}{2}\exp\bigg(\frac{v\cdot x}{D}\bigg)\text{erfc}\bigg(\frac{R\cdot x + v\cdot t}{2\cdot\sqrt{D\cdot R \cdot t}}\bigg)
$$

with $C_0$ = input/source concentration [ML$^{-3}$]

$t$ = time [T]

$v$ = groundwater flow velocity [LT$^{-1}$]

erfc() = represents the complementary error function [See here for details](https://en.wikipedia.org/wiki/Error_function). erfc() can be easily computed using Python Scipy special function library.

```python
# Breakthrough curve using Kinzelbach (1987) analytical solution Main function

# The main function - you may change the value of C_o, lam, R, Dx, v, x
# C_o = input concentration, mg/L
# lam = 0 # 1/d, degradation rate, 1/d
# R = retardation factor, ()
# Dx = dispersion coeff. along x, m^2/d
# v = groundwater velocity, m/d
# x = position where C is to be measured, m

def Cx(t, C_o= 1, lam = 0, R=1, Dx=1, v= 10, x = 20):
    sterm = C_o*np.exp(-lam*t)
    erf_ag1 = (R*x-v*t)/(2*np.sqrt(Dx*R*t))
    erf_ag2 = (R*x+v*t)/(2*np.sqrt(Dx*R*t))

    C = sterm*(1-(0.5*sci.erfc(erf_ag1)-0.5*np.exp((v*x)/Dx)*sci.erfc(erf_ag2)))
    return C

```

```python
# Computing Case 1: Conservative process- R = 1, Lambda = 0
t1 = np.linspace(1e-5,50,1000) # times, d
C1 = Cx(t1, C_o= 1, lam = 0, R=1, Dx=1, v= 1, x = 20)

# Computing Case 2: Conservative system + Retardation - R = 2, Lambda = 0
t2 = np.linspace(1e-5,50,1000) # times, d
C2 = Cx(t2, C_o= 1, lam = 0, R=2, Dx=1, v= 1, x = 20)

# Computing Case 3: Conservative system + Retardation + degradation - R = 2, Lambda = 0.004
t3 = np.linspace(1e-5,50,1000) # times, d
C3 = Cx(t3, C_o= 1, lam = 0.004, R=2, Dx=1, v= 1, x = 20)

# plots - this should be adjusted as required

plt.figure(figsize=(9, 6))

plt.plot(t1, C1, label="Conservative transport")
plt.plot(t2, C2, label = "Reactive transport with sorption")
plt.plot(t3, C3, label = "Reactive transport with sorption and degradation rate")
plt.legend(loc= 3); plt.xlim(0), plt.ylim(0)
plt.xlabel("time (d)"); plt.ylabel(r"Concentration, $C$ (mg/L)")
plt.text(5, 0.2, r"$x= 20$ m")
```

### Mass (Re-)Distribution During Injection / Extraction

**Consider a scenario**:

Water is _injected_ into a certain portion of an aquifer with total volume $V$, bulk density $\rho_b$ and effective porosity $n_e$. Assume that the injected water contains a chemical of total mass $M$, which is adsorbed by the aquifer materials under equilibrium conditions according to Henry isotherm (quantified by $K_d$$).

Bases on the assumption of sorption equilibrium, the total mass $M$ of the chemical is instantaneously(!) split up into a dissolved and a sorbed part. In such case, the mass distribution can be computed as follows (with $R$ = retardation factor, $\rho_b$
= bulk density):

\begin{align}
M &= n_e \cdot V \cdot C + V\cdot\rho_b \cdot C_a \\
&= n_e \cdot V \cdot C + V \cdot \rho_b \cdot K_d \cdot C\\
&= n_e \cdot (1 + \rho_b \cdot K_d/n_e) \cdot V \cdot C\\
&= n_e \cdot R \cdot V \cdot C
\end{align}

In which,

$n_e \cdot V \cdot C$ = dissolved mass

$V\cdot\rho_b \cdot C_a$ = mass of adsorbate

For the dissolved mass we thus have $n_e \cdot V \cdot C = M/R$ and consequently the mass of adsorbate is:

$V\cdot\rho_b \cdot C_a = M- M/R = (1-1/R)\cdot M$

The same approach can be adopted for the **extraction** scenarios, i.e. equilibrium desorption.


```julia
"""Tutorial: Second-Order Moller--Plesset Perturbation Theory (MP2)"""

__author__    = ["D. Menendez", "Dominic A. Sirianni"]
__credit__    = ["D. Menendez", "Dominic A. Sirianni", "Daniel G. A. Smith"]

__copyright__ = "(c) 2014-2020, The Psi4Julia Developers"
__license__   = "BSD-3-Clause"
__date__      = "2020-07-30"
```

    "2020-07-30"

# Second-Order Moller-Plesset Perturbation Theory (MP2)

Moller-Plesset perturbation theory [also referred to as many-body perturbation theory (MBPT)] is an adaptation of the more general Rayleigh-Schrodinger perturbation theory (RSPT), applied to problems in molecular electronic structure theory.  This tutorial will provide a brief overview of both RSPT and MBPT, before walking through an implementation of second-order Moller-Plesset perturbation theory (specifically referred to as MP2) which uses conventional, 4-index ERIs.

### I. Overview of Rayleigh-Schrodinger Perturbation Theory
Given the Hamiltonian operator $\hat{H}$ for a system, perturbation theory solves the Schrodinger equation for that system by rewriting $\hat{H}$ as

\begin{equation}
\hat{H} = \hat{H}{}^{(0)} + \lambda\hat{V},
\tag{[Szabo:1996], pp. 322, Eqn. 6.3}
\end{equation}

were $\hat{H}{}^{(0)}$ is the Hamiltonian operator corresponding to a solved problem which resembles $\hat{H}$, and $\hat{V}$ is the *perturbation* operator, defined as $\hat{V} = \hat{H} - \hat{H}{}^{(0)}$.  Then the Schrodinger equation for the system becomes

\begin{equation}
\hat{H}\mid\Psi_n\rangle = (\hat{H}{}^{(0)} + \lambda\hat{V})\mid\Psi_n\rangle = E_n\mid\Psi_n\rangle.
\tag{[Szabo:1996], pp. 322, Eqn. 6.2}
\end{equation}

The energies $E_n$ and wavefunctions $\mid\Psi_n\rangle$ will both be functions of $\lambda$; they can therefore be written as a Taylor series expansion about $\lambda = 0$ ([Szabo:1996], pp. 322, Eqns. 6.4a & 6.4b):

\begin{align}
E_n &= E_n^{(0)} + \lambda E_n^{(1)} + \lambda^2E_n^{(2)} + \ldots;\tag{[Szabo:1996], pp. 322, Eqn. 6.4a}\\
\mid\Psi_n\rangle &= \mid\Psi_n^{(0)}\rangle + \lambda\mid\Psi_n^{(1)}\rangle + \lambda^2\mid\Psi_n^{(2)}\rangle + \ldots,\tag{[Szabo:1996], pp. 322, Eqn. 6.4b}\\
\end{align}

in practice, these perturbation expansions may be truncated to a given power of $\lambda$.  Substituting the perturbation series above back into the Schrodinger equation yields

\begin{equation*}
(\hat{H}{}^{(0)} + \lambda\hat{V})(\mid\Psi_n^{(0)}\rangle + \lambda\mid\Psi_n^{(1)}\rangle + \lambda^2\mid\Psi_n^{(2)}\rangle + \ldots) = (E_n^{(0)} + \lambda E_n^{(1)} + \lambda^2E_n^{(2)} + \ldots)(\mid\Psi_n^{(0)}\rangle + \lambda\mid\Psi_n^{(1)}\rangle + \lambda^2\mid\Psi_n^{(2)}\rangle + \ldots),
\end{equation*}

which by equating powers of $\lambda$ ([Szabo:1996], pp. 323, Eqns. 6.7a-6.7d) gives expressions for the $E_n^{(i)}$and $\mid\Psi_n^{(i)}\rangle$. Note that for $\lambda^0$, $E_n^{(0)}$ and $\mid\Psi_n^{(0)}\rangle$ are known, as they are the solution to the zeroth-order problem $\hat{H}{}^{(0)}$.  For $\lambda^1$ and $\lambda^2$, the expressions for $E_n^{(1)}$ and $E_n^{(2)}$ are given by

\begin{align}
\lambda^1:\;\;\;\;E_n^{(1)} &= \langle\Psi_n^{(0)}\mid\hat{V}\mid\Psi_n^{(0)}\rangle,\tag{[Szabo:1996], pp. 323, Eqn. 6.8b}\\
\lambda^2:\;\;\;\;E_n^{(2)} &= \sum_{\mu\neq n}\frac{\mid\langle\Psi_{n}^{(0)}\mid\hat{V}\mid\Psi_{\mu}^{(0)}\rangle\mid^2}{E_n^{(0)} - E_{\mu}^{(0)}}\tag{[Szabo:1996], pp. 324, Eqn. 6.12}\\
\end{align}

### II. Overview of Moller-Plesset Perturbation Theory

The exact electronic Hamiltonian for an N-electron molecule with a given atomic configuration is (in atomic units):

\begin{equation}
\hat{H}_{elec} = \sum_i\hat{h}(i) + \sum_{i<j}\frac{1}{r_{ij}}.\tag{[Szabo:1996], pp. 43, Eqn. 2.10}
\end{equation}

Moller-Plesset perturbation theory seeks to solve the molecular electronic Schrodinger equation with the above Hamiltonian using the techniques of Rayleigh-Schroding perturbation theory, by selecting the zeroth-order Hamiltonian and wavefunctions to be those from Hartree-Fock theory:

\begin{align}
\hat{H}{}^{(0)}  &= \sum_i \hat{f}(i) = \sum_i \hat{h}(i) + \upsilon^{HF}(i)\tag{[Szabo:1996], pp. 350, Eqn. 6.59}\\
\hat{V} &= \hat{H} - \hat{H}{}^{(0)} = \sum_{i<j} \frac{1}{r_{ij}} - \upsilon^{HF}(i).\tag{[Szabo:1996], pp. 350, Eqn. 6.60}\\
\end{align}

With these choices of $\hat{H}{}^{(0)}$, $\hat{V}$, and $\mid\Psi_n^{(0)}\rangle$, the ground-state zeroth-order energy is given by $E_0^{(0)} = \sum_i\epsilon_i$ ([Szabo:1996], pp. 351, Eqn. 6.67).  Then, the first-order ground state energy may be computed by $E_0^{(1)} = \langle\Psi_0^{HF}\mid\hat{V}\mid\Psi_0^{HF}\rangle$ to find that the total ground-state energy through first order, $E_0 = E_0^{(0)} + E_0^{(1)}$, is exactly the Hartree-Fock energy. ([Szabo:1996], pp. 351, Eqn. 6.69)  Therefore, the first correction to the Hartree-Fock energy will occur at second-order in the perturbation series, i.e., with $E_0^{(2)}$; truncating the perturbation series at second order is commonly referred to as MP2.  The second order correction to the ground state energy will be given by

\begin{equation}
E_0^{(2)} = \sum_{\mu\neq 0}\frac{\left|\langle\Psi_{0}^{HF}\mid\hat{V}\mid\Psi_{\mu}^{HF}\rangle\right|^2}{E_0^{(0)} - E_{\mu}^{(0)}}\tag{[Szabo:1996], pp. 351, Eqn. 6.70}
\end{equation}

For brevity, we will now drop the "HF" from all zeroth-order wavefunctions.  This summation is over the eigenstate index $\mu$, each associated with a different eigenstate of the zeroth-order Hamiltonian.  For a single-determinant wavefunction constructed from spin orbitals, the summation over the eigenstate index $\mu\neq 0$ therefore refers to determinants which are constructed from *different* spin orbitals than the ground state determinant.  To distinguish such determinants, we will denote MOs occupied in the ground state with indices $i,\,j,\,\ldots$, and MOs which are unoccupied in the ground state with indices $a,\,b,\,\ldots\,$.  Then a determinant where orbital $a$ is substituted for orbital $i$ is denoted $\mid\Psi_i^a\rangle$, and so on.  Before substituting this new notation into the above energy expression, however, we may immediately recognize that many terms $\langle\Psi_{0}\mid\hat{V}\mid\Psi_{\mu}\rangle$ will not contribute to the second order energy:

| Term         | Determinant                               |  Contribution to $E_0^{(2)}$                            |
|--------------|-------------------------------------------|---------------------------------------------------------|
| Singles      | $\mid\Psi_i^a\rangle$                     | 0; Brillouin's Theorem                                  |
| Doubles      | $\mid\Psi_{ij}^{ab}\rangle$               | Survive                                                 |
| Higher-order | $\mid\Psi_{ijk\ldots}^{abc\ldots}\rangle$ | 0; $\hat{V}$ is a two-particle operator            |

Hence we see that only doubly-substituted determinants $\mid\Psi_{ij}^{ab}\rangle$ will contribute to $E_0^{(2)}$.  From Hartree-Fock theory, we know that

\begin{equation}
\langle\Psi_0\mid r_{ij}^{-1}\mid\Psi_{ij}^{ab}\rangle = [ia\| jb],\tag{[Szabo:1996], pp. 72, Tbl. 2.6}
\end{equation}

where $[ia\| jb] = [ia\mid jb] - [ib\mid ja]$ is an antisymmetrized two-electron integral, and the square brackets "$[\ldots ]$" indicate that we are employing chemists' notation.  What about the energies of these doubly substituted determinants?  Recognizing that the difference between the energies of the newly- and formerly-occupied orbitals in each substitution must modify the total energy of the ground state determinant,

\begin{equation}
E_{ij}^{ab} = E_0 - (\epsilon_i - \epsilon_a + \epsilon_j - \epsilon_b).
\end{equation}

Substituting these expressions into the one for the second-order energy, we have that

\begin{equation}
E_0^{(2)} = \sum_{i<j}\sum_{a<b} \frac{\left|\,[ia\|jb]\,\right|^2}{\epsilon_i - \epsilon_a + \epsilon_j - \epsilon_b}.\tag{[Szabo:1996], pp. 352, Eqn. 6.71}
\end{equation}

So far, our discussion has used spin-orbitals instead of the more familiar spatial orbitals.  Indeed, significant speedups are achieved when using spatial orbitals.  Integrating out the spin variable $\omega$ from $E_0^{(2)}$ yields two expressions; one each for the interaction of particles with the same spin (SS) and opposite spin (OS):

\begin{align}
E_{\rm 0,\,SS}^{(2)} = \sum_{ij}\sum_{ab}\frac{(ia\mid jb)[(ia\mid jb) - (ib\mid ja)]}{\epsilon_i - \epsilon_a + \epsilon_j - \epsilon_b}, \quad
E_{\rm 0,\,OS}^{(2)} = \sum_{ij}\sum_{ab}\frac{(ia\mid jb)(ia\mid jb)}{\epsilon_i - \epsilon_a + \epsilon_j - \epsilon_b},
\end{align}

where these spin-free expressions make use of integrals in chemists' notation over spatial orbitals. (Rearranged from [Szabo:1996], pp. 352, Eqn. 6.74)  Note that an exchange integral arises between particles of the same spin; this is because the motions of particles with identical spins are correlated due to the requirement that $\left|\Psi\right|^2$ remain invariant to the exchange of the spatial and spin coordinates of any pair of electrons.  Finally, the total MP2 correction energy $E_0^{(2)} = E_{\rm 0,\,SS}^{(2)} + E_{\rm 0,\,OS}^{(2)}$.

### Implementation of Conventional MP2

Let's begin by importing Psi4, NumPy and TensorOperations, and setting memory and output file options:

```julia
# ==> Import statements & Global Options <==
using PyCall: pyimport
psi4 = pyimport("psi4")
np   = pyimport("numpy") # used only to cast to Psi4 arrays
using TensorOperations: @tensor

psi4.set_memory(Int(2e9))
numpy_memory = 2
psi4.core.set_output_file("output.dat", false)
```

      Memory set to   1.863 GiB by Python driver.

Next, we can define our molecule and Psi4 options.  Notice that we are using `scf_type pk` to indicate that we wish to use conventional, full 4-index ERIs, and that we have specified `mp2_type conv` so that the MP2 algorithm we check against also uses the conventional ERIs.

```julia
# ==> Molecule & Psi4 Options Definitions <==
mol = psi4.geometry("""
O
H 1 1.1
H 1 1.1 2 104
symmetry c1
""")

psi4.set_options(Dict("basis"         => "6-31g",
                      "scf_type"      => "pk",
                      "mp2_type"      => "conv",
                      "e_convergence" => 1e-8,
                      "d_convergence" => 1e-8))
```

Since MP2 is a perturbation on the zeroth-order Hartree-Fock description of a molecular system, all of the relevant information (Fock matrix, orbitals, orbital energies) about the system can be computed using any Hartree-Fock program.  We could use the RHF program that we wrote in tutorial 3a, but we could just as easily use Psi4 to do our dirty work.  In the cell below, use Psi4 to compute the RHF energy and wavefunction, and store them using the `return_wfn=True` keyword argument to `psi4.energy()`:

```julia
# Get the SCF wavefunction & energies
scf_e, scf_wfn = psi4.energy("scf", return_wfn=true)
```

    (-75.95252904632221, PyObject <psi4.core.RHF object at 0x14859e110>)

In the expression for $E_0^{(2)}$, the two summations are over occupied and virtual indices, respectively.  Therefore, we'll need to get the number of occupied orbitals and the total number of orbitals.  Additionally, we must obtain the MO energy eigenvalues; again since the sums are over occupied and virtual orbitals, it is good to separate the occupied orbital energies from the virtual orbital energies.  From the SCF wavefunction you generated above, get the number of doubly occupied orbitals, number of molecular orbitals, and MO energies:

```julia
# ==> Get orbital information & energy eigenvalues <==
# Number of Occupied orbitals & MOs
ndocc = scf_wfn.nalpha()
nmo = scf_wfn.nmo()

# Get orbital energies, cast into NumPy array, and separate occupied & virtual
eps = np.asarray(scf_wfn.epsilon_a())
e_ij = eps[1:ndocc]
e_ab = eps[ndocc+1:end];
```

Unlike the orbital information, Psi4 does not return the ERIs when it does a computation.  Fortunately, however, we can just build them again using the `psi4.core.MintsHelper()` class.  Recall that these integrals will be generated in the AO basis; before using them in the $E_0^{(2)}$ expression, we must transform them into the MO basis.  To do this, we first need to obtain the orbital coefficient matrix, **C**.  In the cell below, generate the ERIs for our molecule, get **C** from the SCF wavefunction, and obtain occupied- and virtual-orbital slices of **C** for future use.

```julia
# ==> ERIs <==
# Create instance of MintsHelper class
mints = psi4.core.MintsHelper(scf_wfn.basisset())

# Memory check for ERI tensor
I_size = nmo^4 * 8.e-9
println("\nSize of the ERI tensor will be $I_size GB.")
memory_footprint = I_size * 1.5
if I_size > numpy_memory
    psi4.core.clean()
    throw(OutOfMemoryError("Estimated memory utilization ($memory_footprint GB) exceeds " *
                           "allotted memory limit of $numpy_memory GB."))
end

# Build ERI Tensor
I = np.asarray(mints.ao_eri())

# Get MO coefficients from SCF wavefunction
C = np.asarray(scf_wfn.Ca())
Cocc = C[:, 1:ndocc]
Cvirt = C[:, ndocc+1:end];
```

    Size of the ERI tensor will be 0.00022848800000000003 GB.

In order to transform the four-index integrals from the AO to the MO basis, we must perform the following contraction:

$$(i\,a\mid j\,b) = C_{\mu i}C_{\nu a}(\mu\,\nu\mid|\,\lambda\,\sigma)C_{\lambda j}C_{\sigma b}$$

Again, here we are using $i,\,j$ as occupied orbital indices and $a,\, b$ as virtual orbital indices.  We could carry out the above contraction all in one step using either `@tensor` or explicit loops:

~~~julia
# Naive Algorithm for ERI Transformation
@tensor I_mo[i,a,j,b] := Cocc[p,i] * Cvirt[q,a] * I[p,q,r,s] * Cocc[r,j] * Cvirt[s,b]
~~~

Notice that the transformation from AO index to occupied (virtual) MO index requires only the occupied (virtual) block of the **C** matrix; this allows for computational savings in large basis sets, where the virtual space can be very large.  This algorithm, while efficient with `@tensor`, has horrendous scaling if the search of an optimal contraction fails.  We will enforce a better contraction. Examining the contraction more closely, we see that there are 8 unique indices, and thus the step above scales as ${\cal O}(N^8)$.  With this algorithm, a twofold increase of the number of MO's would result in $2^8 = 256\times$ expense to perform.  We can, however, refactor the above contraction such that

$$(i\,a\mid j\,b) = \left[C_{\mu i}\left[C_{\nu a}\left[C_{\lambda j}\left[C_{\sigma b}(\mu\,\nu\mid|\,\lambda\,\sigma)\right]\right]\right]\right],$$

where we have now written the transfomation as four ${\cal O}(N^5)$ steps instead of one ${\cal O}(N^8)$ step. This is a savings of $\frac{4}{n^3}$, and is responsible for the feasibility of the MP2 method for application to any but very small systems and/or basis sets.  We may carry out the above ${\cal O}(N^5)$ algorithm by carrying out one index transformation at a time, and storing the result in a temporary array.  In the cell below, transform the ERIs from the AO to MO basis, using our smarter algorithm:

```julia
# ==> Transform I -> I_mo @ O(N⁵) <==
I_mo = @tensor begin
   I_mo[i,q,r,s] := Cocc[p,i]     * I[p,q,r,s]
   I_mo[i,a,r,s] := Cvirt[q,a]    * I_mo[i,q,r,s]
   I_mo[i,a,j,s] :=                 I_mo[i,a,r,s] * Cocc[r,j]
   I_mo[i,a,j,b] :=                 I_mo[i,a,j,s] * Cvirt[s,b]
end
nothing
```

We note here that we can use infrastructure in Psi4 to carry out the above integral transformation; this entails obtaining the occupied and virtual blocks of **C** Psi4-side, and then using the built-in `MintsHelper` function `MintsHelper.mo_eri()` to transform the integrals.  Just to check your work above, execute the next cell to see this tech in action:

```julia
# ==> Compare our Imo to MintsHelper <==
Co = scf_wfn.Ca_subset("AO","OCC")
Cv = scf_wfn.Ca_subset("AO","VIR")
MO = np.asarray(mints.mo_eri(Co, Cv, Co, Cv))
println("Do our transformed ERIs match Psi4's? ", np.allclose(I_mo, np.asarray(MO)))
```

    Do our transformed ERIs match Psi4's? true

Now we have all the pieces needed to compute $E_0^{(2)}$.  This could be done by writing explicit loops over occupied and virtual indices Julia side, e.g.,

```julia
# Compute SS & OS MP2 Correlation
mp2_corr = let mp2_ss_corr = 0.0, mp2_os_corr = 0.0
   nvirt = nmo - ndocc
   for i in 1:ndocc, a in 1:nvirt, j in 1:ndocc, b in 1:nvirt
       numerator = I_mo[i,a,j,b] * (I_mo[i, a, j, b] - I_mo[i, b, j, a])
       mp2_ss_corr += numerator / (e_ij[i] + e_ij[j] - e_ab[a] - e_ab[b])
       mp2_os_corr += I_mo[i,a,j,b]^2 / (e_ij[i] + e_ij[j] - e_ab[a] - e_ab[b])
   end
   mp2_ss_corr + mp2_os_corr
end

# Total MP2 Energy
MP2_E = scf_e + mp2_corr

# ==> Compare to Psi4 <==
psi4.compare_values(psi4.energy("mp2"), MP2_E, 6, "MP2 Energy")
```

    	MP2 Energy........................................................PASSED

    true

In this method it is very clear what is going on and is easy to program. Julia has the distinct advantage loops are as fast as the same block written in a compiled language like C, C++, or Fortran.

However, we will provide an alternative formulation using tensor contractions. It should be clear how to contract the four-index integrals $(i\,a\mid j\,b)$ and $(i\,a\mid j\,b)$ with one another, but what about the energy eigenvalues $\epsilon$?  We can use a Julia trick called *broadcasting* to construct a four-index array of all possible energy denominators, which can then be contracted with the full I_mo arrays.  To do this, we'll use the function `reshape()`:
~~~julia
# Prepare 4d energy denominator array
e_denom  = reshape(e_ij,  1, 1, 1, :)      # Diagonal of 4d array are occupied orbital energies
e_denom -= reshape(e_ab', 1, 1, :)         # all combinations of (e_ij - e_ab)
e_denom += e_ij                            # all combinations of [(e_ij - e_ab) + e_ij]
e_denom -= e_ab'                           # All combinations of full denominator
e_denom  = premutedims(e_denom, (1,2,4,3)) # permute 3rd and 4th dims to have (nocc,nvirt,nocc,nvirt) shape
e_denom  = inv.(e_denom)                   # Take reciprocal for contracting with numerator
~~~
In the cell below, compute the energy denominator using `reshape()` and contract this array with the four-index ERIs to compute the same-spin and opposite-spin MP2 correction using `sum()`. Then, add these quantities to the SCF energy computed above to obtain the total MP2 energy.

Hint: For the opposite-spin correlation, use `permutedims()` to obtain the correct ordering of the indices in the exchange integral.

```julia
#using Einsum: @einsum
# ==> Compute MP2 Correlation & MP2 Energy <==
# Compute energy denominator array
e_denom = reshape(e_ij,1,1,1,:) .- reshape(e_ab',1,1,:) .+ (e_ij .- e_ab')
e_denom = permutedims(e_denom, (1,2,4,3)) # 3 ↔ 4
e_denom = inv.(e_denom)

# check
#using Test
#nvirt = nmo - ndocc
#for i in 1:ndocc, a in 1:nvirt, j in 1:ndocc, b in 1:nvirt
#    @test e_denom[i,a,j,b] ≈ 1 / (e_ij[i] + e_ij[j] - e_ab[a] - e_ab[b])
#end

# Compute SS & OS MP2 Correlation with sum()
bctd_mp2_os_corr = sum(I_mo .* I_mo .* e_denom)
I_mo_swap = permutedims(I_mo,(3,2,1,4)) # 1 ↔ 3
bctd_mp2_ss_corr = sum(I_mo .* (I_mo .- I_mo_swap) .* e_denom)

# Compare broadcasted and loop MP2
@assert bctd_mp2_os_corr + bctd_mp2_ss_corr ≈ mp2_corr

# Total MP2 Energy
MP2_E = scf_e + bctd_mp2_os_corr + bctd_mp2_ss_corr
```

    -76.09464888642944

```julia
# ==> Compare to Psi4 <==
psi4.compare_values(psi4.energy("mp2"), MP2_E, 6, "MP2 Energy")
```

    	MP2 Energy........................................................PASSED

    true

## References

1. Original paper: "Note on an Approximation Treatment for Many-Electron Systems"
	> [[Moller:1934:618](https://journals.aps.org/pr/abstract/10.1103/PhysRev.46.618)] C. Møller and M. S. Plesset, *Phys. Rev.* **46**, 618 (1934)
2. The Laplace-transformation in MP theory: "Minimax approximation for the decomposition of energy denominators in Laplace-transformed Møller–Plesset perturbation theories"
    > [[Takasuka:2008:044112](http://aip.scitation.org/doi/10.1063/1.2958921)] A. Takatsuka, T. Siichiro, and W. Hackbusch, *J. Phys. Chem.*, **129**, 044112 (2008)
3. Equations taken from:
	> [[Szabo:1996](https://books.google.com/books?id=KQ3DAgAAQBAJ&printsec=frontcover&dq=szabo+%26+ostlund&hl=en&sa=X&ved=0ahUKEwiYhv6A8YjUAhXLSCYKHdH5AJ4Q6AEIJjAA#v=onepage&q=szabo%20%26%20ostlund&f=false)] A. Szabo and N. S. Ostlund, *Modern Quantum Chemistry: Introduction to Advanced Electronic Structure Theory*. Courier Corporation, 1996.
4. Algorithms taken from:
	> [Crawford:prog] T. D. Crawford, "The Second-Order Møller–Plesset Perturbation Theory (MP2) Energy."  Accessed via the web at http://github.com/CrawfordGroup/ProgrammingProjects.


```python
# This cell is mandatory in all Dymos documentation notebooks.
missing_packages = []
try:
    import openmdao.api as om
except ImportError:
    if 'google.colab' in str(get_ipython()):
        !python -m pip install openmdao[notebooks]
    else:
        missing_packages.append('openmdao')
try:
    import dymos as dm
except ImportError:
    if 'google.colab' in str(get_ipython()):
        !python -m pip install dymos
    else:
        missing_packages.append('dymos')
try:
    import pyoptsparse
except ImportError:
    if 'google.colab' in str(get_ipython()):
        !pip install -q condacolab
        import condacolab
        condacolab.install_miniconda()
        !conda install -c conda-forge pyoptsparse
    else:
        missing_packages.append('pyoptsparse')
if missing_packages:
    raise EnvironmentError('This notebook requires the following packages '
                           'please install them and restart this notebook\'s runtime: {",".join(missing_packages)}')
```

(examples:the_brachistochrone)=
# The Brachistochrone

```{admonition} Things you'll learn through this example
- How to define a basic Dymos ODE system.
- How to test the partials of your ODE system.
- Adding a Trajectory object with a single Phase to an OpenMDAO Problem.
- Imposing boundary conditions on states with simple bounds via `fix_initial` and `fix_final`.
- Using the Phase.interpolate` method to set a linear guess for state and control values across the Phase.
- Checking the validity of the result through explicit simulation via the `Trajectory.simulate` method.
```

The brachistochrone is one of the most well-known optimal control problems.
It was originally posed as a challenge by Johann Bernoulli.

```{admonition} The brachistochrone problem
_Given two points A and B in a vertical plane, find the path AMB
down which a movable point M must by virtue of its weight fall from
A to B in the shortest possible time._

-  Johann Bernoulli, Acta Eruditorum, June 1696
```

We seek to find the optimal shape of a wire between two points (A and B) such that a bead sliding
without friction along the wire moves from point A to point B in minimum time.

```python
import numpy as np

import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrowPatch, Arc

LW = 2

fig, ax = plt.subplots(1, 1, figsize=(5, 5))
ax.axis('off')

ax.set_xlim(-1, 11)
ax.set_ylim(-1, 11)

circle = plt.Circle((0, 10), radius=0.1, fc='k')
ax.add_patch(circle)
plt.text(0.2, 10.2, 'A')

circle = plt.Circle((10, 5), radius=0.1, fc='k')
ax.add_patch(circle)
plt.text(10.2, 5.2, 'B')

# Choose a to suite, compute b
a = 0.1
b = -0.5 - 10*a
c = 10

def y_wire(x):
    return a*x**2 + b*x + c, 2*a*x + b

x = np.linspace(0, 10, 100)
y, _ = y_wire(x)
plt.plot(x, y, 'b-')

# Add the bead to the wire
x = 3
y, dy_dx = y_wire(x)
plt.plot(x, y, 'ro', ms=10)

# Draw and label the gravity vector
gvec = FancyArrowPatch((x, y), (x, y-2), arrowstyle='->', mutation_scale=10, linewidth=LW, color='k')
lv_line = plt.Line2D((x, x), (y, y-2), visible=False)  # Local vertical
ax.add_patch(gvec)
plt.text(x - 0.5, y-1, 'g')

# Draw and label the velocity vector
dx = 2
dy = dy_dx * dx
vvec = FancyArrowPatch((x, y), (x+dx, y+dy), arrowstyle='->', mutation_scale=10, linewidth=LW, color='k')
ax.add_patch(vvec)
plt.text(x+dx-0.25, y+dy-0.25, 'v')

# Draw angle theta
vvec_line = plt.Line2D((x, x+dx), (y, y+dy), visible=False)
# angle_plot = get_angle_plot(lv_line, vvec_line, color='k', origin=(x, y), radius=3)
# ax.add_patch(angle_plot)
ax.text(x+0.25, y-1.25, r'$\theta$')

# Draw the axes
x = 0
y = 2
dx = 5
dy = 0
xhat = FancyArrowPatch((x, y), (x+dx, y+dy), arrowstyle='->', mutation_scale=10, linewidth=LW, color='k')
ax.add_patch(xhat)
plt.text(x+dx/2.0-0.5, y+dy/2.0-0.5, 'x')

dx = 0
dy = 5
yhat = FancyArrowPatch((x, y), (x+dx, y+dy), arrowstyle='->', mutation_scale=10, linewidth=LW, color='k')
ax.add_patch(yhat)
plt.text(x+dx/2.0-0.5, y+dy/2.0-0.5, 'y')

plt.ylim(1, 11)
plt.xlim(-0.5, 10.5)

plt.savefig('brachistochrone_fbd.png')

plt.show()
```

## State variables

In this implementation, three _state_ variables are used to define the configuration of the system at any given instant in time.

- **x**: The horizontal position of the particle at an instant in time.
- **y**: The vertical position of the particle at an instant in time.
- **v**: The speed of the particle at an instant in time.

## System dynamics

From the free-body diagram above, the evolution of the state variables is given by the following ordinary differential equations (ODE).

\begin{align}
    \frac{d x}{d t} &= v \sin(\theta) \\
    \frac{d y}{d t} &= -v \cos(\theta) \\
    \frac{d v}{d t} &= g \cos(\theta)
\end{align}

## Control variables

This system has a single control variable.

- **$\theta$**: The angle between the gravity vector and the tangent to the curve at the current instant in time.

## The initial and final conditions

In this case, starting point **A** is given as _(0, 10)_.
The point moving along the curve will begin there with zero initial velocity.

The initial conditions are:

\begin{align}
    x_0 &= 0 \\
    y_0 &= 10 \\
    v_0 &= 0
\end{align}

The end point **B** is given as _(10, 5)_.
The point will end there, but the velocity at that point is not constrained.

The final conditions are:

\begin{align}
    x_f &= 10 \\
    y_f &= 5 \\
    v_f &= \mathrm{free}
\end{align}

## Defining the ODE as an OpenMDAO System

In Dymos, the ODE is an OpenMDAO System (a Component, or a Group of components).
The following ExplicitComponent computes the state rates for the brachistochrone problem.

More detail on the workings of an ExplicitComponent can be found in the OpenMDAO documentation.  In summary:

- **initialize**:  Called at setup, and used to define options for the component.  **ALL** Dymos ODE components should have the property `num_nodes`, which defines the number of points at which the outputs are simultaneously computed.
- **setup**: Used to add inputs and outputs to the component, and declare which outputs (and indices of outputs) are dependent on each of the inputs.
- **compute**: Used to compute the outputs, given the inputs.
- **compute_partials**: Used to compute the derivatives of the outputs w.r.t. each of the inputs analytically.  This method may be omitted if finite difference or complex-step approximations are used, though analytic is recommended.

```python
import numpy as np
import openmdao.api as om

class BrachistochroneODE(om.ExplicitComponent):

    def initialize(self):
        self.options.declare('num_nodes', types=int)
        self.options.declare('static_gravity', types=(bool,), default=False,
                             desc='If True, treat gravity as a static (scalar) input, rather than '
                                  'having different values at each node.')

    def setup(self):
        nn = self.options['num_nodes']

        # Inputs
        self.add_input('v', val=np.zeros(nn), desc='velocity', units='m/s')

        if self.options['static_gravity']:
            self.add_input('g', val=9.80665, desc='grav. acceleration', units='m/s/s',
                           tags=['dymos.static_target'])
        else:
            self.add_input('g', val=9.80665 * np.ones(nn), desc='grav. acceleration', units='m/s/s')

        self.add_input('theta', val=np.ones(nn), desc='angle of wire', units='rad')

        self.add_output('xdot', val=np.zeros(nn), desc='velocity component in x', units='m/s',
                        tags=['dymos.state_rate_source:x', 'dymos.state_units:m'])

        self.add_output('ydot', val=np.zeros(nn), desc='velocity component in y', units='m/s',
                        tags=['dymos.state_rate_source:y', 'dymos.state_units:m'])

        self.add_output('vdot', val=np.zeros(nn), desc='acceleration magnitude', units='m/s**2',
                        tags=['dymos.state_rate_source:v', 'dymos.state_units:m/s'])

        self.add_output('check', val=np.zeros(nn), desc='check solution: v/sin(theta) = constant',
                        units='m/s')

        # Setup partials
        arange = np.arange(self.options['num_nodes'])
        self.declare_partials(of='vdot', wrt='theta', rows=arange, cols=arange)

        self.declare_partials(of='xdot', wrt='v', rows=arange, cols=arange)
        self.declare_partials(of='xdot', wrt='theta', rows=arange, cols=arange)

        self.declare_partials(of='ydot', wrt='v', rows=arange, cols=arange)
        self.declare_partials(of='ydot', wrt='theta', rows=arange, cols=arange)

        self.declare_partials(of='check', wrt='v', rows=arange, cols=arange)
        self.declare_partials(of='check', wrt='theta', rows=arange, cols=arange)

        if self.options['static_gravity']:
            c = np.zeros(self.options['num_nodes'])
            self.declare_partials(of='vdot', wrt='g', rows=arange, cols=c)
        else:
            self.declare_partials(of='vdot', wrt='g', rows=arange, cols=arange)

    def compute(self, inputs, outputs):
        theta = inputs['theta']
        cos_theta = np.cos(theta)
        sin_theta = np.sin(theta)
        g = inputs['g']
        v = inputs['v']

        outputs['vdot'] = g * cos_theta
        outputs['xdot'] = v * sin_theta
        outputs['ydot'] = -v * cos_theta
        outputs['check'] = v / sin_theta

    def compute_partials(self, inputs, partials):
        theta = inputs['theta']
        cos_theta = np.cos(theta)
        sin_theta = np.sin(theta)
        g = inputs['g']
        v = inputs['v']

        partials['vdot', 'g'] = cos_theta
        partials['vdot', 'theta'] = -g * sin_theta

        partials['xdot', 'v'] = sin_theta
        partials['xdot', 'theta'] = v * cos_theta

        partials['ydot', 'v'] = -cos_theta
        partials['ydot', 'theta'] = v * sin_theta

        partials['check', 'v'] = 1 / sin_theta
        partials['check', 'theta'] = -v * cos_theta / sin_theta ** 2
```

```{admonition} "Things to note about the ODE system"
- There is no input for the position states ($x$ and $y$).  The dynamics aren't functions of these states, so they aren't needed as inputs.
- While $g$ is an input to the system, since it will never change throughout the trajectory, it can be an option on the system.  This way we don't have to define any partials w.r.t. $g$.
- The output `check` is an _auxiliary_ output, not a rate of the state variables.  In this case, optimal control theory tells us that `check` should be constant throughout the trajectory, so it's a useful output from the ODE.
```

## Testing the ODE

Now that the ODE system is defined, it is strongly recommended to test the analytic partials before using it in optimization.
If the partials are incorrect, then the optimization will almost certainly fail.
Fortunately, OpenMDAO makes testing derivatives easy with the `check_partials` method.
The `assert_check_partials` method in `openmdao.utils.assert_utils` can be used in test frameworks to verify the correctness of the partial derivatives in a model.

The following is a test method which creates a new OpenMDAO problem whose model contains the ODE class.
The problem is setup with the `force_alloc_complex=True` argument to enable complex-step approximation of the derivatives.
Complex step typically produces derivative approximations with an error on the order of 1.0E-16, as opposed to ~1.0E-6 for forward finite difference approximations.

```python
import numpy as np
import openmdao.api as om

num_nodes = 5

p = om.Problem(model=om.Group())

ivc = p.model.add_subsystem('vars', om.IndepVarComp())
ivc.add_output('v', shape=(num_nodes,), units='m/s')
ivc.add_output('theta', shape=(num_nodes,), units='deg')

p.model.add_subsystem('ode', BrachistochroneODE(num_nodes=num_nodes))

p.model.connect('vars.v', 'ode.v')
p.model.connect('vars.theta', 'ode.theta')

p.setup(force_alloc_complex=True)

p.set_val('vars.v', 10*np.random.random(num_nodes))
p.set_val('vars.theta', 10*np.random.uniform(1, 179, num_nodes))

p.run_model()
cpd = p.check_partials(method='cs', compact_print=True)
```

```python
from dymos.utils.testing_utils import assert_check_partials

assert_check_partials(cpd)
```

## Solving the problem with Legendre-Gauss-Lobatto collocation in Dymos

The following script fully defines the brachistochrone problem with Dymos and solves it.  In this section we'll walk through each step.

```python
import openmdao.api as om
import dymos as dm
from dymos.examples.plotting import plot_results
from dymos.examples.brachistochrone import BrachistochroneODE
import matplotlib.pyplot as plt

#
# Initialize the Problem and the optimization driver
#
p = om.Problem(model=om.Group())
p.driver = om.ScipyOptimizeDriver()
p.driver.declare_coloring()

#
# Create a trajectory and add a phase to it
#
traj = p.model.add_subsystem('traj', dm.Trajectory())

phase = traj.add_phase('phase0',
                       dm.Phase(ode_class=BrachistochroneODE,
                                transcription=dm.GaussLobatto(num_segments=10)))

#
# Set the variables
#
phase.set_time_options(fix_initial=True, duration_bounds=(.5, 10))

phase.add_state('x', fix_initial=True, fix_final=True)

phase.add_state('y', fix_initial=True, fix_final=True)

phase.add_state('v', fix_initial=True, fix_final=False)

phase.add_control('theta', continuity=True, rate_continuity=True,
                  units='deg', lower=0.01, upper=179.9)

phase.add_parameter('g', units='m/s**2', val=9.80665)

#
# Minimize time at the end of the phase
#
phase.add_objective('time', loc='final', scaler=10)

p.model.linear_solver = om.DirectSolver()

#
# Setup the Problem
#
p.setup()

#
# Set the initial values
#
p['traj.phase0.t_initial'] = 0.0
p['traj.phase0.t_duration'] = 2.0

p.set_val('traj.phase0.states:x', phase.interp('x', ys=[0, 10]))
p.set_val('traj.phase0.states:y', phase.interp('y', ys=[10, 5]))
p.set_val('traj.phase0.states:v', phase.interp('v', ys=[0, 9.9]))
p.set_val('traj.phase0.controls:theta', phase.interp('theta', ys=[5, 100.5]))

#
# Solve for the optimal trajectory
#
dm.run_problem(p)

# Check the results
print(p.get_val('traj.phase0.timeseries.time')[-1])
```

```python
from openmdao.utils.assert_utils import assert_near_equal

assert_near_equal(p.get_val('traj.phase0.timeseries.time')[-1], 1.8016, tolerance=1.0E-3)
```

```python
# Generate the explicitly simulated trajectory
exp_out = traj.simulate()

plot_results([('traj.phase0.timeseries.states:x', 'traj.phase0.timeseries.states:y',
               'x (m)', 'y (m)'),
              ('traj.phase0.timeseries.time', 'traj.phase0.timeseries.controls:theta',
               'time (s)', 'theta (deg)')],
             title='Brachistochrone Solution\nHigh-Order Gauss-Lobatto Method',
             p_sol=p, p_sim=exp_out)

plt.show()
```

(examples:brachistochrone:explicit_shooting)=
## Solving the problem with single shooting in Dymos

The following script fully defines the brachistochrone problem with Dymos and solves it using a single explicit shooting method.

The code is nearly identical to that using the collocation approach.
Key differences are shown when defining the transcription, specifying how to constraint the final state values of `x` and `y` states, and providing initial guesses for all states.

```python
import openmdao.api as om
import dymos as dm
from dymos.examples.plotting import plot_results
from dymos.examples.brachistochrone import BrachistochroneODE
import matplotlib.pyplot as plt

#
# Initialize the Problem and the optimization driver
#
p = om.Problem(model=om.Group())
p.driver = om.ScipyOptimizeDriver()

# We'll try to use coloring, but OpenMDAO will tell us that it provides no benefit.
p.driver.declare_coloring()

#
# Create a trajectory and add a phase to it
#
traj = p.model.add_subsystem('traj', dm.Trajectory())

phase = traj.add_phase('phase0',
                       dm.Phase(ode_class=BrachistochroneODE,
                                transcription=dm.ExplicitShooting(num_segments=10,
                                                                  num_steps_per_segment=10,
                                                                  method='rk4')))

#
# Set the variables
#
phase.set_time_options(fix_initial=True, duration_bounds=(.5, 10))

# Note, we cannot use fix_final=True with the shooting method
# because the final value of the states are
# not design variables in the transcribed optimization problem.
phase.add_state('x', fix_initial=True)

phase.add_state('y', fix_initial=True)

phase.add_state('v', fix_initial=True)

phase.add_control('theta', continuity=True, rate_continuity=True,
                  units='deg', lower=0.01, upper=179.9)

phase.add_parameter('g', units='m/s**2', val=9.80665)

#
# Minimize time at the end of the phase
#
phase.add_objective('time', loc='final', scaler=10)

#
# Add boundary constraints for x and y since we could not use `fix_final=True`
#
phase.add_boundary_constraint('x', loc='final', equals=10)
phase.add_boundary_constraint('y', loc='final', equals=5)

p.model.linear_solver = om.DirectSolver()

#
# Setup the Problem
#
p.setup()

#
# Set the initial values
#
p['traj.phase0.t_initial'] = 0.0
p['traj.phase0.t_duration'] = 2.0

# Only the initial values of the states are design variables,
# so the phase interp method is not used on states.
p.set_val('traj.phase0.states:x', 0.0)
p.set_val('traj.phase0.states:y', 10.0)
p.set_val('traj.phase0.states:v', 0.0)
p.set_val('traj.phase0.controls:theta', phase.interp('theta', ys=[5, 100.5]))

#
# Solve for the optimal trajectory
#
dm.run_problem(p)

# Check the results
print(p.get_val('traj.phase0.timeseries.time')[-1])
```

```python
assert_near_equal(p.get_val('traj.phase0.timeseries.time')[-1], 1.8016, tolerance=1.0E-3)
```

```python
plot_results([('traj.phase0.timeseries.states:x', 'traj.phase0.timeseries.states:y',
               'x (m)', 'y (m)'),
              ('traj.phase0.timeseries.time', 'traj.phase0.timeseries.controls:theta',
               'time (s)', 'theta (deg)')],
             title='Brachistochrone Solution\nExplicit Shooting Method',
             p_sol=p)

plt.show()
```

Note that the shooting methods provide timeseries values at the start and end of each segment.
Due to plans to enable variable step integration in the future, Dymos cannot provide outputs for each individual step since the number of steps will not be known when setting up the problem.

Because control values and rates may not be continuous across segment bounds, depending on the settings used, there may be discontinuities in ODE outputs and therefore timeseries values are provided on each side of each segment boundary.