```python
%matplotlib inline
```

Adversarial Example Generation
==============================

**Author:** `Nathan Inkawhich <https://github.com/inkawhich>`__

If you are reading this, hopefully you can appreciate how effective some
machine learning models are. Research is constantly pushing ML models to
be faster, more accurate, and more efficient. However, an often
overlooked aspect of designing and training models is security and
robustness, especially in the face of an adversary who wishes to fool
the model.

This tutorial will raise your awareness to the security vulnerabilities
of ML models, and will give insight into the hot topic of adversarial
machine learning. You may be surprised to find that adding imperceptible
perturbations to an image *can* cause drastically different model
performance. Given that this is a tutorial, we will explore the topic
via example on an image classifier. Specifically we will use one of the
first and most popular attack methods, the Fast Gradient Sign Attack
(FGSM), to fool an MNIST classifier.

Threat Model
------------

For context, there are many categories of adversarial attacks, each with
a different goal and assumption of the attacker’s knowledge. However, in
general the overarching goal is to add the least amount of perturbation
to the input data to cause the desired misclassification. There are
several kinds of assumptions of the attacker’s knowledge, two of which
are: **white-box** and **black-box**. A *white-box* attack assumes the
attacker has full knowledge and access to the model, including
architecture, inputs, outputs, and weights. A *black-box* attack assumes
the attacker only has access to the inputs and outputs of the model, and
knows nothing about the underlying architecture or weights. There are
also several types of goals, including **misclassification** and
**source/target misclassification**. A goal of *misclassification* means
the adversary only wants the output classification to be wrong but does
not care what the new classification is. A *source/target
misclassification* means the adversary wants to alter an image that is
originally of a specific source class so that it is classified as a
specific target class.

In this case, the FGSM attack is a *white-box* attack with the goal of
*misclassification*. With this background information, we can now
discuss the attack in detail.

Fast Gradient Sign Attack
-------------------------

One of the first and most popular adversarial attacks to date is
referred to as the *Fast Gradient Sign Attack (FGSM)* and is described
by Goodfellow et. al. in `Explaining and Harnessing Adversarial
Examples <https://arxiv.org/abs/1412.6572>`__. The attack is remarkably
powerful, and yet intuitive. It is designed to attack neural networks by
leveraging the way they learn, *gradients*. The idea is simple, rather
than working to minimize the loss by adjusting the weights based on the
backpropagated gradients, the attack *adjusts the input data to maximize
the loss* based on the same backpropagated gradients. In other words,
the attack uses the gradient of the loss w.r.t the input data, then
adjusts the input data to maximize the loss.

Before we jump into the code, let’s look at the famous
`FGSM <https://arxiv.org/abs/1412.6572>`__ panda example and extract
some notation.

.. figure:: /_static/img/fgsm_panda_image.png
   :alt: fgsm_panda_image

From the figure, $\mathbf{x}$ is the original input image
correctly classified as a “panda”, $y$ is the ground truth label
for $\mathbf{x}$, $\mathbf{\theta}$ represents the model
parameters, and $J(\mathbf{\theta}, \mathbf{x}, y)$ is the loss
that is used to train the network. The attack backpropagates the
gradient back to the input data to calculate
$\nabla_{x} J(\mathbf{\theta}, \mathbf{x}, y)$. Then, it adjusts
the input data by a small step ($\epsilon$ or $0.007$ in the
picture) in the direction (i.e.
$sign(\nabla_{x} J(\mathbf{\theta}, \mathbf{x}, y))$) that will
maximize the loss. The resulting perturbed image, $x'$, is then
*misclassified* by the target network as a “gibbon” when it is still
clearly a “panda”.

Hopefully now the motivation for this tutorial is clear, so lets jump
into the implementation.

```python
from __future__ import print_function
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import numpy as np
import matplotlib.pyplot as plt

# NOTE: This is a hack to get around "User-agent" limitations when downloading MNIST datasets
#       see, https://github.com/pytorch/vision/issues/3497 for more information
from six.moves import urllib
opener = urllib.request.build_opener()
opener.addheaders = [('User-agent', 'Mozilla/5.0')]
urllib.request.install_opener(opener)
```

Implementation
--------------

In this section, we will discuss the input parameters for the tutorial,
define the model under attack, then code the attack and run some tests.

Inputs
~~~~~~

There are only three inputs for this tutorial, and are defined as
follows:

-  **epsilons** - List of epsilon values to use for the run. It is
   important to keep 0 in the list because it represents the model
   performance on the original test set. Also, intuitively we would
   expect the larger the epsilon, the more noticeable the perturbations
   but the more effective the attack in terms of degrading model
   accuracy. Since the data range here is $[0,1]$, no epsilon
   value should exceed 1.

-  **pretrained_model** - path to the pretrained MNIST model which was
   trained with
   `pytorch/examples/mnist <https://github.com/pytorch/examples/tree/master/mnist>`__.
   For simplicity, download the pretrained model `here <https://drive.google.com/drive/folders/1fn83DF14tWmit0RTKWRhPq5uVXt73e0h?usp=sharing>`__.

-  **use_cuda** - boolean flag to use CUDA if desired and available.
   Note, a GPU with CUDA is not critical for this tutorial as a CPU will
   not take much time.

```python
epsilons = [0, .05, .1, .15, .2, .25, .3]
pretrained_model = "data/lenet_mnist_model.pth"
use_cuda=True
```

Model Under Attack
~~~~~~~~~~~~~~~~~~

As mentioned, the model under attack is the same MNIST model from
`pytorch/examples/mnist <https://github.com/pytorch/examples/tree/master/mnist>`__.
You may train and save your own MNIST model or you can download and use
the provided model. The *Net* definition and test dataloader here have
been copied from the MNIST example. The purpose of this section is to
define the model and dataloader, then initialize the model and load the
pretrained weights.

```python
# LeNet Model definition
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

# MNIST Test dataset and dataloader declaration
test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('../data', train=False, download=True, transform=transforms.Compose([
            transforms.ToTensor(),
            ])),
        batch_size=1, shuffle=True)

# Define what device we are using
print("CUDA Available: ",torch.cuda.is_available())
device = torch.device("cuda" if (use_cuda and torch.cuda.is_available()) else "cpu")

# Initialize the network
model = Net().to(device)

# Load the pretrained model
model.load_state_dict(torch.load(pretrained_model, map_location='cpu'))

# Set the model in evaluation mode. In this case this is for the Dropout layers
model.eval()
```

    Downloading http://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz
    Failed to download (trying next):
    HTTP Error 503: Service Unavailable

    Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz
    Downloading https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz to ../data/MNIST/raw/train-images-idx3-ubyte.gz

    HBox(children=(FloatProgress(value=0.0, max=9912422.0), HTML(value='')))

    Extracting ../data/MNIST/raw/train-images-idx3-ubyte.gz to ../data/MNIST/raw

    Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz
    Downloading http://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz to ../data/MNIST/raw/train-labels-idx1-ubyte.gz

    HBox(children=(FloatProgress(value=0.0, max=28881.0), HTML(value='')))

    Extracting ../data/MNIST/raw/train-labels-idx1-ubyte.gz to ../data/MNIST/raw

    Downloading http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz
    Failed to download (trying next):
    HTTP Error 503: Service Unavailable

    Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz
    Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz to ../data/MNIST/raw/t10k-images-idx3-ubyte.gz

    HBox(children=(FloatProgress(value=0.0, max=1648877.0), HTML(value='')))

    Extracting ../data/MNIST/raw/t10k-images-idx3-ubyte.gz to ../data/MNIST/raw

    Downloading http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz
    Failed to download (trying next):
    HTTP Error 503: Service Unavailable

    Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz
    Downloading https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz to ../data/MNIST/raw/t10k-labels-idx1-ubyte.gz

    HBox(children=(FloatProgress(value=0.0, max=4542.0), HTML(value='')))

    Extracting ../data/MNIST/raw/t10k-labels-idx1-ubyte.gz to ../data/MNIST/raw

    Processing...

    /usr/local/lib/python3.7/dist-packages/torchvision/datasets/mnist.py:502: UserWarning: The given NumPy array is not writeable, and PyTorch does not support non-writeable tensors. This means you can write to the underlying (supposedly non-writeable) NumPy array using the tensor. You may want to copy the array to protect its data or make it writeable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at  /pytorch/torch/csrc/utils/tensor_numpy.cpp:143.)
      return torch.from_numpy(parsed.astype(m[2], copy=False)).view(*s)

    Done!
    CUDA Available:  True

    Net(
      (conv1): Conv2d(1, 10, kernel_size=(5, 5), stride=(1, 1))
      (conv2): Conv2d(10, 20, kernel_size=(5, 5), stride=(1, 1))
      (conv2_drop): Dropout2d(p=0.5, inplace=False)
      (fc1): Linear(in_features=320, out_features=50, bias=True)
      (fc2): Linear(in_features=50, out_features=10, bias=True)
    )

FGSM Attack
~~~~~~~~~~~

Now, we can define the function that creates the adversarial examples by
perturbing the original inputs. The ``fgsm_attack`` function takes three
inputs, *image* is the original clean image ($x$), *epsilon* is
the pixel-wise perturbation amount ($\epsilon$), and *data_grad*
is gradient of the loss w.r.t the input image
($\nabla_{x} J(\mathbf{\theta}, \mathbf{x}, y)$). The function
then creates perturbed image as

\begin{align}perturbed\_image = image + epsilon*sign(data\_grad) = x + \epsilon * sign(\nabla_{x} J(\mathbf{\theta}, \mathbf{x}, y))\end{align}

Finally, in order to maintain the original range of the data, the
perturbed image is clipped to range $[0,1]$.

```python
# FGSM attack code
def fgsm_attack(image, epsilon, data_grad):
    # Collect the element-wise sign of the data gradient
    sign_data_grad = data_grad.sign()
    # Create the perturbed image by adjusting each pixel of the input image
    perturbed_image = image + epsilon*sign_data_grad
    # Adding clipping to maintain [0,1] range
    perturbed_image = torch.clamp(perturbed_image, 0, 1)
    # Return the perturbed image
    return perturbed_image
```

Testing Function
~~~~~~~~~~~~~~~~

Finally, the central result of this tutorial comes from the ``test``
function. Each call to this test function performs a full test step on
the MNIST test set and reports a final accuracy. However, notice that
this function also takes an *epsilon* input. This is because the
``test`` function reports the accuracy of a model that is under attack
from an adversary with strength $\epsilon$. More specifically, for
each sample in the test set, the function computes the gradient of the
loss w.r.t the input data ($data\_grad$), creates a perturbed
image with ``fgsm_attack`` ($perturbed\_data$), then checks to see
if the perturbed example is adversarial. In addition to testing the
accuracy of the model, the function also saves and returns some
successful adversarial examples to be visualized later.

```python
def test( model, device, test_loader, epsilon ):

    # Accuracy counter
    correct = 0
    adv_examples = []

    # Loop over all examples in test set
    for data, target in test_loader:

        # Send the data and label to the device
        data, target = data.to(device), target.to(device)

        # Set requires_grad attribute of tensor. Important for Attack
        data.requires_grad = True

        # Forward pass the data through the model
        output = model(data)
        init_pred = output.max(1, keepdim=True)[1] # get the index of the max log-probability

        # If the initial prediction is wrong, don't bother attacking, just move on
        if init_pred.item() != target.item():
            continue

        # Calculate the loss
        loss = F.nll_loss(output, target)

        # Zero all existing gradients
        model.zero_grad()

        # Calculate gradients of model in backward pass
        loss.backward()

        # Collect datagrad
        data_grad = data.grad.data

        # Call FGSM Attack
        perturbed_data = fgsm_attack(data, epsilon, data_grad)

        # Re-classify the perturbed image
        output = model(perturbed_data)

        # Check for success
        final_pred = output.max(1, keepdim=True)[1] # get the index of the max log-probability
        if final_pred.item() == target.item(): # ATTACK FAILED :(
            correct += 1
            # Special case for saving 0 epsilon examples
            if (epsilon == 0) and (len(adv_examples) < 5):
                adv_ex = perturbed_data.squeeze().detach().cpu().numpy()
                adv_examples.append( (init_pred.item(), final_pred.item(), adv_ex) )
        else: # ATTACK SUCCESS! :)
            # Save some adv examples for visualization later
            if len(adv_examples) < 5:
                adv_ex = perturbed_data.squeeze().detach().cpu().numpy()
                adv_examples.append( (init_pred.item(), final_pred.item(), adv_ex) )

    # Calculate final accuracy for this epsilon
    final_acc = correct/float(len(test_loader))
    print("Epsilon: {}\tTest Accuracy = {} / {} = {}".format(epsilon, correct, len(test_loader), final_acc))

    # Return the accuracy and an adversarial example
    return final_acc, adv_examples
```

Run Attack
~~~~~~~~~~

The last part of the implementation is to actually run the attack. Here,
we run a full test step for each epsilon value in the *epsilons* input.
For each epsilon we also save the final accuracy and some successful
adversarial examples to be plotted in the coming sections. Notice how
the printed accuracies decrease as the epsilon value increases. Also,
note the $\epsilon=0$ case represents the original test accuracy,
with no attack.

```python
accuracies = []
examples = []

# Run test for each epsilon
for eps in epsilons:
    acc, ex = test(model, device, test_loader, eps)
    accuracies.append(acc)
    examples.append(ex)
```

    Epsilon: 0	Test Accuracy = 9810 / 10000 = 0.981
    Epsilon: 0.05	Test Accuracy = 9426 / 10000 = 0.9426
    Epsilon: 0.1	Test Accuracy = 8510 / 10000 = 0.851
    Epsilon: 0.15	Test Accuracy = 6826 / 10000 = 0.6826
    Epsilon: 0.2	Test Accuracy = 4301 / 10000 = 0.4301
    Epsilon: 0.25	Test Accuracy = 2082 / 10000 = 0.2082
    Epsilon: 0.3	Test Accuracy = 869 / 10000 = 0.0869

Results
-------

Accuracy vs Epsilon
~~~~~~~~~~~~~~~~~~~

The first result is the accuracy versus epsilon plot. As alluded to
earlier, as epsilon increases we expect the test accuracy to decrease.
This is because larger epsilons mean we take a larger step in the
direction that will maximize the loss. Notice the trend in the curve is
not linear even though the epsilon values are linearly spaced. For
example, the accuracy at $\epsilon=0.05$ is only about 4% lower
than $\epsilon=0$, but the accuracy at $\epsilon=0.2$ is 25%
lower than $\epsilon=0.15$. Also, notice the accuracy of the model
hits random accuracy for a 10-class classifier between
$\epsilon=0.25$ and $\epsilon=0.3$.

```python
plt.figure(figsize=(5,5))
plt.plot(epsilons, accuracies, "*-")
plt.yticks(np.arange(0, 1.1, step=0.1))
plt.xticks(np.arange(0, .35, step=0.05))
plt.title("Accuracy vs Epsilon")
plt.xlabel("Epsilon")
plt.ylabel("Accuracy")
plt.show()
```

Sample Adversarial Examples
~~~~~~~~~~~~~~~~~~~~~~~~~~~

Remember the idea of no free lunch? In this case, as epsilon increases
the test accuracy decreases **BUT** the perturbations become more easily
perceptible. In reality, there is a tradeoff between accuracy
degredation and perceptibility that an attacker must consider. Here, we
show some examples of successful adversarial examples at each epsilon
value. Each row of the plot shows a different epsilon value. The first
row is the $\epsilon=0$ examples which represent the original
“clean” images with no perturbation. The title of each image shows the
“original classification -> adversarial classification.” Notice, the
perturbations start to become evident at $\epsilon=0.15$ and are
quite evident at $\epsilon=0.3$. However, in all cases humans are
still capable of identifying the correct class despite the added noise.

```python
# Plot several examples of adversarial samples at each epsilon
cnt = 0
plt.figure(figsize=(8,10))
for i in range(len(epsilons)):
    for j in range(len(examples[i])):
        cnt += 1
        plt.subplot(len(epsilons),len(examples[0]),cnt)
        plt.xticks([], [])
        plt.yticks([], [])
        if j == 0:
            plt.ylabel("Eps: {}".format(epsilons[i]), fontsize=14)
        orig,adv,ex = examples[i][j]
        plt.title("{} -> {}".format(orig, adv))
        plt.imshow(ex, cmap="gray")
plt.tight_layout()
plt.show()
```

Where to go next?
-----------------

Hopefully this tutorial gives some insight into the topic of adversarial
machine learning. There are many potential directions to go from here.
This attack represents the very beginning of adversarial attack research
and since there have been many subsequent ideas for how to attack and
defend ML models from an adversary. In fact, at NIPS 2017 there was an
adversarial attack and defense competition and many of the methods used
in the competition are described in this paper: `Adversarial Attacks and
Defences Competition <https://arxiv.org/pdf/1804.00097.pdf>`__. The work
on defense also leads into the idea of making machine learning models
more *robust* in general, to both naturally perturbed and adversarially
crafted inputs.

Another direction to go is adversarial attacks and defense in different
domains. Adversarial research is not limited to the image domain, check
out `this <https://arxiv.org/pdf/1801.01944.pdf>`__ attack on
speech-to-text models. But perhaps the best way to learn more about
adversarial machine learning is to get your hands dirty. Try to
implement a different attack from the NIPS 2017 competition, and see how
it differs from FGSM. Then, try to defend the model from your own
attacks.


# scsRPA 的频域与绝热耦合系数积分实现

> 创建日期：2020-08-24

这篇文档中，我们会回顾 scsRPA 的实现过程 Zhang, Xu [^Zhang-Xu.JPCL.2019.10]。对于其理论推导，我们基本不作的讨论。

我们会大量使用 [前一篇文档](dRPA_Comprehense.ipynb) 的结论与程序。但需要注意，这篇文档中，我们所使用的参考态是 PBE0 而非 PBE，并且分子会选用开壳层分子。

```python
%matplotlib notebook

from pyscf import gto, dft, scf, cc, mcscf, df
import numpy as np
import scipy
from scipy.linalg import fractional_matrix_power
from functools import partial
from matplotlib import pyplot as plt

np.set_printoptions(5, linewidth=120, suppress=True)
np.einsum = partial(np.einsum, optimize=True)
```

## 开壳层分子体系与 PBE0 计算

由于算法的特殊性，我们需要使用上下自旋不定的开壳层分子来描述具体过程。因此，我们使用小体系 cc-pVTZ 基组下的 OH 自由基分子。

```python
mol = gto.Mole()
mol.atom = """
O  0. 0. 0.
H  0. 0. 1.
"""
mol.basis = "cc-pVTZ"
mol.verbose = 0
mol.spin = 1
mol.build()
```

    <pyscf.gto.mole.Mole at 0x7fb1e5206e20>

RI 基组选用 cc-pVTZ-ri 预定义基组。

```python
mol_df = mol.copy()
mol_df.basis = "cc-pVTZ-ri"
mol_df.build()
```

    <pyscf.gto.mole.Mole at 0x7fb2185a00d0>

我们首先计算分子的 PBE0 能量，其目的是得到 PBE0 的分子轨道。该计算实例记在 `mf` 中。出于便利，我们使用未经过 DF (Density Fitting) 的自洽场。

```python
mf = dft.UKS(mol, xc="PBE0").run()
mf.e_tot
```

    -75.68100138443273

我们随后需要定义与分子或方法有关的变量。大多数情况下，我们仍然使用 $i, j$ 表示占据分子轨道，$a, b$ 表示非占分子轨道，$p, q, r, s$ 表示全部分子轨道。这种表示方法并不是很严格，因为一般情况下开壳层应当使用类似于 $i_\alpha, j_\beta$ 等记号；但可以比较好地契合程序编写。

$\mu, \nu, \kappa, \lambda$ 仍然表示原子轨道，$P, Q$ 仍然表示 DF 轨道。$\sigma, \gamma$ 在这篇文档中表示 (未指定的) 自旋；$\alpha, \beta$ 表示上、下自旋。

- `nocc` $(n_\mathrm{occ}^\alpha, n_\mathrm{occ}^\beta)$ 占据轨道数

- `nvir` $(n_\mathrm{vir}^\alpha, n_\mathrm{vir}^\beta)$ 未占轨道数

- `dim_ov` $(n_\mathrm{occ}^\alpha n_\mathrm{vir}^\alpha, n_\mathrm{occ}^\beta n_\mathrm{vir}^\beta)$ 占据乘以未占轨道数

- `nmo` $n_\mathrm{MO}$ 分子轨道数，等于占据轨道数 `nao` $n_\mathrm{AO}$

- `naux` $n_\mathrm{aux}$ Density Fitting 基组轨道数 (或者也称辅助基组 Auxiliary)

- `so`, `sv`, `sa` 占据、未占、全轨道分割 (用于程序编写)，变量形式为 `Tuple[slice, slice]`

- `eri0_ao` $(\mu \nu | \kappa \lambda)$ 原子轨道双电子排斥积分

```python
nocc, nmo, nao = mol.nelec, mol.nao, mol.nao
naux = mol_df.nao
nvir = (nmo - nocc[0], nmo - nocc[1])
dim_ov = (nocc[0] * nvir[0], nocc[1] * nvir[1])
so = slice(0, nocc[0]), slice(0, nocc[1])
sv = slice(nocc[0], nmo), slice(nocc[1], nmo)
eri0_ao = mol.intor("int2e")
```

- `e`, `eo`, `ev` $(e_p^\alpha, e_p^\beta)$ 全、占据、未占 PBE 轨道能

- `C`, `Co`, `Cv` $(C_{\mu p}^\alpha, C_{\mu p}^\beta)$ 全、占据、未占 PBE 轨道系数

```python
e, C = mf.mo_energy, mf.mo_coeff
eo, ev = (e[0][so[0]], e[1][so[1]]), (e[0][sv[0]], e[1][sv[1]])
Co, Cv = (C[0][:, so[0]], C[1][:, so[1]]), (C[0][:, sv[0]], C[1][:, sv[1]])
```

- `eng_xc` PBE0 交换相关能中，仅包含其 GGA 而没有杂化部分的能量 $E_\mathrm{xc}^\mathsf{GGA}$

- `eng_exactX` 使用 PBE0 轨道所构建的精确交换能 $E_\mathrm{x}^\mathsf{exact}$

- `eng_HXX` $E^\mathsf{HXX} = E^\mathsf{xxRPA}_\mathrm{tot} - E^\mathsf{xxRPA}_\mathrm{c}$ 即除去 RPA 的总能量，它在杂化泛函中写为

    $$
    E^\mathsf{HXX} = E^\mathsf{hGGA} - E_\mathrm{xc}^\mathsf{GGA} + (1 - c_\mathrm{x}) E_\mathrm{x}^\mathsf{exact}
    $$

```python
ni = dft.numint.NumInt()
eng_xc = ni.nr_uks(mol, mf.grids, "PBE0", mf.make_rdm1())[1]
eng_xc
```

    -6.671932093325463

```python
eng_exactX = - 0.5 * (mf.get_k(mf.make_rdm1()) * mf.make_rdm1()).sum()
eng_exactX
```

    -8.539252705875096

```python
eng_HXX = mf.e_tot - eng_xc + (1 - ni.hybrid_coeff(mf.xc)) * eng_exactX
eng_HXX
```

    -75.41350882051358

- `V_df_ia` $(V_{ia, P}^\alpha, V_{ia, P}^\beta)$ DF 3c-2e 积分的导出结果，其具有下述性质：

    $$
    \sum_{P} V_{ia,P}^\sigma V_{jb,P}^\sigma \simeq (ia|jb)^\sigma
    $$

- `V` 上述张量重塑为矩阵，维度为 $(ia, P)$

```python
int2c2e = mol_df.intor("int2c2e")
int3c2e = df.incore.aux_e2(mol, mol_df)
int2c2e_half = scipy.linalg.cholesky(int2c2e, lower=True)
V_df_mp2 = scipy.linalg.solve_triangular(int2c2e_half, int3c2e.reshape(-1, naux).T, lower=True).reshape(naux, nao, nao).transpose((1, 2, 0))
V_df_ia = (
    np.einsum("uvP, ui, va -> iaP", V_df_mp2, Co[0], Cv[0]),
    np.einsum("uvP, ui, va -> iaP", V_df_mp2, Co[1], Cv[1]))
V = (V_df_ia[0].reshape(dim_ov[0], naux), V_df_ia[1].reshape(dim_ov[1], naux))
```

- `D_ia` $(D_{ia}^\alpha, D_{ia}^\beta)$ 轨道能之差：

    $$
    D_{ia}^\sigma = \varepsilon_i^\sigma - \varepsilon_a^\sigma
    $$

    其是维度为 $(ia, )$ 的向量

```python
D_ia = (
    - eo[0][:, None] + ev[0][None, :],
    - eo[1][:, None] + ev[1][None, :])
D = (D_ia[0].flatten(), D_ia[1].flatten())
```

## dRPA 能量计算回顾

我们在上一篇文档中，提及了 [RI 方法闭壳层 dRPA 相关能计算](dRPA_Comprehense.ipynb#程序实现：$O(N^4)$-的-RI-方法)。开壳层的计算方法也是非常类似的。我们定义函数 `Pi_alpha` 与 `Pi_beta`，它表示两种自旋下的

$$
\Pi_{PQ}^\sigma (\tilde \omega) = - \sum_{ia \in \sigma} \frac{2 V_{ia, P}^\alpha V_{ia, Q}^\sigma D_{ia}^\sigma}{(D_{ia}^\sigma)^2 + \tilde \omega^2}
$$

或者写为矩阵形式

$$
\mathbf{\Pi}^\sigma (\tilde \omega) = - 2 \mathbf{V}^{\sigma \dagger} \mathbf{D}^{\sigma, 1/2} (\mathbf{D}^\sigma + \tilde \omega^2 \mathbf{I})^{-1} \mathbf{D}^{\sigma, 1/2} \mathbf{V}^\sigma
$$

```python
Pi_alpha = lambda omega: - 2 * np.einsum("dP, d, dQ -> PQ", V[0], D[0] / (D[0]**2 + omega**2), V[0])
Pi_beta  = lambda omega: - 2 * np.einsum("dP, d, dQ -> PQ", V[1], D[1] / (D[1]**2 + omega**2), V[1])
```

我们额外定义 `Pi_dRPA` 为

$$
\Pi_{PQ}^\mathsf{dRPA} (\tilde \omega) = \sum_{\sigma \in \{ \alpha, \beta \}} \Pi_{PQ}^\sigma (\tilde \omega) = \Pi_{PQ}^\alpha (\tilde \omega) + \Pi_{PQ}^\beta (\tilde \omega)
$$

或者写为矩阵形式

$$
\mathbf{\Pi}^\mathsf{dRPA} (\tilde \omega) = \mathbf{\Pi}^\alpha (\tilde \omega) + \mathbf{\Pi}^\beta (\tilde \omega)
$$

```python
Pi_dRPA = lambda omega: Pi_alpha(omega) + Pi_beta(omega)
```

在频域上格点积分的格点 $\tilde{\omega}_g$ 与权重 $w (\tilde{\omega}_g)$ 可以通过下述函数获得：

```python
def gen_leggauss_0_inf(ngrid):
    x, w = np.polynomial.legendre.leggauss(ngrid)
    return 0.5 * (1 + x) / (1 - x), w / (1 - x)**2
```

那么，开壳层的 dRPA 能量 `eng_dRPA` 可以表述为

$$
\begin{align}
E_\mathrm{c}^\mathsf{dRPA}
&= \frac{1}{2 \pi} \int_{0}^{+ \infty} \big( \log \det \big( \mathbf{1} - \mathbf{\Pi}^\mathsf{dRPA} (\tilde \omega) \big) + \mathrm{tr} \big( \mathbf{\Pi}^\mathsf{dRPA} (\tilde \omega) \big) \big) \, \mathrm{d} \tilde \omega \\
&= \frac{1}{2 \pi} \sum_{g} w(\tilde{\omega}_g) \big( \log \det \big( \mathbf{1} - \mathbf{\Pi}^\mathsf{dRPA} (\tilde \omega_g) \big) + \mathrm{tr} \big( \mathbf{\Pi}^\mathsf{dRPA} (\tilde \omega_g) \big) \big)
\end{align}
$$

下面的计算中，对 $\tilde{\omega}$ 的积分使用了 40 个格点。

```python
eng_dRPA = 0
for omega, w_omega in zip(*gen_leggauss_0_inf(40)):
    eng_dRPA += 1 / (2 * np.pi) * w_omega * (np.log(np.linalg.det(np.eye(naux) - Pi_RPA(omega))) + Pi_RPA(omega).trace())
eng_dRPA
```

    -0.3287828424653143

因此，dRPA 总能量为

```python
eng_HXX + eng_dRPA
```

    -75.7422916629789

## ssRPA 相关能计算

ssRPA (**S**ame-**S**pin **RPA**) 相关能计算式与 dRPA 能量计算式非常相似：

$$
E_\mathrm{c}^\mathsf{ssRPA}
= \frac{1}{2 \pi} \sum_{g} \sum_\sigma w(\tilde{\omega}_g) \big( \log \det \big( \mathbf{1} - \mathbf{\Pi}^\sigma (\tilde \omega_g) \big) + \mathrm{tr} \big( \mathbf{\Pi}^\sigma (\tilde \omega_g) \big) \big)
$$

```python
eng_ssRPA = 0
for omega, w_omega in zip(*gen_leggauss_0_inf(40)):
    eng_ssRPA += 1 / (2 * np.pi) * w_omega * (np.log(np.linalg.det(np.eye(naux) - Pi_alpha(omega))) + Pi_alpha(omega).trace())
    eng_ssRPA += 1 / (2 * np.pi) * w_omega * (np.log(np.linalg.det(np.eye(naux) - Pi_beta(omega))) + Pi_beta(omega).trace())
eng_ssRPA
```

    -0.19104532917053876

## osRPA1 相关能计算

osRPA1 (**RPA**-type **O**pposite-**S**pin coupling Dyson equation terminated at the **First-order**) 所贡献的相关能计算式包含对耦合系数的积分。其较为基础的实现方式需要额外对绝热路径耦合系数 $\lambda$ 作格点积分 (注意到前一篇文档中使用了 $\alpha$ 作为耦合系数，但在这篇文档中已经被用作自旋向上记号)。

### 耦合系数 $\lambda$ 积分下的 dRPA 相关能

我们再次回顾 dRPA 的计算。事实上，dRPA 在推导过程中，是先有了关于耦合系数 $\lambda$ 的积分表达式，后面才有 log det 的表达式。我们现在就将关于 $\lambda$ 重新写出：

$$
E_\mathrm{c}^\mathsf{dRPA}
= \frac{1}{2 \pi} \int_{0}^{+ \infty} \mathrm{d} \tilde \omega \, \mathrm{tr} \big( \mathbf{\Pi}^\mathsf{dRPA} (\tilde \omega) \big)
- \frac{1}{2 \pi} \int_{0}^{+ \infty} \mathrm{d} \tilde \omega \int_0^1 \mathrm{d} \lambda \, \mathrm{tr} \left( \frac{\mathbf{\Pi}^\mathsf{dRPA} (\tilde \omega)}{\mathbf{1} - \lambda \mathbf{\Pi}^\mathsf{dRPA} (\tilde \omega)} \right)
$$

上式第二项中，尽管矩阵除法的写法并不是很好，但由于我们求的是迹，因此实际写程序的时候我们不需要关心矩阵交换律是否满足。

对于 $\lambda$ 从 0 至 1 的积分可以使用线性缩放后的 Legendre-Gauss 格点实现。下述函数 `gen_leggauss_0_1` 会输出格点 $\lambda_{g'}$ 与其权重 $w(\lambda_{g'})$：

```python
def gen_leggauss_0_1(ngrid):
    x, w = np.polynomial.legendre.leggauss(ngrid)
    return 0.5 * (x + 1), 0.5 * w
```

因此，若将上述积分表达式更换为格点求和，则

$$
E_\mathrm{c}^\mathsf{dRPA}
= \frac{1}{2 \pi} \sum_g w(\tilde{\omega}_g) \, \mathrm{tr} \big( \mathbf{\Pi}^\mathsf{dRPA} (\tilde{\omega}_g) \big)
- \frac{1}{2 \pi} \sum_g w(\tilde{\omega}_g) \sum_{g'} w(\lambda_{g'}) \, \mathrm{tr} \left( \frac{\mathbf{\Pi}^\mathsf{dRPA} (\tilde{\omega}_g)}{\mathbf{1} - \lambda_{g'} \mathbf{\Pi}^\mathsf{dRPA} (\tilde{\omega}_g)} \right)
$$

下面的计算中，对 $\lambda_{g'}$ 的积分使用了 10 个格点。

```python
eng_dRPA_ac = 0
for omega, w_omega in zip(*gen_leggauss_0_inf(40)):
    # Term 1
    Pi_dRPA_matrix = Pi_RPA(omega)
    eng_dRPA_ac += 1 / (2 * np.pi) * w_omega * Pi_dRPA_matrix.trace()
    # Term 2
    W_matrix = np.zeros((naux, naux))
    for lambd, w_lambd in zip(*gen_leggauss_0_1(10)):
        W_matrix += w_lambd * np.linalg.inv(np.eye(naux) - lambd * Pi_dRPA_matrix)
    eng_dRPA_ac -= 1 / (2 * np.pi) * w_omega * (W_matrix @ Pi_dRPA_matrix).trace()
eng_dRPA_ac
```

    -0.3287828424662268

留意到上式中使用了中间变量 `W_matrix`:

$$
\mathbf{W} (\tilde \omega) = \int_0^1 \mathrm{d} \lambda \big( \mathbf{1} - \lambda \mathbf{\Pi}^\mathsf{dRPA} (\tilde \omega) \big)^{-1}
$$

### osRPA1 相关能

osRPA1 相关能可以写为

$$
E_\mathrm{c}^\mathsf{osRPA1}
= \frac{1}{2 \pi} \sum_{\sigma \in \{ \alpha, \beta \}} \int_{0}^{+ \infty} \mathrm{d} \tilde \omega \, \mathrm{tr} \big( \mathbf{\Pi}^\sigma (\tilde \omega) \big)
- \frac{1}{2 \pi} \sum_{\substack{\sigma, \gamma \in \{ \alpha, \beta \} \\ \gamma \neq \sigma}} \int_{0}^{+ \infty} \mathrm{d} \tilde \omega \int_0^1 \mathrm{d} \lambda \, \mathrm{tr} \left( \frac{\mathbf{\Pi}^\sigma (\tilde \omega)}{\mathbf{1} - \lambda \mathbf{\Pi}^\gamma (\tilde \omega)} \right)
$$

程序编写是类似的。对 $\tilde{\omega}_{g}$ 的积分使用 40 个格点，对 $\lambda_{g'}$ 的积分使用 10 个格点。

```python
eng_osRPA1 = 0
for omega, w_omega in zip(*gen_leggauss_0_inf(40)):
    # Term 1
    Pi_alpha_matrix, Pi_beta_matrix = Pi_alpha(omega), Pi_beta(omega)
    eng_osRPA1 += 1 / (2 * np.pi) * w_omega * (Pi_alpha_matrix.trace() + Pi_beta_matrix.trace())
    # Term 2
    W_alpha_matrix, W_beta_matrix = np.zeros((naux, naux)), np.zeros((naux, naux))
    for lambd, w_lambd in zip(*gen_leggauss_0_1(10)):
        W_alpha_matrix += w_lambd * np.linalg.inv(np.eye(naux) - lambd * Pi_alpha_matrix)
        W_beta_matrix  += w_lambd * np.linalg.inv(np.eye(naux) - lambd * Pi_beta_matrix)
    eng_osRPA1 += - 1 / (2 * np.pi) * w_omega * ((W_alpha_matrix @ Pi_beta_matrix).trace())
    eng_osRPA1 += - 1 / (2 * np.pi) * w_omega * ((W_beta_matrix @ Pi_alpha_matrix).trace())
eng_osRPA1
```

    -0.17888391093782954

上述过程的总计算复杂度是 $O(n_g n_\mathrm{aux}^2 n_\mathrm{occ} n_\mathrm{vir} + n_g n_{g'} b n_\mathrm{aux}^3)$，其中第一项是生成 $\mathbf{\Pi}^\sigma (\tilde \omega)$ 的复杂度，第二项是格点积分外加矩阵求逆的复杂度，$b$ 代表矩阵求逆复杂度中的 $O(b n^3)$。

### 闭壳层下 osRPA1 与 ssRPA 相关能等价的说明

留意到对于闭壳层的情况，$\mathbf{\Pi}^\alpha (\tilde \omega) = \mathbf{\Pi}^\beta (\tilde \omega)$，因此可以使用简化式

$$
- \int_0^1 \mathrm{d} \lambda \, \mathrm{tr} \left( \frac{\mathbf{S}}{\mathbf{1} - \lambda \mathbf{S}} \right) = \log \det \big( \mathbf{1} - \mathbf{S} \big)
$$

其中，$\mathbf{S}$ 是正定的对称矩阵。那么 osRPA1 的表达式很容易地会推到与 ssRPA 完全相同的形式。因此，闭壳层下 $E_\mathrm{c}^\mathsf{osRPA1} = E_\mathrm{c}^\mathsf{ssRPA}$。但开壳层下两者的数值会有不同。

## scsRPA 能量结算

随后，我们就可以计算得到各种贡献分项。

$$
\begin{align}
E_\mathrm{c}^\mathsf{osRPA} &= E_\mathrm{c}^\mathsf{dRPA} - E_\mathrm{c}^\mathsf{ssRPA} \\
E_\mathrm{c}^\mathsf{osRPAr} &= E_\mathrm{c}^\mathsf{osRPA} - E_\mathrm{c}^\mathsf{osRPA1} \\
E_\mathrm{c}^\mathsf{ssRPA+} &= E_\mathrm{c}^\mathsf{ssRPA} + E_\mathrm{c}^\mathsf{osRPAr}
\end{align}
$$

```python
eng_osRPA  = eng_dRPA  - eng_ssRPA
eng_osRPAr = eng_osRPA - eng_osRPA1
eng_ssRPAp = eng_ssRPA + eng_osRPAr
```

```python
print("{:6}{:16.8f}".format(" dRPA", eng_dRPA))
print("{:6}{:16.8f}".format("ssRPA", eng_ssRPA))
print("{:6}{:16.8f}".format("osRPA", eng_osRPA))
print("{:6}{:16.8f}".format("osRPA1", eng_osRPA1))
print("{:6}{:16.8f}".format("osRPAr", eng_osRPAr))
print("{:6}{:16.8f}".format("ssRPA+", eng_ssRPAp))
```

     dRPA      -0.32878284
    ssRPA      -0.19104533
    osRPA      -0.13773751
    osRPA1     -0.17888391
    osRPAr      0.04114640
    ssRPA+     -0.14989893

因此，scsRPA 能量可以写为

$$
E_\mathrm{c}^\mathsf{scsRPA} = \frac{6}{5} E_\mathrm{c}^\mathsf{osRPA1} + \frac{3}{4} E_\mathrm{c}^\mathsf{ssRPA+}
$$

```python
eng_scsRPA = 6/5 * eng_osRPA1 + 3/4 * eng_ssRPAp
eng_scsRPA
```

    -0.327084891771009

事实上，上式也等价于

$$
E_\mathrm{c}^\mathsf{scsRPA} = \frac{9}{20} E_\mathrm{c}^\mathsf{osRPA1} + \frac{3}{4} E_\mathrm{c}^\mathsf{dRPA}
$$

```python
eng_scsRPA = 9/20 * eng_osRPA1 + 3/4 * eng_dRPA
eng_scsRPA
```

    -0.327084891771009

那么 scsRPA 总能量则表示为

```python
eng_HXX + eng_scsRPA
```

    -75.74059371228459

## 氮气解离曲线

我们下面将会复现氮气解离曲线 (Zhang, Xu [^Zhang-Xu.JPCL.2019.10] 文献的 Figure 3)。我们会分别使用到闭壳层与开壳层 scsRPA 的计算结果。

```python
def get_UscsRPA(mol, dfbasis="cc-pVTZ-ri"):
    # molecule preparation
    mf = dft.UKS(mol, xc="PBE0").run()
    eng_PBE0 = mf.e_tot
    mol_df = mol.copy()
    mol_df.basis = dfbasis
    mol_df.build()
    # Basic information preparation
    nocc, nmo, nao = mol.nelec, mol.nao, mol.nao
    naux = mol_df.nao
    nvir = (nmo - nocc[0], nmo - nocc[1])
    dim_ov = (nocc[0] * nvir[0], nocc[1] * nvir[1])
    so = slice(0, nocc[0]), slice(0, nocc[1])
    sv = slice(nocc[0], nmo), slice(nocc[1], nmo)
    eri0_ao = mol.intor("int2e")
    e, C = mf.mo_energy, mf.mo_coeff
    eo, ev = (e[0][so[0]], e[1][so[1]]), (e[0][sv[0]], e[1][sv[1]])
    Co, Cv = (C[0][:, so[0]], C[1][:, so[1]]), (C[0][:, sv[0]], C[1][:, sv[1]])
    # eng_HXX calculation
    ni = dft.numint.NumInt()
    eng_xc = ni.nr_uks(mol, mf.grids, mf.xc, mf.make_rdm1())[1]
    eng_exactX = - 0.5 * (mf.get_k(mf.make_rdm1()) * mf.make_rdm1()).sum()
    eng_HXX = mf.e_tot - eng_xc + (1 - ni.hybrid_coeff(mf.xc)) * eng_exactX
    # density fitting
    int2c2e = mol_df.intor("int2c2e")
    int3c2e = df.incore.aux_e2(mol, mol_df)
    int2c2e_half = scipy.linalg.cholesky(int2c2e, lower=True)
    V_df_mp2 = scipy.linalg.solve_triangular(int2c2e_half, int3c2e.reshape(-1, naux).T, lower=True).reshape(naux, nao, nao).transpose((1, 2, 0))
    V_df_ia = (
        np.einsum("uvP, ui, va -> iaP", V_df_mp2, Co[0], Cv[0]),
        np.einsum("uvP, ui, va -> iaP", V_df_mp2, Co[1], Cv[1]))
    V = (V_df_ia[0].reshape(dim_ov[0], naux), V_df_ia[1].reshape(dim_ov[1], naux))
    D_ia = (
        - eo[0][:, None] + ev[0][None, :],
        - eo[1][:, None] + ev[1][None, :])
    D = (D_ia[0].flatten(), D_ia[1].flatten())
    # dRPA Preparation
    Pi_alpha = lambda omega: - 2 * np.einsum("dP, d, dQ -> PQ", V[0], D[0] / (D[0]**2 + omega**2), V[0])
    Pi_beta  = lambda omega: - 2 * np.einsum("dP, d, dQ -> PQ", V[1], D[1] / (D[1]**2 + omega**2), V[1])
    # dRPA, osRPA1 correlation energy calculation
    eng_dRPA, eng_osRPA1, eng_ssRPA = 0, 0, 0
    for omega, w_omega in zip(*gen_leggauss_0_inf(40)):
        Pi_alpha_matrix, Pi_beta_matrix = Pi_alpha(omega), Pi_beta(omega)
        Pi_dRPA_matrix = Pi_alpha_matrix + Pi_beta_matrix
        eng_dRPA   += 1 / (2 * np.pi) * w_omega * (np.log(np.linalg.det(np.eye(naux) - Pi_dRPA_matrix )) + Pi_dRPA_matrix.trace() )
        eng_ssRPA  += 1 / (2 * np.pi) * w_omega * (np.log(np.linalg.det(np.eye(naux) - Pi_alpha_matrix)) + Pi_alpha_matrix.trace())
        eng_ssRPA  += 1 / (2 * np.pi) * w_omega * (np.log(np.linalg.det(np.eye(naux) - Pi_beta_matrix )) + Pi_beta_matrix.trace() )
        eng_osRPA1 += 1 / (2 * np.pi) * w_omega * (Pi_alpha_matrix.trace() + Pi_beta_matrix.trace())
        W_alpha_matrix, W_beta_matrix = np.zeros((naux, naux)), np.zeros((naux, naux))
        for lambd, w_lambd in zip(*gen_leggauss_0_1(10)):
            W_alpha_matrix += w_lambd * np.linalg.inv(np.eye(naux) - lambd * Pi_alpha_matrix)
            W_beta_matrix  += w_lambd * np.linalg.inv(np.eye(naux) - lambd * Pi_beta_matrix)
        eng_osRPA1 += - 1 / (2 * np.pi) * w_omega * ((W_alpha_matrix @ Pi_beta_matrix).trace())
        eng_osRPA1 += - 1 / (2 * np.pi) * w_omega * ((W_beta_matrix @ Pi_alpha_matrix).trace())
    eng_osRPA  = eng_dRPA  - eng_ssRPA
    eng_scsRPA = 9/20 * eng_osRPA1 + 3/4 * eng_dRPA
    return (
        eng_PBE0,
        eng_HXX + eng_osRPA,
        eng_HXX + eng_dRPA,
        eng_HXX + eng_osRPA1,
        eng_HXX + eng_scsRPA,
    )
```

```python
def get_RscsRPA(mol, dfbasis="cc-pVTZ-ri"):
    # molecule preparation
    mf = dft.RKS(mol, xc="PBE0").run()
    eng_PBE0 = mf.e_tot
    mol_df = mol.copy()
    mol_df.basis = dfbasis
    mol_df.build()
    # Basic information preparation
    nocc, nmo, nao = mol.nelec[0], mol.nao, mol.nao
    naux = mol_df.nao
    nvir = nmo - nocc
    dim_ov = nocc * nvir
    so, sv = slice(0, nocc), slice(nocc, nmo)
    eri0_ao = mol.intor("int2e")
    e, C = mf.mo_energy, mf.mo_coeff
    eo, ev = e[so], e[sv]
    Co, Cv = C[:, so], C[:, sv]
    # eng_HXX calculation
    ni = dft.numint.NumInt()
    eng_xc = ni.nr_rks(mol, mf.grids, mf.xc, mf.make_rdm1())[1]
    eng_exactX = - 0.25 * (mf.get_k(mf.make_rdm1()) * mf.make_rdm1()).sum()
    eng_HXX = mf.e_tot - eng_xc + (1 - ni.hybrid_coeff(mf.xc)) * eng_exactX
    # density fitting
    int2c2e = mol_df.intor("int2c2e")
    int3c2e = df.incore.aux_e2(mol, mol_df)
    int2c2e_half = scipy.linalg.cholesky(int2c2e, lower=True)
    V_df_mp2 = scipy.linalg.solve_triangular(int2c2e_half, int3c2e.reshape(-1, naux).T, lower=True).reshape(naux, nao, nao).transpose((1, 2, 0))
    V_df_ia = np.einsum("uvP, ui, va -> iaP", V_df_mp2, Co, Cv)
    V = V_df_ia.reshape(dim_ov, naux)
    D = (- eo[:, None] + ev[None, :]).flatten()
    # dRPA Preparation
    Pi_alpha = lambda omega: - 2 * np.einsum("dP, d, dQ -> PQ", V, D / (D**2 + omega**2), V)
    # dRPA, osRPA1 correlation energy calculation
    eng_dRPA, eng_ssRPA = 0, 0
    for omega, w_omega in zip(*gen_leggauss_0_inf(40)):
        Pi_alpha_matrix = Pi_alpha(omega)
        Pi_dRPA_matrix = 2 * Pi_alpha_matrix
        eng_dRPA   += 1 / (2 * np.pi) * w_omega * (np.log(np.linalg.det(np.eye(naux) - Pi_dRPA_matrix )) + Pi_dRPA_matrix.trace() )
        eng_ssRPA  += 1 / (2 * np.pi) * w_omega * (np.log(np.linalg.det(np.eye(naux) - Pi_alpha_matrix)) + Pi_alpha_matrix.trace()) * 2
    eng_osRPA1 = eng_ssRPA
    eng_osRPA  = eng_dRPA - eng_ssRPA
    eng_scsRPA = 9/20 * eng_osRPA1 + 3/4 * eng_dRPA
    return (
        eng_PBE0,
        eng_HXX + eng_osRPA,
        eng_HXX + eng_dRPA,
        eng_HXX + eng_osRPA1,
        eng_HXX + eng_scsRPA,
    )
```

下述程序是计算处于解离态的氮原子能量。列表中的数值有 5 项，分别为 PBE0, osRPA, dRPA, osRPA1, scsRPA 的结果。

```python
def get_result_N():
    mol = gto.M(atom="N 0 0 0", basis="cc-pVTZ", spin=3, verbose=0)
    return get_UscsRPA(mol)

result_N = np.array(get_result_N())
```

下述程序是计算处于解离态的氮气能量。列表中的数值亦有 5 项。

```python
def get_result_N2(bond_length):
    mol = gto.M(atom="N 0 0 0; N 0 0 {:16.8f}".format(bond_length), basis="cc-pVTZ", verbose=0)
    return get_RscsRPA(mol)

bond_length_list = np.logspace(np.log10(0.8), np.log10(4.0), 30)
results_N2 = np.array([get_result_N2(length) for length in bond_length_list])
```

最后我们就可以对计算所得到的能量作图。

```python
x_list = bond_length_list
y_list = (results_N2 - 2 * result_N).T * 627.5
```

```python
fig, ax = plt.subplots(figsize=(6, 4))
ax.plot(x_list, y_list[0], c="black", linestyle="-.", label="PBE0")
ax.plot(x_list, y_list[1], c="C2", linestyle="-.", label="osRPA")
ax.plot(x_list, y_list[2], c="C0", linestyle="--", label="dRPA")
ax.plot(x_list, y_list[3], c="C1", linestyle="-", label="osRPA1")
ax.plot(x_list, y_list[4], c="C6", linestyle="-.", label="scsRPA")
ax.plot(x_list, np.zeros_like(x_list), c="black", linewidth=0.5)
ax.set_ylim(-300, 400)
ax.legend()
ax.set_xlabel("Bond Length (Angstrom)")
ax.set_ylabel("Relative energy curves (kcal/mol)")
ax.set_title("$\mathsf{N_2}$ dissociation curve without breaking spin symmetry")
fig.tight_layout()
```

    <IPython.core.display.Javascript object>

由于 CCSD 与 CASSCF 的收敛情况似乎容易存在问题，因此这里仅仅计算了这五条曲线。

[^Zhang-Xu.JPCL.2019.10]: Zhang, I. Y.; Xu, X. Simultaneous Attenuation of Both Self-Interaction Error and Nondynamic Correlation Error in Density Functional Theory: A Spin-Pair Distinctive Adiabatic-Connection Approximation. J. Phys. Chem. Lett. 2019, 10 (10), 2617–2623. doi: [10.1021/acs.jpclett.9b00946](https://doi.org/10.1021/acs.jpclett.9b00946).


[Table of Contents](./table_of_contents.ipynb)

# Designing Nonlinear Kalman Filters

```python
%matplotlib inline
from __future__ import division, print_function
```

```python
#format the book
import book_format
book_format.set_style()
```

<style>
.output_wrapper, .output {
    height:auto !important;
    max-height:100000px;
}
.output_scroll {
    box-shadow:none !important;
    webkit-box-shadow:none !important;
}
</style>

## Introduction

** Author's note: I was initially planning to have a design nonlinear chapter that compares various approaches. This may or may not happen, but for now this chapter has no useful content and I suggest not reading it. **

We see that the Kalman filter reasonably tracks the ball. However, as already explained, this is a silly example; we can predict trajectories in a vacuum with arbitrary precision; using a Kalman filter in this example is a needless complication.

### Kalman Filter with Air Drag

I will dispense with the step 1, step 2, type approach and proceed in a more natural style that you would use in a non-toy engineering problem. We have already developed a Kalman filter that does excellently at tracking a ball in a vacuum, but that does not incorporate the effects of air drag into the model. We know that the process model is implemented with $\textbf{F}$, so we will turn our attention to that immediately.

Notionally, the computation that $\textbf{F}$ computes is

$$x' = Fx$$

With no air drag, we had

$$
\mathbf{F} = \begin{bmatrix}
1 & \Delta t & 0 & 0 & 0 \\
0 & 1 & 0 & 0 & 0 \\
0 & 0 & 1 & \Delta t & \frac{1}{2}{\Delta t}^2 \\
0 & 0 & 0 & 1 & \Delta t \\
0 & 0 & 0 & 0 & 1
\end{bmatrix}
$$

which corresponds to the equations

$$
\begin{aligned}
x &= x + v_x \Delta t \\
v_x &= v_x \\
\\
y &= y + v_y \Delta t + \frac{a_y}{2} {\Delta t}^2 \\
v_y &= v_y + a_y \Delta t \\
a_y &= a_y
\end{aligned}
$$

From the section above we know that our new Euler equations must be

$$
\begin{aligned}
x &= x + v_x \Delta t \\
v_x &= v_x \\
\\
y &= y + v_y \Delta t + \frac{a_y}{2} {\Delta t}^2 \\
v_y &= v_y + a_y \Delta t \\
a_y &= a_y
\end{aligned}
$$

## Realistic 2D Position Sensors

The position sensor in the last example are not very realistic. In general there is no 'raw' sensor that provides (x,y) coordinates. We have GPS, but GPS already uses a Kalman filter to create a filtered output; we should not be able to improve the signal by passing it through another Kalman filter unless we incorporate additional sensors to provide additional information. We will tackle that problem later.

Consider the following set up. In an open field we put two transmitters at a known location, each transmitting a signal that we can detect. We process the signal and determine  how far we are from that signal, with some noise. First, let's look at a visual depiction of that.

```python
import matplotlib.pyplot as plt

circle1=plt.Circle((-4, 0), 5, color='#004080',
                   fill=False, linewidth=20, alpha=.7)
circle2=plt.Circle((4, 0), 5, color='#E24A33',
                   fill=False, linewidth=5, alpha=.7)

fig = plt.gcf()
ax = fig.gca()

plt.axis('equal')
plt.xlim((-10, 10))
plt.ylim((-10, 10))

plt.plot ([-4, 0], [0, 3], c='#004080')
plt.plot ([4, 0], [0, 3], c='#E24A33')
plt.text(-4, -.5, "A", fontsize=16, horizontalalignment='center')
plt.text(4, -.5, "B", fontsize=16, horizontalalignment='center')

ax.add_artist(circle1)
ax.add_artist(circle2)
plt.show()
```

Here I have attempted to show transmitter A, drawn in red, at (-4,0) and a second one B, drawn in blue, at (4,0). The red and blue circles show the range from the transmitters to the robot, with the width illustrating the effect of the $1\sigma$ angular error for each transmitter. Here I have given the blue transmitter more error than the red one. The most probable position for the robot is where the two circles intersect, which I have depicted with the red and blue lines. You will object that we have two intersections, not one, but we will see how we deal with that when we design the measurement function.

This is a very common sensor set up. Aircraft still use this system to navigate, where it is called DME (Distance Measuring Equipment). Today GPS is a much more common navigation system, but I have worked on an aircraft where we integrated sensors like this into our filter along with the GPS, INS, altimeters, etc. We will tackle what is called *multi-sensor fusion* later; for now we will just address this simple configuration.

The first step is to design our state variables. We will assume that the robot is traveling in a straight direction with constant velocity. This is unlikely to be true for a long period of time, but is acceptable for short periods of time. This does not differ from the previous problem - we will want to track the values for the robot's position and velocity. Hence,

$$\mathbf{x} =
\begin{bmatrix}x\\v_x\\y\\v_y\end{bmatrix}$$

The next step is to design the state transition function. This also will be the same as the previous problem, so without further ado,

$$
\mathbf{x}' = \begin{bmatrix}1& \Delta t& 0& 0\\0& 1& 0& 0\\0& 0& 1& \Delta t\\ 0& 0& 0& 1\end{bmatrix}\mathbf{x}$$

The next step is to design the control inputs. We have none, so we set ${\mathbf{B}}=0$.

The next step is to design the measurement function $\mathbf{z} = \mathbf{Hx}$. We can model the measurement using the Pythagorean theorem.

$$
z_a = \sqrt{(x-x_A)^2 + (y-y_A)^2} + v_a\\[1em]
z_b = \sqrt{(x-x_B])^2 + (y-y_B)^2} + v_b
$$

where $v_a$ and $v_b$ are white noise.

We see an immediate problem. The Kalman filter is designed for linear equations, and this is obviously nonlinear. In the next chapters we will look at several ways to handle nonlinear problems in a robust way, but for now we will do something simpler. If we know the approximate position of the robot than we can linearize these equations around that point. I could develop the generalized mathematics for this technique now, but instead let me just present the worked example to give context to that development.

Instead of computing $\mathbf{H}$ we will compute the partial derivative of $\mathbf{H}$ with respect to the robot's position $\mathbf{x}$. You are probably familiar with the concept of partial derivative, but if not, it just means how $\mathbf{H}$ changes with respect to the robot's position. It is computed as the partial derivative of $\mathbf{H}$ as follows:

$$\frac{\partial \mathbf{h}}{\partial \mathbf{x}} =
\begin{bmatrix}
\frac{\partial h_1}{\partial x_1} & \frac{\partial h_1}{\partial x_2} &\dots \\
\frac{\partial h_2}{\partial x_1} & \frac{\partial h_2}{\partial x_2} &\dots \\
\vdots & \vdots
\end{bmatrix}
$$

Let's work the first partial derivative. We want to find

$$\frac{\partial }{\partial x} \sqrt{(x-x_A)^2 + (y-y_A)^2}
$$

Which we compute as

$$
\begin{aligned}
\frac{\partial h_1}{\partial x} &= ((x-x_A)^2 + (y-y_A)^2))^\frac{1}{2} \\
&= \frac{1}{2}\times 2(x-x_a)\times ((x-x_A)^2 + (y-y_A)^2))^{-\frac{1}{2}} \\
&= \frac{x_r - x_A}{\sqrt{(x_r-x_A)^2 + (y_r-y_A)^2}}
\end{aligned}
$$

We continue this computation for the partial derivatives of the two distance equations with respect to $x$, $y$, $dx$ and $dy$, yielding

$$\frac{\partial\mathbf{h}}{\partial\mathbf{x}}=
\begin{bmatrix}
\frac{x_r - x_A}{\sqrt{(x_r-x_A)^2 + (y_r-y_A)^2}} & 0 &
\frac{y_r - y_A}{\sqrt{(x_r-x_A)^2 + (y_r-y_A)^2}} & 0 \\
\frac{x_r - x_B}{\sqrt{(x_r-x_B)^2 + (y_r-y_B)^2}} & 0 &
\frac{y_r - y_B}{\sqrt{(x_r-x_B)^2 + (y_r-y_B)^2}} & 0 \\
\end{bmatrix}
$$

That is pretty painful, and these are very simple equations. Computing the Jacobian can be extremely difficult or even impossible for more complicated systems. However, there is an easy way to get Python to do the work for you by using the SymPy module [1]. SymPy is a Python library for symbolic mathematics. The full scope of its abilities are beyond this book, but it can perform algebra, integrate and differentiate equations, find solutions to differential equations, and much more. We will use it to compute our Jacobian!

First, a simple example. We will import SymPy, initialize its pretty print functionality (which will print equations using LaTeX). We will then declare a symbol for NumPy to use.

```python
import sympy
from sympy import init_printing
init_printing(use_latex='png')

phi, x = sympy.symbols('\phi, x')
phi
```

Notice how we use a latex expression for the symbol `phi`. This is not necessary, but if you do it will render as LaTeX when output. Now let's do some math. What is the derivative of $\sqrt{\phi}$?

```python
sympy.diff('sqrt(phi)')
```

We can factor equations.

```python
sympy.factor('phi**3 -phi**2 + phi - 1')
```

SymPy has a remarkable list of features, and as much as I enjoy exercising its features we cannot cover them all here. Instead, let's compute our Jacobian.

```python
from sympy import symbols, Matrix
phi = symbols('\phi')
phi

x, y, xa, xb, ya, yb, dx, dy = symbols('x y x_a x_b y_a y_b dx dy')

H = Matrix([[sympy.sqrt((x-xa)**2 + (y-ya)**2)],
            [sympy.sqrt((x-xb)**2 + (y-yb)**2)]])

state = Matrix([x, dx, y, dy])
H.jacobian(state)
```

In a nutshell, the entry (0,0) contains the difference between the x coordinate of the robot and transmitter A's x coordinate divided by the distance between the robot and A. (2,0) contains the same, except for the y coordinates of the robot and transmitters. The bottom row contains the same computations, except for transmitter B. The 0 entries account for the velocity components of the state variables; naturally the range does not provide us with velocity.

The values in this matrix change as the robot's position changes, so this is no longer a constant; we will have to recompute it for every time step of the filter.

If you look at this you may realize that this is just a computation of x/dist and y/dist, so we can switch this to a trigonometic form with no loss of generality:

$$\frac{\partial\mathbf{h}}{\partial\mathbf{x}}=
\begin{bmatrix}
-\cos{\theta_A} & 0 & -\sin{\theta_A} & 0 \\
-\cos{\theta_B} & 0 & -\sin{\theta_B} & 0
\end{bmatrix}
$$

However, this raises a huge problem. We are no longer computing $\mathbf{H}$, but $\Delta\mathbf{H}$, the change of $\mathbf{H}$. If we passed this into our Kalman filter without altering the rest of the design the output would be nonsense. Recall, for example, that we multiply $\mathbf{Hx}$ to generate the measurements that would result from the given estimate of $\mathbf{x}$ But now that $\mathbf{H}$ is linearized around our position it contains the *change* in the measurement function.

We are forced, therefore, to use the *change* in $\mathbf{x}$ for our state variables. So we have to go back and redesign our state variables.

>Please note this is a completely normal occurrence in designing Kalman filters. The textbooks present examples like this as *fait accompli*, as if it is trivially obvious that the state variables needed to be velocities, not positions. Perhaps once you do enough of these problems it would be trivially obvious, but at that point why are you reading a textbook? I find myself reading through a presentation multiple times, trying to figure out why they made a choice, finally to realize that it is because of the consequences of something on the next page. My presentation is longer, but it reflects what actually happens when you design a filter. You make what seem reasonable design choices, and as you move forward you discover properties that require you to recast your earlier steps. As a result, I am going to somewhat abandon my **step 1**, **step 2**, etc.,  approach, since so many real problems are not quite that straightforward.

If our state variables contain the velocities of the robot and not the position then how do we track where the robot is? We can't. Kalman filters that are linearized in this fashion use what is called a *nominal trajectory* - i.e. you assume a position and track direction, and then apply the changes in velocity and acceleration to compute the changes in that trajectory. How could it be otherwise? Recall the graphic showing the intersection of the two range circles - there are two areas of intersection. Think of what this would look like if the two transmitters were very close to each other - the intersections would be two very long crescent shapes. This Kalman filter, as designed, has no way of knowing your true position from only distance measurements to the transmitters. Perhaps your mind is already leaping to ways of working around this problem. If so, stay engaged, as later sections and chapters will provide you with these techniques. Presenting the full solution all at once leads to more confusion than insight, in my opinion.

So let's redesign our *state transition function*. We are assuming constant velocity and no acceleration, giving state equations of
$$
\dot{x}' = \dot{x} \\
\ddot{x}' = 0 \\
\dot{y}' = \dot{y} \\
\dot{y}' = 0$$

This gives us the the *state transition function* of

$$
\mathbf{F} = \begin{bmatrix}0 &1 & 0& 0\\0& 0& 0& 0\\0& 0& 0& 1\\ 0& 0& 0& 0\end{bmatrix}$$

A final complication comes from the measurements that we pass in. $\mathbf{Hx}$ is now computing the *change* in the measurement from our nominal position, so the measurement that we pass in needs to be not the range to A and B, but the *change* in range from our measured range to our nominal position.

There is a lot here to take in, so let's work through the code bit by bit. First we will define a function to compute $\frac{\partial\mathbf{h}}{\partial\mathbf{x}}$ for each time step.

```python
from math import sin, cos, atan2

def H_of(pos, pos_A, pos_B):
    """ Given the position of our object at 'pos' in 2D, and two
    transmitters A and B at positions 'pos_A' and 'pos_B', return
    the partial derivative of H
    """

    theta_a = atan2(pos_a[1] - pos[1], pos_a[0] - pos[0])
    theta_b = atan2(pos_b[1] - pos[1], pos_b[0] - pos[0])

    return np.array([[0, -cos(theta_a), 0, -sin(theta_a)],
                     [0, -cos(theta_b), 0, -sin(theta_b)]])
```

Now we need to create our simulated sensor.

```python
from numpy.random import randn

class DMESensor(object):
    def __init__(self, pos_a, pos_b, noise_factor=1.0):
        self.A = pos_a
        self.B = pos_b
        self.noise_factor = noise_factor

    def range_of(self, pos):
        """ returns tuple containing noisy range data to A and B
        given a position 'pos'
        """

        ra = math.sqrt((self.A[0] - pos[0])**2 + (self.A[1] - pos[1])**2)
        rb = math.sqrt((self.B[0] - pos[0])**2 + (self.B[1] - pos[1])**2)

        return (ra + randn()*self.noise_factor,
                rb + randn()*self.noise_factor)
```

Finally, we are ready for the Kalman filter code. I will position the transmitters at x=-100 and 100, both with y=-20. This gives me enough space to get good triangulation from both as the robot moves. I will start the robot at (0,0) and move by (1,1) each time step.

```python
import kf_book.book_plots as bp
from filterpy.kalman import KalmanFilter
import math
import numpy as np

pos_a = (100, -20)
pos_b = (-100, -20)

f1 = KalmanFilter(dim_x=4, dim_z=2)

f1.F = np.array ([[0, 1, 0, 0],
                  [0, 0, 0, 0],
                  [0, 0, 0, 1],
                  [0, 0, 0, 0]], dtype=float)

f1.R *= 1.
f1.Q *= .1

f1.x = np.array([[1, 0, 1, 0]], dtype=float).T
f1.P = np.eye(4) * 5.

# initialize storage and other variables for the run
count = 30
xs, ys = [], []
pxs, pys = [], []

# create the simulated sensor
d = DMESensor(pos_a, pos_b, noise_factor=3.)

# pos will contain our nominal position since the filter does not
# maintain position.
pos = [0, 0]

for i in range(count):
    # move (1,1) each step, so just use i
    pos = [i, i]

    # compute the difference in range between the nominal track
    # and measured ranges
    ra,rb = d.range_of(pos)
    rx,ry = d.range_of((pos[0] + f1.x[0, 0], pos[1] + f1.x[2, 0]))
    z = np.array([[ra - rx], [rb - ry]])

    # compute linearized H for this time step
    f1.H = H_of (pos, pos_a, pos_b)

    # store stuff so we can plot it later
    xs.append(f1.x[0, 0]+i)
    ys.append(f1.x[2, 0]+i)
    pxs.append(pos[0])
    pys.append(pos[1])

    # perform the Kalman filter steps
    f1.predict()
    f1.update(z)

bp.plot_filter(xs, ys)
bp.plot_track(pxs, pys)
plt.legend(loc=2)
plt.show()
```

## Linearizing the Kalman Filter

Now that we have seen an example of linearizing the Kalman filter we are in a position to better understand the math.

We start by assuming some function $\mathbf f$

## Example: A falling Ball

**author's note: ignore this section for now. **

In the **Designing Kalman Filters** chapter I first considered tracking a ball in a vacuum, and then in the atmosphere. The Kalman filter performed very well for vacuum, but diverged from the ball's path in the atmosphere. Let us look at the output; to avoid littering this chapter with code from that chapter I have placed it all in the file `ekf_internal.py'.

```python
import kf_book.ekf_internal as ekf
ekf.plot_ball()
```

We can artificially force the Kalman filter to track the ball by making $Q$ large. That would cause the filter to mistrust its prediction, and scale the kalman gain $K$ to strongly favor the measurments. However, this is not a valid approach. If the Kalman filter is correctly predicting the process we should not 'lie' to the filter by telling it there are process errors that do not exist. We may get away with that for some problems, in some conditions, but in general the Kalman filter's performance will be substandard.

Recall from the **Designing Kalman Filters** chapter that the acceleration is

$$a_x = (0.0039 + \frac{0.0058}{1+\exp{[(v-35)/5]}})*v*v_x \\
a_y = (0.0039 + \frac{0.0058}{1+\exp{[(v-35)/5]}})*v*v_y- g
$$

These equations will be *very* unpleasant to work with while we develop this subject, so for now I will retreat to a simpler one dimensional problem using this simplified equation for acceleration that does not take the nonlinearity of the drag coefficient into account:

$$\begin{aligned}
\ddot{y} &= \frac{0.0034ge^{-y/20000}\dot{y}^2}{2\beta} - g \\
\ddot{x} &= \frac{0.0034ge^{-x/20000}\dot{x}^2}{2\beta}
\end{aligned}$$

Here $\beta$ is the ballistic coefficient, where a high number indicates a low drag.

This is still nonlinear, so we need to linearize this equation at the current state point. If our state is position and velocity, we need an equation for some arbitrarily small change in $\mathbf{x}$, like so:

$$ \begin{bmatrix}\Delta \dot{x} \\ \Delta \ddot{x} \\ \Delta \dot{y} \\ \Delta \ddot{y}\end{bmatrix} =
\large\begin{bmatrix}
\frac{\partial \dot{x}}{\partial x} &
\frac{\partial \dot{x}}{\partial \dot{x}} &
\frac{\partial \dot{x}}{\partial y} &
\frac{\partial \dot{x}}{\partial \dot{y}} \\
\frac{\partial \ddot{x}}{\partial x} &
\frac{\partial \ddot{x}}{\partial \dot{x}}&
\frac{\partial \ddot{x}}{\partial y}&
\frac{\partial \dot{x}}{\partial \dot{y}}\\
\frac{\partial \dot{y}}{\partial x} &
\frac{\partial \dot{y}}{\partial \dot{x}} &
\frac{\partial \dot{y}}{\partial y} &
\frac{\partial \dot{y}}{\partial \dot{y}} \\
\frac{\partial \ddot{y}}{\partial x} &
\frac{\partial \ddot{y}}{\partial \dot{x}}&
\frac{\partial \ddot{y}}{\partial y}&
\frac{\partial \dot{y}}{\partial \dot{y}}
\end{bmatrix}\normalsize
\begin{bmatrix}\Delta x \\ \Delta \dot{x} \\ \Delta \dot{y} \\ \Delta \ddot{y}\end{bmatrix}$$

The equations do not contain both an x and a y, so any partial derivative with both in it must be equal to zero. We also know that $\large\frac{\partial \dot{x}}{\partial x}\normalsize = 0$ and that  $\large\frac{\partial \dot{x}}{\partial \dot{x}}\normalsize = 1$, so our matrix ends up being

$$\mathbf{F} = \begin{bmatrix}0&1&0&0 \\
\frac{0.0034e^{-x/22000}\dot{x}^2g}{44000\beta}&0&0&0
\end{bmatrix}$$

$$\begin{aligned}\ddot{x} &= -\frac{1}{2}C_d\rho A \dot{x}\\
\ddot{y} &= -\frac{1}{2}C_d\rho A \dot{y}-g\end{aligned}$$

```python
from sympy.abc import *
from sympy import *

init_printing(pretty_print=True, use_latex='mathjax')

x1 = (0.0034*g*exp(-x/22000)*((x)**2))/(2*b) - g

x2 = (a*g*exp(-x/c)*(Derivative(x)**2))/(2*b) - g

#pprint(x1)
#pprint(Derivative(x)*Derivative(x,n=2))
#pprint(diff(x2, x))
```

** orphan text
This approach has many issues. First, of course, is the fact that the linearization does not produce an exact answer. More importantly, we are not linearizing the actual path, but our filter's estimation of the path. We linearize the estimation because it is statistically likely to be correct; but of course it is not required to be. So if the filter's output is bad that will cause us to linearize an incorrect estimate, which will almost certainly lead to an even worse estimate. In these cases the filter will quickly diverge. This is where the 'black art' of Kalman filter comes in. We are trying to linearize an estimate, and there is no guarantee that the filter will be stable. A vast amount of the literature on Kalman filters is devoted to this problem. Another issue is that we need to linearize the system using analytic methods. It may be difficult or impossible to find an analytic solution to some problems. In other cases we may be able to find the linearization, but the computation is very expensive. **

## References

[1] http://sympy.org


# Analyze runs

# Purpose
Analyze the batch simulations from MLFlow

# Methodology
* Load data about runs from MLFlow API
* Analyze this data

```python
# %load ../imports.py
## Local packages:

%matplotlib inline
%load_ext autoreload
%autoreload 2
%config Completer.use_jedi = False  ## (To fix autocomplete)

## External packages:
import pandas as pd
pd.options.display.max_rows = 999
pd.options.display.max_columns = 999
pd.set_option("display.max_columns", None)
import numpy as np
np.set_printoptions(linewidth=150)

import numpy as np
import os
import matplotlib.pyplot as plt
#if os.name == 'nt':
#    plt.style.use('presentation.mplstyle')  # Windows

import plotly.express as px
import plotly.graph_objects as go

import seaborn as sns
import sympy as sp
from sympy.physics.mechanics import (dynamicsymbols, ReferenceFrame,
                                      Particle, Point)
from sympy.physics.vector.printing import vpprint, vlatex
from IPython.display import display, Math, Latex
from src.substitute_dynamic_symbols import run, lambdify

import pyro

import sklearn
import pykalman
from statsmodels.sandbox.regression.predstd import wls_prediction_std
import statsmodels.api as sm

from scipy.integrate import solve_ivp

## Local packages:
from src.data import mdl

from src.symbols import *
from src.parameters import *
import src.symbols as symbols
from src import prime_system
from src.models import regression
from src.visualization.regression import show_pred
from src.visualization.plot import track_plot

## Load models:
# (Uncomment these for faster loading):
import src.models.vmm_abkowitz  as vmm
from src.models.vmm import ModelSimulator

```

```python
from src.models.vmm import ModelSimulator
from src.data.mdl import load_test
```

```python
from mlflow.tracking.client import MlflowClient
from mlflow.entities import ViewType
from src.to_mlflow import tracking_uri
mlflow_client = MlflowClient(tracking_uri=tracking_uri)
```

```python
experiment = mlflow_client.get_experiment_by_name(name='VCT_linear')
```

```python
experiment.name
```

```python
runs = MlflowClient().search_runs(
  experiment_ids=experiment.experiment_id,
  filter_string="",
  #run_view_type=ViewType.ACTIVE_ONLY,
  max_results=1000,
  order_by=["metrics.accuracy DESC"]
)
len(runs)
```

```python
df = pd.DataFrame()

for run in runs:

    if run.info.status!='FINISHED':
        continue  # Unfinnished simulation

    row = pd.Series(run.data.params, dtype='object', name=run.info.run_id)
    row = row.append(pd.Series(run.data.metrics))
    row['status'] = run.info.status
    row['artifact_uri'] = run.info.artifact_uri
    row['experiment_id'] = run.info.experiment_id
    row['experiment'] = experiment.name

    df = df.append(row, ignore_index=True)

df.index.name='run_id'
```

```python
df.head()
```

```python
df.sort_values(by='r2', ascending=False).head()
```

```python
row = df.sort_values(by='r2', ascending=False).iloc[0]
```

```python
from IPython.display import Image
Image(os.path.join(row.artifact_uri.replace(r'file:///',''),'track_plot.png'))

```

```python
row
```

## Load and resimulate

```python
run_params = row
df, meta_data = load_test(**run_params)

model = ModelSimulator.load(run_params['model'])
result = model.simulate(df)
```

```python

```