```python
from resources.workspace import *
from IPython.display import display
from scipy.integrate import odeint
import copy

%matplotlib inline
```

# Lyapunov vectors and ensemble based covariances

We return now to the discussion of perturbations in nonlinear models.  Recall the equations for the evolution of perturbations.  We may define *linear* dynamics, generated by the Jacobian equation,

$$\begin{align}
\frac{{\rm d} \mathbf{x}}{{\rm d} t} = \nabla f_{\rvert_{\mathbf{x}c}},
\end{align}$$

computed along some control trajectory  $\mathbf{x}^c(t)$.  This system of equations is known as the tangent-linear model.  The tangent space can be understood as the space of perturbations at a point (or along a trajectory).  It can be used to define a vector field on a "space" describing the time evolution of trajectories.

By derivative work: McSush (talk)Tangentialvektor.png: TNThe original uploader was TN at German Wikipedia (Tangentialvektor.png) [Public domain], via Wikimedia Commons

As we have seen with the [breeding of errors](#breeding), we can compute the log-average growth rate of the dominant growing mode of the tangent-linear model to approximate the leading Lyapunov exponent.  In a linear system, with fixed matrix  $\mathbf{M}$ this corresponds exactly to taking the log-average growth rate in the power method.  Recall the generalization of the power method to the QR algorithm.  We might hope to find **all of the log-average growth rates** in the tangent-linear model by computing the log averages of the diagonal elements in the QR factors  $\mathbf{U}_k$.

In this case, the Gram-Schmidt factors form a basis for the tangent linear model where:

 the leading vector aligns with the dominant growth mode;

 the subsequent vectors separate out the lower order growth rates.

Suppose the tangent-linear model is computed discretely in time, where  $\textbf{M}_k$ takes the perturbations at time $t_{k-1}$ to time $t_k$.  Suppose that in matrix form, $\mathbf{M}_k$, produces the following QR factorization,

$$\begin{align}
\mathbf{M}_k \mathbf{E}_{k-1} &= \mathbf{E}_k \mathbf{U}_k
\end{align}$$

**Exc 4.36**: Suggest a constructive definition for:

 the Lyapunov exponents of the nonlinear model

 the "Lyapunov vectors"

Note: We will not stress yet what kind of Lyapunov vector we are constructing.  It turns out that Lyapunov exponents are generally globally defined, but there are several types of Lyapunov vectors that are locally defined.

```python
# Answer

# show_answer('lyapunov_vs_es')
```

The type of vector this construction leads to is the **"backward" Lyapunov vectors**.  They are denoted "backward" because they contain information from the in the past, leading to the current time, i.e., the "errors of the day" described by Toth and Kalnay.

For completeness, we remark that the forcing singular vectors mentioned above can be shown to converge to the "forward" Lyapunov vectors.  These are denoted "forward" because they contain information about the future evolution of pertubations, pulled back to the current time.  A third type of Lyapunov vector is also often studied, which are called "covariant".  Similar to how eigenvectors always are mapped to a scaled copy of themselves, covariant vectors share an analogous property with respect to time-varying dynamics.

A full discussion of the significance of forward, backward, and covariant vectors goes well beyond the scope of this tutorial.  For a comprehensive discussion of Lyapunov theory, aimed at practicioners, it is recommended to read the work of [Legras & Vautard](http://www.lmd.ens.fr/legras/publis/liapunov.pdf) and [Kutpsov & Parlitz](https://arxiv.org/abs/1105.5228).

## Ensemble based covariances

When we restrict the forecasting problem to the situation where we assume:

 we can perfectly model and compute the purely deterministic dynamics; and

 prediction error originates soley from the uncertainty in initial conditions,

it is realistic to think of the data assimilation problem as tracking the evolution of perturbations along a control trajectory.  This control trajectory is the "true" state of the dynamical system we are studying, which we receive observations of, and try to predict the behavior of at future times.

In this case, we suppose we have a prior distribution for the state of the dynamical system.  Suppose we sample for an ensemble of "nearby" initial conditions.  If the ensemble remains sufficiently close to the control trajectory, and we can use the [approximation for the evolution of perturbations](#perturbation_equation) to accurately model the forecast errors, then the ensemble spread will be characterized by the backward Lyapunov vectors and their growth and decay rates.

Consider the conceptual image below.  Suppose the initial prior covariance is given by the unit disk, centered on the "true" state of a dynamical system.  If the evolution of uncertainty can be well approximated by the tangent linear model along the "truth", we will expect the covariance to stretch and deform into an ellipse according to the directions of **growth** and **decay**.

By Mrocklin (original creation) [CC BY-SA 3.0 or GFDL], via Wikimedia Commons

In the following, we will implement a simple, "square-root" ensemble Kalman filter in the Lorenz-63 model.  The ensemble Kalman filter will be given **noisy observations** of a control trajectory.   Along the control trajectory, we will compute the **QR factorization** of the tangent-linear model.

We will define the **projection coefficient** of the covariance matrix  $\mathbf{P}_k$ into the $j$-th QR factor to be the quantity,

$$\begin{align}
\left(\mathbf{Q}_k^j\right)^{\rm T} \mathbf{P}_k \mathbf{Q}^j
\end{align}$$

We will compute the projection coefficients of the ensemble based covariance into each of the QR factors, as defined above.

Exc 4.38: Can you conjecture how the projection coefficients will vary in the index $j$?

**Hint**: consider **Exc 4.34**.

Exc 4.40: Test your conjecture from Exc 4.40.  Use the code below to investigate the relationship between the ensemble based covariance and its projeciton into each of the QR factors.  We plot the average projection coefficient for the EnKF covariance into each of the QR factors over the number of analyses.  Similarly, we plot the EnKF mean square error, to vefify that the ensemble mean lies within the variance of the error in these directions.

**Note**: we may consider the QR factorizations to give approximately the "true" Lyapunov vectors when the log-average growth rate approaches the Lyapunov exponents for the system.  The Lyapunov exponents are approximately given by,

$$\begin{align}
\lambda_1 &\approx 0.905 \\
\lambda_2 & = 0 \\
\lambda_3 & \approx -14.571.
\end{align}$$

```python
SIGMA = 10.0
BETA  = 8/3
RHO   = 28.0

sigma = SIGMA
beta = BETA
rho = RHO

def dxdt(xyz, t0, sigma=SIGMA, beta=BETA, rho=RHO):
    """Compute the time-derivative of the Lorenz-63 system."""
    x, y, z = xyz
    return array([
        sigma * (y - x),
        x * (rho - z) - y,
        x * y - beta * z
    ])

def l63_jac(x):
    jac = np.array([
        [-sigma,  sigma,     0 ],
        [rho - x[2], -1,  -x[0]],
        [x[1],      x[0], -beta]
        ]
    )

    return jac

def l63_step_TLM(x, Y, h):

    h_mid = h/2

    # calculate the evolution of x to the midpoint
    x_mid = l63_rk4_step(x, h_mid)

    # calculate x to the next time step
    x_next = l63_rk4_step(x_mid, h_mid)

    k_y_1 = l63_jac(x).dot(Y)
    k_y_2 = l63_jac(x_mid).dot(Y + k_y_1 * (h / 2.0))
    k_y_3 = l63_jac(x_mid).dot(Y + k_y_2 * (h / 2.0))
    k_y_4 = l63_jac(x_next).dot(Y + k_y_3 * h)

    Y_next = Y + (h / 6.0) * (k_y_1 + 2 * k_y_2 + 2 * k_y_3 + k_y_4)

    return [x_next, Y_next]

def l63_rk4_step(xyz, h):
    """ calculate the evolution of Lorenz-63 one step forward via RK-4"""

    k_xyz_1 = dxdt(xyz, h, sigma=SIGMA, beta=BETA, rho=RHO)
    k_xyz_2 = dxdt(xyz + k_xyz_1 * (h / 2.0), h, sigma=SIGMA, beta=BETA, rho=RHO)
    k_xyz_3 = dxdt(xyz + k_xyz_2 * (h / 2.0), h, sigma=SIGMA, beta=BETA, rho=RHO)
    k_xyz_4 = dxdt(xyz + k_xyz_3 * h, h, sigma=SIGMA, beta=BETA, rho=RHO)

    xyz_step = xyz + (h / 6.0) * (k_xyz_1 + 2 * k_xyz_2 + 2 * k_xyz_3 + k_xyz_4)

    return xyz_step

def animate_enkf_covariance(nanl=0):

    # Initial conditions: perturbations around some control state
    tanl=0.005
    h = 0.001
    tl_steps = int(tanl / h)
    N = 4
    obs_un = 0.25
    obs_dim = 3
    R = np.eye(obs_dim) * obs_un
    H = np.eye(3, M=obs_dim).T
    proj_traj = np.zeros([nanl, 3])
    err = np.zeros([nanl])

    seed(1)
    x_0 = array([-6.1, 1.2, 32.5])               # define the control

    # define the perturbations, randomly generated but of fixed norm epsilon
    A_f = randn([3, N])
    a_m = np.mean(A_f, axis=1)
    A_f = A_f.T - a_m
    del a_m
    A_f = (x_0 + A_f).T

    lam = np.zeros(3)
    Q = np.eye(3)

    # for each analysis cycle
    for kk in range(nanl):
        for j in range(tl_steps):
            x_0, Q = l63_step_TLM(x_0, Q, h)
            for l in range(2):
                for j in range(N):
                    A_f[:, j] = l63_rk4_step(A_f[:, j], h / 2)

        # perform QR step and find the local Lyapunov exponents
        Q, U = np.linalg.qr(Q)
        lam += np.log(np.abs(np.diag(U))) / tanl

        # define an observation
        y_0 = H @ x_0 + np.random.multivariate_normal(np.zeros([obs_dim]), R)

        # forecast mean, and rmse
        x_f = np.mean(A_f, axis=1)

        A_f = (A_f.T - x_f).T

        # forecast covaraince
        P_f = (N-1) ** (-1) * A_f @ A_f.T

        # form the Kalman gain
        K = P_f @ H.T @ np.linalg.inv(H @ P_f @ H.T + R)

        # analysis mean
        x_a = x_f + K @ (y_0 - H @ x_f)
        err[kk] = np.mean((x_a - x_0)**2)

        # analyze the ensemble
        T = np.eye(N) - (N-1)**(-1) * (H @ A_f).T @ np.linalg.inv(H @ P_f @ H.T + R) @ (H @ A_f)
        U, S, V_h = np.linalg.svd(T)
        T = U @ np.diag(np.sqrt(S)) @ U.T
        A_f = A_f @ T

        P_a = (N-1)**(-1) * A_f @ A_f.T
        # find the forecast projection coefficients
        for i in range(3):
            proj_traj[kk, i] = Q[:, i].T @ P_a @ Q[:, i]

        A_f = (x_a + A_f.T).T

    # PLOTTING
    avg_proj = np.zeros([nanl, 3])

    # we plot the average projection coefficient into each blv
    for i in range(1,nanl):
        avg_proj[i, :] = np.mean(proj_traj[:i, :], axis=0)

    lam = lam / nanl
    fig = plt.figure(figsize=(16,8))
    ax = plt.subplot(111)

    for i in range(3):
        ax.plot(range(nanl), avg_proj[:, i], label='BLV projection ' + str(i + 1) + ', log-avg growth rate ' +
                str(np.round(lam[i],decimals=3)).zfill(3))

    ax.axhline(y=np.mean(err), color='k', label='Mean square error')
    plt.legend(fontsize=20)
    ax.set_xbound([1, nanl])
    ax.set_yscale('log')
    ax.tick_params(labelsize=20)
    plt.show()

w = interactive(animate_enkf_covariance,nanl=(10,20010,2000))
w
```

Exc 4.42: Answer the following questions.

How do the projection coefficients relate to the log-average growth rates in each direction?

What is significant about the effective rank of the covariance?

Can you conjecture what this means about the necessary number of ensemble members to prevent filter divergence?

 Hint: the ensemble should capture the effective spread of the uncertainty.

tutorials/DA and the Dynamics of Ensemble Based Forecasting/T4 - Lyapunov vectors and ensemble based covariances.ipynb


# **The Reformer - Pushing the limits of language modeling**

***How the Reformer uses less than 8GB of RAM to train on sequences of half a million tokens***

The Reformer model as introduced by [Kitaev, Kaiser et al. (2020)](https://arxiv.org/pdf/2001.04451.pdf) is one of the most memory-efficient transformer models for long sequence modeling as of today.

Recently, long sequence modeling has experienced a surge of interest as can be seen by the many submissions from this year alone - [Beltagy et al. (2020)](https://arxiv.org/abs/2004.05150), [Roy et al. (2020)](https://arxiv.org/abs/2003.05997), [Tay et al.](https://arxiv.org/abs/2002.11296), [Wang et al.](https://arxiv.org/abs/2006.04768) to name  a few.
The motivation behind long sequence modeling is that many tasks in NLP, *e.g.* summarization, question answering, require the model to process longer input sequences than models, such as BERT, are able to handle. In tasks that require the model to process a large input sequence, long sequence models do not have to cut the input sequence to avoid memory overflow and thus have been shown to outperform standard "BERT"-like models *cf.* [Beltagy et al. (2020)](https://arxiv.org/abs/2004.05150).

The Reformer pushes the limit of longe sequence modeling by its ability to process up to half a million tokens at once as shown in this [demo](https://github.com/patrickvonplaten/notebooks/blob/master/PyTorch_Reformer.ipynb). As a comparison, a conventional `bert-base-uncased` model limits the input length to only 512 tokens. In Reformer, each part of the standard transformer architecture is re-engineered to optimize for minimal memory requirement without a significant drop in performance.

The memory improvements can be attributed to **4** features which the Reformer authors introduced to the transformer world:

1.   **Reformer Self-Attention Layer** - *How to efficiently implement self-attention without being restricted to a local context?* => see [this colab](https://colab.research.google.com/drive/15oP52_7W5dRcAnbgX3tYADsu4R3cjMIf?usp=sharing)
2.  **Chunked Feed Forward Layers** - *How to get a better time-memory trade-off for large feed forward layers?* => see [this colab](https://colab.research.google.com/drive/1xKK32Yhda-iYgtoA3eCrnCVuy_lraQR9?usp=sharing)
3.   **Reversible Residual Layers**  - *How to drastically reduce memory consumption in training by a smart residual architecture?* => see [this colab](https://colab.research.google.com/drive/1BLffcRt9LXmM7nKU2UXhtm0PqAG0UE7J#scrollTo=mk1ETLlfEMGA)
4.   **Axial Positional Encodings** - *How to make positional encodings usable for extremely large input sequences?*

The goal of this blog post is to give the reader an **in-depth** understanding of each of the four Reformer features mentioned above. While the explanations are focussed on the Reformer, the reader should get a better intuition under which circumstances each of the four features can be effective for other transformer models as well.
The four sections are only loosely connected, so they can very well be read individually.

Reformer is part of the 🤗Transformers library. For all users of the Reformer, it is advised to go through this very detailed blog post to better understand how the model works and how to correctly set its configuration. All equations are accompanied by their equivalent name for the Reformer config, *e.g.* `config.<param_name>`, so that the reader can quickly relate to the official docs and configuration file.

**Note**: *Axial Positional Encodings* are not explained in the official Reformer paper, but are extensively used in the official codebase. This blog post gives the first in-depth explanation of Axial Positional Encodings.

## **4. Axial Positional Encodings**

Reformer makes it possible to process huge input sequences. However, for such long input sequences standard positional encoding weight matrices alone would use more than 1GB to store its weights.
To prevent such large positional encoding matrices, the official Reformer code makes use of *Axial Position Encodings*.

**Important:** *Axial Position Encodings were not explained in the official paper, but can be well understood from looking into the code and talking to the authors*

### **Axial Positional Encodings in Reformer**

Transformers need positional encodings to account for the order of words in the input because self-attention layers have *no notion of order*.
Positional encodings are usually defined by a simple look-up matrix $\mathbf{E} = \left[\mathbf{e}_1, \ldots, \mathbf{e}_{n_\text{max}}\right]$ The positional encoding vector $\mathbf{e}_i$ is then simply added to the *ith* input vector $\mathbf{x}_i + \mathbf{e}_i$ so that the model can distinguish if an input vector (*a.k.a* token) is at position $i$ or $j$.
For every input position, the model needs to be able to look up the corresponding positional encoding vector so that the dimension of $\mathbf{E}$ is defined by the maximum length of input vectors the model can process `config.max_position_embeddings`, *i.e.* $n_\text{max}$, and the `config.hidden_size`, *i.e.* $d_h$ of the input vectors.

Assuming $d_h=4$ and $n_\text{max}=49$, such a positional encoding matrix can be visualized as follows:

Here, we showcase only the positional encodings $\mathbf{e}_1$, $\mathbf{e}_2$, and $\mathbf{e}_{49}$ each of dimension, *a.k.a* height 4.

Let's imagine, we want to train a Reformer model on sequences of a length of up to 0.5M tokens and an input vector `config.hidden_size` of 1024 (see notebook [here](https://github.com/patrickvonplaten/notebooks/blob/master/PyTorch_Reformer.ipynb)). The corresponding positional embeddings have a size of $0.5M \times 1024 \sim 512M$ parameters, which corresponds to a size of 2GB.

Such positional encodings would use an unnecessarily large amount of memory both when loading the model in memory and when saving the model on a hard drive.

The Reformer authors managed to drastically shrink the positional encodings in size by cutting the `config.hidden_size` dimension in two and smartly factorizing the $n_\text{max}$ dimension. In Transformer, the user can decide into which shape $n_\text{max}$ can be factorized into by setting `config.axial_pos_shape` to an appropriate list of two values $n_\text{max}^1$ and $n_\text{max}^2$ so that $n_\text{max}^1 \times n_\text{max}^2 = n_\text{max}$. By setting `config.axial_pos_embds_dim` to an appropriate list of two values $d_h^1$ and $d_h^2$ so that $d_h^1 + d_h^2 = d_h$, the user can decide how the hidden size dimension should be cut.
Now, let's visualize and explain more intuitively.

One can think of factorizing $n_\text{max}$ as folding the dimension into a third axis, which is shown in the following for the factorization `config.axial_pos_shape = [7, 7]`:

Each of the three standing rectangular prisms corresponds to one of the encoding vectors $\mathbf{e}_1, \mathbf{e}_2, \mathbf{e}_{49}$, but we can see that the 49 encoding vectors are divided into 7 rows of 7 vectors each.
Now the idea is to use only one row of 7 encoding vectors and expand those vectors to the other 6 rows, essentially reusing their values.
Because it is discouraged to have the same values for different encoding vectors, each vector of dimension (*a.k.a* height) `config.hidden_size=4` is cut into the lower encoding vector $\mathbf{e}_\text{down}$ of size $1$ and $\mathbf{e}_\text{up}$ of size $3$, so that the lower part can be expanded along the row dimension and the upper part can be expanded along the column dimension.
Let's visualize for more clarity.

We can see that we have cut the embedding vectors into $\mathbf{e}_\text{down}$ (*in blue*) and $\mathbf{e}_\text{up}$ (*in yellow*).
Now for the "sub"-vectors $\mathbf{E}_\text{down} = \left[\mathbf{e}_{\text{down},1}, \ldots, \mathbf{e}_{\text{down},49}\right]$ only the first row, *a.k.a.* the width in the graphic, of $7$ is kept and expanded along the column dimension, *a.k.a.* the depth of the graphic. Inversely, for the "sub"-vectors $\mathbf{E}_\text{up} = \left[\mathbf{e}_{\text{up},1}, \ldots, \mathbf{e}_{\text{up},49}\right]$ only the first column of $7$ is kept and expanded along the row dimension.
The resulting embedding vectors $\mathbf{e'}_i$ then correspond to

  \begin{align}
    \mathbf{e'}_i &= \begin{bmatrix}
           \mathbf{e}_{\text{down, } i \% n_\text{max}^1} \\
           \mathbf{e}_{\text{up, } \left \lfloor{\frac{i}{n_\text{max}^2}}\right \rceil}
         \end{bmatrix}
  \end{align}
whereas $n_\text{max}^1 = 7$ and $n_\text{max}^2 = 7$ in our example.
These new encodings $\mathbf{E'} = \left[\mathbf{e'}_1, \ldots, \mathbf{e'}_{n_\text{max}}\right]$ are called **Axial Position Encodings**.

In the following, these axial position encodings are illustrated in more detail for our example.

Now it should be more understandable how the final positional encoding vectors $\mathbf{E'}$ are calculated only from $\mathbf{E}_{\text{down}}$ of dimension $d_h^1 \times n_\text{max}^1$ and $\mathbf{E}_{\text{up}}$ of dimension $d_h^2 \times n_\text{max}^2$.

The crucial aspect to see here is that Axial Positional Encodings make sure that none of the vectors $\left[\mathbf{e'}_1, \ldots, \mathbf{e'}_{n_\text{max}}\right]$ are equal to each other by design and that the overall size of the encoding matrix is reduced from $n_\text{max} \times d_h$ to $n_\text{max}^1 \times d_h^1 + n_\text{max}^2 \times d_h^2$.
By allowing each axial positional encoding vector to be different by design the model is given much more flexibility to learn efficient positional representations if axial positional encodings are learned by the model.

To demonstrate the drastic reduction in size,
let's assume we would have set `config.axial_pos_shape = [1024, 512]` and `config.axial_pos_embds_dim = [512, 512]` for a Reformer model that can process inputs up to a length of 0.5M tokens. The resulting axial positional encoding matrix would have had a size of only $1024 \times 512 + 512 \times 512 \sim 800K$ parameters which corresponds to roughly 3MB. This is a drastic reduction from the 2GB a standard positional encoding matrix would require in this case.

For a more condensed and math-heavy explanation please refer to the 🤗Transformers docs [here](https://huggingface.co/transformers/model_doc/reformer.html#axial-positional-encodings).

### **Benchmark**

Lastly, let's also compare the peak memory consumption of conventional positional embeddings to *axial positional embeddings*.

```
#@title Installs and Imports
# pip installs
!pip -qq install git+https://github.com/huggingface/transformers.git
!pip install -qq py3nvml

from transformers import ReformerConfig, PyTorchBenchmark, PyTorchBenchmarkArguments, ReformerModel
```

         |████████████████████████████████| 3.0MB 7.2MB/s
         |████████████████████████████████| 1.1MB 32.1MB/s
         |████████████████████████████████| 890kB 42.5MB/s
      Building wheel for transformers (setup.py) ... done
      Building wheel for sacremoses (setup.py) ... done
         |████████████████████████████████| 61kB 3.7MB/s

Positional embeddings depend only on two configuration parameters: The maximum allowed length of input sequences `config.max_position_embeddings` and `config.hidden_size`. Let's use a model that pushes the maximum allowed length of input sequences to half a million tokens, called `google/reformer-crime-and-punishment`, to see the effect of using axial positional embeddings.

To begin with, we will compare the shape of axial position encodings with standard positional encodings and the number of parameters in the model.

```
config_no_pos_axial_embeds = ReformerConfig.from_pretrained("google/reformer-crime-and-punishment", axial_pos_embds=False)  # disable axial positional embeddings
config_pos_axial_embeds = ReformerConfig.from_pretrained("google/reformer-crime-and-punishment", axial_pos_embds=True, axial_pos_embds_dim=(64, 192), axial_pos_shape=(512, 1024))  # enable axial positional embeddings

print("Default Positional Encodings")
print(20 * '-')
model = ReformerModel(config_no_pos_axial_embeds)
print(f"Positional embeddings shape: {model.embeddings.position_embeddings}")
print(f"Num parameters of model: {model.num_parameters()}")
print(20 * '-' + '\n\n')

print("Axial Positional Encodings")
print(20 * '-')
model = ReformerModel(config_pos_axial_embeds)
print(f"Positional embeddings shape: {model.embeddings.position_embeddings}")
print(f"Num parameters of model: {model.num_parameters()}")
print(20 * '-' + '\n\n')
```

    HBox(children=(FloatProgress(value=0.0, description='Downloading', max=1151.0, style=ProgressStyle(description…

    Default Positional Encodings
    --------------------
    Positional embeddings shape: PositionEmbeddings(
      (embedding): Embedding(524288, 256)
    )
    Num parameters of model: 136572416
    --------------------

    Axial Positional Encodings
    --------------------
    Positional embeddings shape: AxialPositionEmbeddings(
      (weights): ParameterList(
          (0): Parameter containing: [torch.FloatTensor of size 512x1x64]
          (1): Parameter containing: [torch.FloatTensor of size 1x1024x192]
      )
    )
    Num parameters of model: 2584064
    --------------------

Having read the theory, the shape of the axial positional encoding weights should not be a surprise to the reader.

Regarding the results, it can be seen that for models being capable of processing such long input sequences, it is not practical to use default positional encodings.
In the case of `google/reformer-crime-and-punishment`, standard positional encodings alone contain more than 100M parameters.
Axial positional encodings reduce this number to just over 200K.

Lastly, let's also compare the required memory at inference time.

```
benchmark_args = PyTorchBenchmarkArguments(sequence_lengths=[512], batch_sizes=[8], models=["Reformer-No-Axial-Pos-Embeddings", "Reformer-Axial-Pos-Embeddings"], no_speed=True, no_env_print=True)
benchmark = PyTorchBenchmark(configs=[config_no_pos_axial_embeds, config_pos_axial_embeds], args=benchmark_args)
result = benchmark.run()
```

    1 / 2
    2 / 2

    ====================      INFERENCE - MEMORY - RESULT       ====================
    --------------------------------------------------------------------------------
              Model Name             Batch Size     Seq Length    Memory in MB
    --------------------------------------------------------------------------------
    Reformer-No-Axial-Pos-Embeddin       8              512             959
    Reformer-Axial-Pos-Embeddings        8              512             447
    --------------------------------------------------------------------------------

It can be seen that using axial positional embeddings reduces the memory requirement to approximately half in the case of `google/reformer-crime-and-punishment`.


# CHSH不等式の破れを確認する

この最初の実習では、量子コンピュータにおいて**量子力学的状態、特に「エンタングルメント」**が実現しているか検証してみましょう。実習を通じて量子力学の概念と量子コンピューティングの基礎を紹介していきます。

## 目次
1. {ref}`sec:probing_quantumness`
1. {ref}`sec:chsh_inequality`
1. {ref}`sec:qiskit`
  1. {ref}`subsec:qubits`
  1. {ref}`subsec:circuit`
  1. {ref}`subsec:common_gates`
  1. {ref}`subsec:circuit_visualization`
1. {ref}`sec:writing_chsh`
1. {ref}`sec:using_ibmq`
1. {ref}`sec:analysis`

$\newcommand{ket}[1]{|#1\rangle}$

(sec:probing_quantumness)=
## 本当に量子コンピュータなのか？

このワークブックの主旨が量子コンピュータ（QC）を使おう、ということですが、QCなんて数年前までSFの世界の存在でした。それが今やクラウドの計算リソースとして使えるというわけですが、ではそもそも私たちがこれから使おうとしている機械は本当にQCなのでしょうか。どうしたらそれが調べられるでしょうか。

QCの基本的な仕組みは、**「何らの物理的な系（超電導共振器や冷却原子など）をうまく操作して、求める計算の結果がその系の量子状態に表現されるようにする」**ということです。つまり、量子状態が長く保たれてかつ思うように操作できる対象と、「計算」という実体のなさそうなものを具体的な「量子操作」に対応させるアルゴリズムの両方があって初めてQCが成り立ちます。アルゴリズムの部分はこのワークブックを通じてちょっとずつ紹介していくので、今回は「量子状態が保たれ、それを操作できる」ということを確認してみましょう。

(sec:chsh_inequality)=
## CHSH不等式とは

CHSH不等式やその根本であるBell不等式の意義については、{doc}`quantum_basics`を参照してください。かいつまんで言うと、CHSH不等式とは「二体系の特定の観測量について、エンタングルメントなど量子力学固有の現象がなければ保たれる不等式」です。やや回りくどいロジックですが、つまりQC（だと考えられる機械）で測ったこの観測量の値がCHSH不等式を破っていれば、その機械は実際に量子現象を利用していると言えることになります。

通常このような実験を行うには高度なセットアップ（レーザーと非線形結晶、冷却原子など）が必要ですが、クラウドQCではブラウザ一つしか要りません。このワークブックではJupyter NotebookでPythonのプログラムを書き、[IBM Quantum Experience](https://quantum-computing.ibm.com/)を通じて[IBM Q System One](https://www.ibm.com/quantum-computing/systems/)を利用します。

(sec:qiskit)=
## Qiskitの基本構造

IBM Q System OneのQCで量子計算を実行するには、IBMの提供する[Qiskit](https://qiskit.org/)というPythonライブラリを利用します。Qiskitの基本的な使い方は

1. 空っぽの量子回路を定義する。
1. 回路に量子計算の操作（ゲート）を足していく。
1. 回路を実行して計算結果を得る。ここでは二通りのオプションがあり、
  - 回路をQCの実機に送り、実行させる。
  - 回路をシミュレートする。
1. 計算結果を解析する。

です。以下でこの流れを一通り、重要な概念の説明を混ぜながら実行してみましょう。

(subsec:qubits)=
### 量子ビット

　**量子ビット**（qubit=キュビット）とは量子コンピュータの基本構成要素のことで、量子情報の入れ物の最小単位です。量子ビットや量子レジスタ（量子ビットを並べたもの）は量子コンピュータ中で常に一つの「状態」にあります。

　量子レジスタの状態を物理学の習わしに従ってしばしば「ケット」という$\ket{\psi}$のような記号で表しますが、これ自体はただの記号なので特に意味はありません。別に「枠」なしで$\psi$と書いても、絵文字を使って🔱と書いても、何でも構いません。

　重要なのは各量子ビットに対して2つの**基底状態**が定義できることで、量子計算の習わしではそれらを$\ket{0}$と$\ket{1}$で表し、「計算基底」とも呼びます[^basis]。そして、量子ビットの任意の状態は、2つの複素数$\alpha, \beta$を使って
$$
\alpha \ket{0} + \beta \ket{1}
$$
と2つの基底の「重ね合わせ」で表せます。ここで$\alpha, \beta$を確率振幅や単に**振幅**（amplitude）と呼ぶこともあります。繰り返しですが別に表記法自体に深い意味はなく、例えば同じ状態を$[\alpha, \beta]$と書いてもいいわけです[^complexarray]。

　量子ビットの任意の状態が2つの複素数で表せるということは、逆に言えば一つの量子ビットには2つの複素数に相当する情報を記録できるということになります。ただこれには少し注釈があって、量子力学の決まりごとから、$\alpha$と$\beta$は
$$
|\alpha|^2 + |\beta|^2 = 1
$$
という関係を満たさなければならず、かつ全体の位相（global phase）は意味を持たない、つまり
$$
\alpha \ket{0} + \beta \ket{1} \sim e^{i\theta} (\alpha \ket{0} + \beta \ket{1})
$$
（ここで $\sim$ は「同じ量子状態を表す」という意味）という制約があります。

　複素数1つは実数2つで書けるので$\alpha$と$\beta$をあわせて実数4つ分の情報が入っているようですが、2つの拘束条件があるため、実際の自由度は$4 - 2 = 2$個です。自由度の数をあらわにして量子ビットの状態を記述するときは、
$$
e^{i\phi/2}\cos\frac{\theta}{2}\ket{0} + e^{-i\phi/2}\sin\frac{\theta}{2}\ket{1}
$$
と書いたりもします。この表記法をブロッホ球表現と呼ぶこともあります（詳細は{doc}`quantum_basics`にて）。

　面白くなるのは量子ビットが複数ある場合です。例えば量子ビット2つなら、それぞれに$\ket{0}, \ket{1}$の計算基底があるので、任意の状態は
$$
\alpha \ket{0}\ket{0} + \beta \ket{0}\ket{1} + \gamma \ket{1}\ket{0} + \delta \ket{1}\ket{1}
$$
と4つの複素数を使った重ね合わせになります。2つの量子ビットの基底を並べた$\ket{0}\ket{0}$のような状態が、このレジスタの計算基底ということになります。$\ket{00}$と略したりもします。

　上で登場した量子力学の決まりごとはこの場合
$$
|\alpha|^2 + |\beta|^2 + |\gamma|^2 + |\delta|^2 = 1
$$
と
$$
\alpha \ket{00} + \beta \ket{01} + \gamma \ket{10} + \delta \ket{11} \sim e^{i\theta} (\alpha \ket{00} + \beta \ket{01} + \gamma \ket{10} + \delta \ket{11})
$$
となります。量子ビットがいくつあっても拘束条件は2つだけです。

　つまり、量子ビット$n$個レジスタでは、基底の数が$2^n$個で、それぞれに複素数の振幅がかかるので、実数$2 \times 2^n - 2$個分の情報が記録できることになります。これが量子計算に関して「指数関数的」という表現がよく用いられる所以です。

　量子レジスタの計算基底状態の表記法としては、上に書いたようにケットを$n$個並べたり$n$個の0/1を一つのケットの中に並べたりする方法がありますが、さらにコンパクトなのが、0/1の並び（ビット列）を二進数とみなして、対応する（十進数の）数字で表現する方法です。例えば4量子ビットのレジスタで状態$\ket{0000}$と$\ket{1111}$はそれぞれ$\ket{0}$と$\ket{15}$と書けます。

　ただし、ここで注意すべきなのは、左右端のどちらが「1の位」なのか事前に約束しないといけないことです。$\ket{0100}$を$\ket{4}$（右端が1の位）とするか$\ket{2}$（左端が1の位）とするかは約束次第です。このワークブックでは、Qiskitでの定義に従って、右端を1の位とします。同時に、レジスタの最初の量子ビットが1の位に対応するようにしたいので、ケットや0/1を並べて計算基底を表現するときは、右から順にレジスタの量子ビットを並べていくことにします。

　Qiskitには量子レジスタオブジェクトがあり、
```python
from qiskit import QuantumRegister
register = QuantumRegister(4, 'myregister')
```
のように量子ビット数（この場合4）と名前（`'myregister'`）を指定して初期化します。初期状態では、量子ビットはすべて$\ket{0}$状態にあります。レジスタオブジェクトはこのままではあまり使い道がなく、基本的には次に紹介する量子回路の一部として利用します。

(subsec:circuit)=
### ゲート、回路、測定

　量子計算とは、端的に言えば、量子レジスタに特定の状態を生成し、その振幅を利用することと言えます。

　とは言っても、いきなり「えいや」と好きな量子状態を作れるわけではなく、パターンの決まった単純操作（$\ket{0}$と$\ket{1}$の振幅を入れ替える、ブロッホ球表現での位相角度$\phi$を増減させる、など）を順番に組み合わせて複雑な状態を作っていきます。この単純操作のオペレーションのことを一般に量子**ゲート**といい、ゲートの種類や順番を指定したプログラムに相当するものを量子**回路**と呼びます。

　Qiskitでは、量子回路を`QuantumCircuit`オブジェクトで表します。
```python
from qiskit import QuantumCircuit, QuantumRegister
register = QuantumRegister(4, 'myregister')
circuit = QuantumCircuit(register)
```
という具合です。

　作られた量子回路は、量子ビットの数は決まっているもののゲートが一つもない「空っぽ」の状態なので、そこにゲートをかけていきます。例えば下で説明するアダマールゲートをレジスタの2個目の量子ビットに作用させるには
```python
circuit.h(register[1])
```
とします。

　上で「振幅を利用する」という曖昧な表現をしましたが、それはいろいろな利用の仕方があるからです。しかし、どんな方法であっても、必ず量子レジスタの**測定**という操作を行います。量子コンピュータから何かしらの情報を得るための唯一の方法が測定です。

　測定は量子レジスタの状態を「覗き見る」ような操作ですが、一回の測定操作で具体的に起きることは、各量子ビットに対して0もしくは1という値が得られるというだけです。つまり、量子状態が$2^n$個の計算基底の複雑な重ね合わせであったとしても、測定をすると一つの計算基底に対応するビット列が出てくるだけということになります。しかも、一度測定してしまった量子ビットはもう状態を変えてしまっていて、複雑な重ね合わせは失われてしまいます。

　ではこの「一つの計算基底」がどの基底なのかというと、実は特殊な場合を除いて決まっていません。全く同じ回路を何度も実行して測定を繰り返すと、その都度ランダムにビット列が決まります。ただし、このランダムさには法則があって、**特定のビット列が得られる確率は、対応する計算基底の振幅の絶対値自乗**となっています。つまり、$n$長レジスタの状態$\sum_{j=0}^{2^n-1} c_j \ket{j}$があるとき、測定でビット列$k$が得られる確率は$|c_k|^2$です。

　したがって、回路の実行と測定を何度も繰り返して、それぞれのビット列が現れる頻度を記録すれば、だんだん$|c_j|^2$の値がわかっていき、レジスタがどんな量子状態にあるのか見えてきます。しかしわかるのは$c_j$の絶対値だけで、複素位相については知る術なしです。どうもすっきりしませんが、これが量子コンピュータから情報を得る方法です。

　逆に、指数関数的な内部の情報量をうまく使って計算を行いつつ、測定という限定的な方法でも答えが読み出せるように工夫するのが、量子アルゴリズム設計の真髄ということになります。例えば理想的には、何か計算の答えが整数$k$であり、それを計算する回路の終状態が単純に$\ket{k}$となるようであれば、一度の測定で答えがわかる（上でいった特殊な場合に相当）わけです。単純に$\ket{k}$でなくても、重ね合わせ$\sum_{j=0}^{2^n-1} c_j \ket{j}$において$|c_k| \gg |c_{j \neq k}|$を実現できれば、数回の測定で答えが高確率でわかります。{doc}`shor`で紹介する位相推定アルゴリズムはその好例です。

(subsec:common_gates)=
### よく使うゲート

　IBM Q System Oneのような超電導振動子を利用した量子コンピュータでは、実際に使用できるゲートは量子ビット1つにかかるものと2つにかかるものに限定されます。しかし、それらを十分な数組み合わせれば、$n$量子ビットレジスタにおいてどのような状態も実現できることが、数学的に証明できます。

　1量子ビットの操作でよく使われるゲートには、以下のようなものがあります。

- $X$: $\ket{0}$と$\ket{1}$の振幅を交換する
- $Z$: $\ket{1}$の振幅を$-1$倍する
- $H$（アダマールゲート）: 計算基底それぞれに対して、以下の変形をする。
$$
H\ket{0} = \frac{1}{\sqrt{2}} (\ket{0} + \ket{1}) \\
H\ket{1} = \frac{1}{\sqrt{2}} (\ket{0} - \ket{1})
$$
（「量子状態にゲートを作用させる」ことをケットの記法で書くときは、ゲートに対応する記号をケットに左からかけます。）

例えば状態$\ket{\psi} = \alpha\ket{0} + \beta\ket{1}$に対しては、
$$
\begin{align}
H\ket{\psi} & = \alpha \frac{1}{\sqrt{2}} (\ket{0} + \ket{1}) + \beta \frac{1}{\sqrt{2}} (\ket{0} - \ket{1}) \\
            & = \frac{1}{\sqrt{2}} (\alpha + \beta) \ket{0} + \frac{1}{\sqrt{2}} (\alpha - \beta) \ket{1}
\end{align}
$$
となります。
- $R_{y}$: パラメータ$\theta$を取り、計算基底それぞれに対して、以下の変形をする。
$$
R_{y}(\theta)\ket{0} = \cos\frac{\theta}{2}\ket{0} - \sin\frac{\theta}{2}\ket{1} \\
R_{y}(\theta)\ket{1} = \sin\frac{\theta}{2}\ket{0} + \cos\frac{\theta}{2}\ket{1}
$$
- $R_{z}$: パラメータ$\phi$を取り、計算基底それぞれに対して、以下の変形をする。
$$
R_{z}(\phi)\ket{0} = e^{-i\phi/2}\ket{0} \\
R_{z}(\phi)\ket{1} = e^{i\phi/2}\ket{1}
$$

　それでは、2量子ビットレジスタの第1ビットに$H, R_y, X$の順にゲートをかけて、最後に測定をする回路をQiskitで書いてみましょう。

```python
from qiskit import QuantumCircuit
import numpy as np

circuit = QuantumCircuit(2) # レジスタを介さずビット数を指定して回路を作成することもできます
circuit.h(0) # その場合、ゲートにはregister[0]ではなく直接量子ビットの番号を指定します
circuit.ry(-np.pi / 2., 0) #　θ = -π/2
circuit.x(0)
circuit.measure_all()

print('This circuit has', circuit.num_qubits, 'qubits and', circuit.size(), 'operations')
```

    This circuit has 2 qubits and 5 operations

　最後のプリント文で、ゲートが3つなのにも関わらず"5 operations"と出てくるのは、各量子ビットの測定も一つのオペレーションとして数えられるからです。

　量子計算に慣れる目的で、この$H, R_y(-\pi/2), X$という順の操作で第1ビットに何が起こるかを辿ってみましょう。初期状態は$\ket{0}$で、ケット記法では操作は左からかけていくので、$X R_y(-\pi/2) H \ket{0}$を計算することになります。
$$
\begin{align}
X R_y\left(\frac{\pi}{2}\right) H \ket{0} & = X R_y\left(-\frac{\pi}{2}\right) \frac{1}{\sqrt{2}}(\ket{0} + \ket{1}) \\
& = \frac{1}{\sqrt{2}} X \left[\left(\cos\left(-\frac{\pi}{4}\right)\ket{0} - \sin\left(-\frac{\pi}{4}\right)\ket{1}\right) + \left(\sin\left(-\frac{\pi}{4}\right)\ket{0} + \cos\left(-\frac{\pi}{4}\right)\ket{1}\right)\right] \\
& = \frac{1}{\sqrt{2}} X \frac{1}{\sqrt{2}} \left[\left(\ket{0} + \ket{1}\right) + \left(-\ket{0} + \ket{1}\right)\right] \\
& = \frac{1}{2} X \ket{1} \\
& = \ket{0}
\end{align}
$$
なので、結局$\ket{0}$状態に戻る操作でした。

　2量子ビットの操作は、量子ビットの超電導素子での実装の都合上、全て一つの方式で行われます。この方式は制御ゲート（controlled gates）と呼ばれ、2つのビットのうち片方を制御（control）、もう片方を標的（target）として、制御ビットが1の時だけ標的ビットに何らかの操作がかかるというものです。例として、ある1ビットゲート$U$を制御ゲート化した$C^i_j[U]$を考えます。ここで$i$が制御、$j$が標的ビットとします。2ビットレジスタの第2ビットを制御、第1ビットを標的とすると（並べて書くときは左が第2で右が第1です）
$$
\begin{align}
C^1_0[U](\ket{0}\ket{0}) & = \ket{0}\ket{0} \\
C^1_0[U](\ket{0}\ket{1}) & = \ket{0}\ket{1} \\
C^1_0[U](\ket{1}\ket{0}) & = \ket{1}U\ket{0} \\
C^1_0[U](\ket{1}\ket{1}) & = \ket{1}U\ket{1}
\end{align}
$$
です。

　上で紹介した頻出する1ビットゲート$X, Z, H, R_y, R_z$のうち、$H$以外は制御ゲート化バージョンもよく使われます。

　Qiskitで2ビットレジスタに制御ゲートを用い、計算基底$\ket{0}, \ket{1}, \ket{2}, \ket{3}$の振幅の絶対値自乗が$1:2:3:4$の比になるような状態を作ってみましょう。さらに$C^0_1[Z]$ゲートを使って$\ket{3}$だけ振幅の符号が他と異なるようにします。

```python
theta1 = np.arctan(np.sqrt(7. / 3.))
theta2 = np.arctan(np.sqrt(2.))
theta3 = np.arctan(np.sqrt(4. / 3))

circuit = QuantumCircuit(2)
circuit.ry(-theta1, 1)
circuit.ry(-theta2, 0)
circuit.cry(-(theta3 - theta2), 1, 0) # C[Ry]　1が制御で0が標的
circuit.cz(0, 1) # C[Z] 0が制御で1が標的（実はC[Z]ではどちらが制御でも結果は同じ）
circuit.measure_all()
```

　やや複雑ですが、また計算を追って見ましょう。まず角度$\theta_1, \theta_2, \theta_3$は定義から以下の関係を満たします。
$$
\begin{align}
R_y(-\theta_1)\ket{0} & = \sqrt{\frac{3}{10}} \ket{0} + \sqrt{\frac{7}{10}} \ket{1} \\
R_y(-\theta_2)\ket{0} & = \sqrt{\frac{1}{3}} \ket{0} + \sqrt{\frac{2}{3}} \ket{1} \\
R_y(-(\theta_3 - \theta_2))R_y(-\theta_2)\ket{0} & = R_y(-\theta_3)\ket{0} = \sqrt{\frac{3}{7}} \ket{0} + \sqrt{\frac{4}{7}} \ket{1}.
\end{align}
$$
したがって、
$$
\begin{align}
C^1_0[R_y(-(\theta_3 - \theta_2))]R_y(-\theta_1)\ket{0}R_y(-\theta_2)\ket{0} & = C^1_0[R_y(-(\theta_3 - \theta_2))]\left(\sqrt{\frac{3}{10}} \ket{0} + \sqrt{\frac{7}{10}} \ket{1}\right) R_y(-\theta_2)\ket{0}\\
& = \sqrt{\frac{3}{10}} \ket{0} R_y(-\theta_2)\ket{0} + \sqrt{\frac{7}{10}} \ket{1} R_y(-\theta_3)\ket{0} \\
& = \sqrt{\frac{3}{10}} \ket{0} \left(\sqrt{\frac{1}{3}} \ket{0} + \sqrt{\frac{2}{3}} \ket{1}\right) + \sqrt{\frac{7}{10}} \ket{1} \left(\sqrt{\frac{3}{7}} \ket{0} + \sqrt{\frac{4}{7}} \ket{1}\right) \\
& = \sqrt{\frac{1}{10}} \ket{00} + \sqrt{\frac{2}{10}} \ket{01} + \sqrt{\frac{3}{10}} \ket{10} + \sqrt{\frac{4}{10}} \ket{11}
\end{align}
$$
最後に$C[Z]$をかけると、$\ket{11}$だけ符号が反転します。

(subsec:circuit_visualization)=
### 回路図の書き方と読み方

　量子回路を可視化する方法として、「回路図」の標準的な描き方が決まっています。Qiskitでは`QuantumCircuit`オブジェクトの`draw()`というメソッドを使って自動描画できます。

```python
circuit.draw('mpl')
```

　ここで`draw()`の引数`'mpl'`はmatplotlibライブラリを使ってカラーで描くことを指定しています。実行環境によっては対応していないこともあるので、その場合は引数なしの`draw()`を使います。結果は`mpl`の場合に比べて見劣りしますが、内容は同じです。

```python
circuit.draw()
```

        ┌──────────────┐┌──────────────┐    ░ ┌─┐
   q_0: ┤ RY(-0.95532) ├┤ RY(0.098245) ├─■──░─┤M├───
        ├──────────────┤└──────┬───────┘ │  ░ └╥┘┌─┐
   q_1: ┤ RY(-0.99116) ├───────■─────────■──░──╫─┤M├
        └──────────────┘                    ░  ║ └╥┘
meas: 2/═══════════════════════════════════════╩══╩═
                                               0  1

　回路図は左から右に読んでいきます。水平の2本の実線が上からそれぞれ第1、第2量子ビットに対応し、その上にかぶさっている四角がゲート、最後にある矢印が下に伸びている箱が測定を表します。1ビットゲートから伸びている先端の丸い縦線は制御を表します。一番下の二重線は「古典レジスタ」に対応し、測定結果の0/1が記録される部分です。

(sec:writing_chsh)=
## CHSH不等式を計算する回路を書く

　それではいよいよ本題に入りましょう。

(sec:using_ibmq)=
## IBMQを使う

(sec:analysis)=
## 量子測定結果の解析

[^basis]: ここで言う「基底」は線形代数での意味（basis）で、「線形空間中の任意のベクトルを要素の線形和で表せる最小の集合」です。基底となる量子状態だから「基底状態」と呼びます。化学や量子力学で言うところのエネルギーの最も低い状態「基底状態」（ground state）とは関係ありません。
[^complexarray]: 実際に量子計算のシミュレーションをコンピュータ上で行う時などは、量子レジスタの状態を複素数の配列で表すので、この表記の方がよく対応します。


```python
from IPython.core.display import display_html
from urllib.request import urlopen

display_html(urlopen('http://j.mp/1Ye5iWA').read(), raw=True)
```

<link href='http://fonts.googleapis.com/css?family=EB+Garamond' rel='stylesheet' type='text/css'>
<link href='http://fonts.googleapis.com/css?family=Lato:100,100italic' rel='stylesheet' type='text/css'>
<link rel="stylesheet" href="//cdn.jsdelivr.net/font-hack/2.013/css/hack-extended.min.css">

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

        div.cell{
	        max-width:    800px;
	        margin-left:  auto;
	        margin-right: auto;}

	h1 {
        	font-family: 'Lato', sans-serif;}

	h2 {
        	font-family: 'Lato', serif;}

	h3{
		font-family:   'Lato', serif;
        	margin-top:    12px;
		margin-bottom: 3px;}

        h4{
		font-family: 'Lato', serif;}

	h5 {
        	font-family: 'Lato', serif;}

	div.text_cell_render{
        	font-family:  'EB Garamond', 'Computer Modern', 'Helvetica Neue', Arial, Helvetica, Geneva, sans-serif;
	        line-height:  140%;
	        font-size:    14pt;
		max-width:    650px;
		margin-left:  auto;
        	margin-right: auto;}

	.CodeMirror{
        	font-family: 'Hack', 'Consolas', 'Inconsolata', monospace;}

	.text_cell_render h1 {
	        font-weight:   100;
	        font-size:     48pt;
		line-height:   100%;
	        color:         #CD2305;
	        margin-bottom: 0.5em;
	        margin-top:    0.5em;
	        display:       block;}

	.text_cell_render h2 {
        	font-weight:   100;
        	font-size:     36pt;}

	.text_cell_render h3 {
        	font-weight:   100;
        	font-size:     21pt;}

	.text_cell_render h4 {
        	font-weight:   100;
        	font-size:     18pt;}

	.text_cell_render h5 {
	        font-weight:   100;
	        font-size:     16pt;
	        color:         #CD2305;
	        font-style:    italic;
	        margin-bottom: 0.5em;
	        margin-top:    0.5em;
	        display:       block;}

	.text_cell_render h6 {
	        font-weight:   100;
	        font-size:     14pt;
	        color:         #CD2305;
	        font-style:    italic;
	        margin-bottom: 0.5em;
	        margin-top:    0.5em;
	        display:       block;}

	.warning{
	    	background-color: #FFFFE0;
	    	border-color: #FFCC99;
	    	border-left: 5px solid #FFCC99;
	    	padding: 0.5em;}

	.error{
	    	background-color: #fcf2f2;
	    	border-color: #dFb5b4;
	    	border-left: 5px solid #dfb5b4;
	    	padding: 0.5em;}

</style>

# Dinámica del Robot Manipulador SCARA

Se tiene un robot manipulador tipo SCARA, como el de la siguiente figura, del cual se quiere obtener sus ecuaciones de movimiento, asi como una simulación bajo cierta ley de  control.

Empezaremos obteniendo una representación simple de las posiciones de los centros de masa para los eslabones del robot.

Primero tenemos que importar las librerias de computo simbolico:

```python
from sympy import var, sin, cos, Matrix, Integer, eye, Function, Rational, exp, Symbol, I, solve, pi, trigsimp, dsolve, sinh, cosh, simplify
from sympy.physics.mechanics import mechanics_printing
mechanics_printing()
```

Y declaramos todas las constantes involucradas en este calculo simbolico:

```python
var("m1 m2 m3 J1 J2 J3 l1 l2 L1 L2 L0 t g")
```

Asi como algunas de las variables de nuestro problema:

```python
q1 = Function("q1")(t)
q2 = Function("q2")(t)
q3 = Function("q3")(t)
```

Y empezamos con el calculo de la posicion de los centros de masa de los eslabones, asi como su derivada y el cuadrado de la velocidad lineal de cada eslabon:

```python
x1 = l1*cos(q1)
y1 = l1*sin(q1)
z1 = L0
v1 = x1.diff("t")**2 + y1.diff("t")**2 + z1.diff("t")**2
v1.trigsimp()
```

```python
x2 = L1*cos(q1) + l2*cos(q1 + q2)
y2 = L1*sin(q1) + l2*sin(q1 + q2)
z2 = L0
v2 = x2.diff("t")**2 + y2.diff("t")**2 + z2.diff("t")**2
v2.trigsimp()
```

```python
x3 = L1*cos(q1) + L2*cos(q1 + q2)
y3 = L1*sin(q1) + L2*sin(q1 + q2)
z3 = L0 - q3
v3 = x3.diff("t")**2 + y3.diff("t")**2 + z3.diff("t")**2
v3.trigsimp()
```

Declaramos $\omega_i$ como la velocidad angular de cada eslabon:

```python
ω1 = q1.diff("t")
ω2 = q2.diff("t")
ω3 = 0
```

Y procedemos al calculo de la energía cinética de cada eslabon:

```python
K1 = Rational(1, 2)*m1*v1 + Rational(1, 2)*J1*ω1**2
K1
```

```python
K2 = Rational(1, 2)*m1*v2 + Rational(1, 2)*J2*ω2**2
K2
```

```python
K3 = Rational(1, 2)*m1*v3 + Rational(1, 2)*J3*ω3**2
K3
```

Calculamos tambien la energía potencial de cada eslabon:

```python
U1 = m1*g*z1
U1
```

```python
U2 = m2*g*z2
U2
```

```python
U3 = m3*g*z3
U3
```

Por lo que ya podemos calcular tanto la energía cinética de nuestro sistema como la potencial:

```python
K = K1 + K2 + K3
K
```

```python
U = U1 + U2 + U3
U
```

y el Lagrangiano queda:

```python
L = (K - U).expand().simplify()
L
```

Por lo que ahora solo tenemos que aplicar la ecuación de Euler-Lagrange para obtener las ecuaciones de movimiento del sistema:

```python
τ1 = (L.diff(q1.diff(t)).diff(t) - L.diff(q1)).simplify().expand().collect(q1.diff(t).diff(t)).collect(q2.diff(t).diff(t))
```

```python
τ2 = (L.diff(q2.diff(t)).diff(t) - L.diff(q2)).simplify().expand().collect(q1.diff(t).diff(t)).collect(q2.diff(t).diff(t))
```

```python
τ3 = (L.diff(q3.diff(t)).diff(t) - L.diff(q3)).simplify().expand().collect(q1.diff(t).diff(t)).collect(q2.diff(t).diff(t))
```

```python
τ1
```

```python
τ2
```

```python
τ3
```

Una vez que tenemos las ecuaciones de movimiento, debemos simular el comportamiento del sistema por medio de la función ```odeint```, y obtener una gráfica de la trayectoria del sistema:

```python
from scipy.integrate import odeint
from numpy import linspace
```

Para utilizar la función ```odeint```, debemos crear una función, que describa la dinámica del sistema:

```python
def scara(estado, tiempo):
    # Se importan funciones necesarias
    from numpy import sin, cos, matrix
    # Se desenvuelven variables del estado y tiempo
    q1, q2, q3, q̇1, q̇2, q̇3 = estado
    t = tiempo

    # Se declaran constantes del sistema
    m1, m2, m3 = 1, 1, 1
    J1, J2, J3 = 1, 1, 1
    l1, l2 = 0.5, 0.5
    L1, L2 = 1, 1
    L = 1
    g = 9.81

    # Se declaran constantes del control
    kp1, kp2, kp3 = -30, -60, -60
    kv1, kv2, kv3 = -20, -20, -18

    # Señales de control nulas
    #tau1, tau2, tau3 = 0, 0, 0

    # Posiciones a alcanzar
    qd1, qd2, qd3 = 1, 1, 1

    # Se declaran señales de control del sistema
    tau1 = kp1*(q1 - qd1) + kv1*q̇1
    tau2 = kp2*(q2 - qd2) + kv2*q̇2
    tau3 = kp3*(q3 - qd3) + kv3*q̇3 + m3*g

    # Se calculan algunos terminos comunes
    λ1 = m1*L1*(l2 + L2)
    λ2 = m1*(l2**2 + L2**2)
    λ3 = m1*(l1**2 + L1**2)

    # Se calculan las matrices de masas, Coriolis,
    # y vectores de gravedad, control, posicion y velocidad
    M = matrix([[J1 + 2*λ1*cos(q2) + m1*L1**2 + λ2 + λ3, λ1*cos(q2) + λ2, 0],
                [λ1*cos(q2) + λ2, J2 + λ2, 0],
                [0, 0, m1]])
    C = matrix([[-2*q̇1, -q̇2, 0], [q̇1, 0, 0], [0, 0, 0]])
    G = matrix([[0], [0], [-m3*g]])
    Tau = matrix([[tau1], [tau2], [tau3]])
    q = matrix([[q1], [q2], [q3]])
    q̇ = matrix([[q̇1], [q̇2], [q̇3]])

    # Se calcula la derivada del estado del sistema
    qp1 = q̇1
    qp2 = q̇2
    qp3 = q̇3

    qpp = M.I*(Tau - C*q̇ - G)
    qpp1, qpp2, qpp3 = qpp.tolist()

    return [qp1, qp2, qp3, qpp1[0], qpp2[0], qpp3[0]]
```

Y declarar un arreglo con todos los tiempos a simular, mandar a llamar a la función ```odeint```, y listo!

```python
t = linspace(0, 10, 1000)
estados_simulados = odeint(func = scara, y0  = [0, 0, 0, 0, 0, 0], t = t)
```

Desempacamos los elementos que nos entrega ```odeint```:

```python
q1, q2, q3, q̇1, q̇2, q̇3 = list(zip(*estados_simulados.tolist()))
```

Importamos la libreria para graficar:

```python
%matplotlib notebook
from matplotlib.pyplot import plot, style, figure
from mpl_toolkits.mplot3d import Axes3D
style.use("ggplot")
```

Hacemos la grafica de las trayectorias del sistema, $q_1$, $q_2$ y $q_3$:

```python
fig1 = figure(figsize=(12, 8))

ax1 = fig1.gca()

p1, = ax1.plot(t, q1)
p2, = ax1.plot(t, q2)
p3, = ax1.plot(t, q3)
ax1.legend([p1, p2, p3],[r"$q_1$", r"$q_2$", r"$q_3$"])
ax1.set_ylim(-0.1, 1.2)
ax1.set_xlim(-0.1, 10);
```

    <IPython.core.display.Javascript object>

Realizamos la cinemática del manipulador para poder graficar en 3D:

```python
def tras_x(x):
    from numpy import matrix
    A = matrix([[1, 0, 0, x],
                [0, 1, 0, 0],
                [0, 0, 1, 0],
                [0, 0, 0, 1]])
    return A

def tras_z(z):
    from numpy import matrix
    A = matrix([[1, 0, 0, 0],
                [0, 1, 0, 0],
                [0, 0, 1, z],
                [0, 0, 0, 1]])
    return A

def rot_z(θ):
    from numpy import matrix, sin, cos
    A = matrix([[cos(θ), -sin(θ), 0, 0],
                [sin(θ), cos(θ), 0, 0],
                [0, 0, 1, 0],
                [0, 0, 0, 1]])
    return A

def cinematica_scara(q1, q2, q3, l1, l2, L):
    from numpy import matrix
    p0 = matrix([[0], [0], [0], [1]])
    p1 = tras_z(L)*p0
    p2 = rot_z(q1)*tras_x(l1)*p1
    p3 = rot_z(q2)*tras_x(l2)*p2
    p4 = tras_z(-q3)*p3

    return [[p0.tolist()[0][0], p1.tolist()[0][0], p2.tolist()[0][0], p3.tolist()[0][0], p4.tolist()[0][0]],
             [p0.tolist()[1][0], p1.tolist()[1][0], p2.tolist()[1][0], p3.tolist()[1][0], p4.tolist()[1][0]],
             [p0.tolist()[2][0], p1.tolist()[2][0], p2.tolist()[2][0], p3.tolist()[2][0], p4.tolist()[2][0]]]
```

```python
from numpy import pi
τ = 2*pi
xs, ys, zs = cinematica_scara(τ/12, τ/9, 0.5, 1, 1, 1)
```

```python
fig2 = figure(figsize=(8, 8))
ax2 = fig2.gca(projection='3d')
ax2.plot(xs, ys, zs, "-o")
ax2.set_xlim(-2, 2)
ax2.set_ylim(-2, 2)
ax2.set_zlim(0, 1.5);
```

    <IPython.core.display.Javascript object>

Cambiamos el ambiente de matplotlib, para poder graficar interactivamente, y declaramos una función que tome la cinematica del robot y grafique los puntos:

```python
%matplotlib inline
```

```python
def grafica_scara(q1, q2, q3, l1, l2, L):
    xs, ys , zs = cinematica_scara(q1, q2, q3, l1, l2, L)

    fig = figure(figsize=(8, 8))

    ax = fig.gca(projection='3d')

    ax.plot(xs, ys, zs, "-o")

    ax.set_xlim(-2, 2)
    ax.set_ylim(-2, 2)
    ax.set_zlim(0, 1.5);
```

```python
grafica_scara(τ/12, τ/9, 0.2, 1, 1, 1)
```

Importamos la libreria para interactuar con los datos y le pasamos la función que grafica:

```python
# Se importan widgets de IPython para interactuar con la funcion
from IPython.html.widgets import interact, fixed
```

    /Users/roberto/miniconda3/lib/python3.5/site-packages/IPython/html.py:14: ShimWarning: The `IPython.html` package has been deprecated. You should import from `notebook` instead. `IPython.html.widgets` has moved to `ipywidgets`.
      "`IPython.html.widgets` has moved to `ipywidgets`.", ShimWarning)

```python
# Se llama a la funcion interactiva
interact(grafica_scara, q1=(0, τ), q2=(0, τ), q3=(0, 1.0), l1=fixed(1), l2=fixed(1), L=fixed(1))
```

Para realizar una animación con los datos de la simulación primero importamos la libreria necesaria y creamos la misma grafica dentro del ambiente de animación:

```python
from matplotlib import animation
from numpy import arange
```

```python
# Se define el tamaño de la figura
fig = figure(figsize=(8, 8))

# Se define una sola grafica en la figura y se dan los limites de los ejes x y y
axi = fig.add_subplot(111, autoscale_on=False, xlim=(-2, 2), ylim=(-2, 2), projection='3d')

# Se utilizan graficas de linea para el resorte y amortiguador
robot, = axi.plot([], [], [], "-o", lw=2)

def init():
    # Esta funcion se ejecuta una sola vez y sirve para inicializar el sistema
    robot.set_data([], [])
    return robot

def animate(i):
    # Esta funcion se ejecuta para cada cuadro del GIF

    # Se obtienen las coordenadas del robot y se meten los datos en su grafica de linea
    xs, ys, zs = cinematica_scara(q1[i], q2[i], q3[i], 1, 1, 1)
    robot.set_data(xs, ys)
    robot .set_3d_properties(zs)

    return robot

# Se hace la animacion dandole el nombre de la figura definida al principio, la funcion que
# se debe ejecutar para cada cuadro, el numero de cuadros que se debe de hacer, el periodo
# de cada cuadro y la funcion inicial
ani = animation.FuncAnimation(fig, animate, arange(1, len(q1)), interval=25,
                              blit=True, init_func=init)

# Se guarda el GIF en el archivo indicado
ani.save('./imagenes/simulacion-scara.gif', writer='imagemagick');
```


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

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

__copyright__ = "(c) 2014-2018, The Psi4NumPy Developers"
__license__   = "BSD-3-Clause"
__date__      = "2017-05-23"
```

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

### Implementation of Conventional MP2

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

```python
# ==> Import statements & Global Options <==
import psi4
import numpy as np

psi4.set_memory(int(2e9))
numpy_memory = 2
psi4.core.set_output_file('output.dat', False)
```

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

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

psi4.set_options({'basis':        '6-31g',
                  'scf_type':     'pk',
                  'mp2_type':     'conv',
                  'e_convergence': 1e-8,
                  'd_convergence': 1e-8})
```

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

```python
# Get the SCF wavefunction & energies
scf_e, scf_wfn = psi4.energy('scf', return_wfn=True)
```

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

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

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

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

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

# Memory check for ERI tensor
I_size = (nmo**4) * 8.e-9
print('\nSize of the ERI tensor will be %4.2f GB.' % I_size)
memory_footprint = I_size * 1.5
if I_size > numpy_memory:
    psi4.core.clean()
    raise Exception("Estimated memory utilization (%4.2f GB) exceeds allotted memory \
                     limit of %4.2f GB." % (memory_footprint, numpy_memory))

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

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

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

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

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

~~~python
# Naive Algorithm for ERI Transformation
Imo = np.einsum('pi,qa,pqrs,rj,sb->iajb', Cocc, Cvirt, I, Cocc, Cvirt)
~~~

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

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

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

```python
# ==> Transform I -> I_mo @ O(N^5) <==
tmp = np.einsum('pi,pqrs->iqrs', Cocc, I)
tmp = np.einsum('qa,iqrs->iars', Cvirt, tmp)
tmp = np.einsum('iars,rj->iajs', tmp, Cocc)
I_mo = np.einsum('iajs,sb->iajb', tmp, Cvirt)
```

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

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

Now we have all the pieces needed to compute $E_0^{(2)}$.  This could be done by writing explicit loops over occupied and virtual indices Python side, e.g.,
~~~python
# Compute SS MP2 Correlation
mp2_ss_corr = 0.0
for i in xrange(ndocc):
    for a in xrange(nmo - ndocc):
        for j in xrange(ndocc):
            for b in xrange(nmo - ndocc):
                numerator = I_mo[i,a,j,b] * (I_mo[i, a, j, b] - I_mo[i, b, j, a])
                mp2_ss_corr += numerator / (e_ij[i] + e_ij[j] - e_ab[a] - e_ab[b])
~~~
This method, while very clear what is going on and easy to program, has the distinct disadvantage that Python loops are much slower than the same block written in a compiled language like C, C++, or Fortran.  For this reason, it is better to use infrastructure available in NumPy like `np.einsum()`, `np.dot`, etc. to explicitly compute the above quantity C-side.  It should be clear how to contract the four-index integrals $(i\,a\mid j\,b)$ and $(i\,a\mid j\,b)$ with one another, but what about the energy eigenvalues $\epsilon$?  We can use a NumPy trick called *broadcasting* to construct a four-index array of all possible energy denominators, which can then be contracted with the full I_mo arrays.  To do this, we'll use the function `np.reshape()`:
~~~python
# Prepare 4d energy denominator array
e_denom = e_ij.reshape(-1, 1, 1, 1) # Diagonal of 4d array are occupied orbital energies
e_denom -= e_ab.reshape(-1, 1, 1)   # all combinations of (e_ij - e_ab)
e_denom += e_ij.reshape(-1, 1)      # all combinations of [(e_ij - e_ab) + e_ij]
e_denom -= e_ab                     # All combinations of full denominator
e_denom = 1 / e_denom               # Take reciprocal for contracting with numerator
~~~
In the cell below, compute the energy denominator using `np.reshape()` and contract this array with the four-index ERIs to compute the same-spin and opposite-spin MP2 correction using `np.einsum()`.  Then, add these quantities to the SCF energy computed above to obtain the total MP2 energy.

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

```python
# ==> Compute MP2 Correlation & MP2 Energy <==
# Compute energy denominator array
e_denom = 1 / (e_ij.reshape(-1, 1, 1, 1) - e_ab.reshape(-1, 1, 1) + e_ij.reshape(-1, 1) - e_ab)

# Compute SS & OS MP2 Correlation with Einsum
mp2_os_corr = np.einsum('iajb,iajb,iajb->', I_mo, I_mo, e_denom)
mp2_ss_corr = np.einsum('iajb,iajb,iajb->', I_mo, I_mo - I_mo.swapaxes(1,3), e_denom)

# Total MP2 Energy
MP2_E = scf_e + mp2_os_corr + mp2_ss_corr
```

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

## References

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

```python

```


# CTIPe

The kamodofied CTIPe requires `netcdf4` (and numpy, scipy, and pandas, but those are required by kamodo)

    pip install netcdf4

!!! warning
    This is a work in progress!

```python
%load_ext autoreload
%autoreload 2
```

```python
from netCDF4 import Dataset
import numpy as np
import os
import scipy
from scipy import interpolate
from scipy.interpolate import RegularGridInterpolator
import math
import time
import kamodo
from datetime import datetime,timezone
import pandas as pd
from plotly.offline import init_notebook_mode, iplot, plot
init_notebook_mode(connected = True)
```

```python
#filename = '/Users/lrastaet/Kamodo_data/CTIPe/2020-05-04-plot-density.nc' # each CTIPe run has plot-density, plot-height, plot-neutral and plot-plasma files for each date
#filename = '/Users/lrastaet/Kamodo_data/CTIPe//Sujay_Pal_042620_IT_1/2015-03-20-plot-neutral.nc'
filename = '/Users/lrastaet/Kamodo_data/CTIPe//Sujay_Pal_042620_IT_1/2015-03-20-plot-plasma.nc'
#filename = '/Users/lrastaet/Kamodo_data/CTIPe//Sujay_Pal_042620_IT_1/2015-03-20-plot-height.nc'
#filename= '/Users/lrastaet/Kamodo_data/CTIPe/Sujay_Pal_042620_IT_1/2015-03-20-plot-density.nc'
```

```python
dt=datetime.strptime("2015/06/12 05:07:01",'%Y/%m/%d %H:%M:%S') #,tzinfo=timezone.utc)
```

```python
delta_t=dt-datetime(1970,1,1)
print(delta_t.total_seconds())
```

    1434085621.0

```python
from readers.ctipe_singletime import CTIPe,ctipe_kamodo_variable_names
from readers.geopack_transform import seconds_from_19700101
def seconds_from_19700101_to_DateTime(seconds):
    from datetime import datetime,timezone,timedelta
    time0=datetime(1970,1,1,0,0,0,tzinfo=timezone.utc)+timedelta(seconds=seconds)
    return time0
from tiegcm.util import datetime_to_epoch
```

By default, the kamodofied CTIPe reader will load all 3d and 4d variables from the model. For performance reasons, you may want to only include certain variables for analysis. To do so, pass a list of variables in addition to the filename:

```python
ctipe = CTIPe(filename,variables_requested=['T'],date="2015/03/20",time="01:15:00",debug=False)
#ctipe = CTIPe(filename,debug=False)
#ctipe = CTIPe_Kamodo(filename, ['TN', 'UN', 'VN', 'EFLUX', 'Z'])
```

    /Users/lrastaet/Kamodo_data/CTIPe/Sujay_Pal_042620_IT_1/2015-03-20-plot-density.nc
    /Users/lrastaet/Kamodo_data/CTIPe/Sujay_Pal_042620_IT_1/2015-03-20-plot-height.nc
    /Users/lrastaet/Kamodo_data/CTIPe/Sujay_Pal_042620_IT_1/2015-03-20-plot-neutral.nc
    /Users/lrastaet/Kamodo_data/CTIPe/Sujay_Pal_042620_IT_1/2015-03-20-plot-plasma.nc
    [1.4268105e+09 1.4268114e+09 1.4268123e+09 1.4268132e+09 1.4268141e+09
     1.4268150e+09 1.4268159e+09 1.4268168e+09 1.4268177e+09 1.4268186e+09
     1.4268195e+09 1.4268204e+09 1.4268213e+09 1.4268222e+09 1.4268231e+09
     1.4268240e+09 1.4268249e+09 1.4268258e+09 1.4268267e+09 1.4268276e+09
     1.4268285e+09 1.4268294e+09 1.4268303e+09 1.4268312e+09 1.4268321e+09
     1.4268330e+09 1.4268339e+09 1.4268348e+09 1.4268357e+09 1.4268366e+09
     1.4268375e+09 1.4268384e+09 1.4268393e+09 1.4268402e+09 1.4268411e+09
     1.4268420e+09 1.4268429e+09 1.4268438e+09 1.4268447e+09 1.4268456e+09
     1.4268465e+09 1.4268474e+09 1.4268483e+09 1.4268492e+09 1.4268501e+09
     1.4268510e+09 1.4268519e+09 1.4268528e+09 1.4268537e+09 1.4268546e+09
     1.4268555e+09 1.4268564e+09 1.4268573e+09 1.4268582e+09 1.4268591e+09
     1.4268600e+09 1.4268609e+09 1.4268618e+09 1.4268627e+09 1.4268636e+09
     1.4268645e+09 1.4268654e+09 1.4268663e+09 1.4268672e+09 1.4268681e+09
     1.4268690e+09 1.4268699e+09 1.4268708e+09 1.4268717e+09 1.4268726e+09
     1.4268735e+09 1.4268744e+09 1.4268753e+09 1.4268762e+09 1.4268771e+09
     1.4268780e+09 1.4268789e+09 1.4268798e+09 1.4268807e+09 1.4268816e+09
     1.4268825e+09 1.4268834e+09 1.4268843e+09 1.4268852e+09 1.4268861e+09
     1.4268870e+09 1.4268879e+09 1.4268888e+09 1.4268897e+09 1.4268906e+09
     1.4268915e+09 1.4268924e+09 1.4268933e+09 1.4268942e+09 1.4268951e+09
     1.4268960e+09]
    1426814100.0
    3
    Norm_distance:  1.0
    odict_keys(['time', 'ht', 'lat', 'lon', 'electron_density', 'atomic_oxygen_ion_density', 'atomic_hydrogen_ion_density', 'ion_temperature', 'electron_temperature'])
    opening /Users/lrastaet/Kamodo_data/CTIPe/Sujay_Pal_042620_IT_1/2015-03-20-plot-density.nc
    odict_keys(['time', 'plev', 'lat', 'lon', 'doy', 'min', 'density', 'height', 'temperature', 'mean_molecular_mass'])
    odict_keys(['time', 'ht', 'lat', 'lon', 'electron_density', 'atomic_oxygen_ion_density', 'atomic_hydrogen_ion_density', 'ion_temperature', 'electron_temperature'])
    odict_keys(['time', 'plev', 'lat', 'lon', 'height', 'elat', 'elon', 'meridional_neutral_wind', 'zonal_neutral_wind', 'vertical_neutral_wind', 'neutral_temperature', 'mean_molecular_mass', 'electron_density', 'neutral_density', 'solar_heating', 'joule_heating', 'radiation_heat_cool', 'atomic_oxygen_density', 'molecular_oxygen_density', 'molecular_nitrogen_density', 'nitric_oxide_density', 'nitric_oxide_ion_density', 'molecular_nitrogen_ion_density', 'molecular_oxygen_ion_density', 'atomic_nitrogen_ion_density', 'atomic_oxygen_ion_density', 'atomic_hydrogen_ion_density', 'pedersen_conductivity', 'hall_conductivity', 'meridional_ion_velocity', 'zonal_ion_velocity', 'height_integrated_joule_heating', 'energy_influx', 'mean_energy', 'theta_electric_field_at_140km', 'theta_electric_field_at_300km', 'lambda_electric_field_at_140km', 'lambda_electric_field_at_300km', 'total_electron_content'])
    [['temperature', 'height'], [], ['height']]
    temperature
    height
    height
    registered 2 variables

```python
ctipe
```

\begin{equation}\rho{\left(\vec{x} \right)} = \lambda{\left(\vec{x} \right)}\end{equation}\begin{equation}H{\left(\vec{x} \right)} [m] = \lambda{\left(\vec{x} \right)}\end{equation}

```python
ctipe.set_plot('LonLat',cutV=15,latrange=dict(min=-20,max=70,n=31),lonrange=dict(min=0,max=120,n=61))
plot=ctipe.get_plot('H',colorscale="Rainbow")

iplot(plot)
```

    Time resetting plot and precomputing interpolations: 0.0003 seconds

```python
ctipe.set_plot('LatIP',cutV=50,latrange=dict(min=80,max=90),hrange=dict(min=5))
plot=ctipe.get_plot('rho',colorscale="Rainbow")
iplot(plot)
```

    Time resetting plot and precomputing interpolations: 0.0003 seconds
    xrange:  {'min': 80.0, 'max': 90.0, 'n': 31}
    yrange:  {'min': 5.0, 'max': 15.0, 'n': 15}
    aspect ratio: x= 1  y= 1.0

ctipe.set_plot('LonIP')
plot=ctipe.get_plot('rho',colorscale="Rainbow")
iplot(plot)

```python
ctipe.set_plot('LonIP',cutV=0)
plot=ctipe.get_plot('rho',colorscale="Rainbow")
iplot(plot)
```

    Time resetting plot and precomputing interpolations: 0.0003 seconds

```python
print(ctipe['H'].data[0,:,0,0])
```

```python
H=ctipe['H'].data
#hintp=RegularGridInterpolator((ctipe._time,H[0,:,0,0].flatten(),ctipe._lat_density,ctipe._lon_density),H)
hintp=HeightInterpolator((ctipe._time,H[0,:,0,0].flatten(),ctipe._lat_density,ctipe._lon_density),H)

```

```python
print(hintp([ctipe._time[0],125000.,10.3,50.3],method='linear'))

```

```python
## comparison test with IDL in run
## IDL> print,get_pointdata([50.,0.,10.])
## % Compiled module: GET_POINTDATA_3D.
##  5.07806e-10      188.940      19.6322     -6.30557      28.9232      1.83336      595.492  7.52705e+10  1.99925e-05   0.00973094
##    -0.985651  1.08981e+16  1.08948e+14  4.56949e+15  1.68988e+12  1.35428e+10      6169.57  7.52248e+09      29051.7  5.42052e+10
##  1.73583e+06  3.07283e-05  1.89448e-06     -3.32263      2.88644  1.39743e-06      0.00000      0.00000 -0.000311317 -0.000303866
## -0.000346761 -0.000338463      36.2960      0.00000  1.38551e+12      260.000
##IDL> print,qusntities
##% PRINT: Variable is undefined: QUSNTITIES.
##% Execution halted at: $MAIN$              1 /ccmcshare/IDL/3DView/ak3me.pro
##IDL> print,quantities
##rho H Rmt Vn_lat Vn_lon Vn_IP T_n N_e Psolar Pjoule Prad N_O N_O2 N_N2 N_NO N_NO+ N_N2+ N_O2+ N_N+ N_O+ N_H+ sigma_P sigma_H Vi_lat
##Vi_lon Wjoule Win En_avg E140_theta E300_theta E140_lambda E300_lambda TEC_above_140km Vi_IP NmF2 HmF2
##IDL> print,year,month,time,put_time(time)
##        2015           3   9.9999999e-0900:00:00.0000
```

```python
testdata=dict()
testdata['ip']=5.5
testdata['Height']='125'
testdata['date']='2015/3/20'
testdata['time']='00:15:00'
testdata['lon']=50.
testdata['lat']=10.
testdata['rho']='1.49067e-07  1.64051e-08'
testdata['H']='105.230      125.000'
testdata['Rmt']='26.3395      23.9059'
testdata['Vn_lat']='-0.0601084      43.4968'
testdata['Vn_lon']='-23.2656     -3.15513'
testdata['Vn_IP']='-0.0434272   -0.0818311'
testdata['T']='298.432      430.817'  #T_n
testdata['N_e']='1.06982e+08  5.90486e+08'
testdata['Psolar']='1.10576e-06  2.72283e-06'
testdata['Pjoule']='4.78753e-05  0.000398958'
testdata['Prad']=' -0.210089     -1.06214'
testdata['N_O']='4.70951e+17  1.42260e+17'
testdata['N_O2']='3.72280e+17  1.27961e+16'
testdata['N_N2']='2.51143e+18  2.56911e+17'
testdata['N_NO']='6.02975e+13  3.22197e+13'
testdata['N_NOplus']='1.06908e+08  5.04418e+08'
testdata['N_N2plus']='4.83005     0.987384'
testdata['N_O2plus']='74194.3  8.11062e+07'
testdata['N_Nplus']='0.615991     0.795714'
testdata['N_Oplus']='3.41415  4.96166e+06' #N_O+
testdata['N_Hplus']='0.00000      32.4158'
testdata['sigma_P']='9.66247e-07  5.87935e-06'
testdata['sigma_H']='9.17236e-06  8.91819e-06'
testdata['Vi_lat']='-0.723088      40.6683'
testdata['Vi_lon']='-23.9126     -9.70373'
testdata['Wjoule']='5.02903e-07  5.02903e-07'
testdata['Win ']='0.00000      0.00000'
testdata['En_avg']='0.00000      0.00000'
testdata['E_theta140km']='-0.000191656 -0.000191656'
testdata['E_theta300km']='-0.000187069 -0.000187069'
testdata['E_lambda140km']='-0.000340932 -0.000340932'
testdata['E_lambda300km']='-0.000332773 -0.000332773'
testdata['TEC']='40.6796      40.6796'  # 'TEC_above_140km'
testdata['Vi_IP']='0.00000      0.00000'
testdata['NmF2']='1.57999e+12  1.57999e+12'
testdata['HmF2']='300.000      300.000'
```

```python
checked=dict()
pos_ip=[ctipe._time[0],5.5,10,50]
# the Python algortihm calculates IP and uses that for regular interpolation in grid
# resulting in a slightly smaller height value thatn input.
# the required input height value resulting in 125 km in the output is 126.2 km.
# THis shift in altitude / height may vary over the longitude-latitude grid.
# the IDL implemetnation uses the height value and interpolates at the 2 nearest (lon,lat)
# grid stacks using the height values before applying the weights in longitude and latitude
# to arrive at the final interpolatied value for each quantity
# numerical value for efield parameters are different between IDL and Python
# The efield values on the 37x73 grid are re-interpolated to the 91x20 plasma grid, a step not performed in Python
#
#pos_h=[ctipe._time[0],126200.,10,50]
pos_h=[ctipe._time[0],125000.,10,50]
print('pos:    %12.4g %12.4g %12.4g %12.4g' % (pos_ip[0],pos_ip[1],pos_ip[2],pos_ip[3]))
print('pos_h:  %12.4g %12.4g %12.4g %12.4g' % (pos_h[0],pos_h[1],pos_h[2],pos_h[3]))
ilev=ctipe.lev_density(ctipe,t=pos_h[0],z=pos_h[1],lon=pos_h[2],lat=pos_h[3])
#ilev=ctipe.lev_density_log(ctipe,t=pos_h[0],z=pos_h[1],lon=pos_h[2],lat=pos_h[3])
pos_h_ip=pos_h.copy() # lists may be assigned by reference unless using copy()
pos_h_ip[1]=ilev
print(pos_h_ip)
for var in ctipe:
    if var.name not in checked:  # prevent duplication
        checked[var.name]=1   #

        try:
            val_ip=np.squeeze(np.array(ctipe[var](pos_ip[0],pos_ip[1],pos_ip[2],pos_ip[3])))
        except:
            val_ip=np.squeeze(np.array(ctipe[var](pos_ip[0],pos_ip[2],pos_ip[3])))
        try:
            val_h=np.squeeze(np.array(ctipe[var](pos_h[0],ilev,pos_h[2],pos_h[3])))
        except:
            val_h=np.squeeze(np.array(ctipe[var](pos_h[0],pos_h[2],pos_h[3])))
#        try:
        val_h2=np.squeeze(np.array(ctipe.vert_interp(var,pos_h[0],pos_h[1],pos_h[2],pos_h[3])))
#        except:
#            val_h2=np.NaN # np.squeeze(np.array(ctipe[var](pos_h[0],pos_h[2],pos_h[3])))
        if var.name in testdata:
            print("%s %12.4g %12.4g %12.4g IDL: %s" % (var,val_ip,val_h,val_h2,testdata[var.name]))
        else:
            print("%s %12.4g %12.4g  IDL: %s" % (var,val_ip,val_h,"---"))
#          print(var,np.squeeze(np.array(ctipe[var](positions[0],positions[2],positions[3]))))
#            except:
#                val_ip=np.squeeze(np.array(ctipe[var](pos_ip[0],pos_ip[1],pos_ip[2],pos_ip[3])))
#                val_h=np.squeeze(np.array(ctipe[var](positions[0],ilev,positions[2],positions[3])))
#                print(var,val_ip,val_h)
##                print(var,np.squeeze(np.array(ctipe[var]([positions[0],positions[1],positions[2],positions[3]]))))
```

```python
ctipe.lev_density(ctipe,ctipe._time[0],125000.,30.,20.)
```

```python
ctipe._ctipe_density.variables['height'].data.shape
```

```python
print(pos_h)
ctipe.vert_interp('rho',pos_h[0],pos_h[1],pos_h[2],pos_h[3])
```

```python
lon=pos_h[2]
lat=pos_h[3]
lon_rep=np.repeat(lon,len(ctipe._ilev_density))
lat_rep=np.repeat(lat,len(ctipe._ilev_density))
xvec=np.vstack((ctipe._ilev_density,lat_rep,lon_rep)).T
print(lon_rep)
print(lat_rep)
print(ctipe._ilev_density)
print(xvec)
```

```python
z_levels=np.squeeze(np.log(ctipe.H(t = ctipe._time[0], lat = 30., lon = 20.)))
```

```python
ctipe._ilev_neutral
```

```python
from scipy.interpolate import RegularGridInterpolator, interp1d
level = interp1d(z_levels, ctipe._ilev_density, bounds_error=False)
```

```python
z_levels_abs=np.exp(z_levels)
level_abs = interp1d(z_levels_abs, ctipe._ilev_density, bounds_error=False)
```

```python
np.array(np.where(ctipe._lon_density <= 0.)).max()
```

```python
from scipy.interpolate import interp1d
import numpy as np
#@np.vectorize
def ctipe_vert_interp(self,varname,time,z,lat,lon):
    data=self[varname].data
    lons=self._lon_density
    lats=self._lat_density
    times=self._time
    #H=np.log(self['H'].data)
    H=self['H'].data

    ilon=np.array(np.where(lons[0:-1] <= lon)).max()
    if ilon == len(lons):
        ilon=ilon-1
    w0_lon=(lons[ilon+1]-lon)/(lons[ilon+1]-lons[ilon])
    print('ILON=%d W0=%f' %( ilon, w0_lon))

    ilat=np.array(np.where(lats[0:-1] <= lat)).max()
    if ilat == len(lats):
        ilat=ilat-1
    w0_lat=(lats[ilat+1]-lat)/(lats[ilat+1]-lats[ilat])
    print('ILAT=%d W0=%f' %( ilat, w0_lat) )

    itime=np.array(np.where(times <= time)).max()
# need to load next day's data and use itime=0 from there
    if itime == len(times):
        itime=ilime-1
    w0_time=(times[itime+1]-time)/(times[itime+1]-times[itime])
    print('ITIME=%d W0=%f' %( itime, w0_time) )
    h_000=np.squeeze(H[itime,:,ilat,ilon])
    d_000=np.squeeze(data[itime,:,ilat,ilon])
    int000=interp1d(h_000,d_000,bounds_error=False)
    print('Lons: ',lons.shape)
    print('Lats: ',lats.shape)
    print('Height: ',H.shape)
    print('Data:' ,data.shape)

    ilon1=ilon+1
    ilat1=ilat+1
    itime1=itime+1

    w1_lon=1.-w0_lon
    w1_lat=1-w0_lat
    w1_time=1-w0_time

    h_000=np.squeeze(H[itime,:,ilat,ilon])
    d_000=np.squeeze(data[itime,:,ilat,ilon])
    int000=interp1d(h_000,d_000,bounds_error=False)

    h_001=np.squeeze(H[itime,:,ilat,ilon1])
    d_001=np.squeeze(data[itime,:,ilat,ilon+1])
    int001=interp1d(h_000,d_000,bounds_error=False)

    h_010=np.squeeze(H[itime,:,ilat1,ilon])
    d_010=np.squeeze(data[itime,:,ilat+1,ilon])
    int010=interp1d(h_000,d_000,bounds_error=False)

    h_011=np.squeeze(H[itime,:,ilat1,ilon1])
    d_011=np.squeeze(data[itime,:,ilat+1,ilon+1])
    int011=interp1d(h_000,d_000,bounds_error=False)

    h_100=np.squeeze(H[itime+1,:,ilat,ilon])
    d_100=np.squeeze(data[itime+1,:,ilat,ilon])
    int100=interp1d(h_000,d_000,bounds_error=False)

    h_101=np.squeeze(H[itime1,:,ilat,ilon1])
    d_101=np.squeeze(data[itime1,:,ilat,ilon1])
    int101=interp1d(h_000,d_000,bounds_error=False)

    h_110=np.squeeze(H[itime1,:,ilat1,ilon])
    d_110=np.squeeze(data[itime1,:,ilat1,ilon])
    int110=interp1d(h_000,d_000,bounds_error=False)

    h_111=np.squeeze(H[itime1,:,ilat1,ilon1])
    d_111=np.squeeze(data[itime1,:,ilat1,ilon1])
    int111=interp1d(h_000,d_000,bounds_error=False)
    #zlog=np.log(z)
    zlog=z
    return(int000(zlog)*w0_lon*w0_lat*w0_time
        +int001(zlog)*w1_lon*w0_lat*w0_time
        +int010(zlog)*w0_lon*w1_lat*w0_time
        +int011(zlog)*w1_lon*w1_lat*w0_time
        +int100(zlog)*w0_lon*w0_lat*w1_time
        +int101(zlog)*w1_lon*w0_lat*w1_time
        +int110(zlog)*w0_lon*w1_lat*w1_time
        +int111(zlog)*w1_lon*w1_lat*w1_time
          )

```

```python
pos_h=[ctipe._time[0],125000.,10,50]
print(ctipe.vert_interp('H',ctipe._time[0],125000.,10,50.))
```

```python
print(ctipe_vert_interp(ctipe,'rho',ctipe._time[0],125000.,10,50.))

```

```python
H=np.array([[1.,2.,3.],[4.,5.,6.]])
```

```python
print(H.shape)
```

```python

```