# Autograd (3)：梯度下降法解量子化学 RHF 自洽场能量

> 创建时间：2019-12-17

这一篇文档将会简单地介绍量子化学计算中经常使用的 RHF 方法，在自动求导下的一种应用。

RHF (Restricted Hartree-Fock) 方法是量化计算的基本方法；对于更高精度的分子体系或晶体等周期性体系的计算，譬如 DFT、Post-HF 方法如 CC 和 CI，通常都是以 HF 方法为原型开发。但 RHF 方法的求解方程通常都是通过自洽场方法 (self-consistent field, SCF) 迭代计算；一种经典的迭代计算算法是 DIIS (见 [简单理解 SCF 中的 DIIS](../../QC_Notes/SCF_Series/diis_comprehen.ipynb))。关于 SCF 求解 RHF 的入门文档可以参考 Szabo, Ostlund [^Szabo-Ostlund.Dover.1996]。关于自洽场求解 RHF 的程序相关问题，可以参考 [pyxdh 文档](https://py-xdh.readthedocs.io/zh_CN/latest/qcbasic/basic_rhf.html)。

而事实上，更经典但现在不太用于 RHF 方程求解的方法是将其化为局域凸优化问题。尽管通常这类凸优化问题属于类牛顿法的范畴，但如果将 Hessian 矩阵始终视为单位矩阵的常数倍，这就化为了普通的梯度下降问题，也恰好是自动求导可以解决的问题范畴。这篇文档就用 PyTorch 简单地实现这样一个功能。

在继续这篇文档之前，我们先引入程序库：

```python
import torch
import math
import numpy as np
import scipy
from pyscf import gto, scf
import warnings

torch.set_printoptions(precision=5, sci_mode=False, linewidth=120)
np.set_printoptions(precision=5, suppress=True, linewidth=120)
warnings.filterwarnings("ignore")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
```

:::{attention}

这篇文档是找了一个作者稍微熟悉的优化问题，用自动求导来解决。但这不是通常处理这类问题的做法，况且并不高效。

同时需要指出的是，一般来说 RHF 并不是一个凸优化问题；但在比较合理的初猜下，可以将问题 **近似地** 看成凸优化问题。后面尽管默认了凸优化这个条件，但许多现实的 RHF 收敛问题实际上来源于非凸优化的复杂性。作者也不是搞优化问题的，因此很多术语可能用得不对。

因此，这篇文档是 **娱乐向** 的文档。

:::

## 分子体系的定义

### 分子体系能量

我们使用的分子体系是不对称的双氧水分子，基组为 6-31G，用 PySCF 进行计算。分子体系的总能量为 `energy_tot` $E_\mathrm{tot}$ -150.5850338 Hartree。这篇文档的目标就是计算得到这个数值。

```python
mol = gto.Mole()
mol.atom = """
O  0.0  0.0  0.0
O  0.0  0.0  1.5
H  1.0  0.0  0.0
H  0.0  0.7  1.0
"""
mol.basis = "6-31G"
mol.verbose = 0
mol.build()
```

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

```python
scf_eng = scf.RHF(mol).run()
energy_tot = scf_eng.e_tot
energy_tot
```

    -150.5850337808369

### 恒定物理量与矩阵

这里列举的是随分子不同而不同的量。这些量可能计算耗时小，也可能是计算过程中的重要构成部分；但它们与凸优化问题的变量无关，在这篇文档中姑且称为恒定量。

分子体系总能量中，一部分能量是无需通过量化方法，而只需要经典物理的库伦排斥力计算得到。这部分能量是原子核排斥能 `energy_nuc` $E_\mathrm{nuc}$ 37.8846744 Hartree。为了简化叙述，我们就直接借用了 PySCF 的程序的结果。

```python
energy_nuc = scf_eng.energy_nuc()
energy_nuc
```

    37.88467440864127

下面是原子轨道角标的矩阵或张量。无电子相互作用的 Hamiltonian 积分 (Hamiltonian Core 积分) $h_{\mu \nu}$ 记为 `H`，重叠积分 $S_{\mu \nu}$ 记为 `S`，双电子积分 $(\mu \nu | \kappa \lambda)$ 记为 `eri_ao`。

```python
T = mol.intor("int1e_kin")
Vnuc = mol.intor("int1e_nuc")
H = T + Vnuc
S = mol.intor("int1e_ovlp")
eri_ao = mol.intor("int2e")
```

上述矩阵或张量都是通过 PySCF 给出的；PySCF 使用的默认引擎是 numpy。我们需要将它们转为 PyTorch 引擎的矩阵或张量。

```python
H = torch.tensor(H, device=device)
S = torch.tensor(S, device=device)
eri_ao = torch.tensor(eri_ao, device=device)
```

### 凸优化自变量与密度矩阵

在凸优化问题中，RHF 问题可以看作是对下述原子轨道空间下 $\mathbf{X}$ 矩阵 (或写作矩阵元的形式，$X_{\mu \nu}$) 作为自变量的函数优化问题 (Helgaker et al. [^Helgaker-Jorgensen.Wiley.2013], eq 10.7.69，但这篇文档对密度矩阵 $R_{\mu \nu}$ 的定义恰好是 Helgaker 课本中的两倍)：

$$
E_\mathrm{tot} [\mathbf{X}] = \mathrm{tr} (\mathbf{R}[\mathbf{X}]) \mathbf{h} + \frac{1}{4} \mathrm{tr} (\mathbf{R}[\mathbf{X}] \mathbf{G} \mathbf{R}[\mathbf{X}]) + E_\mathrm{nuc}
$$

其中，根据 Helgaker, eq 10.7.32，有

$$
\mathbf{R}[\mathbf{X}] = \exp (- \mathbf{X} \mathbf{S}) \mathbf{R}_0 \exp (\mathbf{S} \mathbf{X})
$$

关于这些公式的意义与计算过程，我们后面还会有更详细的描述。这里需要留意的是，$\mathbf{G}$ 是一个四维张量；其它地方的记号与矩阵乘法没有区别。

我们现在需要知道的是，其一，$E_\mathrm{tot} [\mathbf{X}]$ 是我们的凸优化目标；对应于统计或机器学习中的概念，就类似于损失函数 (loss function)。其二，尽管 $\mathbf{X}$ 是我们的学习目标，但真正有物理意义的、可以被实验观测的量是 $\mathbf{R}[\mathbf{X}]$；该量的意义是电子云密度 (简称密度)。换一种说法的话，RHF 方法作为量子化学方法的核心目标是给出基态的电子云密度；有了电子云密度之后，许多其它的分子性质是可以从中导出的。因此，对于 RHF 方法而言，其电子云密度的价值近乎等价于作为 Schrodinger 方程的求解目的的波函数。

方才我们已经用 PySCF 计算了分子的 RHF 的能量。PySCF 也可以给出分子基态下的 RHF 密度 `D` $D_{\mu \nu}$：

```python
D = scf_eng.make_rdm1()
D = torch.tensor(D, device=device)
```

对于 RHF 问题而言，我们会说，$\mathbf{X}$ 取到函数 $E_\mathrm{tot} [\mathbf{X}]$ 处于最小值的极小值几乎等价于问题

$$
\mathbf{D} = \mathbf{R} [\mathbf{X}]
$$

我们等下会验证这个结论。

需要补充的是，PySCF 给出电子密度的方式与上述过程完全不相同。在一般的自洽场迭代过程中，电子态密度是由分子轨道系数得到的：

$$
\mathbf{D} = 2 \mathbf{C}_\mathrm{occ} \mathbf{C}_\mathrm{occ}^\mathrm{T}
$$

或用 Einstein Summation Convention 的语言，

$$
D_\mathrm{\mu \nu} = 2 C_{\mu i} C_{\nu i}
$$

这与 $\mathbf{D} = \mathbf{R} [\mathbf{X}]$ 看似是完全无关的公式，但至少从结果上，两者确实相同。

## 损失函数的定义

### PyTorch 矩阵幂运算

这一段只是讨论一个技术细节。尽管 PyTorch 功能确实很强大，但一个比较尴尬的地方是 PyTorch 不支持不少需要一些数学技巧才能保证效率与稳定性的算法。其中一个算法是矩阵幂 (作者 Bing 了一刻钟发现实在找不到，于是作出了断言，但高 zi 贵 bi 的作者是不会提 issue 或者 stackoverflow 的 >.>)。

为了解决这样一个量化的问题，我们无奈需要自己造轮子，目标是只要能用就行。

下面的程序 `calc_exp` 是根据矩阵幂的定义给出的：

$$
\exp(\mathbf{X}) = \sum_{k=0}^{K} \frac{\mathbf{X}^k}{k!}
$$

其中用到了矩阵指数运算 $\mathbf{X}^k$ 的函数 `mat_power`；该函数通过递归定义。

我们知道，这种函数的求和上界应当是 $K \rightarrow \infty$；但实际的程序不能允许这种操作，必须让 $K$ 作有界的截断。这种有界的截断的判据是，若通过截断计算得到的 $\exp(\mathbf{X})$ 和 numpy 给出的不截断的 $\exp(\mathbf{X})$ 近乎相等 (`np.allclose` 的判标下)，那么就允许截断。

```python
def mat_power(X, order):
    if order > 0:
        return X @ mat_power(X, order - 1)
    return torch.eye(X.shape[0], dtype=X.dtype, device=X.device)

def calc_exp(X, debug=False):
    target = scipy.linalg.expm(X.cpu().clone().detach().numpy())
    result = mat_power(X, 0) + mat_power(X, 1)
    order = 1
    while not np.allclose(result.cpu().clone().detach().numpy(), target):
        order += 1
        result += mat_power(X, order) / math.factorial(order)
        if debug: print(np.linalg.norm(result.cpu().clone().detach().numpy() - target))
    if debug: print("Current order is", order)
    return result
```

读者可以尝试几个矩阵幂运算的例子。但需要留意，这个程序的鲁棒性显然是不高的；遇到上述定义下 $\mathbf{X}$ 超出收敛域的情形下，这个程序就失效了。

### 分子能量函数 $E_\mathrm{tot} [\mathbf{X}]$

下面我们回到分子能量函数 (损失函数) $E_\mathrm{tot} [\mathbf{X}]$ 的定义中。事实上让程序来做这件事其实是相当容易的，只需要几行代码就能搞定：

```python
def calc_RX(X, S, R0):
    return calc_exp(-X @ S) @ R0 @ calc_exp(S @ X)
```

```python
def calc_EX(X, S, R0, H, eri_ao, energy_nuc):
    RX = calc_RX(X, S, R0)
    return (
        torch.einsum("uv, uv ->", RX, H)
        + 0.5 * torch.einsum("uv, uvkl, kl ->", RX, eri_ao, RX)
        - 0.25 * torch.einsum("uv, ukvl, kl ->", RX, eri_ao, RX)
        + energy_nuc
    )
```

但我们需要对其中的过程作一些说明。首先是关于密度的导出式 `calc_RX`：

$$
\mathbf{R}[\mathbf{X}] = \exp (- \mathbf{X} \mathbf{S}) \mathbf{R}_0 \exp (\mathbf{S} \mathbf{X})
$$

这里的 $\mathbf{S}$ 是重叠积分。$\mathbf{R}_0$ 原则上是任意的、满足下述三式的矩阵 (Helgaker, eq 10.7.25-27)：

$$
\begin{align}
\mathbf{R}_0^\mathrm{T} &= \mathbf{R}_0 \\
\mathrm{tr} (\mathbf{R}_0 \mathbf{S}) &= N \\
\mathbf{R}_0 \mathbf{S} \mathbf{R}_0 &= \mathbf{R}_0
\end{align}
$$

上述三式对应的性质分别称为对称性、电子数守恒、基于重叠积分 $\mathbf{S}$ 的幂等性。其中，$N$ 代表体系的电子数。

但实际上，从数值稳定性的角度上讲，我们不希望让 $\mathbf{X}$ 太大以至于超出收敛域；因此 $\mathbf{R}_0$ 总是需要尽可能接近真实密度 $\mathbf{D}$。在这篇文档中，我们称 $\mathbf{R}_0$ 为密度初猜。

同时，$\mathbf{X}$ 是一种旋转矩阵；它必须要满足反对称性，即 $\mathbf{X}^\mathrm{T} = - \mathbf{X}$。这是相当重要且有用的性质，但我们这里不作更多分析。

而对于能量函数 `calc_EX`：

$$
E_\mathrm{tot} [\mathbf{X}] = \mathrm{tr} (\mathbf{R}[\mathbf{X}]) \mathbf{h} + \frac{1}{4} \mathrm{tr} (\mathbf{R}[\mathbf{X}] \mathbf{G} \mathbf{R}[\mathbf{X}]) + E_\mathrm{nuc}
$$

由于其中的四脚标张量的存在稍有麻烦，在程序编写便利、以及借用强大的 `einsum` 功能的角度上，我们采用 Einstein Summation Convention 可以将上式写作

$$
E_\mathrm{tot} [\mathbf{X}] = R_{\mu \nu} [\mathbf{X}] h_{\mu \nu} + \frac{1}{2} R_{\mu \nu} [\mathbf{X}] (\mu \nu | \kappa \lambda) R_{\kappa \lambda} [\mathbf{X}] - \frac{1}{4} R_{\mu \nu} [\mathbf{X}] (\mu \kappa | \nu \lambda) R_{\kappa \lambda} [\mathbf{X}] + E_\mathrm{nuc}
$$

各个矩阵元都已经在上文中被定义了，因此程序的书写也并不困难。

### 最小值的极小点分析

这一段我们分析极小点附近的性质。我们之前提到过，$\mathbf{R}_0$ 总是需要尽可能接近真实密度 $\mathbf{D}$。如果这两者相等是什么情形？

现在我们假定 $E_\mathrm{tot} [\mathbf{X}]$ 已经取到了处于最小值的极小点。我们说起过这种情况下与 $\mathbf{D} = \mathbf{R} [\mathbf{X}]$ 几乎等价。但我们并没有验证过这个结论。下面我们就来验证之。

若同时有 $\mathbf{R}_0 = \mathbf{D}$ 与 $\mathbf{D} = \mathbf{R} [\mathbf{X}]$，那么

$$
\exp (- \mathbf{X} \mathbf{S}) \mathbf{D} \exp (\mathbf{S} \mathbf{X}) = \mathbf{D}
$$

一个显然满足上式的 $\mathbf{X}$ 是零矩阵。在这种情况下，我们应当会期待分子体系总能量 `energy_autograd` 就是使用现成量化软件给出的总能量 `energy_tot`：

```python
X = torch.zeros_like(S, requires_grad=True, device=S.device)
R0 = D.clone()
```

```python
energy_autograd = calc_EX(X, S, R0, H, eri_ao, energy_nuc)
float(energy_autograd)
```

    -150.58503378083833

并且，我们会期待 $\mathbf{X}$ 确实取到了极小值，通过自动求导得到的导数是不会发生任何变化的零矩阵：

```python
energy_autograd.backward()
gX = X.grad
```

```python
torch.norm(gX)
```

    tensor(    0.00000, device='cuda:0', dtype=torch.float64)

从这个意义上，我们可以说，这一套工作流程是基本合理的。

### 真实密度初猜的分析

但上面一段是假设我们已经知道体系密度的理想情况。现实情况是，我们并不知道体系的真正密度，从而需要求自变量 $\mathbf{X}$ 的结果。

现实中有许多给出密度初猜 $\mathbf{R}_0$ 的方式。其中最简单的密度初猜是将 Fock 矩阵当作 Hamiltonian Core 矩阵，进行对角化后得到的密度矩阵。在 PySCF 中，它通过 `init_guess_by_1e` 方法给出。同时，初始的自变量 (或称旋转矩阵) $\mathbf{X}$ 定为零矩阵。

```python
X = torch.zeros_like(S, requires_grad=True, device=S.device)
R0 = torch.tensor(scf_eng.init_guess_by_1e(), device=S.device)
```

以这个初始密度代入能量的计算公式中，会得到一个比刚才的 -150 Hartree 高出不少的能量：

```python
energy_autograd = calc_EX(X, S, R0, H, eri_ao, energy_nuc)
energy_autograd
```

    tensor(-140.54132, device='cuda:0', dtype=torch.float64, grad_fn=<AddBackward0>)

得到这样的结果也是情理之中：这完全满足变分原理的要求，即任意非基态的波函数 (在 RHF 中，也可以称为其对应的基态密度 $\mathbf{D}$)，那么其给出的能量会高于最低能量。这也满足优化问题的要求，即任何偏离作为最小值的极小点附近的变量 $\mathbf{X}$，一定给出高于最小值的结果来。

下面我们来考虑关于 $\mathbf{X}$ 的梯度。首先，由于 $\mathbf{X}$ 在最小值附近，我们假设其附近是凸性的，因此其梯度必然不为零：

```python
energy_autograd.backward()
gX = X.grad
torch.norm(gX)
```

    tensor(6.70115, device='cuda:0', dtype=torch.float64)

其次，我们说起过，$\mathbf{X}$ 具有反对称的性质。如果让能量函数 $E_\mathrm{tot} [\mathbf{X}]$ 作反向传播，那么 $\mathbf{X}$ 的梯度也应当是满足这种反对称性质。

```python
torch.allclose(-gX.T, gX)
```

    True

## RHF 的梯度下降求取

### 学习率递降 class `Scheduler`

这其实是一个非常细节的问题，但却需要不少代码来解决。

我们都知道在深度学习中，梯度下降算法与学习率递降都是非常有必要的。但由于我们需要在训练的过程中始终保证作为自变量的 $\mathbf{X}$ 具有反对称性，因此它不能像 PyTorch 深度网络中参数一样，使用较为自由的优化方式；而是需要在优化 (机器学习中类似于训练) 过程中始终重新将 $\mathbf{X}$ 反对称化。

因此，这里暂时不使用 PyTorch 提供的 optim 类的 `Adam` 或 `ReduceLROnPlateau` 类来解决问题，而需要手动造与学习率有关的轮子；至于梯度下降法只能用 Naive 的版本，而不使用 Adam 等更高级的版本。至于是不是一定不能用 PyTorch 提供的自动化工具，作者暂时还不清楚。

```python
class Sheduler:

    def __init__(self, init_lr=0.05, min_lr=1e-6, factor=0.95, patience=3, stop_train=30, stop_thresh=1e-6, debug=True):
        self.lr = init_lr
        self.min_lr = min_lr
        self.factor = factor
        self.patience = patience
        self.stop_train = stop_train
        self.stop_thresh = stop_thresh
        self.debug = debug

        self.loss = 1e10
        self.epoch = 0

        self.flag_min = False
        self.end_train = False
        self.loss_list = []

    def step(self, loss, epoch_id=-1):
        self.epoch += 1
        self.loss_list.append(float(loss))
        if loss < self.loss:
            self.loss = loss
            self.epoch = 0
        if not self.flag_min:
            if self.epoch >= self.patience:
                self.lr *= self.factor
                self.epoch = 0
                print("learning rate decrease to {:7.4f} on epoch {:5d}.".format(self.lr, epoch)) if self.debug else None
            if self.lr < self.min_lr:
                self.flag_min = True
                print("Hit minimum learning rate.") if self.debug else None
        if self.flag_min:
            if self.epoch >= self.stop_train:
                print("End learning at epoch {:5d} since {:5d} epoch loss is larger than lowest loss.".format(epoch, self.stop_train)) if self.debug else None
                self.end_train = True
        if np.all(np.abs(np.array(self.loss_list[-self.stop_train:]) - float(self.loss)) < self.stop_thresh) and len(self.loss_list) > 30:
            print("End learning at epoch {:5d} since last {:5d} epoch loss is allclose to lowest loss".format(epoch, self.stop_train) +
                  " within threshold {:7.2e}.".format(self.stop_thresh)) if self.debug else None
            self.end_train = True
```

这个类的优化过程稍微有些一点点复杂 (但比起现成的程序已经简单太多)，一些简单的说明是：

- 优化过程的学习率从 `init_lr` (默认 0.05) 渐渐降低为 `min_lr` (默认 1e-6)，每次降低 `factor` (默认 0.95) 倍，但降到最低后就不再下降；
- 学习率降低的判据是，如果连续 `patience` (默认 3) 次的损失函数大于最低的损失函数值，就降低学习率
- 上述过程比通常深度学习中的要求高很多；这是因为通常的深度学习不是凸优化问题，在损失函数势能面上游走不见得是坏事；但 RHF 问题我们假设在合理的密度初猜下是凸优化，损失函数不应该经常上升；
- 停止优化的第一判据是，若学习率 `self.lr` 降到最低后连续 `stop_train` (默认 30) 次损失函数的值大于最低的那一次，那么停止优化；
- 停止优化的第二判据是，若最后连续 `stop_train` (默认 30) 次损失函数都与最低的那一次之间差不超过 `stop_thresh` (默认 1e-6) Hartree 大小，那么停止优化。

### 优化过程

有了上面文字的铺垫，后面的优化过程若读者有任何一点深度学习的脚本经验，应该是非常容易理解的。这里用到的初始条件是

- 密度初猜 $\mathbf{R}_0$：采用 Hamiltonian Core 作为 Fock 矩阵对角化所得密度
- 初始自变量 $\mathbf{X}$：采用零矩阵

```python
X = torch.zeros_like(S, requires_grad=True, device=S.device)
R0 = torch.tensor(scf_eng.init_guess_by_1e(), device=S.device)
```

随后学习率从 0.04 开始递降，进行损失函数 (即能量函数 $E_\mathrm{tot} [\mathbf{X}]$) 的优化。

```python
sheduler = Sheduler(init_lr=0.04)

for epoch in range(0, 5000):
    # Calculate loss function
    energy_autograd = calc_EX(X, S, R0, H, eri_ao, energy_nuc)
    energy_autograd.backward()
    # Make optimization step
    t = X - X.grad * sheduler.lr
    # Force anti-symmetrize X
    X = ((t - t.T) / 2).clone().detach().requires_grad_(True)
    # Update learning rate or stop training
    sheduler.step(energy_autograd, epoch)
    if sheduler.end_train:
        print("End at epoch {:5d}.".format(epoch))
        break
print(float(energy_autograd))
```

    learning rate decrease to  0.0380 on epoch     7.
    learning rate decrease to  0.0361 on epoch    10.
    learning rate decrease to  0.0343 on epoch   182.
    learning rate decrease to  0.0326 on epoch   196.
    End learning at epoch   859 since last    30 epoch loss is allclose to lowest loss within threshold 1.00e-06.
    End at epoch   859.
    -150.5850332855298

这就完成了借助于自动求导的 RHF 的能量计算问题！我们可以拿它与 PySCF 作为量化程序计算得到的能量作比较：

```python
energy_tot
```

    -150.5850337808369

## 小结

我们在这份文档中简单地讨论了在给定电子积分的情况下，使用 PyTorch 的自动求导功能，以 Naive 梯度下降法给出 RHF 能量。这样的程序既可以在 CPU 下运行，也能在 GPU 下运行。

简单的代码总结可以是，在定义了与量子化学问题无关的矩阵指数 `mat_power`、矩阵幂 `calc_exp`、学习率递降管理器 `Sheduler` 后，我们可以进行 RHF 的能量计算，其代码也非常简单。

**分子定义**

```python
mol = gto.Mole()
mol.atom = """
O  0.0  0.0  0.0
O  0.0  0.0  1.5
H  1.0  0.0  0.0
H  0.0  0.7  1.0
"""
mol.basis = "6-31G"
mol.verbose = 0
mol.build()
```

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

**电子积分定义**

```python
H = torch.tensor(mol.intor("int1e_kin") + mol.intor("int1e_nuc"), device=device)
S = torch.tensor(mol.intor("int1e_ovlp"), device=device)
eri_ao = torch.tensor(mol.intor("int2e"), device=device)
```

**核库伦排斥能定义**

```python
natm = mol.natm
Z_A, A_t = mol.atom_charges(), mol.atom_coords()
r_AB = np.linalg.norm(A_t[:, None, :] - A_t[None, :, :], axis=-1)
r_AB += np.diag(np.ones(natm) * np.inf)
energy_nuc = 0.5 * (Z_A[None, :] * Z_A[:, None] / r_AB).sum()
```

**密度矩阵与能量损失函数定义**

```python
def calc_RX(X, S, R0):
    return calc_exp(-X @ S) @ R0 @ calc_exp(S @ X)
```

```python
def calc_EX(X, S, R0, H, eri_ao, energy_nuc):
    RX = calc_RX(X, S, R0)
    return (
        torch.einsum("uv, uv ->", RX, H)
        + 0.5 * torch.einsum("uv, uvkl, kl ->", RX, eri_ao, RX)
        - 0.25 * torch.einsum("uv, ukvl, kl ->", RX, eri_ao, RX)
        + energy_nuc
    )
```

**密度初猜与自变量定义**

```python
X = torch.zeros_like(S, requires_grad=True, device=S.device)
R0 = torch.tensor(scf_eng.init_guess_by_1e(), device=S.device)
```

**梯度下降得到 RHF 能量**

```python
sheduler = Sheduler(init_lr=0.04, debug=False)

for epoch in range(0, 5000):
    energy_autograd = calc_EX(X, S, R0, H, eri_ao, energy_nuc)
    energy_autograd.backward()
    t = X - X.grad * sheduler.lr
    X = ((t - t.T) / 2).clone().detach().requires_grad_(True)
    sheduler.step(energy_autograd, epoch)
    if sheduler.end_train: break
```

**最终能量**

```python
float(energy_autograd)
```

    -150.5850332855298

[^Szabo-Ostlund.Dover.1996]: Szabo, A.; Ostlund, N. S. *Modern Quantum Chemistry: Introduction to Advanced Electronic Structure Theory*; Dover Publications, 1996.

[^Helgaker-Jorgensen.Wiley.2013]: Helgaker, T.; Olsen, J.; Jorgensen, P. *Molecular Electronic-Structure Theory*; Wiley-Blackwell, 2013.


# Tutorial 1: Learning to Predict
**Week 3, Day 4: Reinforcement Learning**

**By Neuromatch Academy**

__Content creators:__ Marcelo Mattar and Eric DeWitt with help from Matt Krause

__Content reviewers:__ Ella Batty, Byron Galbraith and Michael Waskom

**Our 2021 Sponsors, including Presenting Sponsor Facebook Reality Labs**

---

# Tutorial objectives

*Estimated timing of tutorial: 50 min*

In this tutorial, we will learn how to estimate state-value functions in a classical conditioning paradigm using Temporal Difference (TD) learning and examine TD-errors at the presentation of the conditioned and unconditioned stimulus (CS and US) under different CS-US contingencies. These exercises will provide you with an understanding of both how reward prediction errors (RPEs) behave in classical conditioning and what we should expect to see if Dopamine represents a "canonical" model-free RPE.

At the end of this tutorial:
* You will learn to use the standard tapped delay line conditioning model
* You will understand how RPEs move to CS
* You will understand how variability in reward size effects RPEs
* You will understand how differences in US-CS timing effect RPEs

```python
# @title Tutorial slides

# @markdown These are the slides for all videos in this tutorial.
from IPython.display import IFrame
IFrame(src=f"https://mfr.ca-1.osf.io/render?url=https://osf.io/2jzdu/?direct%26mode=render%26action=download%26mode=render", width=854, height=480)
```

---
# Setup

```python
# Imports

import numpy as np
import matplotlib.pyplot as plt
```

```python
# @title Figure Settings
import ipywidgets as widgets       # interactive display
%config InlineBackend.figure_format = 'retina'
plt.style.use("https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle")
```

```python
# @title Plotting Functions
from matplotlib import ticker

def plot_value_function(V, ax=None, show=True):
  """Plot V(s), the value function"""
  if not ax:
    fig, ax = plt.subplots()

  ax.stem(V, use_line_collection=True)
  ax.set_ylabel('Value')
  ax.set_xlabel('State')
  ax.set_title("Value function: $V(s)$")

  if show:
    plt.show()

def plot_tde_trace(TDE, ax=None, show=True, skip=400):
  """Plot the TD Error across trials"""
  if not ax:
    fig, ax = plt.subplots()

  indx = np.arange(0, TDE.shape[1], skip)
  im = ax.imshow(TDE[:,indx])
  positions = ax.get_xticks()
  # Avoid warning when setting string tick labels
  ax.xaxis.set_major_locator(ticker.FixedLocator(positions))
  ax.set_xticklabels([f"{int(skip * x)}" for x in positions])
  ax.set_title('TD-error over learning')
  ax.set_ylabel('State')
  ax.set_xlabel('Iterations')
  ax.figure.colorbar(im)
  if show:
    plt.show()

def learning_summary_plot(V, TDE):
  """Summary plot for Ex1"""
  fig, (ax1, ax2) = plt.subplots(nrows = 2, gridspec_kw={'height_ratios': [1, 2]})

  plot_value_function(V, ax=ax1, show=False)
  plot_tde_trace(TDE, ax=ax2, show=False)
  plt.tight_layout()
```

```python
# @title Helper Functions and Classes

def reward_guesser_title_hint(r1, r2):
  """"Provide a mildly obfuscated hint for a demo."""
  if (r1==14 and r2==6) or (r1==6 and r2==14):
    return "Technically correct...(the best kind of correct)"

  if  ~(~(r1+r2) ^ 11) - 1 == (6 | 24): # Don't spoil the fun :-)
    return "Congratulations! You solved it!"

  return "Keep trying...."

class ClassicalConditioning:

    def __init__(self, n_steps, reward_magnitude, reward_time):

        # Task variables
        self.n_steps = n_steps
        self.n_actions = 0
        self.cs_time = int(n_steps/4) - 1

        # Reward variables
        self.reward_state = [0,0]
        self.reward_magnitude = None
        self.reward_probability = None
        self.reward_time = None

        self.set_reward(reward_magnitude, reward_time)

        # Time step at which the conditioned stimulus is presented

        # Create a state dictionary
        self._create_state_dictionary()

    def set_reward(self, reward_magnitude, reward_time):

        """
        Determine reward state and magnitude of reward
        """
        if reward_time >= self.n_steps - self.cs_time:
            self.reward_magnitude = 0

        else:
            self.reward_magnitude = reward_magnitude
            self.reward_state = [1, reward_time]

    def get_outcome(self, current_state):

        """
        Determine next state and reward
        """
        # Update state
        if current_state < self.n_steps - 1:
            next_state = current_state + 1
        else:
            next_state = 0

        # Check for reward
        if self.reward_state == self.state_dict[current_state]:
            reward = self.reward_magnitude
        else:
            reward = 0

        return next_state, reward

    def _create_state_dictionary(self):

        """
        This dictionary maps number of time steps/ state identities
        in each episode to some useful state attributes:

        state      - 0 1 2 3 4 5 (cs) 6 7 8 9 10 11 12 ...
        is_delay   - 0 0 0 0 0 0 (cs) 1 1 1 1  1  1  1 ...
        t_in_delay - 0 0 0 0 0 0 (cs) 1 2 3 4  5  6  7 ...
        """
        d = 0

        self.state_dict = {}
        for s in range(self.n_steps):
            if s <= self.cs_time:
                self.state_dict[s] = [0,0]
            else:
                d += 1 # Time in delay
                self.state_dict[s] = [1,d]

class MultiRewardCC(ClassicalConditioning):
  """Classical conditioning paradigm, except that one randomly selected reward,
    magnitude, from a list, is delivered of a single fixed reward."""
  def __init__(self, n_steps, reward_magnitudes, reward_time=None):
    """"Build a multi-reward classical conditioning environment
      Args:
        - nsteps: Maximum number of steps
        - reward_magnitudes: LIST of possible reward magnitudes.
        - reward_time: Single fixed reward time
      Uses numpy global random state.
      """
    super().__init__(n_steps, 1, reward_time)
    self.reward_magnitudes = reward_magnitudes

  def get_outcome(self, current_state):
    next_state, reward = super().get_outcome(current_state)
    if reward:
      reward=np.random.choice(self.reward_magnitudes)
    return next_state, reward

class ProbabilisticCC(ClassicalConditioning):
  """Classical conditioning paradigm, except that rewards are stochastically omitted."""
  def __init__(self, n_steps, reward_magnitude, reward_time=None, p_reward=0.75):
    """"Build a multi-reward classical conditioning environment
      Args:
        - nsteps: Maximum number of steps
        - reward_magnitudes: Reward magnitudes.
        - reward_time: Single fixed reward time.
        - p_reward: probability that reward is actually delivered in rewarding state
      Uses numpy global random state.
      """
    super().__init__(n_steps, reward_magnitude, reward_time)
    self.p_reward = p_reward

  def get_outcome(self, current_state):
    next_state, reward = super().get_outcome(current_state)
    if reward:
      reward*= int(np.random.uniform(size=1)[0] < self.p_reward)
    return next_state, reward
```

---
# Section 1: Temporal difference learning

```python
# @title Video 1: Introduction
from ipywidgets import widgets

out2 = widgets.Output()
with out2:
  from IPython.display import IFrame
  class BiliVideo(IFrame):
      def __init__(self, id, page=1, width=400, height=300, **kwargs):
          self.id=id
          src = 'https://player.bilibili.com/player.html?bvid={0}&page={1}'.format(id, page)
          super(BiliVideo, self).__init__(src, width, height, **kwargs)

  video = BiliVideo(id="BV13f4y1d7om", width=854, height=480, fs=1)
  print('Video available at https://www.bilibili.com/video/{0}'.format(video.id))
  display(video)

out1 = widgets.Output()
with out1:
  from IPython.display import YouTubeVideo
  video = YouTubeVideo(id="YoNbc9M92YY", width=854, height=480, fs=1, rel=0)
  print('Video available at https://youtube.com/watch?v=' + video.id)
  display(video)

out = widgets.Tab([out1, out2])
out.set_title(0, 'Youtube')
out.set_title(1, 'Bilibili')

display(out)
```

    Tab(children=(Output(), Output()), _titles={'0': 'Youtube', '1': 'Bilibili'})

__Environment:__

- The agent experiences the environment in episodes or trials.
- Episodes terminate by transitioning to the inter-trial-interval (ITI) state and they are initiated from the ITI state as well. We clamp the value of the terminal/ITI states to zero.
- The classical conditioning environment is composed of a sequence of states that the agent deterministically transitions through. Starting at State 0, the agent moves to State 1 in the first step, from State 1 to State 2 in the second, and so on.  These states represent time in the tapped delay line representation
- Within each episode, the agent is presented a CS and US (reward).
- The CS is always presented at 1/4 of the total duration of the trial. The US (reward) is then delivered after the CS. The interval between the CS and US is specified by `reward_time`.
- The agent's goal is to learn to predict expected rewards from each state in the trial.

**General concepts**

* Return $G_{t}$: future cumulative reward, which can be written in a recursive form
\begin{align}
G_{t} &= \sum \limits_{k = 0}^{\infty} \gamma^{k} r_{t+k+1} \\
&= r_{t+1} + \gamma G_{t+1}
\end{align}
where $\gamma$ is discount factor that controls the importance of future rewards, and $\gamma \in [0, 1]$. $\gamma$ may also be interpreted as probability of continuing the trajectory.
* Value funtion $V_{\pi}(s_t=s)$: expecation of the return
\begin{align}
V_{\pi}(s_t=s) &= \mathbb{E} [ G_{t}\; | \; s_t=s, a_{t:\infty}\sim\pi] \\
& = \mathbb{E} [ r_{t+1} + \gamma G_{t+1}\; | \; s_t=s, a_{t:\infty}\sim\pi]
\end{align}
With an assumption of **Markov process**, we thus have:
\begin{align}
V_{\pi}(s_t=s) &= \mathbb{E} [ r_{t+1} + \gamma V_{\pi}(s_{t+1})\; | \; s_t=s, a_{t:\infty}\sim\pi] \\
&= \sum_a \pi(a|s) \sum_{r, s'}p(s', r)(r + V_{\pi}(s_{t+1}=s'))
\end{align}

**Temporal difference (TD) learning**

* With a Markovian assumption, we can use $V(s_{t+1})$ as an imperfect proxy for the true value $G_{t+1}$ (Monte Carlo bootstrapping), and thus obtain the generalised equation to calculate TD-error:
\begin{align}
\delta_{t} = r_{t+1} + \gamma V(s_{t+1}) - V(s_{t})
\end{align}

* Value updated by using the learning rate constant $\alpha$:
\begin{align}
V(s_{t}) \leftarrow V(s_{t}) + \alpha \delta_{t}
\end{align}

  (Reference: https://web.stanford.edu/group/pdplab/pdphandbook/handbookch10.html)

__Definitions:__

* TD-error:
\begin{align}
\delta_{t} = r_{t+1} + \gamma V(s_{t+1}) - V(s_{t})
\end{align}

* Value updates:
\begin{align}
V(s_{t}) \leftarrow V(s_{t}) + \alpha \delta_{t}
\end{align}

## Coding Exercise 1: TD-learning with guaranteed rewards

Implement TD-learning to estimate the state-value function in the classical-conditioning world with guaranteed rewards, with a fixed magnitude, at a fixed delay after the conditioned stimulus, CS. Save TD-errors over learning (i.e., over trials) so we can visualize them afterwards.

In order to simulate the effect of the CS, you should only update $V(s_{t})$ during the delay period after CS. This period is indicated by the boolean variable `is_delay`. This can be implemented by multiplying the expression for updating the value function by `is_delay`.

Use the provided code to estimate the value function. We will use helper class `ClassicalConditioning`.

```python
def td_learner(env, n_trials, gamma=0.98, alpha=0.001):
  """ Temporal Difference learning

  Args:
    env (object): the environment to be learned
    n_trials (int): the number of trials to run
    gamma (float): temporal discount factor
    alpha (float): learning rate

  Returns:
    ndarray, ndarray: the value function and temporal difference error arrays
  """
  V = np.zeros(env.n_steps) # Array to store values over states (time)
  TDE = np.zeros((env.n_steps, n_trials)) # Array to store TD errors

  for n in range(n_trials):

    state = 0 # Initial state

    for t in range(env.n_steps):

      # Get next state and next reward
      next_state, reward = env.get_outcome(state)

      # Is the current state in the delay period (after CS)?
      is_delay = env.state_dict[state][0]

      ########################################################################
      ## TODO for students: implement TD error and value function update
      # Fill out function and remove
      #raise NotImplementedError("Student excercise: implement TD error and value function update")
      #################################################################################
      # Write an expression to compute the TD-error
      TDE[state, n] = reward + gamma * V[next_state] - V[state]

      # Write an expression to update the value function
      V[state] += alpha * TDE[state, n]  * is_delay

      # Update state
      state = next_state

  return V, TDE

# Initialize classical conditioning class
env = ClassicalConditioning(n_steps=40, reward_magnitude=10, reward_time=10)

# Perform temporal difference learning
V, TDE = td_learner(env, n_trials=20000)

# Visualize
learning_summary_plot(V, TDE)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D4_ReinforcementLearning/solutions/W3D4_Tutorial1_Solution_adeb004b.py)

*Example output:*

## Interactive Demo 1.1: US to CS Transfer

During classical conditioning, the subject's behavioral response (e.g., salivating) transfers from the unconditioned stimulus (US; like the smell of tasty food) to the conditioned stimulus (CS; like Pavlov ringing his bell) that predicts it. Reward prediction errors play an important role in this process by adjusting the value of states according to their expected, discounted return.

Use the widget below to examine how reward prediction errors change over time.

Before training (orange line), only the reward state has high reward prediction error. As training progresses (blue line, slider), the reward prediction errors shift to the conditioned stimulus, where they end up when the trial is complete (green line).

Dopamine neurons, which are thought to carry reward prediction errors _in vivo_, show exactly the same behavior!

```python
#@title

#@markdown Make sure you execute this cell to enable the widget!

n_trials = 20000

@widgets.interact
def plot_tde_by_trial(trial = widgets.IntSlider(value=5000, min=0, max=n_trials-1 , step=1, description="Trial #")):
  if 'TDE' not in globals():
    print("Complete Exercise 1 to enable this interactive demo!")
  else:

    fig, ax = plt.subplots()
    ax.axhline(0, color='k') # Use this + basefmt=' ' to keep the legend clean.
    ax.stem(TDE[:, 0], linefmt='C1-', markerfmt='C1d', basefmt=' ',
            label="Before Learning (Trial 0)",
            use_line_collection=True)
    ax.stem(TDE[:, -1], linefmt='C2-', markerfmt='C2s', basefmt=' ',
            label="After Learning (Trial $\infty$)",
            use_line_collection=True)
    ax.stem(TDE[:, trial], linefmt='C0-', markerfmt='C0o', basefmt=' ',
            label=f"Trial {trial}",
            use_line_collection=True)

    ax.set_xlabel("State in trial")
    ax.set_ylabel("TD Error")
    ax.set_title("Temporal Difference Error by Trial")
    ax.legend()
```

    interactive(children=(IntSlider(value=5000, description='Trial #', max=19999), Output()), _dom_classes=('widge…

## Interactive Demo 1.2: Learning Rates and Discount Factors

Our TD-learning agent has two parameters that control how it learns: $\alpha$, the learning rate, and $\gamma$, the discount factor. In Exercise 1, we set these parameters to $\alpha=0.001$ and $\gamma=0.98$ for you. Here, you'll investigate how changing these parameters alters the model that TD-learning learns.

Before enabling the interactive demo below, take a moment to think about the functions of these two parameters. $\alpha$ controls the size of the Value function updates produced by each TD-error. In our simple, deterministic world, will this affect the final model we learn? Is a larger $\alpha$ necessarily better in more complex, realistic environments?

The discount rate $\gamma$ applies an exponentially-decaying weight to returns occuring in the future, rather than the present timestep. How does this affect the model we learn? What happens when $\gamma=0$ or $\gamma \geq 1$?

Use the widget to test your hypotheses.

```python
#@title

#@markdown Make sure you execute this cell to enable the widget!

@widgets.interact
def plot_summary_alpha_gamma(alpha = widgets.FloatSlider(value=0.0001, min=0.0001, max=0.1, step=0.0001, readout_format='.4f', description="alpha"),
                             gamma = widgets.FloatSlider(value=0.980, min=0, max=1.1, step=0.010, description="gamma")):
  env = ClassicalConditioning(n_steps=40, reward_magnitude=10, reward_time=10)
  try:
    V_params, TDE_params = td_learner(env, n_trials=20000, gamma=gamma, alpha=alpha)
  except NotImplementedError:
    print("Finish Exercise 1 to enable this interactive demo")

  learning_summary_plot(V_params,TDE_params)
```

    interactive(children=(FloatSlider(value=0.0001, description='alpha', max=0.1, min=0.0001, readout_format='.4f'…

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D4_ReinforcementLearning/solutions/W3D4_Tutorial1_Solution_d20f92ed.py)

---
# Section 2: TD-learning with varying reward magnitudes

*Estimated timing to here from start of tutorial: 30 min*

In the previous exercise, the environment was as simple as possible. On every trial, the CS predicted the same reward, at the same time, with 100% certainty. In the next few exercises, we will make the environment more progressively more complicated and examine the TD-learner's behavior.

## Interactive Demo 2: Match the Value Functions

First, will replace the environment with one that dispenses one of several rewards, chosen at random. Shown below is the final value function $V$ for a TD learner that was trained in an enviroment where the CS predicted a reward of 6 or 14 units; both rewards were equally likely).

Can you find another pair of rewards that cause the agent to learn the same value function? Assume each reward will be dispensed 50% of the time.

Hints:
* Carefully consider the definition of the value function $V$. This can be solved analytically.
* There is no need to change $\alpha$ or $\gamma$.
* Due to the randomness, there may be a small amount of variation.

```python
# @markdown Make sure you execute this cell to enable the widget! Please allow some time for the new figure to load

n_trials = 20000
np.random.seed(2020)
rng_state = np.random.get_state()
env = MultiRewardCC(40, [6, 14], reward_time=10)
V_multi, TDE_multi = td_learner(env, n_trials, gamma=0.98, alpha=0.001)

@widgets.interact
def reward_guesser_interaction(r1 = widgets.IntText(value=0, min=0, max=50, description="Reward 1"),
                               r2 = widgets.IntText(value=0, min=0, max=50, description="Reward 2")):
  try:
    env2 = MultiRewardCC(40, [r1, r2], reward_time=10)
    V_guess, _ = td_learner(env2, n_trials, gamma=0.98, alpha=0.001)
    fig, ax = plt.subplots()
    m, l, _ = ax.stem(V_multi, linefmt='y-', markerfmt='yo', basefmt=' ', label="Target",
            use_line_collection=True)
    m.set_markersize(15)
    m.set_markerfacecolor('none')
    l.set_linewidth(4)
    m, _, _ = ax.stem(V_guess, linefmt='r', markerfmt='rx', basefmt=' ', label="Guess",
                      use_line_collection=True)
    m.set_markersize(15)

    ax.set_xlabel("State")
    ax.set_ylabel("Value")
    ax.set_title("Guess V(s)\n" + reward_guesser_title_hint(r1, r2))
    ax.legend()
  except NotImplementedError:
    print("Please finish Exercise 1 first!")
```

## Think! 2: Examining the TD Error

Run the cell below to plot the TD errors from our multi-reward environment. A new feature appears in this plot? What is it? Why does it happen?

```python
plot_tde_trace(TDE_multi)
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D4_ReinforcementLearning/solutions/W3D4_Tutorial1_Solution_9fd42e7d.py)

---
# Section 3: TD-learning with probabilistic rewards
*Estimated timing to here from start of tutorial: 40 min*

## Think! 3: Probabilistic rewards

In this environment, we'll return to delivering a single reward of ten units. However, it will be delivered intermittently: on 20 percent of trials, the CS will be shown but the agent will not receive the usual reward; the remaining 80% will proceed as usual.

 Run the cell below to simulate. How does this compare with the previous experiment?

Earlier in the notebook, we saw that changing $\alpha$ had little effect on learning in a deterministic environment. What happens if you set it to an large value, like 1, in this noisier scenario? Does it seem like it will _ever_ converge?

```python
# @markdown Execute this cell to see visualization
np.random.set_state(rng_state) # Resynchronize everyone's notebooks
n_trials = 20000
try:
  env = ProbabilisticCC(n_steps=40, reward_magnitude=10, reward_time=10,
                        p_reward=0.8)
  V_stochastic, TDE_stochastic = td_learner(env, n_trials*2, alpha=1)
  learning_summary_plot(V_stochastic, TDE_stochastic)
except NotImplementedError:
  print("Please finish Exercise 1 first")
```

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D4_ReinforcementLearning/solutions/W3D4_Tutorial1_Solution_162d105b.py)

---
# Summary

*Estimated timing of tutorial: 50 min*

In this notebook, we have developed a simple TD Learner and examined how its state representations and reward prediction errors evolve during training. By manipualting its environment and parameters ($\alpha$, $\gamma$), you developed an intuition for how it behaves.

This simple model closely resembles the behavior of subjects undergoing classical conditioning tasks and the dopamine neurons that may underlie that behavior. You may have implemented TD-reset or used the model to recreate a common experimental error. The update rule used here has been extensively studied for [more than 70 years](https://www.pnas.org/content/108/Supplement_3/15647) as a possible explanation for artificial and biological learning.

However, you may have noticed that something is missing from this notebook. We carefully calculated the value of each state, but did not use it to actually do anything. Using values to plan _**Actions**_ is coming up next!

---
# Bonus

## Bonus Think! 1: Removing the CS

In Coding Exercise 1, you (should have) included a term that depends on the conditioned stimulus. Remove it and see what happens. Do you understand why?
This phenomena often fools people attempting to train animals--beware!

[*Click for solution*](https://github.com/NeuromatchAcademy/course-content/tree/master//tutorials/W3D4_ReinforcementLearning/solutions/W3D4_Tutorial1_Solution_5e1c3185.py)


# Notebook

**Roadmap**

* Data Generating Process
* Objects of Interest
* Processing Specification
* Setting up Simulation
* Conducting Estimation
* Inspection of Results

Over the next couple of lectures, we will then constantly refine the basic code and explore elements of software engineering such as Object-Oriented Programming, Unit Testing, Debugging, and Profiling.

Before we get started, let us import some basic libraries.

```python
%%capture
# Notebook metods
from IPython.core.display import HTML, Image

# Unix Pattern Extensions
import glob

# Operating System Interfaces
import os

# Lexical Analysis
import shlex

# Copy operations
import copy

# Encoders and Decoders
import codecs

# Statistical Modeling and Econometrics
import statsmodels.api as sm

# Plotting
import matplotlib.pyplot as plt
%pylab inline --no-import-all

# Scientific Computing
import numpy as np
from scipy.stats    import norm
from scipy.optimize import minimize
```

## Processing of Model Specification

We manage the model specification in an external text file, which is called *init.ini*. This file will turn out to be useful to provide the parameters for a simulation of a synthetic sample or the initialization of starting values for an estimation.

```python
Image(filename='material/images/init.png', width=1000)
```

Now we will develop a set of function that processes the initialization file.

```python
def _process_bene(list_, dict_, keyword):
    """ This function processes the BENE part of the initialization file.
    """
    # Distribute information
    name, val_treated, val_untreated = list_[0], list_[1], list_[2]

    # Initialize dictionary
    if 'TREATED' not in dict_.keys():
        for subgroup in ['TREATED', 'UNTREATED']:
            dict_[subgroup] = {}
            dict_[subgroup]['coeff'] = []
            dict_[subgroup]['int'] = None
            dict_[subgroup]['sd'] = None

    # Type conversion
    val_treated = float(val_treated)
    val_untreated = float(val_untreated)

    # Collect information
    if name in ['coeff']:
        dict_['TREATED'][name] += [val_treated]
        dict_['UNTREATED'][name] += [val_untreated]
    else:
        dict_['TREATED'][name] = val_treated
        dict_['UNTREATED'][name] = val_untreated

    # Finishing
    return dict_

def _process_not_bene(list_, dict_, keyword):
    """ This function processes all of the initialization file, but the
        BENE section.
    """
    # Distribute information
    name, val = list_[0], list_[1]

    # Prepare container.
    if name not in dict_[keyword].keys():
        if name in ['coeff']:
            dict_[keyword][name] = []

    # Type conversion
    if name in ['agents', 'maxiter']:
        val = int(val)
    elif name in ['source', 'algorithm', 'start', 'version']:
        val = str(val)
    else:
        val = float(val)

    # Collect information
    if name in ['coeff']:
        dict_[keyword][name] += [val]
    else:
        dict_[keyword][name] = val

    # Finishing.
    return dict_

def _check_integrity_process(dict_):
    """ Check integrity of initFile dict.
    """
    # Antibugging
    assert (isinstance(dict_, dict))

    # Check number of agents
    assert (dict_['BASICS']['agents'] > 0)
    assert (isinstance(dict_['BASICS']['agents'], int))

    # Check optimizer
    assert (dict_['ESTIMATION']['algorithm'] in ['bfgs', 'nm'])

    # Check starting values
    assert (dict_['ESTIMATION']['start'] in ['random', 'init'])

    # Maximum iterations
    assert (dict_['ESTIMATION']['maxiter'] >= 0)

    # Finishing
    return True

def _add_auxiliary(dict_):
    """ Add some auxiliary objects.
    """
    # Antibugging
    assert (isinstance(dict_, dict))

    # Initialize container
    dict_['AUX'] = {}

    # Full set of coefficients.
    for key_ in ['TREATED', 'UNTREATED', 'COST']:
        dict_[key_]['all'] = [dict_[key_]['int']]
        dict_[key_]['all'] += dict_[key_]['coeff']
        dict_[key_]['all'] = np.array(dict_[key_]['all'])

    # Number of covariates
    num_covars_out = len(dict_['TREATED']['all'])
    num_covars_cost = len(dict_['COST']['all'])

    dict_['AUX']['num_covars_out'] = num_covars_out
    dict_['AUX']['num_covars_cost'] = num_covars_cost

    # Number of parameters
    dict_['AUX']['num_paras'] = 2 * num_covars_out + num_covars_cost + 2 + 2

    # Starting values
    dict_['AUX']['init_values'] = []

    for key_ in ['TREATED', 'UNTREATED', 'COST']:
        dict_['AUX']['init_values'] += dict_[key_]['all'].tolist()

    dict_['AUX']['init_values'] += [dict_['TREATED']['sd']]
    dict_['AUX']['init_values'] += [dict_['UNTREATED']['sd']]
    dict_['AUX']['init_values'] += [dict_['DIST']['rho1']]
    dict_['AUX']['init_values'] += [dict_['DIST']['rho0']]

    # Finishing
    return dict_

def _process_cases(list_):
    """ Process cases and determine whether keyword or empty
        line.
    """
    # Antibugging
    assert (isinstance(list_, list))

    # Get information
    is_empty = (len(list_) == 0)

    if not is_empty:
        is_keyword = list_[0].isupper()
    else:
        is_keyword = False

    # Antibugging
    assert (is_keyword in [True, False])
    assert (is_empty in [True, False])

    # Finishing
    return is_empty, is_keyword
```

```python
def process(file_):
    """ Process initialization file.
    """
    # Initialization
    dict_ = {}

    for line in open(file_).readlines():

        # Remove UTF-3 marker
        if line.startswith(codecs.BOM_UTF8):
            line = line[3:]

        # Split line
        list_ = shlex.split(line)

        # Determine special cases
        is_empty, is_keyword = _process_cases(list_)

        # Applicability
        if is_empty:
            continue

        if is_keyword:
            keyword = list_[0]
            dict_[keyword] = {}
            continue

        if keyword not in ['BENE']:
            dict_ = _process_not_bene(list_, dict_, keyword)

        else:
            dict_ = _process_bene(list_, dict_, keyword)

    # Remove BENE
    del dict_['BENE']

    # Add auxiliary objects
    dict_ = _add_auxiliary(dict_)

    # Check quality.
    _check_integrity_process(dict_)

    # Finishing.
    return dict_
```

Let us check if it is all working.

```python
init_dict = process('material/msc/init.ini')
```

# Setting up the Simulation

**Distributinal Assumptions**

*Observables*

\begin{align}
X & \sim \mathbb{N}(0, 1) \\
Z & \sim \mathbb{N}(0, 1) \\
\end{align}

*Unobservables*

\begin{eqnarray*}
\begin{pmatrix}U_{1}\\
U_{0}\\
V
\end{pmatrix} & \sim & \mathbb{N}\left[\left(\begin{array}{c}
0\\
0\\
0
\end{array}\right),\left(\begin{array}{ccc}
\sigma_{U_1}^2 & 0 & \sigma_{U_1,V}\\
0 & \sigma_{U_0}^2 & \sigma_{U_0,V}\\
\sigma_{U_1,V} & \sigma_{U_0,V} & \sigma_{V}^2
\end{array}\right)\right]\\
\end{eqnarray*}

```python
init_dict = process('material/msc/init.ini')
```

```python
def _check_integrity_simulate(Y1, Y0, Y, D):
    """ Check quality of simulated sample.
    """
    assert (np.all(np.isfinite(Y1)))
    assert (np.all(np.isfinite(Y0)))

    assert (np.all(np.isfinite(Y)))
    assert (np.all(np.isfinite(D)))

    assert (Y1.dtype == 'float')
    assert (Y0.dtype == 'float')

    assert (Y.dtype == 'float')
    assert (D.dtype == 'float')

    assert (D.all() in [1.0, 0.0])

def _write_out(Y, D, X, Z, source, unobserved=False, Y1=None, Y0=None):
    """ Write out simulated data to file.
    """

    if not unobserved:

        np.savetxt(source, np.column_stack((Y, D, X, Z)), fmt='%8.3f')

    else:

        assert (isinstance(Y1, np.ndarray))
        assert (isinstance(Y0, np.ndarray))

        np.savetxt(source, np.column_stack((Y, D, X, Z, Y1, Y0)),
                   fmt='%8.3f')
```

```python
def simulate(init_dict, unobserved=False):
    """ Simulate a model based on the initialization file.
    """

    # Antibugging
    assert (isinstance(init_dict, dict))
    assert (unobserved in [True, False])

    # Ensure recomputability
    np.random.seed(123)

    # Distribute information
    num_agents = init_dict['BASICS']['agents']
    source = init_dict['BASICS']['source']

    Y1_coeffs = init_dict['TREATED']['all']
    Y0_coeffs = init_dict['UNTREATED']['all']

    C_coeffs = np.array(init_dict['COST']['all'])

    U1_sd = init_dict['TREATED']['sd']
    U0_sd = init_dict['UNTREATED']['sd']

    V_sd = init_dict['COST']['sd']

    U1V_rho = init_dict['DIST']['rho1']
    U0V_rho = init_dict['DIST']['rho0']

    # Auxiliary objects
    U1V_cov = U1V_rho * U1_sd * V_sd
    U0V_cov = U0V_rho * U0_sd * V_sd

    num_covars_out = Y1_coeffs.shape[0]
    num_covars_cost = C_coeffs.shape[0]

    # Simulate observables
    means = np.tile(0.0, num_covars_out)
    covs = np.identity(num_covars_out)

    X = np.random.multivariate_normal(means, covs, num_agents)

    means = np.tile(0.0, num_covars_cost)
    covs = np.identity(num_covars_cost)

    Z = np.random.multivariate_normal(means, covs, num_agents)

    # Add intercepts. The first column of the X and Z matrix always contains
    # the intercept term. This is exploited throughout the code.
    Z[:,0], X[:, 0] = 1.0, 1.0

    # Construct index of observable characteristics
    Y1_level = np.dot(Y1_coeffs, X.T)
    Y0_level = np.dot(Y0_coeffs, X.T)
    C_level = np.dot(C_coeffs, Z.T)

    # Simulate unobservables
    means = np.tile(0.0, 3)
    vars_ = [U1_sd**2, U0_sd**2, V_sd**2]
    covs = np.diag(vars_)

    covs[0, 2] = U1V_cov
    covs[2, 0] = covs[0, 2]

    covs[1, 2] = U0V_cov
    covs[2, 1] = covs[1, 2]

    U = np.random.multivariate_normal(means, covs, num_agents)

    # Simulate endogenous variables
    Y1 = np.tile(np.nan, num_agents)
    Y0 = np.tile(np.nan, num_agents)
    Y = np.tile(np.nan, num_agents)

    D = np.tile(np.nan, num_agents)

    for i in range(num_agents):

        # Select individual unobservables and observables
        u1, u0, v = U[i, 0], U[i, 1], U[i, 2]

        y1_idx, y0_idx, c_idx = Y1_level[i], Y0_level[i], C_level[i]

        # Decision Rule
        expected_benefits = y1_idx - y0_idx
        cost = c_idx + v

        d = np.float((expected_benefits - cost > 0))

        # Potential outcomes
        y1, y0 = y1_idx + u1, y0_idx + u0

        # Observed outcomes
        y = d * y1 + (1.0 - d) * y0

        # Collect data matrices
        Y[i], Y0[i], Y1[i], D[i] = y, y1, y0, d

    # Check integrity of simulated data
    _check_integrity_simulate(Y1, Y0, Y, D)

    # Save to disk
    _write_out(Y, D, X, Z, source, unobserved, Y1, Y0)

    # Return selected features of data
    return Y1, Y0, D
```

Let us check if it is all working.

```python
init_dict = process('material/msc/init.ini')

Y1, Y0, D = simulate(init_dict)

```

Given our parametrization, let us revisit our objects of interest. We start with the individual-specific benefits. Please note the use of the [StatsModels](http://statsmodels.sourceforge.net/devel/index.html) library,
that provides classes and functions for the estimation of many different statistical models, as well as for conducting statistical tests, and statistical data exploration. Think of it as a replacement for using [R](http://www.r-project.org/). In general, [rpy2](http://rpy.sourceforge.net/rpy2.html) provides a low-level interface to R from Python.

```python
# Auxiliary variables
B = Y1 - Y0

# Create histogram and density estimate of benefits.
kde = sm.nonparametric.KDEUnivariate(B)
kde.fit()

# Initialize canvas
ax = plt.figure(figsize=(12,8)).add_subplot(111)

# Plot histogram and density
ax.hist(B, bins=50, normed=True, color='blue')
ax.plot(kde.support, kde.density, lw=2, color='black')

# Set axis labels
ax.set_xlabel('Individual-Specific Benefits', fontsize=18)
ax.set_ylabel('Density and Histogram', fontsize=18)

# Change background color
ax.set_axis_bgcolor('white')

# Remove first element on y-axis
ax.yaxis.get_major_ticks()[0].set_visible(False)

# Calcuate the average treatment effects
ate, tt, tut = np.mean(B), np.mean(B[D==1]), np.mean(B[D==0])

# Pretty formatting of strings and output
fmt = '     {0:<5}{1:10.2f}\n'

print '\nAverage Treatment Effects\n'
print fmt.format('ATE ', ate)
print fmt.format('TT', tt)
print fmt.format('TUT ', tut)

# Let us add them to our plot.
plt.axvline(x=ate, ymin=0, ymax=5, linewidth=2, color='g')
plt.axvline(x=tt, ymin=0, ymax=5, linewidth=2, color='b')
plt.axvline(x=tut, ymin=0, ymax=5, linewidth=2, color='y')

# Add title
plt.suptitle('Distribution of Individual-Specific Benefits', fontsize=20)
```

## Estimation

Now, we will perform Maximum Likelihood Estimation using alternative optimization algorithms. Here is the likelihood function:

\begin{align}
\mathcal{L}(\theta; X, Z) =\sum^N_{i=1} D\mathcal{L_{i,1}} + (1 - D)\mathcal{L_{i,0}},
\end{align}

where

\begin{align}
\mathcal{L_1} = & \log\left(\frac{1}{\sigma_{U_1}}\phi\left(\frac{Y_i - X_i\beta_1}{\sigma_{U_1}}\right)\Phi\left(\frac{Z_i\gamma - \sigma_V/\sigma_{U_1}(Y_i - X_i\beta_1)}{\sqrt{(1 - \rho_{U_1,V})*\sigma^2_{V}}}\right)\right) \\
\mathcal{L_0} =  &\log\left(\frac{1}{\sigma_{U_0}}\phi\left(\frac{Y_i - X_i\beta_0}{\sigma_{U_0}}\right)\Phi\left(\frac{Z_i\gamma - \sigma_V/\sigma_{U_0}(Y_i - X_i\beta_0)}{\sqrt{(1 - \rho_{U_0,V})*\sigma^2_{V}}}\right)\right) \\
\end{align}

```python
def _distribute_parameters(x, init_dict, num_covars_out):
    """ Distribute the parameters.
    """
    # Antibugging
    assert (isinstance(x, np.ndarray))
    assert (isinstance(num_covars_out, int))
    assert (num_covars_out > 0)

    # Initialize containers
    rslt = dict()

    rslt['TREATED'] = dict()
    rslt['UNTREATED'] = dict()
    rslt['COST'] = dict()
    rslt['DIST'] = dict()

    # Distribute parameters
    rslt['TREATED']['all'] = x[:num_covars_out]
    rslt['UNTREATED']['all'] = x[num_covars_out:(2 * num_covars_out)]

    rslt['COST']['all'] = x[(2 * num_covars_out):(-4)]
    rslt['COST']['sd'] = init_dict['COST']['sd']

    rslt['TREATED']['sd'] = np.exp(x[(-4)])
    rslt['UNTREATED']['sd'] = np.exp(x[(-3)])

    rslt['DIST']['rho1'] = -1.0 + 2.0 / (1.0 + float(np.exp(-x[-2])))
    rslt['DIST']['rho0'] = -1.0 + 2.0 / (1.0 + float(np.exp(-x[-1])))

    # Update auxiliary versions
    rslt['AUX'] = dict()

    rslt['AUX']['x_internal'] = x.copy()
    rslt['AUX']['x_internal'][-4] = np.exp(x[(-4)])
    rslt['AUX']['x_internal'][-3] = np.exp(x[(-3)])
    rslt['AUX']['x_internal'][-2] = -1.0 + 2.0 / (1.0 + float(np.exp(-x[-2])))
    rslt['AUX']['x_internal'][-1] = -1.0 + 2.0 / (1.0 + float(np.exp(-x[-1])))

    rslt['AUX']['init_values'] = init_dict['AUX']['init_values']

    # Finishing.
    return rslt

def _max_interface(x, Y, D, X, Z, version, init_dict):
    """ Interface to the SciPy maximization routines.
    """
    # Auxiliary objects
    num_covars_out = X.shape[1]

    # Collect maximization arguments
    rslt = _distribute_parameters(x, init_dict, num_covars_out)

    # Calculate likelihood
    likl = _negative_log_likelihood(rslt, Y, D, X, Z)

    # Finishing.
    return likl

def _negative_log_likelihood(args, Y, D, X, Z):
    """ Negative Log-likelihood function of the generalized Roy model.
    """
    # Distribute arguments
    Y1_coeffs, Y0_coeffs, C_coeffs, choice_coeffs, U1_sd, U0_sd, U1V_rho, \
    U0V_rho, V_sd = _distribute_arguments(args)

    # Auxiliary objects
    num_agents = Y.shape[0]

    # Initialize containers
    likl = np.tile(np.nan, num_agents)
    choice_idx = np.tile(np.nan, num_agents)

    # Likelihood construction.
    for i in range(num_agents):

        g = np.concatenate((X[i, :], Z[i,:]))
        choice_idx[i] = np.dot(choice_coeffs, g)

        # Select outcome information
        if D[i] == 1.00:
            coeffs, rho, sd = Y1_coeffs, U1V_rho, U1_sd
        else:
            coeffs, rho, sd = Y0_coeffs, U0V_rho, U0_sd

        arg_one = (Y[i] - np.dot(coeffs, X[i, :])) / sd
        arg_two = (choice_idx[i] - rho * V_sd * arg_one) / \
                  np.sqrt((1.0 - rho ** 2) * V_sd**2)

        pdf_evals, cdf_evals = norm.pdf(arg_one), norm.cdf(arg_two)

        if D[i] == 1.0:
            contrib = (1.0 / float(sd)) * pdf_evals * cdf_evals
        else:
            contrib = (1.0 / float(sd)) * pdf_evals * (1.0 - cdf_evals)

        likl[i] = contrib

    # Transformations.
    likl = -np.mean(np.log(np.clip(likl, 1e-20, np.inf)))

    # Quality checks.
    assert (isinstance(likl, float))
    assert (np.isfinite(likl))

    # Finishing.
    return likl

def _get_start(which, init_dict, Y, D, X, Z):
    """ Get different kind of starting values.
    """
    # Antibugging.
    assert (which in ['random', 'init', 'auto'])

    # Distribute auxiliary objects
    num_paras = init_dict['AUX']['num_paras']
    num_covars_cost = init_dict['AUX']['num_covars_cost']

    # Construct auxiliary objects
    G = np.concatenate((X, Z[:, 1:]), axis=1)

    # Select relevant values.
    if which == 'random':
        x0 = np.random.uniform(size=num_paras)

        # Variances
        x0[(-4)] = max(x0[(-4)], 0.01)
        x0[(-3)] = max(x0[(-3)], 0.01)

        # Correlations
        x0[(-2)] -= 0.5
        x0[(-1)] -= 0.5

    elif which == 'init':
        x0 = np.array(init_dict['AUX']['init_values'][:])

    elif which == 'auto':

        # Subsetting
        Y1, X1 = Y[D == 1], X[(D == 1), :]
        olsRslt = sm.OLS(Y1, X1).fit()

        # Extract results
        coeffs_treated = olsRslt.params
        sd_treated = np.array(np.sqrt(olsRslt.scale))

        # Subsetting
        Y0, X0 = Y[D == 0], X[(D == 0), :]
        olsRslt = sm.OLS(Y0, X0).fit()

        # Extract results
        coeffs_untreated = olsRslt.params
        sd_untreated = np.array(np.sqrt(olsRslt.scale))

        # Estimate choice model
        probitRslt = sm.Probit(D, G).fit()
        sd = init_dict['COST']['sd']
        coeffs = probitRslt.params*sd

        # Special treatment of cost intercept
        cost_int = coeffs_treated[0] - coeffs_untreated[0] - coeffs[0]

        # Collect results
        x0 = np.concatenate((coeffs_treated, coeffs_untreated))
        x0 = np.concatenate((x0, [cost_int], -coeffs[-(num_covars_cost - 1):]))
        x0 = np.concatenate((x0, [sd_treated, sd_untreated]))
        x0 = np.concatenate((x0, [0.00, 0.00]))

    else:
        raise AssertionError

    # Document starting values
    init_dict['AUX']['start_values'] = x0.copy()

    # Transform to real line
    x0 = _transform_start(x0)

    # Type conversion
    x0 = np.array(x0)

    # Quality assurance.
    assert (np.all(np.isfinite(x0)))

    # Finishing.
    return x0

def _transform_start(x):
    """ Transform starting values to cover the whole real line.
    """

    # Coefficients
    x[:(-4)] = x[:(-4)]

    # Variances
    x[(-4)] = np.log(x[(-4)])
    x[(-3)] = np.log(x[(-3)])

    # Correlations
    transform = (x[(-2)] + 1) / 2
    x[(-2)] = np.log(transform / (1.0 - transform))

    transform = (x[(-1)] + 1) / 2
    x[(-1)] = np.log(transform / (1.0 - transform))

    # Finishing
    return x

def _load_data(init_dict):
    """ Load dataset.
    """
    # Auxiliary objects
    num_covars_out = init_dict['AUX']['num_covars_out']
    num_covars_cost = init_dict['AUX']['num_covars_cost']
    num_agents = init_dict['BASICS']['agents']

    # Read dataset
    data = np.genfromtxt(init_dict['BASICS']['source'])

    # Reshaping, this ensure that the program also runs with just one agent
    # as otherwise only an vector is created. This creates problems for the
    # subsetting of the overall data into the components.
    data = np.array(data, ndmin=2)

    # Distribute data
    Y, D = data[:, 0], data[:, 1]

    X, Z = data[:, 2:(num_covars_out + 2)], data[:, -num_covars_cost:]

    # Finishing
    return Y, D, X, Z

def _distribute_arguments(args):
    """ Distribute arguments for evaluation of criterion function and some
        auxiliary parameters.
    """
    Y1_coeffs = np.array(args['TREATED']['all'])
    Y0_coeffs = np.array(args['UNTREATED']['all'])

    C_coeffs = np.array(args['COST']['all'])

    U1_sd = args['TREATED']['sd']
    U0_sd = args['UNTREATED']['sd']

    U1V_rho = args['DIST']['rho1']
    U0V_rho = args['DIST']['rho0']
    V_sd = args['COST']['sd']

    choice_coeffs = np.concatenate((Y1_coeffs - Y0_coeffs, - C_coeffs))

    # Finishing
    return Y1_coeffs, Y0_coeffs, C_coeffs, choice_coeffs, U1_sd, U0_sd, \
           U1V_rho, U0V_rho, V_sd
```

Let us evaluate the criterion function and see if that part is working.

```python
# Load model information
init_dict  = process('material/msc/init.ini')

Y, D, X, Z = _load_data(init_dict)

likl = _negative_log_likelihood(init_dict, Y, D, X, Z)

print 'Evaluation of criterion function', likl
```

    Evaluation of criterion function -2.44668227139

Finally, we create a function for the estimation. It builds on all our previously defined elements.

```python
def estimate(init_dict):
    """ Estimate our version of the generalized Roy model.
    """
    # Antibugging
    assert (isinstance(init_dict, dict))

    # Load dataset
    Y, D, X, Z = _load_data(init_dict)

    # Create auxiliary objects
    start = init_dict['ESTIMATION']['start']
    maxiter = init_dict['ESTIMATION']['maxiter']

    optimizer = init_dict['ESTIMATION']['algorithm']
    num_covars_out = init_dict['AUX']['num_covars_out']

    # Initialize different starting values
    x0 = _get_start(start, init_dict, Y, D, X, Z)

    # Select optimizer
    if optimizer == 'nm':

        optimizer = 'Nelder-Mead'

    elif optimizer == 'bfgs':

        optimizer = 'BFGS'

    # Provide additional arguments to the optimizer
    opts = dict()

    opts['maxiter'] = maxiter

    # Run optimization or just evaluate function at starting values
    if maxiter == 0:

        # Collect maximization arguments.
        rslt = _distribute_parameters(np.array(x0), init_dict, num_covars_out)

        # Calculate likelihood according to user's request
        likl = _negative_log_likelihood(rslt, Y, D, X, Z)

        # Compile results
        x_rslt, fun, success = x0, likl, False

    else:

        # Check out the SciPy documentation for details about the interface
        # to the `minimize' function that provides a convenient interface to
        # a variety of alternative maximization algorithms. You will also
        # find information about the return information.
        opt_rslt = minimize(_max_interface, x0,
                            args=(Y, D, X, Z, init_dict),
                            method=optimizer, options=opts)

        # Compile results
        x_rslt, fun = opt_rslt['x'], opt_rslt['fun']
        success = opt_rslt['success']

    # Tranformation to internal parameters
    rslt = _distribute_parameters(x_rslt, init_dict, num_covars_out)

    rslt['fval'], rslt['success'] = fun, success

    # Finishing
    return rslt
```

Now, let us actually run some estimations. Our current implementation allow to easily switch between alternative optimzation algorithms and starting values.

```python
# Process model specification
init_dict = process('material/msc/init.ini')

# Simulate a synthetic sample
simulate(init_dict)

# Estimate the generalized Roy model
for algorithm in ['bfgs', 'nm']:

    init_dict['ESTIMATION']['algorithm'] = algorithm

    for start in ['random', 'init']:

        init_dict['ESTIMATION']['start'] = start

        # Monitoring
        print '\n\n Current Request \n'

        print ' Algorithm: ', algorithm

        print ' Start:     ', start

        # Run estimation
        rslt = estimate(init_dict)

        # Inspect subset of results
        print ' Variances: ', rslt['TREATED']['sd'], rslt['UNTREATED']['sd']
```

     Current Request

     Algorithm:  bfgs
     Start:      random
     Variances:  0.0854101015863 0.580645990592

     Current Request

     Algorithm:  bfgs
     Start:      init
     Variances:  0.02 0.02

     Current Request

     Algorithm:  nm
     Start:      random
     Variances:  0.981374779551 0.694354634872

     Current Request

     Algorithm:  nm
     Start:      init
     Variances:  0.02 0.02

## Inspection of Results

After conduction an estimation, the focus shifts back to the economic interpretation of our results. As we often have to study the results from hundreds of different estimation runs, it is convenient to have a function set up that produces the main objects of interest.

```python
def inspect(rslt, init_dict):
    """ This function simulates a sample from the estimates of the model
        and reports the average effects of treatment in a file.
    """
    # Antibugging
    assert (isinstance(rslt, dict))
    assert (isinstance(init_dict, dict))

    # Update results
    modified_init = copy.deepcopy(init_dict)

    for key_ in rslt.keys():

        if key_ in ['fval', 'success']:
            continue

        for subkey in rslt[key_].keys():

            modified_init[key_][subkey] = rslt[key_][subkey]

    # Modified dataset
    modified_init['BASICS']['file'] = 'simulated.grm.txt'

    # Simulate from estimation results
    Y1, Y0, D = simulate(modified_init, True)

    # Calculate the average treatment effectsa
    B = Y1 - Y0

    effects = []
    effects += [np.mean(B)]
    effects += [np.mean(B[D == 1])]
    effects += [np.mean(B[D == 0])]

    # Print selected results to file
    with open('results.grm.txt', 'w') as file_:

        file_.write('\n softEcon: Generalized Roy Model')
        file_.write('\n -------------------------------\n')

        # Average effects of treatment
        fmt = '     {0:<5}{1:10.2f}\n\n'

        file_.write('\n Average Treatment Effects\n\n')

        for i, label in enumerate(['ATE', 'TT', 'TUT']):

            str_ = fmt.format(label, effects[i])

            file_.write(str_)

        file_.write('\n Parameters\n\n')
        file_.write('     Start    Finish\n\n')

        num_paras = init_dict['AUX']['num_paras']

        # Structural parameters
        x0, x = init_dict['AUX']['start_values'], rslt['AUX']['x_internal']

        fmt = '{0:10.2f}{1:10.2f}\n'

        for i in range(num_paras):

            str_ = fmt.format(x0[i], x[i])

            file_.write(str_)

```

Let us try it out.

```python
# Process initialization file
init_dict = process('material/msc/init.ini')

# Simulate synthetic sample
simulate(init_dict)

# Estimate model
rslt = estimate(init_dict)

# Write results
inspect(rslt, init_dict)

# Inspect the results
%cat results.grm.txt
```

     softEcon: Generalized Roy Model
     -------------------------------

     Average Treatment Effects

         ATE       -0.40

         TT        -0.46

         TUT       -0.29

     Parameters

         Start    Finish

          0.50      0.50
          0.20      0.20
          0.52      0.52
          0.52      0.52
          0.10      0.10
          0.12      0.12
          0.30      0.30
          0.52      0.52
          0.20      0.20
          0.12      0.12
          0.25      0.25
          0.35      0.35
          0.02      0.02
          0.02      0.02
          0.10      0.10
          0.12      0.12

## Cleanup

```python
# Create list of all files generated the notebook
files = glob.glob('*.grm.*')

# Remove files
for file_ in files:

    os.remove(file_)
```

**Formatting**

```python
import urllib; from IPython.core.display import HTML
HTML(urllib.urlopen('http://bit.ly/1OKmNHN').read())
```

<link href='http://fonts.googleapis.com/css?family=Alegreya+Sans:100,300,400,500,700,800,900,100italic,300italic,400italic,500italic,700italic,800italic,900italic' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Arvo:400,700,400italic' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=PT+Mono' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Shadows+Into+Light' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Philosopher:400,700,400italic,700italic' rel='stylesheet' type='text/css'>

<style>

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

#notebook_panel { /* main background */
    background: #888;
    color: #f6f6f6;
}

div.cell { /* set cell width to about 80 chars */
    width: 1000px;
}

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

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

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

/* Put a solid color box around each cell and its output, visually linking them together */
div.cell.code_cell {
    background-color: rgba(171,165,131,0.3);
    border-radius: 10px; /* rounded borders */
    padding: 1em;
    margin-top: 1em;
}

div.text_cell_render{
    font-family: 'Arvo' sans-serif;
    line-height: 130%;
    font-size: 150%;
    width:900px;
    margin-left:auto;
    margin-right:auto;
}

/* Formatting for header cells */
.text_cell_render h1 {
    font-family: 'Philosopher', sans-serif;
    font-weight: 400;
    font-size: 32pt;
    line-height: 100%;
    color: rgb(12,85,97);
    margin-bottom: 0.1em;
    margin-top: 0.1em;
    display: block;
}

.text_cell_render h2 {
    font-family: 'Philosopher', serif;
    font-weight: 700;
    font-size: 24pt;
    line-height: 100%;
    color: rgb(171,165,131);
    margin-bottom: 0.1em;
    margin-top: 0.1em;
    display: block;
}

.text_cell_render h3 {
    font-family: 'Philosopher', serif;
    margin-top:12px;
    margin-bottom: 3px;
    font-style: italic;
    color: rgb(95,92,72);
}

.text_cell_render h4 {
    font-family: 'Philosopher', serif;
}

.text_cell_render h5 {
    font-family: 'Alegreya Sans', sans-serif;
    font-weight: 300;
    font-size: 16pt;
    color: grey;
    font-style: italic;
    margin-bottom: .1em;
    margin-top: 0.1em;
    display: block;
}

.text_cell_render h6 {
    font-family: 'PT Mono', sans-serif;
    font-weight: 300;
    font-size: 10pt;
    color: grey;
    margin-bottom: 1px;
    margin-top: 1px;
}

.CodeMirror{
        font-family: "PT Mono";
        font-size: 120%;
}

</style>


```python
import numpy as np
import matplotlib.pyplot as plt
import xlrd
import sympy as sp
from matplotlib.ticker import MultipleLocator, FormatStrFormatter ,FuncFormatter
import pandas as pd
from scipy import optimize
from scipy.interpolate import spline
import os

strt=["1tally  14        nps =   ","1tally  24        nps =   ",
      "1tally  34        nps =   ","1tally  44        nps =   ",
      "1tally  54        nps =   ","1tally  64        nps =   ",]#14 24
un=["g/kg","g/kg","mm","mm","g/cm$^3$","g/cm$^3$","mm","$^{\circ}$$C$"]
st=["o","h","H","s","d","^","<",">","1","2","p"]
c=["r","b","g","c","m","y","k","blueviolet","lightgreen","darkmagenta","lightgreen"]
la=["地层水矿化度($g/kg$)","井眼矿化度($g/kg$)","厚度($mm$)","厚度($mm$)","泥浆密度($g/cm^3$)","泥浆密度($g/cm^3$)","井径($mm$)",
   "井温($^{\circ}$$C$)"]
col=["孔隙度（pu）","近中子计数器-1","近中子计数器-2","远中子计数器-1","远中子计数器-2","远中子计数器-1","远中子计数器-2",
     "第1组远中子计数和","第2组远中子计数和","第1组近远中子比","第2组近远中子比","总近中子计数和","误差","总远中子计数和","误差",
     "总近远中子比"]
def linersmooth7(a): #七点平滑
    b=[]
    N=len(a)-1
    for i in range(3,N-2): #不要前三个和最后三个点
        b.append((130*a[i]+140*(a[i-1]+a[i+1])+65*(a[i-2]+a[i+2])+40*(a[i-3]+a[i+3]))/610) #公式形式
    return b
def dele(a): #y轴去掉了前三个和后三个点，x轴数据也要去掉前三个和后三个
    b=a.copy()
    for i in range(0,3):
        del b[0]
        b.pop()
    return b
```

```python
 #用flag给指示是否要将标准情况放入其中
#存在用到标准曲线的情况
```

```python
data=pd.read_excel("E:\王馨奕\随钻-19-4\MCNP程序及结果\\02-中子图版\\00-标准曲线\\标准曲线.xlsx",index_col=0)
y0=data["总近远中子比"].values#取得标准刻度井计算结果
x0=data["孔隙度（pu）"].values
```

```python

```

```python
def rf(flag,x0,excelname):
    path=os.getcwd()
    tn=[[] for i in range(7)]#每个中子计数器的计数值,第一个是记孔隙度值
    jx=[]#文件名-环境变量值
    aimf=[]#目标文件
    ind=[]
    err=[[] for i in range(7)]
    files= os.listdir(path)#读文件名
    for file in files: #遍历文件夹
         if os.path.isdir(file):
                if (file!=".ipynb_checkpoints"):#判断是否是文件夹，是文件夹才打开
                    jx.append(file)
    jx=np.array(jx)#将字符转换为数字，去掉引号
    jx=jx.astype(np.int)#环境因素
    jx.sort()
    for i in range(len(jx)):
        path_dir1="%s"%(jx[i])
        path_list=os.listdir(path_dir1)
        for filename in path_list:
            i=0
            if(i<11):
                if os.path.splitext(filename)[1] == '.o':
                    name=os.path.splitext(filename)[0]
                    aimf.append(name)
                    i+=1
        aimf=np.array(aimf)#将字符转换为数字，去掉引号
        aimf=aimf.astype(np.int)#环境因素
        aimf.sort()
#         print(aimf)
        for i in range(11):#读到的文件文件名做index将中子计数值存储下来
            ind.append(path_dir1)
            tn[0].append(x0[i])#
            err[0].append(x0[i])
            path_dir=path+"\\"+path_dir1+"\\"+"%s.o"%(aimf[i])#读到了含有输出文件的文件路径
            i=0
            j=0
            f=open(path_dir,"r")#读取该文件
            for n in range(6):
                for line in f:
                    if(line.find(strt[n])!=-1):
                        i=1
                    if(i==1):
                        j+=1
                    if(j==11):
                        m=line.strip()
                        tn[n+1].append(m.split()[0])
                        err[n+1].append(m.split()[1])
                        i=0
                        j=0
                        break
        aimf=[]
    nu=len(tn[0])
#     for i in range(len(tn)):
#         print(len(tn[i]))
    tn=np.array(tn)#将字符转换为数字，去掉引号
    tn=tn.astype(np.float)
    err=np.array(err)#将字符转换为数字，去掉引号
    err=err.astype(np.float)
    t0=tn.T#nu*6维矩阵转换
    t0=t0.tolist()#numpy to list
    print("已读取到路径%s:文件%s"%(path,jx))
    print("读到%s份输出文件"%(nu))
    for i in range(nu):
#         t0[i].append(tn[0][i])#参考组1
#         t0[i].append(tn[1][i])#参考组2
        t0[i].append(tn[3][i]+tn[5][i])#1组远
        t0[i].append(tn[4][i]+tn[6][i])#2组远
        t0[i].append((tn[1][i])/(tn[3][i]+tn[5][i]))#1组比
        t0[i].append((tn[2][i])/(tn[4][i]+tn[6][i]))#2组比
        t0[i].append(tn[2][i]+tn[1][i])#总近和
        t0[i].append(err[2][i]+err[1][i])#总近和误差
        t0[i].append(tn[6][i]+tn[3][i]+tn[4][i]+tn[5][i])#总远和
        t0[i].append(err[6][i]+err[3][i]+err[4][i]+err[5][i])#总远和
        t0[i].append((tn[2][i]+tn[1][i])/(tn[6][i]+tn[3][i]+tn[4][i]+tn[5][i]))#总比
        R=(tn[2][i]+tn[1][i])/(tn[6][i]+tn[3][i]+tn[4][i]+tn[5][i])
    print("已计算三组近远中子比")
    print("数据存入%s×%s矩阵"%(len(t0),len(t0[0])))
    data=pd.DataFrame(data=t0,index=ind,columns=col)
    data.to_excel("%s.xlsx"%(excelname),encoding="utf_8_sig")
    print("存文件:%s.xlsx"%(excelname))
    if (len(tn[0])==len(t0)):
        print("正确录入数据")
    else:
        print("未正确录入数据")
    return (t0,jx,len(jx),err)
```

```python
t0,jx,f,err=rf(0,x0,"间隙300")
```

    已读取到路径E:\王馨奕\随钻-19-4\MCNP程序及结果\02-中子图版\4间隙校正\300:文件[ 0  5 10 15 20 25 30 40 50]
    读到99份输出文件
    已计算三组近远中子比
    数据存入99×16矩阵
    存文件:间隙300.xlsx
    正确录入数据

```python
#画图
```

```python
def f2(x,A,B,C):
    return A*x*x+B*x+C
def pic1(name_pic,flag):
    x=np.arange(0,55,5)
    D=jx
    d=len(D)
    un=["g/kg","g/kg","mm","mm","g/cm$^3$","g/cm$^3$","mm","$^{\circ}$$C$"]
    st=["o","h","H","s","d","^","<",">","1","2","p"]
    c=["r","b","g","c","m","y","k","blueviolet","lightgreen","darkmagenta","lightgreen"]
    plt.figure(figsize=(6.67, 5), dpi=100) #图片大小：6.75cm×5cm
    plt.rcParams['font.sans-serif']=['SimHei'] #中文默认黑体
    plt.rcParams['axes.unicode_minus']=False #显示负号
    for i in range(6):
        A2,B2,C2=optimize.curve_fit(f2,x,y1[i])[0]
        x2=np.arange(0,50,0.5)
        y2=A2*x2*x2+B2*x2+C2
        plt.plot(x2,y2,color="%s"%(c[i]))
        plt.plot(x,y1[i],"%s"%(st[i]),color="%s"%(c[i]),label="%.1f%s:%.4f×x$^2$%+.2f×x%+.2f"%(D[i],un[flag],A2,B2,C2))
    plt.legend()
    plt.xticks(fontproperties='Times New Roman',fontsize=13)
    plt.yticks(fontproperties='Times New Roman',fontsize=13)
    plt.xlabel("孔隙度($pu$)",fontproperties='SimHei',fontsize=13)
    plt.ylabel("近远热中子比值",fontproperties='SimHei',fontsize=13)
    plt.grid(linewidth='1.5' ,linestyle='-',color="black",alpha=0.5)
    plt.savefig('%s趋势1.png'%(name_pic), format='png', dpi=100,pad_inches=0.0,bbox_inches = 'tight')
    plt.savefig('%s趋势1.svg'%(name_pic), format='svg', dpi=100,pad_inches=0.0,bbox_inches = 'tight')
    plt.show()
    plt.figure(figsize=(6.67, 5), dpi=100) #图片大小：6.75cm×5cm
    plt.rcParams['font.sans-serif']=['SimHei'] #中文默认黑体
    plt.rcParams['axes.unicode_minus']=False #显示负号
    x=np.arange(0,55,5)
    for i in range(0,11):
        A,B,C=optimize.curve_fit(f2,D,[x[i] for x in y1])[0]# l=[0---250] 取不同空隙度对应的比值 横行
        x2=np.arange(min(D),max(D)+0.01,0.01)
        y2=f2(x2,A,B,C)
        plt.plot(x2,y2,color="%s"%(c[i])) #先画拟合的直线，再画实际结果
        plt.plot(D,[x[i] for x in y1],"%s"%(st[i]),color="%s"%(c[i]),label="%spu"%(x[i]),markersize=7)
    plt.xticks(fontproperties='Times New Roman',fontsize=17)
    plt.yticks(fontproperties='Times New Roman',fontsize=17)
    plt.xlabel("%s"%(la[flag]),fontproperties='SimHei',fontsize=17)
    plt.ylabel("近远热中子比值",fontproperties='SimHei',fontsize=17)
    plt.legend(bbox_to_anchor=(1,1.001),borderaxespad=0.1,edgecolor="black",handletextpad=-0.3)
    plt.grid(linewidth='1.5' ,color="black",alpha=0.5)
    plt.savefig('%s趋势2.png'%(name_pic), format='png', dpi=100,pad_inches=0.0,bbox_inches = 'tight')
    plt.savefig('%s趋势2.svg'%(name_pic), format='svg', dpi=100,pad_inches=0.0,bbox_inches = 'tight')
    plt.show()
```

```python
len(t0[0])
```

    16

```python
y1=[[] for i in range(f)]
for m in range(f):
    for i in range(len(x0)):
        y1[m].append(t0[i+m*len(x0)][15])#读到的y[0]为标准曲线的值
#     print(len(y1[m]))
pic1("间隙300",3)
```

```python
def dom(name_pic,flag):
    #计算视孔隙度
    ynew=[]
    x=np.arange(0,55,5)
    for i in range(0,len(jx)):
        co = np.polyfit(y1[i],x,2) # 用多项式拟合
        A2=co[0]
        B2=co[1]
        C2=co[2]
        x2=np.arange(1,20,0.2)
        re = np.polyval(co, x2)
        X=sp.Symbol('X')
        temp=[]
        for j in range(len(x)): #解方程 从近远中子比解孔隙度
            y3=sp.solve(A2*X*X+B2*X+C2-x[j],X) #求出来比值 solve函数
            temp.append(y3[1])
        ynew.append(temp)
    for i in range(len(jx)):
        ynew[i][0]=sp.re(ynew[i][0])+sp.im(ynew[i][0])
    ynew=np.array(ynew)
    pnew=[]
    A,B,C=optimize.curve_fit(f2,y0,x)[0]
    for i in range(len(jx)):
        pnew.append(A*ynew[i]*ynew[i]+B*ynew[i]+C)#带入标准曲线得到视孔隙度7个井径
    plt.figure(figsize=(6.67, 2.5), dpi=100) #图片大小：6.75cm×5cm
    plt.rcParams['font.sans-serif']=['SimHei'] #中文默认黑体
    plt.rcParams['axes.unicode_minus']= False #显示负号
    st=["o","h","H","s","d","^","<",">","1","2","p"]
    c=["r","b","g","c","m","y","k","blueviolet","lightgreen","darkmagenta","lightgreen"]
    pnew=np.array(pnew)
    l=jx
    for i in range(0,11):
        k=[x[i] for x in pnew]
        k=np.array(k,dtype='float')
        plt.plot(k,l,"%s"%(st[i]),color="%s"%(c[i]),markersize=5)#,"%s"%(st[i]),color="%s"%(c[i]),markersize=5
        z1 = np.polyfit(l,k,2) # 用3次多项式拟合
        z = np.polyval(z1, l)
        plt.plot(z,l,"-",color="%s"%(c[i]))
    plt.xticks(fontproperties='Times New Roman',fontsize=13)
    plt.yticks(fontproperties='Times New Roman',fontsize=13)
    plt.ylabel("%s"%(la[flag]),fontproperties='SimHei',fontsize=13)
    plt.xlabel("孔隙度($pu$)[视石灰岩]",fontproperties='SimHei',fontsize=13)
    #plt.legend(prop=font1,bbox_to_anchor=(1,1.001),borderaxespad=0.1,edgecolor="black",handletextpad=-0.3)
    plt.xlim(0,50)#x轴范围
    plt.ylim(min(l),max(l))
    plt.grid(linewidth='1.5' ,color="black",alpha=0.5)
    # plt.tight_layout()
    plt.savefig('%s校正.png'%(name_pic), format='png', dpi=100,bbox_inches = 'tight')
    plt.savefig('%s校正.svg'%(name_pic), format='svg', dpi=100,bbox_inches = 'tight')
    plt.show()
    plt.figure(figsize=(6.67, 2.5), dpi=100) #图片大小：6.75cm×5cm
    plt.rcParams['font.sans-serif']=['SimHei'] #中文默认黑体
    plt.rcParams['axes.unicode_minus']= False #显示负号
    st=["o","h","H","s","d","^","<",">","1","2","p"]
    c=["r","b","g","c","m","y","k","blueviolet","lightgreen","darkmagenta","lightgreen"]
    pnew=np.array(pnew)
    l=jx
    for i in range(0,11):
        k=[x[i] for x in pnew]
        k=np.array(k,dtype='float')
        z1 = np.polyfit(l,k,2) # 用3次多项式拟合
        z = np.polyval(z1, l)
        plt.plot(z,l,"-",color="#054E9F")
    plt.xticks(fontproperties='Times New Roman',fontsize=13)
    plt.yticks(fontproperties='Times New Roman',fontsize=13)
    plt.ylabel("%s"%(la[flag]),fontproperties='SimHei',fontsize=13)
    plt.xlabel("孔隙度($pu$)[视石灰岩]",fontproperties='SimHei',fontsize=13)
    #plt.legend(prop=font1,bbox_to_anchor=(1,1.001),borderaxespad=0.1,edgecolor="black",handletextpad=-0.3)
    plt.xlim(0,50)#x轴范围
    plt.ylim(min(l),max(l))
    plt.grid(linewidth='1.5' ,color="black",alpha=0.5)
    # plt.tight_layout()
    plt.savefig('%s校正2.png'%(name_pic), format='png', dpi=100,bbox_inches = 'tight')
    plt.savefig('%s校正2.svg'%(name_pic), format='svg', dpi=100,bbox_inches = 'tight')
    plt.show()
```

```python
dom("间隙250",1)
```

```python

```