```python
%matplotlib inline
```

심화 과정 : Bi-LSTM CRF와 동적 결정
======================================================

동적, 정적 딥 러닝 툴킷(toolkits) 비교
--------------------------------------------

Pytorch는 *동적* 신경망 툴킷입니다. 다른 동적 신경망 툴킷으로는
`Dynet <https://github.com/clab/dynet>`_ 이 있습니다.(이 툴킷을
예로 든 이유는 사용하는 법이 Pytorch와 비슷하기 때문입니다. Dynet의 예제를 보면
Pytorch로 구현할 때도 도움이 될 것입니다.) 반대로 *정적* 툴킷들로
Theano, Keras, TensorFlow 등이 있습니다. 주요 차이점은 다음과 같습니다:

* 정적 툴킷을 쓸 때는 계산 그래프를 한 번만 정의하고, 컴파일 한 후,
  데이터를 계산 그래프에 넘깁니다.
* 동적 툴킷에서는 *각 데이터* 의 계산 그래프를 정의하며 컴파일하지
  않고 즉각 실행됩니다.

경험이 많지 않다면 두 방식의 차이를 알기 어렵습니다. 딥 러닝 기반의
구구조 분석기(constituent parser)를 예로 들어보겠습니다. 모델은 대략
다음과 같은 과정을 수행합니다:

* 트리를 상향식(bottom-up)으로 만들어 나갑니다.
* 최상위 노드를 태깅합니다. (문장의 각 단어)
* 거기서부터 신경망과 단어들의 임베딩을 이용해 구구조를 이루는 조합을
  찾아냅니다. 새로운 구구조를 생성할 때마다 구구조의 임베딩을 얻기 위한
  어떤 기술이 필요합니다. 지금은 신경망이 오직 입력 문장만 참고할
  것입니다. "The green cat scratched the wall"이란 문장에서, 모델의 어느 시점에
  $(i,j,r) = (1, 3, \text{NP})$ 범위 (단어 1에서부터 단어 3까지가
  NP 구구조라는 뜻이며, 이 문장에서는 "The green cat") 를 결합하길 원할
  것입니다.

그런데, 또다른 문장 "Somewhere, the big fat cat scratched the wall" 에서는
어느 시점에 $(2, 4, NP)$ 구구조를 만들기를 원할 것입니다. 우리가
만들기 원하는 구구조들은 문장에 따라 다릅니다. 만약 정적 툴킷에서처럼
계산 그래프를 한 번만 컴파일한다면, 이 과정을 프로그래밍하기 매우 어렵거나
불가능할 것입니다. 하지만 동적 툴킷에서는 하나의 계산 그래프만 있지
않습니다. 각 문장들마다 새로운 계산 그래프가 있을 수 있기 때문에 이런
문제가 없습니다.

동적 틀킷은 디버깅 하기 더 쉽고, 코드가 기반 언어와 더 비슷합니다
(Pytorch와 Dynet이 Keras 또는 Theano 보다 Python 코드와 더 비슷합니다).

Bi-LSTM Conditional Rnadom Field 설명
-------------------------------------------

이 영역에서는 개체명 인식을 수행하는 완성된 Bi-LSTM Conditional Random
Field 예시를 살펴보겠습니다. 위에 나온 LSTM 태거(tagger)는 일반적으로
품사 태깅을 하기에 충분합니다. 하지만 CRF 같은 연속된 데이터를 다루는
모델은 좋은 개체명 인식 모델(NER)에 꼭 필요합니다. 여러분이 CRF를 잘 알고
있다고 가정하겠습니다. 이름이 무섭게 들릴 수도 있지만, LSTM이 특징을
제공하는 점을 제외하면 이 모델은 CRF 입니다. 하지만 더 발전된 모델이며,
이 튜토리얼의 앞부분에 나왔던 모델보다 훨씬 복잡합니다. 넘어가고 싶다면
넘어가도 괜찮습니다. 이해할 수 있다고 생각한다면, 아래를 읽어보세요:

-  태그 k에 대한 i번째 단계의 비터비(viterbi) 변수를 위해 순환 흐름을 만든다.
-  순방향 변수를 계산하기 위해 위의 순한 흐름을 조정한다.
-  순방향 변수를 로그 공간에서 계산하기 위해 다시 한 번 조정한다.
   (힌트 : 로그-합-지수승)

위의 세가지를 할 수 있다면, 아래의 코드를 이해할 수 있을 것입니다.
CRF는 조건부 확률을 계산한다는 점을 기억하세요. $y$ 를 연속된
태그라 하고, $x$ 를 연속된 입력 단어라 하겠습니다. 그러면 아래의
식을 계산할 수 있습니다.

\begin{align}P(y|x) = \frac{\exp{(\text{Score}(x, y)})}{\sum_{y'} \exp{(\text{Score}(x, y')})}\end{align}

점수(score) 함수는 아래와 같이 정의된 로그 포텐셜(potential) $\log \psi_i(x,y)$
함수에 의해 결정됩니다.

\begin{align}\text{Score}(x,y) = \sum_i \log \psi_i(x,y)\end{align}

분배 함수(partition function)를 단순화하기 위해서, 포텐셜이 주변의
특징들만 반영한다고 하겠습니다.

Bi-LSTM CRF 안에 배출(emission), 전이(transition) 두 종류의 포텐셜을
정의합니다. $i$ 번째 단어에 대한 배출 포텐셜은 Bi-LSTM의
$i$ 번째 시점의 은닉 상태가 결정합니다. 전이 점수는 $|T|x|T|$
형태인 행렬 $\textbf{P}$ 에 저장되어 있습니다. $T$ 는
태그의 집합입니다. 이 구현에서, $\textbf{P}_{j,k}$ 는 tag $j$ 에서
tag $k$ 로의 전이 점수를 의미합니다. 따라서:

\begin{align}\text{Score}(x,y) = \sum_i \log \psi_\text{EMIT}(y_i \rightarrow x_i) + \log \psi_\text{TRANS}(y_{i-1} \rightarrow y_i)\end{align}

\begin{align}= \sum_i h_i[y_i] + \textbf{P}_{y_i, y_{i-1}}\end{align}

두 번째 식에서 고유하고 음수가 아닌 인덱스에 의해 태그가 부여됐다고
간주합니다.

위의 설명이 너무 간단하다고 생각한다면 CRF에 대한 Michael Collins의
글을 `여기 <http://www.cs.columbia.edu/%7Emcollins/crf.pdf>`__ 에서
읽어보세요.

구현 문서
--------------------

아래의 예시는 로그 공간에서 분배 함수를 계산하기 위한 순방향 알고리즘과
복호화하기 위한 비터비 알고리즘을 구현한 것입니다. 역전파
단계에서 변화도는 자동으로 계산될 것입니다. 우리가 직접 할 일은
없습니다.

이 구현은 최적의 상태가 아닙니다. 과정을 이해했다면, 순방향 알고리즘
상에서 다음 태그를 순차적으로 처리하는 과정을 하나의 큰 연산으로 줄일
수 있다는 것을 아마 빠르게 알 수 있을 것입니다. 이 코드는 가능한 읽기
쉽게 작성했습니다. 적절하게 수정하면, 이 태거를 실제 문제들에 사용할
수도 있을 것입니다.

```python
# 작성자: Robert Guthrie

import torch
import torch.autograd as autograd
import torch.nn as nn
import torch.optim as optim

torch.manual_seed(1)
```

코드 가독성을 높여주는 보조 함수들

```python
def argmax(vec):
    # argmax를 파이썬 정수형으로 반환합니다.
    _, idx = torch.max(vec, 1)
    return idx.item()

def prepare_sequence(seq, to_ix):
    idxs = [to_ix[w] for w in seq]
    return torch.tensor(idxs, dtype=torch.long)

# 순방향 알고리즘을 위해 수치적으로 안정적인 방법으로 로그 합 지수승을 계산합니다.
def log_sum_exp(vec):
    max_score = vec[0, argmax(vec)]
    max_score_broadcast = max_score.view(1, -1).expand(1, vec.size()[1])
    return max_score + \
        torch.log(torch.sum(torch.exp(vec - max_score_broadcast)))
```

모델 생성

```python
class BiLSTM_CRF(nn.Module):

    def __init__(self, vocab_size, tag_to_ix, embedding_dim, hidden_dim):
        super(BiLSTM_CRF, self).__init__()
        self.embedding_dim = embedding_dim
        self.hidden_dim = hidden_dim
        self.vocab_size = vocab_size
        self.tag_to_ix = tag_to_ix
        self.tagset_size = len(tag_to_ix)

        self.word_embeds = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim // 2,
                            num_layers=1, bidirectional=True)

        # LSTM의 출력을 태그 공간으로 대응시킵니다.
        self.hidden2tag = nn.Linear(hidden_dim, self.tagset_size)

        # 전이 매개변수 행렬. i, j 성분은 i에서 j로 변할 때의 점수입니다.
        self.transitions = nn.Parameter(
            torch.randn(self.tagset_size, self.tagset_size))

        # 이 두 코드는 시작 태그로 전이하지 않고, 정지 태그에서부터
        # 전이하지 않도록 강제합니다.
        self.transitions.data[tag_to_ix[START_TAG], :] = -10000
        self.transitions.data[:, tag_to_ix[STOP_TAG]] = -10000

        self.hidden = self.init_hidden()

    def init_hidden(self):
        return (torch.randn(2, 1, self.hidden_dim // 2),
                torch.randn(2, 1, self.hidden_dim // 2))

    def _forward_alg(self, feats):
        # 분배 함수를 계산하기 위해 순방향 알고리즘을 수행합니다.
        init_alphas = torch.full((1, self.tagset_size), -10000.)
        # START_TAG는 모든 점수를 갖고 있습니다.
        init_alphas[0][self.tag_to_ix[START_TAG]] = 0.

        # 자동으로 역전파 되도록 변수로 감쌉니다.
        forward_var = init_alphas

        # 문장의 각 성분을 반복 처리합니다.
        for feat in feats:
            alphas_t = []  # 현재 시점의 순방향 텐서
            for next_tag in range(self.tagset_size):
                # 이전의 태그와 상관없이 배출 점수를 전파합니다.
                emit_score = feat[next_tag].view(
                    1, -1).expand(1, self.tagset_size)
                # trans_score의 i번째 성분은 i로부터 next_tag로 전이할 점수입니다.
                trans_score = self.transitions[next_tag].view(1, -1)
                # next_tag_var의 i번째 성분은 로그-합-지수승을 계산하기 전
                # i에서 next_tag로 가는 간선의 값입니다.
                next_tag_var = forward_var + trans_score + emit_score
                # 이 태그의 순방향 변수는 모든 점수들의 로그-합-지수승 계산
                # 결과입니다.
                alphas_t.append(log_sum_exp(next_tag_var).view(1))
            forward_var = torch.cat(alphas_t).view(1, -1)
        terminal_var = forward_var + self.transitions[self.tag_to_ix[STOP_TAG]]
        alpha = log_sum_exp(terminal_var)
        return alpha

    def _get_lstm_features(self, sentence):
        self.hidden = self.init_hidden()
        embeds = self.word_embeds(sentence).view(len(sentence), 1, -1)
        lstm_out, self.hidden = self.lstm(embeds, self.hidden)
        lstm_out = lstm_out.view(len(sentence), self.hidden_dim)
        lstm_feats = self.hidden2tag(lstm_out)
        return lstm_feats

    def _score_sentence(self, feats, tags):
        # 주어진 태그 순열에 점수를 매깁니다.
        score = torch.zeros(1)
        tags = torch.cat([torch.tensor([self.tag_to_ix[START_TAG]], dtype=torch.long), tags])
        for i, feat in enumerate(feats):
            score = score + \
                self.transitions[tags[i + 1], tags[i]] + feat[tags[i + 1]]
        score = score + self.transitions[self.tag_to_ix[STOP_TAG], tags[-1]]
        return score

    def _viterbi_decode(self, feats):
        backpointers = []

        # 비터비 변수를 로그 공간 상에 초기화합니다.
        init_vvars = torch.full((1, self.tagset_size), -10000.)
        init_vvars[0][self.tag_to_ix[START_TAG]] = 0

        # i 단계의 forward_var는 i-1 단계의 비터비 변수를 갖고 있습니다.
        forward_var = init_vvars
        for feat in feats:
            bptrs_t = []  # 현재 단계의 backpointer를 갖고 있습니다.
            viterbivars_t = []  # 현재 단계의 비터비 변수를 갖고 있습니다.

            for next_tag in range(self.tagset_size):
                # next_tag_var[i]는 이전 단계의 태그 i에 대한 비터비 변수와,
                # 태그 i에서 next_tag로 전이할 점수를 더한 값을 갖고 있습니다.
                # 배출 점수는 argmax와 상관 없기 때문에(아래 코드에서 추가할 것입니다)
                # 여기에 포함되지 않습니다.
                next_tag_var = forward_var + self.transitions[next_tag]
                best_tag_id = argmax(next_tag_var)
                bptrs_t.append(best_tag_id)
                viterbivars_t.append(next_tag_var[0][best_tag_id].view(1))
            # 이제 배출 점수를 더합니다. 그리고 방금 계산한 비터비 변수의
            # 집합을 forward_var에 할당합니다.
            forward_var = (torch.cat(viterbivars_t) + feat).view(1, -1)
            backpointers.append(bptrs_t)

        # STAP_TAG로의 전이
        terminal_var = forward_var + self.transitions[self.tag_to_ix[STOP_TAG]]
        best_tag_id = argmax(terminal_var)
        path_score = terminal_var[0][best_tag_id]

        # 최적의 경로를 구하기 위해 back pointer를 따라갑니다.
        best_path = [best_tag_id]
        for bptrs_t in reversed(backpointers):
            best_tag_id = bptrs_t[best_tag_id]
            best_path.append(best_tag_id)
        # 시작 태그를 빼냅니다 (시작 태그는 반환된 필요가 없습니다)
        start = best_path.pop()
        assert start == self.tag_to_ix[START_TAG]  # 완결성 검사 (Sanity check)
        best_path.reverse()
        return path_score, best_path

    def neg_log_likelihood(self, sentence, tags):
        feats = self._get_lstm_features(sentence)
        forward_score = self._forward_alg(feats)
        gold_score = self._score_sentence(feats, tags)
        return forward_score - gold_score

    def forward(self, sentence):  # 이 함수와 위의 _forward_alg를 헷갈리지 마세요.
        # Bi-LSTM으로부터 배출 점수를 얻습니다.
        lstm_feats = self._get_lstm_features(sentence)

        # 주어진 특징(배출 점수)들로 최적의 경로를 찾아냅니다.
        score, tag_seq = self._viterbi_decode(lstm_feats)
        return score, tag_seq
```

훈련 실행

```python
START_TAG = "<START>"
STOP_TAG = "<STOP>"
EMBEDDING_DIM = 5
HIDDEN_DIM = 4

# 훈련용 데이터를 만듭니다.
training_data = [(
    "the wall street journal reported today that apple corporation made money".split(),
    "B I I I O O O B I O O".split()
), (
    "georgia tech is a university in georgia".split(),
    "B I O O O O B".split()
)]

word_to_ix = {}
for sentence, tags in training_data:
    for word in sentence:
        if word not in word_to_ix:
            word_to_ix[word] = len(word_to_ix)

tag_to_ix = {"B": 0, "I": 1, "O": 2, START_TAG: 3, STOP_TAG: 4}

model = BiLSTM_CRF(len(word_to_ix), tag_to_ix, EMBEDDING_DIM, HIDDEN_DIM)
optimizer = optim.SGD(model.parameters(), lr=0.01, weight_decay=1e-4)

# 훈련 전 예측 결과를 확인합니다.
with torch.no_grad():
    precheck_sent = prepare_sequence(training_data[0][0], word_to_ix)
    precheck_tags = torch.tensor([tag_to_ix[t] for t in training_data[0][1]], dtype=torch.long)
    print(model(precheck_sent))

# Make sure prepare_sequence from earlier in the LSTM section is loaded
# 위의 보조 함수 영역에 있는 prepare_sequence 함수가 불러와 졌는지 확인합니다.
for epoch in range(
        300):  # 다시 말하지만, 아마 300 에폭을 실행하진 않을 것입니다. 이것은 연습용 데이터입니다.
    for sentence, tags in training_data:
        # 1단계. Pytorch가 변화도를 누적한다는 것을 기억하세요.
        # 그것들을 제거합니다.
        model.zero_grad()

        # 2단계. 입력 데이터를 신경망에 사용될 수 있도록 단어
        # 인덱스들의 텐서로 변환합니다.
        sentence_in = prepare_sequence(sentence, word_to_ix)
        targets = torch.tensor([tag_to_ix[t] for t in tags], dtype=torch.long)

        # 3단계. 순방향 계산을 수행합니다.
        loss = model.neg_log_likelihood(sentence_in, targets)

        # 4단계. 손실값, 변화도를 계산하고 optimizer.step()을 호출하여
        # 매개변수들을 갱신합니다.
        loss.backward()
        optimizer.step()

# 훈련이 끝난 후 예측 결과를 확인합니다.
with torch.no_grad():
    precheck_sent = prepare_sequence(training_data[0][0], word_to_ix)
    print(model(precheck_sent))
# 다 했습니다!
```

연습 : 판별적(discriminative) 태깅을 위한 새로운 손실 함수
------------------------------------------------------------

사실 복호화 할 때는 비터비 경로 점수로 역전파를 하지 않기 때문에 계산
그래프를 만들 필요가 없었습니다. 그러나 이미 만들었으니, 비터비 경로
점수와 실제 정답 경로 점수의 차이를 손실 함수로 사용해서 태거를
학습시켜 보세요. 손실 함수의 값은 음수가 아니어야 하며, 예측된 태그
순열이 정답이라면 손실 함수의 값은 0이어야 합니다. 이것은 본질적으로
*구조화된 퍼셉트론* 입니다.

이미 비터비와 score_sentence 함수가 구현되어 있기 때문에 간단히 수정할
수 있습니다. 이 모델은 *학습 데이터에 따라 변하는* 계산 그래프의 한
예시입니다. 이 모델을 정적 툴킷에서 구현해 보지는 않았는데, 구현이
가능하지만 덜 직관적일 수 있습니다.

실제 데이터를 사용해보고 비교해보세요!


# Question duplicates

- In this, you will use Siamese networks to natural language processing to build a classifier that will allow you to identify whether two questions asked are same or not.

- You will process the data first and then pad in a similar way you have done in the previous assignment. Your model will take in the two question embeddings, run them through an LSTM, and then compare the outputs of the two sub networks using cosine similarity.

# Part 1: Importing the Data

### 1.1 Loading in the data

You will be using the Quora question answer dataset to build a model that could identify similar questions. This is a useful task because you don't want to have several versions of the same question posted.

```python
import os
import nltk
import trax
from trax import layers as tl
from trax.supervised import training
from trax.fastmath import numpy as fastnp
import numpy as np
import pandas as pd
import random as rnd

# set random seeds
trax.supervised.trainer_lib.init_random_number_generators(34)
rnd.seed(34)
```

    INFO:tensorflow:tokens_length=568 inputs_length=512 targets_length=114 noise_density=0.15 mean_noise_span_length=3.0

```python
data = pd.read_csv("questions.csv")
N=len(data)
print('Number of question pairs: ', N)
data.head()
```

    Number of question pairs:  404351

      id
      qid1
      qid2
      question1
      question2
      is_duplicate

      0
      0
      1
      2
      What is the step by step guide to invest in sh...
      What is the step by step guide to invest in sh...
      0

      1
      1
      3
      4
      What is the story of Kohinoor (Koh-i-Noor) Dia...
      What would happen if the Indian government sto...
      0

      2
      2
      5
      6
      How can I increase the speed of my internet co...
      How can Internet speed be increased by hacking...
      0

      3
      3
      7
      8
      Why am I mentally very lonely? How can I solve...
      Find the remainder when [math]23^{24}[/math] i...
      0

      4
      4
      9
      10
      Which one dissolve in water quikly sugar, salt...
      Which fish would survive in salt water?
      0

We first split the data into a train and test set. The test set will be used later to evaluate our model.

```python
N_train = 300000
N_test  = 10*1024
data_train = data[:N_train]
data_test  = data[N_train:N_train+N_test]
print("Train set:", len(data_train), "Test set:", len(data_test))
del(data) # remove to free memory
```

    Train set: 300000 Test set: 10240

We select only the question pairs that are duplicate to train the model.

We build two batches as input for the Siamese network and we assume that question $q1_i$ (question $i$ in the first batch) is a duplicate of $q2_i$ (question $i$ in the second batch), but all other questions in the second batch are not duplicates of $q1_i$.
The test set uses the original pairs of questions and the status describing if the questions are duplicates.

```python
td_index = (data_train['is_duplicate'] == 1).to_numpy()
td_index = [i for i, x in enumerate(td_index) if x]
print('number of duplicate questions: ', len(td_index))
print('indexes of first ten duplicate questions:', td_index[:10])
```

    number of duplicate questions:  111486
    indexes of first ten duplicate questions: [5, 7, 11, 12, 13, 15, 16, 18, 20, 29]

```python
print(data_train['question1'][5])  #  Example of question duplicates (first one in data)
print(data_train['question2'][5])
print('is_duplicate: ', data_train['is_duplicate'][5])
```

    Astrology: I am a Capricorn Sun Cap moon and cap rising...what does that say about me?
    I'm a triple Capricorn (Sun, Moon and ascendant in Capricorn) What does this say about me?
    is_duplicate:  1

```python
Q1_train_words = np.array(data_train['question1'][td_index])
Q2_train_words = np.array(data_train['question2'][td_index])

Q1_test_words = np.array(data_test['question1'])
Q2_test_words = np.array(data_test['question2'])
y_test  = np.array(data_test['is_duplicate'])
```

```python
print('TRAINING QUESTIONS:\n')
print('Question 1: ', Q1_train_words[0])
print('Question 2: ', Q2_train_words[0], '\n')
print('Question 1: ', Q1_train_words[5])
print('Question 2: ', Q2_train_words[5], '\n')

print('TESTING QUESTIONS:\n')
print('Question 1: ', Q1_test_words[0])
print('Question 2: ', Q2_test_words[0], '\n')
print('is_duplicate =', y_test[0], '\n')
```

    TRAINING QUESTIONS:

    Question 1:  Astrology: I am a Capricorn Sun Cap moon and cap rising...what does that say about me?
    Question 2:  I'm a triple Capricorn (Sun, Moon and ascendant in Capricorn) What does this say about me?

    Question 1:  What would a Trump presidency mean for current international master’s students on an F1 visa?
    Question 2:  How will a Trump presidency affect the students presently in US or planning to study in US?

    TESTING QUESTIONS:

    Question 1:  How do I prepare for interviews for cse?
    Question 2:  What is the best way to prepare for cse?

    is_duplicate = 0

```python
#create arrays
Q1_train = np.empty_like(Q1_train_words)
Q2_train = np.empty_like(Q2_train_words)

Q1_test = np.empty_like(Q1_test_words)
Q2_test = np.empty_like(Q2_test_words)
```

```python
# Building the vocabulary with the train set
from collections import defaultdict

vocab = defaultdict(lambda: 0)
vocab['<PAD>'] = 1

for idx in range(len(Q1_train_words)):
    Q1_train[idx] = nltk.word_tokenize(Q1_train_words[idx])
    Q2_train[idx] = nltk.word_tokenize(Q2_train_words[idx])
    q = Q1_train[idx] + Q2_train[idx]
    for word in q:
        if word not in vocab:
            vocab[word] = len(vocab) + 1
print('The length of the vocabulary is: ', len(vocab))
```

    The length of the vocabulary is:  36268

```python
print(vocab['<PAD>'])
print(vocab['Astrology'])
print(vocab['Astronomy'])  #not in vocabulary, returns 0
```

    1
    2
    0

```python
for idx in range(len(Q1_test_words)):
    Q1_test[idx] = nltk.word_tokenize(Q1_test_words[idx])
    Q2_test[idx] = nltk.word_tokenize(Q2_test_words[idx])
```

```python
print('Train set has reduced to: ', len(Q1_train) )
print('Test set length: ', len(Q1_test) )
```

    Train set has reduced to:  111486
    Test set length:  10240

### 1.2 Converting a question to a tensor

```python
# Converting questions to array of integers
for i in range(len(Q1_train)):
    Q1_train[i] = [vocab[word] for word in Q1_train[i]]
    Q2_train[i] = [vocab[word] for word in Q2_train[i]]

for i in range(len(Q1_test)):
    Q1_test[i] = [vocab[word] for word in Q1_test[i]]
    Q2_test[i] = [vocab[word] for word in Q2_test[i]]
```

```python
print('first question in the train set:\n')
print(Q1_train_words[0], '\n')
print('encoded version:')
print(Q1_train[0],'\n')

print('first question in the test set:\n')
print(Q1_test_words[0], '\n')
print('encoded version:')
print(Q1_test[0])
```

    first question in the train set:

    Astrology: I am a Capricorn Sun Cap moon and cap rising...what does that say about me?

    encoded version:
    [2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21]

    first question in the test set:

    How do I prepare for interviews for cse?

    encoded version:
    [32, 38, 4, 107, 65, 1015, 65, 11509, 21]

You will now split your train set into a training/validation set so that you can use it to train and evaluate your Siamese model.

```python
# Splitting the data
cut_off = int(len(Q1_train)*.8)
train_Q1, train_Q2 = Q1_train[:cut_off], Q2_train[:cut_off]
val_Q1, val_Q2 = Q1_train[cut_off: ], Q2_train[cut_off:]
print('Number of duplicate questions: ', len(Q1_train))
print("The length of the training set is:  ", len(train_Q1))
print("The length of the validation set is: ", len(val_Q1))
```

    Number of duplicate questions:  111486
    The length of the training set is:   89188
    The length of the validation set is:  22298

### 1.3 Understanding the iterator

Most of the time in Natural Language Processing, and AI in general we use batches when training our data sets. If you were to use stochastic gradient descent with one example at a time, it will take you forever to build a model. In this, we build a data generator that takes in $Q1$ and $Q2$ and returns a batch of size `batch_size`  in the following format $([q1_1, q1_2, q1_3, ...]$, $[q2_1, q2_2,q2_3, ...])$. The tuple consists of two arrays and each array has `batch_size` questions. Again, $q1_i$ and $q2_i$ are duplicates, but they are not duplicates with any other elements in the batch.

The command ```next(data_generator)```returns the next batch. This iterator returns the data in a format that you could directly use in your model when computing the feed-forward of your algorithm. This iterator returns a pair of arrays of questions.

```python
def data_generator(Q1, Q2, batch_size, pad=1, shuffle=True):
    """Generator function that yields batches of data

    Args:
        Q1 (list): List of transformed (to tensor) questions.
        Q2 (list): List of transformed (to tensor) questions.
        batch_size (int): Number of elements per batch.
        pad (int, optional): Pad character from the vocab. Defaults to 1.
        shuffle (bool, optional): If the batches should be randomnized or not. Defaults to True.
    Yields:
        tuple: Of the form (input1, input2) with types (numpy.ndarray, numpy.ndarray)
        NOTE: input1: inputs to your model [q1a, q2a, q3a, ...] i.e. (q1a,q1b) are duplicates
              input2: targets to your model [q1b, q2b,q3b, ...] i.e. (q1a,q2i) i!=a are not duplicates
    """

    input1 = []
    input2 = []
    idx = 0
    len_q = len(Q1)
    question_indexes = [*range(len_q)]

    if shuffle:
        rnd.shuffle(question_indexes)

    ### START CODE HERE (Replace instances of 'None' with your code) ###
    while True:
        if idx >= len_q:
            # if idx is greater than or equal to len_q, set idx accordingly
            # (Hint: look at the instructions above)
            idx = 0
            # shuffle to get random batches if shuffle is set to True
            if shuffle:
                rnd.shuffle(question_indexes)

        # get questions at the `question_indexes[idx]` position in Q1 and Q2
        q1 = Q1[question_indexes[idx]]
        q2 = Q2[question_indexes[idx]]

        # increment idx by 1
        idx += 1
        # append q1
        input1.append(q1)
        # append q2
        input2.append(q2)
        if len(input1) == batch_size:
            # determine max_len as the longest question in input1 & input 2
            max_len = max(len(max(input1, key=len)), len(max(input2, key=len)))
            # pad to power-of-2 (Hint: look at the instructions above)
            max_len = 2**int(np.ceil(np.log2(max_len)))
            b1 = []
            b2 = []
            for q1, q2 in zip(input1, input2):
                # add [pad] to q1 until it reaches max_len
                q1 = q1 + [pad] * (max_len - len(q1))
                # add [pad] to q2 until it reaches max_len
                q2 = q2 + [pad] * (max_len - len(q2))
                # append q1
                b1.append(q1)
                # append q2
                b2.append(q2)
            # use b1 and b2
            yield np.array(b1), np.array(b2)
    ### END CODE HERE ###
            # reset the batches
            input1, input2 = [], []  # reset the batches
```

```python
batch_size = 2
res1, res2 = next(data_generator(train_Q1, train_Q2, batch_size))
print("First questions  : ",'\n', res1, '\n')
print("Second questions : ",'\n', res2)
```

    First questions  :
     [[  30   87   78  134 2132 1981   28   78  594   21    1    1    1    1
         1    1]
     [  30   55   78 3541 1460   28   56  253   21    1    1    1    1    1
         1    1]]

    Second questions :
     [[  30  156   78  134 2132 9508   21    1    1    1    1    1    1    1
         1    1]
     [  30  156   78 3541 1460  131   56  253   21    1    1    1    1    1
         1    1]]

# Part 2: Defining the Siamese model

### 2.1 Understanding Siamese Network
A Siamese network is a neural network which uses the same weights while working in tandem on two different input vectors to compute comparable output vectors.The Siamese network you are about to implement looks like this:

You get the question embedding, run it through an LSTM layer, normalize $v_1$ and $v_2$, and finally use a triplet loss (explained below) to get the corresponding cosine similarity for each pair of questions. As usual, you will start by importing the data set. The triplet loss makes use of a baseline (anchor) input that is compared to a positive (truthy) input and a negative (falsy) input. The distance from the baseline (anchor) input to the positive (truthy) input is minimized, and the distance from the baseline (anchor) input to the negative (falsy) input is maximized. In math equations, you are trying to maximize the following.

$$\mathcal{L}(A, P, N)=\max \left(\|\mathrm{f}(A)-\mathrm{f}(P)\|^{2}-\|\mathrm{f}(A)-\mathrm{f}(N)\|^{2}+\alpha, 0\right)$$

$A$ is the anchor input, for example $q1_1$, $P$ the duplicate input, for example, $q2_1$, and $N$ the negative input (the non duplicate question), for example $q2_2$.

$\alpha$ is a margin; you can think about it as a safety net, or by how much you want to push the duplicates from the non duplicates.

```python
def Siamese(vocab_size=len(vocab), d_model=128, mode='train'):
    """Returns a Siamese model.

    Args:
        vocab_size (int, optional): Length of the vocabulary. Defaults to len(vocab).
        d_model (int, optional): Depth of the model. Defaults to 128.
        mode (str, optional): 'train', 'eval' or 'predict', predict mode is for fast inference. Defaults to 'train'.

    Returns:
        trax.layers.combinators.Parallel: A Siamese model.
    """

    def normalize(x):  # normalizes the vectors to have L2 norm 1
        return x / fastnp.sqrt(fastnp.sum(x * x, axis=-1, keepdims=True))

    ### START CODE HERE (Replace instances of 'None' with your code) ###
    q_processor = tl.Serial(  # Processor will run on Q1 and Q2.
        tl.Embedding(vocab_size=vocab_size, d_feature=d_model), # Embedding layer
        tl.LSTM(n_units=d_model), # LSTM layer
        tl.Mean(axis=1), # Mean over columns
        tl.Fn('Normalize', lambda x: normalize(x))  # Apply normalize function
    )  # Returns one vector of shape [batch_size, d_model].

    ### END CODE HERE ###

    # Run on Q1 and Q2 in parallel.
    model = tl.Parallel(q_processor, q_processor)
    return model

```

Setup the Siamese network model

```python
# check your model
model = Siamese()
print(model)
```

    Parallel_in2_out2[
      Serial[
        Embedding_41699_128
        LSTM_128
        Mean
        Normalize
      ]
      Serial[
        Embedding_41699_128
        LSTM_128
        Mean
        Normalize
      ]
    ]

### 2.2 Hard  Negative Mining

You will now implement the `TripletLoss`.

TripletLoss is composed of two terms. One term utilizes the mean of all the non duplicates, the second utilizes the *closest negative*. Our loss expression is then:

\begin{align}
 \mathcal{Loss_1(A,P,N)} &=\max \left( -cos(A,P)  + mean_{neg} +\alpha, 0\right) \\
 \mathcal{Loss_2(A,P,N)} &=\max \left( -cos(A,P)  + closest_{neg} +\alpha, 0\right) \\
\mathcal{Loss(A,P,N)} &= mean(Loss_1 + Loss_2) \\
\end{align}

```python
def TripletLossFn(v1, v2, margin=0.25):
    """Custom Loss function.

    Args:
        v1 (numpy.ndarray): Array with dimension (batch_size, model_dimension) associated to Q1.
        v2 (numpy.ndarray): Array with dimension (batch_size, model_dimension) associated to Q2.
        margin (float, optional): Desired margin. Defaults to 0.25.

    Returns:
        jax.interpreters.xla.DeviceArray: Triplet Loss.
    """
    ### START CODE HERE (Replace instances of 'None' with your code) ###

    # use fastnp to take the dot product of the two batches (don't forget to transpose the second argument)
    scores = fastnp.dot(v1, v2.T)  # pairwise cosine sim
    # calculate new batch size
    batch_size = len(scores)
    # use fastnp to grab all postive `diagonal` entries in `scores`
    positive = fastnp.diagonal(scores)  # the positive ones (duplicates)
    # multiply `fastnp.eye(batch_size)` with 2.0 and subtract it out of `scores`
    negative_without_positive = scores - fastnp.eye(batch_size) * 2.0
    # take the row by row `max` of `negative_without_positive`.
    # Hint: negative_without_positive.max(axis = [?])
    closest_negative = negative_without_positive.max(axis=1)
    # subtract `fastnp.eye(batch_size)` out of 1.0 and do element-wise multiplication with `scores`
    negative_zero_on_duplicate = fastnp.multiply((1.0 - fastnp.eye(batch_size)), scores)
    # use `fastnp.sum` on `negative_zero_on_duplicate` for `axis=1` and divide it by `(batch_size - 1)`
    mean_negative = fastnp.sum(negative_zero_on_duplicate, axis=1) / (batch_size - 1)
    # compute `fastnp.maximum` among 0.0 and `A`
    # A = subtract `positive` from `margin` and add `closest_negative`
    triplet_loss1 = fastnp.maximum(margin - positive + closest_negative, 0)
    # compute `fastnp.maximum` among 0.0 and `B`
    # B = subtract `positive` from `margin` and add `mean_negative`
    triplet_loss2 = fastnp.maximum(margin - positive + mean_negative, 0)
    # add the two losses together and take the `fastnp.mean` of it
    triplet_loss = fastnp.mean(triplet_loss1 + triplet_loss2)

    ### END CODE HERE ###

    return triplet_loss
```

```python
v1 = np.array([[0.26726124, 0.53452248, 0.80178373],[0.5178918 , 0.57543534, 0.63297887]])
v2 = np.array([[ 0.26726124,  0.53452248,  0.80178373],[-0.5178918 , -0.57543534, -0.63297887]])
TripletLossFn(v2,v1)
print("Triplet Loss:", TripletLossFn(v2,v1))
```

    Triplet Loss: 0.5

```python
from functools import partial
def TripletLoss(margin=0.25):
    triplet_loss_fn = partial(TripletLossFn, margin=margin)
    return tl.Fn('TripletLoss', triplet_loss_fn)
```

# Part 3: Training

```python
batch_size = 256
train_generator = data_generator(train_Q1, train_Q2, batch_size, vocab['<PAD>'])
val_generator = data_generator(val_Q1, val_Q2, batch_size, vocab['<PAD>'])
print('train_Q1.shape ', train_Q1.shape)
print('val_Q1.shape   ', val_Q1.shape)
```

    train_Q1.shape  (89188,)
    val_Q1.shape    (22298,)

```python
lr_schedule = trax.lr.warmup_and_rsqrt_decay(400, 0.01)

def train_model(Siamese, TripletLoss, lr_schedule, train_generator=train_generator, val_generator=val_generator, output_dir='model/'):
    """Training the Siamese Model

    Args:
        Siamese (function): Function that returns the Siamese model.
        TripletLoss (function): Function that defines the TripletLoss loss function.
        lr_schedule (function): Trax multifactor schedule function.
        train_generator (generator, optional): Training generator. Defaults to train_generator.
        val_generator (generator, optional): Validation generator. Defaults to val_generator.
        output_dir (str, optional): Path to save model to. Defaults to 'model/'.

    Returns:
        trax.supervised.training.Loop: Training loop for the model.
    """
    output_dir = os.path.expanduser(output_dir)

    ### START CODE HERE (Replace instances of 'None' with your code) ###

    train_task = training.TrainTask(
        labeled_data=train_generator,       # Use generator (train)
        loss_layer=TripletLoss(),         # Use triplet loss. Don't forget to instantiate this object
        optimizer=trax.optimizers.Adam(learning_rate=0.01),          # Don't forget to add the learning rate parameter
        lr_schedule=lr_schedule, # Use Trax multifactor schedule function
    )

    eval_task = training.EvalTask(
        labeled_data=val_generator,       # Use generator (val)
        metrics=[TripletLoss()],          # Use triplet loss. Don't forget to instantiate this object
    )

    ### END CODE HERE ###

    training_loop = training.Loop(Siamese(),
                                  train_task,
                                  eval_task=eval_task,
                                  output_dir=output_dir)

    return training_loop
```

```python
train_steps = 5
training_loop = train_model(Siamese, TripletLoss, lr_schedule)
training_loop.run(train_steps)
```

    Step      1: train TripletLoss |  0.49954823
    Step      1: eval  TripletLoss |  0.49950948

The model was only trained for 5 steps due to the constraints of this environment.

# Part 4:  Evaluation

### 4.1 Evaluating your siamese network

You will first start by loading a pretrained model and then you will use it to predict.

```python
# Loading in the saved model
model = Siamese()
model.init_from_file('model.pkl.gz')
```

### 4.2 Classify
To determine the accuracy of the model, we will utilize the test set that was configured earlier. While in training we used only positive examples, the test data, Q1_test, Q2_test and y_test, is setup as pairs of questions, some of which are duplicates some are not.
This routine will run all the test question pairs through the model, compute the cosine simlarity of each pair, threshold it and compare the result to  y_test - the correct response from the data set. The results are accumulated to produce an accuracy.

```python
def classify(test_Q1, test_Q2, y, threshold, model, vocab, data_generator=data_generator, batch_size=64):
    """Function to test the accuracy of the model.

    Args:
        test_Q1 (numpy.ndarray): Array of Q1 questions.
        test_Q2 (numpy.ndarray): Array of Q2 questions.
        y (numpy.ndarray): Array of actual target.
        threshold (float): Desired threshold.
        model (trax.layers.combinators.Parallel): The Siamese model.
        vocab (collections.defaultdict): The vocabulary used.
        data_generator (function): Data generator function. Defaults to data_generator.
        batch_size (int, optional): Size of the batches. Defaults to 64.

    Returns:
        float: Accuracy of the model.
    """
    accuracy = 0
    ### START CODE HERE (Replace instances of 'None' with your code) ###
    for i in range(0, len(test_Q1), batch_size):
        # Call the data generator (built in Ex 01) with shuffle=False using next()
        # use batch size chuncks of questions as Q1 & Q2 arguments of the data generator. e.g x[i:i + batch_size]
        # Hint: use `vocab['<PAD>']` for the `pad` argument of the data generator
        q1, q2 = next(data_generator(test_Q1[i:i+batch_size], test_Q2[i:i+batch_size], batch_size=batch_size, pad=vocab['PAD'], shuffle=False))
        # use batch size chuncks of actual output targets (same syntax as example above)
        y_test = y[i:i+batch_size]
        # Call the model
        v1, v2 = model([q1, q2])

        for j in range(batch_size):
            # take dot product to compute cos similarity of each pair of entries, v1[j], v2[j]
            # don't forget to transpose the second argument
            d = np.dot(v1[j], v2[j].T)
            # is d greater than the threshold?
            res = d > threshold
            # increment accurancy if y_test is equal `res`
            accuracy += (y_test[j] == res)
    # compute accuracy using accuracy and total length of test questions
    accuracy = accuracy / len(test_Q1)
    ### END CODE HERE ###

    return accuracy
```

```python
# this takes around 1 minute
accuracy = classify(Q1_test,Q2_test, y_test, 0.7, model, vocab, batch_size = 512)
print("Accuracy", accuracy)
```

    Accuracy 0.3796875

# Part 5: Testing with your own questions

In this section you will test the model with your own questions. You will write a function `predict` which takes two questions as input and returns $1$ or $0$ depending on whether the question pair is a duplicate or not.

```python
def predict(question1, question2, threshold, model, vocab, data_generator=data_generator, verbose=False):
    """Function for predicting if two questions are duplicates.

    Args:
        question1 (str): First question.
        question2 (str): Second question.
        threshold (float): Desired threshold.
        model (trax.layers.combinators.Parallel): The Siamese model.
        vocab (collections.defaultdict): The vocabulary used.
        data_generator (function): Data generator function. Defaults to data_generator.
        verbose (bool, optional): If the results should be printed out. Defaults to False.

    Returns:
        bool: True if the questions are duplicates, False otherwise.
    """
    ### START CODE HERE (Replace instances of 'None' with your code) ###
    # use `nltk` word tokenize function to tokenize
    q1 = nltk.word_tokenize(question1)  # tokenize
    q2 = nltk.word_tokenize(question2)  # tokenize
    Q1, Q2 = [], []
    for word in q1:  # encode q1
        # increment by checking the 'word' index in `vocab`
        Q1 += [vocab[word]]
    for word in q2:  # encode q2
        # increment by checking the 'word' index in `vocab`
        Q2 += [vocab[word]]

    # Call the data generator (built in Ex 01) using next()
    # pass [Q1] & [Q2] as Q1 & Q2 arguments of the data generator. Set batch size as 1
    # Hint: use `vocab['<PAD>']` for the `pad` argument of the data generator
    Q1, Q2 = next(data_generator([Q1], [Q2], 1, vocab['<PAD>']))
    # Call the model
    v1, v2 = model([Q1, Q2])
    # take dot product to compute cos similarity of each pair of entries, v1, v2
    # don't forget to transpose the second argument
    d = np.dot(v1[0], v2[0].T)
    # is d greater than the threshold?
    res = d > threshold

    ### END CODE HERE ###

    if(verbose):
        print("Q1  = ", Q1, "\nQ2  = ", Q2)
        print("d   = ", d)
        print("res = ", res)

    return res
```

```python
# Feel free to try with your own questions
question1 = "When will I see you?"
question2 = "When can I see you again?"
# 1 means it is duplicated, 0 otherwise
predict(question1 , question2, 0.7, model, vocab, verbose = True)
```

    Q1  =  [[585  76   4  46  53  21   1   1]]
    Q2  =  [[ 585   33    4   46   53 7280   21    1]]
    d   =  0.88113236
    res =  True

    True

```python
# Feel free to try with your own questions
question1 = "Do they enjoy eating the dessert?"
question2 = "Do they like hiking in the desert?"
# 1 means it is duplicated, 0 otherwise
predict(question1 , question2, 0.7, model, vocab, verbose=True)
```

    Q1  =  [[  443  1145  3159  1169    78 29017    21     1]]
    Q2  =  [[  443  1145    60 15302    28    78  7431    21]]
    d   =  0.477536
    res =  False

    False

You can see that the Siamese network is capable of catching complicated structures. Concretely it can identify question duplicates although the questions do not have many words in common.

###   On Siamese networks

Siamese networks are important and useful. Many times there are several questions that are already asked in quora, or other platforms and you can use Siamese networks to avoid question duplicates.


## Teoría de perturbaciones

Consiste en resolver un sistema perturbado(se conoce la solución al no perturbado), y donde el interés es conocer la contribución de la parte perturbada $H'$ al nuevo sistema total.

$$ H = H^{0} + H'$$

Para sistemas no degenerados, la corrección a la energía a primer orden se calcula como

$$ E_{n}^{(1)} = \int\psi_{n}^{(0)*} H' \psi_{n}^{(0)}d\tau$$

** Tarea 1 : Programar esta ecuación si conoces $H^{0}$ y sus soluciones. **

Y la corrección a la función de onda, también a primer orden, se obtiene como:

 $$ \psi_{n}^{(1)} = \sum_{m\neq n} \frac{\langle\psi_{m}^{(0)} | H' | \psi_{n}^{(0)} \rangle}{E_{n}^{(0)} - E_{m}^{(0)}} \psi_{m}^{(0)}$$

**Tarea 2: Programar esta ecuación si conoces $H^{0}$ y sus soluciones. **

**Tarea 3: Investigue las soluciones a segundo orden y también programe las soluciones. **

**Tarea 4.  Resolver el átomo de helio aplicando los programas anteriores.**

**Tarea 5: Método variacional-perturbativo. **

Este método nos permite estimar de forma precisa $E^{(2)}$ y correcciones perturbativas de la energía de órdenes más elevados para el estado fundamental del sistema, sin evaluar sumas infinitas. Ver ecuación 9.38 del libro.

**Resolver el átomo de helio, considerando este método (sección 9.4), como mejor le parezca. **

**Tarea 6. Revisar sección 9.7. **

Inicialmente a mano, y sengunda instancia favor de intentar programar sección del problema, i.e. integral de Coulomb  e integral de intercambio.

```python
from sympy import *
from sympy import init_printing; init_printing(use_latex = 'mathjax')
import numpy as np
import sympy.physics.hydrogen as hyd
import mpmath as mp
import sympy.functions.special.spherical_harmonics as sp
import scipy.integrate as integ
import scipy.constants as scc
```

```python
var('x r r_1 r_2 phi_1 phi_2 theta_1 theta_2 phi theta r_r zeta')
var('n l m', positive = 'True', integer = 'True')
var('m_e hbar Z omega a_0 e L c d Z', real = 'True', constant = 'True')
```

$$\left ( m_{e}, \quad \hbar, \quad Z, \quad \omega, \quad a_{0}, \quad e, \quad L, \quad c, \quad d, \quad Z\right )$$

### Átomo de hidrógeno

```python
Yhyd = hyd.R_nl(n, l, r)
Yhyd
```

$$\frac{2 \left(\frac{2 r}{n}\right)^{l} e^{- \frac{r}{n}}}{n^{2} \sqrt{\left(l + n\right)!}} L_{- l + n - 1}^{\left(2 l + 1\right)}\left(\frac{2 r}{n}\right) \sqrt{\left(- l + n - 1\right)!}$$

```python
Ehyd = hyd.E_nl(n)
Ehyd
```

$$- \frac{1}{2 n^{2}}$$

```python
#Corrección energía primer orden
```

```python
Ehyd_1 = integrate(Yhyd.subs([(n, 1), (l, 0)])*(hbar/I)**4*(-1/(8*m**3*c**2))*\
          diff(Yhyd.subs([(n, 1), (l, 0)]), r, 4)*r**2, (r, 0, oo))
Ehyd_1
```

$$- \frac{\hbar^{4}}{8 c^{2} m^{3}}$$

```python
Ehyd_1.subs([(hbar,scc.hbar),(c, scc.c),(m,scc.m_e)]) /(4*scc.epsilon_0)**4
```

$$-1.44634324025469 \cdot 10^{-22}$$

```python
#Correción función primer orden
```

```python
Yhyd_num = [simplify(integrate(Yhyd.subs([(n, i), (l, 0)])*(hbar/I)**4*(-1/(8*m**3*c**2))*\
          diff(Yhyd.subs([(n, 1), (l, 0)]), r, 4)*sin(theta)*r**2,\
                (r, 0, 1), (theta, 0, pi), (phi, 0, 2*pi)).args[0][0]) for i in range(2,5)]
Yhyd_num
```

$$\left [ - \frac{5 \sqrt{2} \pi \hbar^{4}}{48 c^{2} m^{3}}, \quad - \frac{49 \sqrt{3} \pi \hbar^{4}}{1215 c^{2} m^{3}}, \quad - \frac{979 \pi \hbar^{4}}{23040 c^{2} m^{3}}\right ]$$

```python
Yhyd_den = [(Ehyd.subs([(n, 1)]) - Ehyd.subs([(n, i)])) for i in range(2, 5)]
Yhyd_den
```

$$\left [ - \frac{3}{8}, \quad - \frac{4}{9}, \quad - \frac{15}{32}\right ]$$

```python
Yhyd_1 = sum([(Yhyd_num[i]*Yhyd.subs([(n, i+1), (l, 0)]))/Yhyd_den[i] for i in range(3)])
Yhyd_1
```

$$\frac{49 \sqrt{6} \pi \hbar^{4} e^{- \frac{r}{2}}}{2160 c^{2} m^{3}} \left(- r + 2\right) + \frac{979 \sqrt{3} \pi \hbar^{4} e^{- \frac{r}{3}}}{145800 c^{2} m^{3}} \left(\frac{2 r^{2}}{9} - 2 r + 3\right) + \frac{5 \sqrt{2} \pi \hbar^{4} e^{- r}}{9 c^{2} m^{3}}$$

```python
#Corrección energía segundo orden
```

```python
Ehyd_2_num = [simplify(integrate(Yhyd.subs([(n, i), (l, 0)])*(-1/(8*m**3*c**2))*\
          diff(Yhyd.subs([(n, 1), (l, 0)])*sin(theta)*r**2, r, 4),\
                (r, 0, 1), (theta, 0, pi), (phi, 0, 2*pi)).args[0][0]**2) for i in range(2,5)]
Ehyd_2_num
```

$$\left [ \frac{24649 \pi^{2}}{1152 c^{4} m^{6}}, \quad \frac{2719201 \pi^{2}}{492075 c^{4} m^{6}}, \quad \frac{1159879249 \pi^{2}}{530841600 c^{4} m^{6}}\right ]$$

```python
Ehyd_2_den = [(Ehyd.subs([(n, 1)]) - Ehyd.subs([(n, i)])) for i in range(2, 5)]
Ehyd_2_den
```

$$\left [ - \frac{3}{8}, \quad - \frac{4}{9}, \quad - \frac{15}{32}\right ]$$

```python
Ehyd_2 = sum([Ehyd_2_num[i] / Ehyd_2_den[i] for i in range(3)])
Ehyd_2 *(27.12*scc.eV)
```

$$- \frac{3.22200858259767 \cdot 10^{-16} \pi^{2}}{c^{4} m^{6}}$$

### Oscilador armónico

```python
import sympy.physics.qho_1d as qho
```

```python
Yqho = qho.psi_n(n, x, m, omega)
Yqho
```

$$\frac{\sqrt[4]{m} \sqrt[4]{\omega} e^{- \frac{m \omega x^{2}}{2 \hbar}}}{2^{\frac{n}{2}} \sqrt[4]{\hbar} \sqrt[4]{\pi} \sqrt{n!}} H_{n}\left(\frac{\sqrt{\omega} x}{\sqrt{\hbar}} \sqrt{m}\right)$$

```python
Eqho = qho.E_n(n, omega)
Eqho
```

$$\hbar \omega \left(n + \frac{1}{2}\right)$$

```python
#Corrección energía primer orden
```

```python
H_1 = c*x**3 + d*x**4
```

```python
Eqho_1 = integrate(Yqho.subs([(n, 0)])*H_1*Yqho.subs([(n, 0)]), (x, -oo, oo))
Eqho_1.args[0][0]
```

$$\frac{3 \hbar^{2} d}{4 m^{2} \omega^{2}}$$

```python
#Corrección función primer orden
```

```python
Yqho_num = [integrate(Yqho.subs([(n, i)])*H_1*Yqho.subs([(n, 0)]), (x, -oo, oo)) for i in range(4) if i != 0]
Yqho_num = [Yqho_num[i].args[0][0] for i in range(3)]
Yqho_num
```

$$\left [ \frac{3 \sqrt{2} \hbar^{\frac{3}{2}} c}{4 m^{\frac{3}{2}} \omega^{\frac{3}{2}}}, \quad \frac{3 \sqrt{2} \hbar^{2} d}{2 m^{2} \omega^{2}}, \quad \frac{\sqrt{3} \hbar^{\frac{3}{2}} c}{2 m^{\frac{3}{2}} \omega^{\frac{3}{2}}}\right ]$$

```python
Yqho_den = [(Eqho.subs([(n, 0)]) - Eqho.subs([(n, i)])) for i in range(4) if i != 0]
Yqho_den
```

$$\left [ - \hbar \omega, \quad - 2 \hbar \omega, \quad - 3 \hbar \omega\right ]$$

```python
Yqho_1 = sum([(Yqho_num[i]*Yqho.subs([(n, i+1)]))/Yqho_den[i] for i in range(3)])
simplify(Yqho_1)
```

$$\frac{e^{- \frac{m \omega x^{2}}{2 \hbar}}}{\hbar^{\frac{5}{4}} \sqrt[4]{\pi} m^{\frac{7}{4}} \omega^{\frac{11}{4}}} \left(- \frac{c x^{3}}{3} m^{2} \omega^{2} - \hbar c m \omega x - \frac{3 d}{2} \hbar m \omega x^{2} + \frac{3 d}{4} \hbar^{2}\right)$$

```python
#Corrección energía segundo orden
```

```python
Eqho_2_num = [integrate(Yqho.subs([(n, i)])*H_1*Yqho.subs([(n, 0)]), (x, -oo, oo)) for i in range(4) if i != 0]
Eqho_2_num = [Eqho_2_num[i].args[0][0]**2 for i in range(3)]
Eqho_2_num
```

$$\left [ \frac{9 \hbar^{3} c^{2}}{8 m^{3} \omega^{3}}, \quad \frac{9 \hbar^{4} d^{2}}{2 m^{4} \omega^{4}}, \quad \frac{3 \hbar^{3} c^{2}}{4 m^{3} \omega^{3}}\right ]$$

```python
Eqho_2_den = [(Eqho.subs([(n, 0)]) - Eqho.subs([(n, i)])) for i in range(4) if i != 0]
Eqho_2_den
```

$$\left [ - \hbar \omega, \quad - 2 \hbar \omega, \quad - 3 \hbar \omega\right ]$$

```python
Eqho_2 = sum([Eqho_2_num[i]/Eqho_2_den[i] for i in range(3)])
Eqho_2
```

$$- \frac{11 \hbar^{2} c^{2}}{8 m^{3} \omega^{4}} - \frac{9 \hbar^{3} d^{2}}{4 m^{4} \omega^{5}}$$

### Partícula en caja

```python
Ypiab = sqrt(2/L) * sin(n*pi * x/L)
Ypiabm = Ypiab.subs(n,m)
Ypiab
```

$$\sqrt{2} \sqrt{\frac{1}{L}} \sin{\left (\frac{\pi x}{L} n \right )}$$

```python
Epiab = (hbar*pi*n)**2/(2*m_e*L**2)
Epiab
```

$$\frac{\pi^{2} \hbar^{2} n^{2}}{2 L^{2} m_{e}}$$

```python
#Correción energía primer orden
```

```python
H_1 = x**2
Epiab_1 = integrate(Ypiab.subs([(n, 1)])*H_1*Ypiab.subs([(n, 1)]), (x, 0, L)).simplify()
Epiab_1
```

$$- \frac{L^{2}}{2 \pi^{2}} + \frac{L^{2}}{3}$$

```python
#Correción función de onda primer orden
```

```python
def ev(x, eL, en, em, epi, H):
    return Ypiab.subs([(L, eL), (n, en), (pi, epi)])*H*Ypiabm.subs([(L, eL), (m, em), (pi, epi)])
```

```python
en = 1
eL = 1e-9
Ypiab_num = np.array([integ.quad(lambdify(x, ev(x, eL, en, i, np.pi, H_1), 'numpy'), 0, eL)[0] for i in range(1,5) if i != en])
Ypiab_den = [Epiab.subs([(n, en), (L, eL)]) - Epiab.subs([(n, i), (L, eL)]) for i in range(1,5) if i != en]
```

```python
Ampiab = Ypiab_num / Ypiab_den
Ampiab[0], Ampiab[1], Ampiab[2]
```

$$\left ( \frac{1.20084365798327 \cdot 10^{-37} m_{e}}{\pi^{2} \hbar^{2}}, \quad - \frac{9.49886096646936 \cdot 10^{-39} m_{e}}{\pi^{2} \hbar^{2}}, \quad \frac{1.92134985277342 \cdot 10^{-39} m_{e}}{\pi^{2} \hbar^{2}}\right )$$

```python
Ypiab_1 = sum([Ampiab[i-2] * Ypiab.subs([(n, i), (L, eL)]) for i in range(1,5) if i != 1])
Ypiab_1
```

$$\frac{3.79740107299537 \cdot 10^{-33} m_{e}}{\pi^{2} \hbar^{2}} \sqrt{2} \sin{\left (2000000000.0 \pi x \right )} - \frac{3.00380358313115 \cdot 10^{-34} m_{e}}{\pi^{2} \hbar^{2}} \sqrt{2} \sin{\left (3000000000.0 \pi x \right )} + \frac{6.07584171679318 \cdot 10^{-35} m_{e}}{\pi^{2} \hbar^{2}} \sqrt{2} \sin{\left (4000000000.0 \pi x \right )}$$

```python
#Corrección energía segundo orden
```

```python
Epiab_2_num = np.array([integ.quad(lambdify(x, ev(x, eL, en, i, np.pi, H_1), 'numpy'), 0, eL)[0]**2 for i in range(1,5) if i != en])
Epiab_2_num
```

    array([  3.24455735e-38,   1.44365375e-39,   2.07651671e-40])

```python
Epiab_2_den = [Epiab.subs([(n, en), (L, eL)]) - Epiab.subs([(n, i), (L, eL)]) for i in range(1,5) if i != en]
Epiab_2_den
```

$$\left [ - \frac{\pi^{2} \hbar^{2}}{m_{e}} 1.5 \cdot 10^{18}, \quad - \frac{\pi^{2} \hbar^{2}}{m_{e}} 4.0 \cdot 10^{18}, \quad - \frac{\pi^{2} \hbar^{2}}{m_{e}} 7.5 \cdot 10^{18}\right ]$$

```python
Epiab_2 = sum([Epiab_2_num[i]/Epiab_2_den[i] for i in range(3)])
Epiab_2
```

$$- \frac{2.20189826918465 \cdot 10^{-56} m_{e}}{\pi^{2} \hbar^{2}}$$

### Átomo de helio

```python
#Corrección energía primer orden
```

```python
Yhe_1 = hyd.R_nl(n, l, r_1, Z = 2)
Yhe_2 = Yhe_1.subs([(r_1, r_2)])
Yhe_1
```

$$\frac{4 \sqrt{2} \left(\frac{4 r_{1}}{n}\right)^{l} e^{- \frac{2 r_{1}}{n}}}{n^{2} \sqrt{\left(l + n\right)!}} L_{- l + n - 1}^{\left(2 l + 1\right)}\left(\frac{4 r_{1}}{n}\right) \sqrt{\left(- l + n - 1\right)!}$$

```python
Ehe = hyd.E_nl(n, Z=2)
Ehe
```

$$- \frac{2}{n^{2}}$$

```python
def Ehe_r(nn, mm):
    return Yhe_1.subs([(n, nn), (l, 0)])*Yhe_2.subs([(n, mm), (l, 0)])*r_1**2*r_2**2*(1/r_r)
Ehe_r(1, 1)
```

$$\frac{32 r_{1}^{2}}{r_{r}} r_{2}^{2} e^{- 2 r_{1}} e^{- 2 r_{2}}$$

```python
def Ehe_ev(nn, mm):
    return integrate(integrate(Ehe_r(1, nn).subs([(r_r, r_2)]), (r_1, 0, r_2)) + \
                 integrate(Ehe_r(1, mm).subs([(r_r, r_1)]), (r_1, r_2, oo)), (r_2, 0, oo))

print(Ehe_ev(1, 1) * (27.2114), 'eV' )
```

    34.0142500000000 eV

## Variacional-perturbativo

```python
Yhe_vp = hyd.R_nl(1, 0, r_1, Z=zeta) * hyd.R_nl(1, 0, r_2, Z=zeta)
simplify(Yhe_vp)
```

$$4 \zeta^{3} e^{- \zeta \left(r_{1} + r_{2}\right)}$$

```python
Ehe_vp = integrate(integrate(Yhe_vp*(zeta - Z)*((1/r_1)+(1/r_2))*Yhe_vp*r_1**2*r_2**2, (r_1, 0, oo)).args[0][0]\
                   , (r_2, 0, oo)).args[0][0] - zeta**2
Ehe_vp
```

$$- 2 Z \zeta + \zeta^{2}$$

```python
Ehe_vp_1 = Ehe_vp + Ehe_ev(1, 1)*zeta/2
Ehe_vp_1
```

$$- 2 Z \zeta + \zeta^{2} + \frac{5 \zeta}{8}$$

```python
Ehe_vp_deriv = diff(Ehe_vp_1, zeta, 1)
Ehe_vp_deriv
```

$$- 2 Z + 2 \zeta + \frac{5}{8}$$

```python
zeta_min = solve(Ehe_vp_deriv, zeta)[0]
zeta_min
```

$$Z - \frac{5}{16}$$

```python
Ehe_vp_corr = simplify(Ehe_vp_1.subs([(zeta, zeta_min)]))
Ehe_vp_corr
```

$$- Z^{2} + \frac{5 Z}{8} - \frac{25}{256}$$

```python
print(Ehe_vp_corr.subs([(Z, 2)]) * (27.2114), 'eV')
```

    -77.4887132812500 eV

```python

```


```python
# Standard libraries
import numpy as np
import scipy as sp
from scipy import signal

from tqdm.notebook import tqdm
from multiprocess import pool

import matplotlib as mpl
import matplotlib.pyplot as plt
cmap = plt.get_cmap("tab10")
cmap2 = plt.get_cmap("tab20")
```

```python
# Special libraries
import qutip as qt

# https://pypi.org/project/dmdlab/
import dmdlab as dmd
```

```python
import importlib
hf = importlib.import_module('helper_functions')
```

# Bilinear Dynamic Mode Decomposition for Quantum Control

https://arxiv.org/abs/2010.14577

Data-driven methods for establishing quantum optimal control (QOC) using time-dependent control pulses tailored to specific quantum dynamical systems and desired control objectives are critical for many emerging quantum technologies. We develop a data-driven regression procedure, bilinear dynamic mode decomposition (biDMD), that leverages time-series measurements to establish quantum system identification for QOC. The biDMD optimization framework is a physics-informed regression that makes use of the known underlying Hamiltonian structure. Further, the biDMD can be modified to model both fast and slow sampling of control signals, the latter by way of stroboscopic sampling strategies. The biDMD method provides a flexible, interpretable, and adaptive regression framework for real-time, online implementation in quantum systems. Further, the method has strong theoretical connections to Koopman theory, which approximates non-linear dynamics with linear operators. In comparison with many machine learning paradigms, it requires minimal data and the biDMD model is easily updated as new data is collected. We demonstrate the efficacy and performance of the approach on a number of representative quantum systems, showing that it also matches experimental results.

Andy Goldschmidt, Eurika Kaiser, Jonathan L. Dubois, Steven L. Brunton, J. Nathan Kutz

## Data storage

First we define a utilty class for storing experiment data.

```python
class QubitExperiment:
    ''' A container for data saved during a set of experiments. '''
    def __init__(self):
        self.X = []
        self.Y = []
        self.Uy = []
        self.Ty = []
        self.Args_List = []
        self._finalized = False

    def save(self, iX, iY, iUy, iTy, iArgs):
        if self._finalized:
            raise ValueError('Cannot save data. QubitExperiment has been finalized.')

        self.X.append(iX)
        self.Y.append(iY)
        self.Uy.append(iUy)
        self.Ty.append(iTy)
        self.Args_List.append(iArgs)

    def dump(self):
        if self._finalized:
            raise ValueError('Cannot dump data. QubitExperiment has been finalized.')

        for i in range(len(self.X)):
            yield self.X[i], self.Y[i], self.Uy[i], self.Ty[i], self.Args_List[i]

    def finalize(self):
        self._finalized = True
        self.X_flat = np.hstack(self.X)
        self.Y1_flat = np.hstack([y[:, :-1] for y in self.Y])
        self.Y2_flat = np.hstack([y[:, 1:] for y in self.Y])
        self.U1y_flat = np.hstack([u[:, :-1] for u in self.Uy])
        self.T1y_flat = np.hstack([t[:-1] for t in self.Ty])

def join_qubit_experiments(qbExpermList):
    joinedQbExperm = QubitExperiment()
    for aQbExperm in qbExpermList:
        for data in aQbExperm.dump():
            joinedQbExperm.save(*data)
    joinedQbExperm.finalize()
    return joinedQbExperm
```

# Example 1

We start with the system described by the Hamiltonian,
\begin{equation}
    H(t) = \pi \sigma_3 + V_0 \cos(\Omega t) \sigma_1
\end{equation}

This description gives us time units that are in terms of the bare oscillation. In these units, can make a physical estimate of $V_0$. To do this, we make some assumptions on our experiment. Suppose that our qubit has a resonance frequency of 4.1 GHz. Further, let's say that a Pi pulse (the control signal that maps $\sigma_z=0$ to $\sigma_z=1$, or 1/4 of the period of a Rabi oscillation) takes 86ns when the input drive amplitude is 100mV. So far these numbers imply $V_0\approx 0.004$. If we have a maximum drive amplitude of 500 mV, our coupling can be increased by a factor of 5 to $V_0 \approx 0.02$. Finally, suppose we have a sample rate around 32 Ghz. That means we can sample approximately 8 times in one period of the bare frequency.

The experiment we will run below won't really match these values. We increase the coupling strength to $V_0 = 1$ because we want to apply a pulse that is off resonance and get some action in $\sigma_z$ in not too many timesteps. If we want to use physical values, we could decrease the amount we are off of resonace and increase the timesteps to accommodate but this is perhaps less  useful visually.

We will run multiple experiments to make sure that our model doesn't overfit a single training pulse. You don't need many experiments for take care of this.

This first block of code sets a lot of the experimental parameters we can adjust.

```python
# Qubit frequency
w0 = 1
Delta = 2*np.pi*w0

# Time series parameters
# -- Pick a simulation discretization and a measurement discretization (measured per bare oscillation)
# -- Make sure your discretizations divide each other...
sim_discretization = 32
m_discretization = 16
sim_dt =  2*np.pi/Delta/sim_discretization
dt = 2*np.pi/Delta/m_discretization
t0 = 0
tf = 5
ts = np.linspace(t0, tf, (tf-t0)//sim_dt, endpoint=True)
ts_test = np.linspace(t0, 5*tf, (5*tf-t0)//sim_dt, endpoint=True)

# Hamiltonian
# -- Lab frame (Linear Polarization)
H0_lab = Delta/2*qt.sigmaz()
V0 = 1
fc_lab = lambda t, args: args['u'](t)*np.cos(args['Omega']*t)
H_lab = [H0_lab, [qt.sigmax(), fc_lab]]

# Dissipation
dissipation = [[.04*qt.sigmax()], [.1*qt.sigmam()]]

# Choose measurements and initial state
measure_list = [qt.sigmax(), qt.sigmay(), qt.sigmaz()]
measure_name = ['<{}>'.format(s) for s in ['x','y','z']]
n_measure = len(measure_list)
psi0 = qt.basis(2,1)

# Set measurement certainty
N_measure = 1e4
sigma = 1/np.sqrt(N_measure)
```

We run experiments for a random assortment of frequencies that we think might be in the ballpark of the resonance frequency.

```python
np.random.seed(1)

# Run experiments at some off-resonance frequencies and miss resonance by at least +/- 0.05
n_experiments = 10
frequency_list = np.hstack([np.random.rand(n_experiments//2)*.1+.85,
                            np.random.rand(n_experiments - n_experiments//2)*.1+1.05])
print(frequency_list)
QbData = QubitExperiment()
for w1 in frequency_list:
    # Simulate the training dynamics
    c_args = {'Omega': 2*np.pi*w1, 'u': lambda t: V0*np.ones_like(t)}
    res = qt.mesolve(H_lab, psi0, ts, dissipation, measure_list, args=c_args)

    # Measure only between some skipped points
    skip = sim_discretization//m_discretization
    X = np.array(res.expect)
    Y = X[:,::skip]
    Ty = ts[::skip]
    Uy = fc_lab(Ty, c_args).reshape(1,-1)

    # Add noise to everything but the prepared state
    Y[:,1:] += np.random.randn(*Y[:,1:].shape)*sigma

    # Save the experiment data
    QbData.save(X, Y, Uy, Ty, c_args)

# Finish running experiments
QbData.finalize()
```

    [0.8917022  0.92203245 0.85001144 0.88023326 0.86467559 1.05923386
     1.06862602 1.08455607 1.08967675 1.10388167]

We fit the bilinear DMD model and back out a resonance frequency from the observed eigenvalues.

```python
# Fit the model
model = dmd.biDMD(QbData.Y2_flat, QbData.Y1_flat, QbData.U1y_flat, QbData.T1y_flat, threshold_type='count', threshold=6)

# Estimate the bare frequency
est_bare = np.abs(np.angle(model.eigs[1])/dt/2/np.pi)

# Plot the eigenvvalues
fig, ax = plt.subplots(1,figsize=[4,4])
ax.set_aspect('equal'), ax.set_xlim([-1.15, 1.15]), ax.set_ylim([-1.15, 1.15])
ax.set_xticks([-1,0,1])
ax.set_yticks([-1,0,1])
ax.set_xticklabels([-1,0,1],fontsize=24)
ax.set_yticklabels([-1,0,1],fontsize=24)
for eigs in model.eigs:
    ax.plot(eigs.real, eigs.imag, ms=12, marker='x', markeredgewidth=3, ls='', color='k')
ax.add_artist(plt.Circle((0, 0), 1, color='k', linestyle='--', fill=False))
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
print('Estimated bare frequency:', est_bare)
print('Eigs: ', model.eigs)
```

```python
# Plot some of the training data and the training fit
for j in range(min(n_experiments, 3)):
    fig, axes = plt.subplots(3, 1, figsize=[2,2.75])
    X_predict = model.predict_dst(control=QbData.Uy[j], x0=QbData.Y[j][:,0])
    print(QbData.Args_List[j]['Omega']/2/np.pi)
    for m in range(n_measure):
        axes[m].plot(QbData.Ty[j][:-1], X_predict[m,:-1], color=cmap2(2*m+1), marker='o', markeredgecolor=cmap(m), markerfacecolor='None', ms=3)
        axes[m].plot(QbData.Ty[j], QbData.Y[j][m], color=cmap2(2*m+1))
        axes[m].set_ylim([-1.4,1.4])
    [ax.set_xticks([]) for ax in axes.flatten()[:-1]]
    fig.subplots_adjust(wspace=0, hspace=0)
```

Now use the model and the estimated resonance frequency to see what our prediction is for an example pulse that is on resonance. The prediction is made using only the model, a given initial state and the new control pulse.

You might also consider running a new experiment where you put the estimated bare frequency into your list and try to get advantages from this feedback.

```python
# Simulate the testing dynamics using the est_bare from the previous section
c_args_test = {'Omega': 2*np.pi*est_bare, 'u': lambda t: V0*np.ones_like(t)}
result_lab_test = qt.mesolve(H_lab, psi0, ts_test, dissipation, measure_list, args=c_args_test)

# Collect test data
skip = sim_discretization//m_discretization
result = np.array(result_lab_test.expect)
X = result[:,::skip]
T1 = ts_test[::skip]
U1 = fc_lab(T1, c_args_test).reshape(1,-1)
X_predict = model.predict_dst(control=U1, x0=X[:,0])

# Plot the prediction
fig, axes = plt.subplots(3,1, figsize=[4,2.75])
for m in range(n_measure):
    axes[m].plot(ts_test, result[m], color=cmap2(2*m+1),label=measure_name[m])
    axes[m].plot(T1, X_predict[m],color=cmap2(2*m))
    axes[m].set_ylim([-1.4,1.4])
    axes[m].set_xlim([-.1,10])

for m in range(n_measure-1):
    axes[m].set_xticks([])
axes[-1].set_xticks([0,5,10])
fig.subplots_adjust(wspace=0, hspace=0)
```

## A variation on Example 1

Suppose you already know the resonance frequency of your qubit and you are only interested in checking up on the model for your qubit in the interaction frame. For a weakly coupled dipole interaction, the rotating wave approximation will allow us to treat the model control as the envelope on some carrier at (an approximate) resonance value.

The Hamiltonian in the lab frame is
\begin{equation}
    H_{\mathrm{LP}}(t) = \frac{\Delta}{2} \sigma_z + V \cos(\Omega t) \sigma_x \,.
\end{equation}

The Hamiltonian in the rotating frame is
\begin{equation}
    \tilde{H}_{\mathrm{LP}}(t)
    =\underbrace{\frac{\Omega}{2} \mathbb{1}+\frac{\left(\Delta-\Omega\right)}{2} \sigma_z+\frac{V}{2} \sigma_x}_{\text {Rotating Wave Approximation }}
    +\frac{V}{2}\left( \cos (2 \Omega t) \sigma_x - \sin (2 \Omega t) \sigma_y  \right) \,.
\end{equation}

The tilde refers to the fact that this new Hamiltonian describes the dynamics of a state $|\tilde{\psi}\rangle$ where some of the time dependence has been split off to act on the operators such that $|\psi\rangle = P(t,0)|\tilde{\psi}\rangle $.

Recall that changing to the rotating frame means applying a transformation to the operators:
\begin{equation}
    \tilde{H}_{\mathrm{LP}}(t) = P^{\dagger} H_{\mathrm{LP}}(t) P - i P^{\dagger} \dot{P}_{t} \,.
\end{equation}

Here, we use $P(t,0) = \exp{i \frac{\Omega t}{2} \left(\mathbb{1} - \sigma_z \right)}$ because it turns out to be equal to the Floquet operator (it's the obvious rotation around the z-axis we'd use to go to a rotating or interaction frame, adjusted for periodicity).

Note that for a Hamiltonian with a circularly polarized drive, the rotating wave approximation is actually exact and time-independent. For reference, a circularly polarized system has a the Hamiltonian in the lab frame given by
\begin{equation}
    H_{\mathrm{CP}}(t) = \frac{\Delta}{2} \sigma_z + \frac{V}{2} \left( \cos(\Omega t) \sigma_x + \sin(\Omega t) \sigma_y \right)
    \Rightarrow  \tilde{H}_{\mathrm{CP}} = \frac{\Omega}{2} \mathbb{1}+\frac{\left(\Delta-\Omega\right)}{2} \sigma_z+\frac{V}{2} \sigma_x \,.
\end{equation}

The conclusion of our discussion is that we can treat the envelope $V(t)=V_0 f(t)$ as the control input for a model of measurements in the rotating frame. We'll explore later the connections between the rotating frame and Floquet theory.

In the code below we explore the differences between the frames. Consider changing the coupling $V_0$ between $1$ and $0.1$ to see the differences between the RWA and the rotating frame emerge at stronger coupling strength (but you should also consider descreasing or increasing the final time by a similar order of magnitude, respectively).

```python
# Qubit frequency
w0 = 1
Delta = 2*np.pi*w0

# Carrier frequency
w1 = 1
Omega = 2*np.pi*w1

# Coupling strength of control
# -- Stronger coupling makes the RWA less accurate (compare .1 and 1)
V0 = .1

# Time series parameters
sim_discretization = 16
sim_dt =  2*np.pi/Omega/sim_discretization
dt =  2*np.pi/Omega/m_discretization
t0 = 0
tf = 10/V0
ts = np.linspace(t0, tf, (tf-t0)//sim_dt, endpoint=True)

# Choose measurements
measure_list = [qt.sigmax(), qt.sigmay(), qt.sigmaz()]
measure_name = ['<{}>'.format(s) for s in ['x','y','z']]
n_measure = len(measure_list)

# Hamiltonians
# -- Rotating frame
H0 = Omega/2*qt.identity(2) + (Delta - Omega)/2*qt.sigmaz()
pulse_minus = lambda t, args: (1 + np.exp(-2j*t*args['Omega']))/2*args['u'](t)
pulse_plus = lambda t, args: (1 + np.exp(2j*t*args['Omega']))/2*args['u'](t)
H_rot = [H0, [V0*qt.sigmam(), pulse_minus], [V0*qt.sigmap(), pulse_plus]]

# -- Rotating wave approximation in the rotating frame
pulse_rwa = lambda t, args: args['u'](t)/2
H_rwa = [H0, [V0*qt.sigmax(), pulse_rwa]]

# -- Lab frame (Linear Polarization)
H0_lab = Delta/2*qt.sigmaz()
pulse_lab = lambda t, args: args['u'](t)*np.cos(args['Omega']*t)
H_lab = [H0_lab, [V0*qt.sigmax(), pulse_lab]]

# Simulate some dynamics
psi0 = qt.basis(2,1)
c_args = {'Omega': Omega, 'u': lambda t: np.ones_like(t)}
result_rot = qt.mesolve(H_rot, psi0, ts, [], measure_list, args=c_args)
result_rwa = qt.mesolve(H_rwa, psi0, ts, [], measure_list, args=c_args)
result_lab = qt.mesolve(H_lab, psi0, ts, [], measure_list, args=c_args)

# Plot
fig, axes = plt.subplots(1,3,figsize=[12,2.5])
for i in range(n_measure):
    axes[1].plot(ts, result_rot.expect[i], color=cmap(i), label=measure_name[i])
    axes[2].plot(ts, result_rwa.expect[i], color=cmap(i))
    axes[0].plot(ts, result_lab.expect[i], color=cmap(i))

axes[1].legend(ncol=3)
axes[0].set_title('Lab frame')
axes[1].set_title('Rotating frame')
axes[2].set_title('Rotating wave approximation')
```

We will take data from the rotating frame as our training set and fit a bilinear model assuming the rotating wave approximation so the control envelope is our input.

We then see how well we can predict with a different control envelope. Also, we are using reasonable experimental values in this example (so reasonable statistical noise and errors in the carrier, and some Markovian dissipation).

```python
# -- Rotating frame is all we need
V0 = .05
H0 = Omega/2*qt.identity(2) + (Delta - Omega)/2*qt.sigmaz()
pulse_minus = lambda t, args: (1 + np.exp(-2j*t*args['Omega']))/2*args['u'](t)
pulse_plus = lambda t, args: (1 + np.exp(2j*t*args['Omega']))/2*args['u'](t)
H_rot = [H0, [V0*qt.sigmam(), pulse_minus], [V0*qt.sigmap(), pulse_plus]]

# Set measurement certainty
N_measure = 1e3
sigma = 1/np.sqrt(N_measure)

# Time series parameters
sim_discretization = 16
sim_dt =  2*np.pi/Omega/sim_discretization
dt =  2*np.pi/Omega/m_discretization
t0 = 0
tf = 150
ts = np.linspace(t0, tf, (tf-t0)//sim_dt, endpoint=True)
m_discretization = 1
# -- Measuremetns occur only between skipped points
skip = int(sim_discretization/m_discretization)

# Dissipation
dissipation = [[.04*qt.sigmax()], [.01*qt.sigmam()]]

# Train data
QbData = QubitExperiment()
train_envelope_1 = lambda t: 1.5*np.sin(np.pi*(t-t0)/(tf-t0)) - 1.5*np.cos(1.7*np.pi*(t-t0)/(tf-t0))
train_envelope_2 = lambda t: 2.5*np.cos(2.3*np.pi*(t-t0)/(tf-t0)) - 1*np.cos(1.5*np.pi*(t-t0)/(tf-t0))

for te in [train_envelope_1, train_envelope_2]:
    c_args = {'Omega': Omega, 'u': te}
    result_train = qt.mesolve(H_rot, psi0, ts, dissipation, measure_list, args=c_args)
    X = np.array(result_train.expect)
    Y = X[:, ::skip]
    Y[:,1:] = Y[:,1:] + np.random.randn(*Y[:,1:].shape)*sigma
    Ty = ts[::skip]
    Uy = pulse_rwa(Ty, c_args).reshape(1, -1)
    QbData.save(X, Y, Uy, Ty, c_args)
QbData.finalize()

# Fit a model
model = dmd.biDMD(QbData.Y2_flat, QbData.Y1_flat, QbData.U1y_flat, QbData.T1y_flat, threshold=1e-3, threshold_type='percent')

# Test data
test_envelope = lambda t: .25-3*np.cos(1.5*np.pi/(tf-t0)*(t-t0)) + np.cos(4.1*np.pi/(tf-t0)*(t-t0))
c_args_test = {'Omega': Omega, 'u': test_envelope}
result_test = qt.mesolve(H_rot, psi0, ts, dissipation, measure_list, args=c_args_test)
Uy_test = pulse_rwa(Ty, c_args_test).reshape(1, -1)
predict = model.predict_dst(control=Uy_test)

# Plot control
fig,axes = plt.subplots(1,3,figsize=[12,2.5])
for i in range(2):
    ax=axes[i]
    ax.plot(Ty, QbData.Uy[i][0], marker='+', color='k')
    ax.set_title('Training input {}'.format(i+1))
ax=axes[2]
ax.plot(Ty, Uy_test[0], marker='+', color='k')
ax.set_title('Test input')
[iax.set_ylim([-2.1,2.1]) for iax in axes]

# Plot test prediction
fig, ax = plt.subplots(1, figsize=[8,2.5])
for m in range(n_measure):
    ax.fill_between(ts, result_test.expect[m]-sigma, result_test.expect[m]+sigma, color=cmap2(2*m+1))
    ax.plot(Ty, predict[m], color=cmap(m))
ax.set_ylim([-1.5,1.5])
ax.set_title('Light color is the test value $\pm \sigma$. Dark color is the the model prediction.')

# In the rotating frame only the B operator has oscillating modes.
print('Control eigenvalues: ', np.linalg.eigvals(model.B))

fig, axes = plt.subplots(1, 2, figsize=[7.5,4])
qt.hinton(model.A/dt, ax=axes[0])
axes[0].set_title('Drift operator A')
qt.hinton(model.B/dt, ax=axes[1])
axes[1].set_title('Control operator B')
```

There are a number of things we could consider doing to improve this fit. We might want to constrain the eigenvalues to be on the unit circle or give our operators a certain structure using an optimization-based DMD. We might also want to look into higher-order expansions of the discretization; that is, we may look into the way the control appears if you expand the continuous dynamics to $\mathcal{O}(\Delta t^3)$.

# Example 2

## Floquet reshape

If we reshape DMD, we can learn a Floquet Hamiltonian for multiple measurments within a single strobe. We reshape the data such that
    \begin{align}
        \mathbf{X} =&
        \begin{bmatrix}
            x(t_0) & x(t_0 + T) & \dots & x(t_0 + (n-2) T) \\
            x(t_1) & x(t_1 + T) & \dots & x(t_1 + (n-2) T) \\
            \vdots & \vdots & & \vdots  \\
            x(t_{m-1}) & x(t_{m-1} + T) & \dots & x(t_{m-1} + (n-2) T)
        \end{bmatrix} \\
        \mathbf{X'} =&
        \begin{bmatrix}
            x(t_0 + T) & x(t_0 + 2 T) \dots & x(t_0 + (n-1) T) \\
            x(t_1 + T) & x(t_1 + 2 T) \dots & x(t_1 + (n-1) T) \\
            \vdots & & \vdots  \\
            x(t_{m-1} + T) & x(t_{m-1} + 2 T) \dots & x(t_{m-1} + (n-1) T)
        \end{bmatrix}
    \end{align}
where we take $m$ measurements inside a single period. The DMD algorithm is looking for an operator that takes a step size of 1 period, but it does not matter that we are simultaneously learning this operator for multiple times within a single strobe.

```python
def fshape(data, n_measure, invert=False):
    '''
    Floquet reshape.

    data.shape = data_dim, n_measure*n_periods

    The idea of the code is the following:
    Initial   ->(data_dim, n_measure*n_periods)
    Transpose ->(n_measure*n_periods, data_dim)
    Reshape   ->(n_periods, n_measure, data_dim)
              ->(n_periods, n_measure*data_dim)
    Transpose ->(n_measure*data_dim, n_periods)

    Paramters:
        data (`ndarray` of float): Data is spatial dimension by time dimension.
        n_measure (int): Number of measurements taken witin one period.
        invert (bool): Undo the reshape operation. fshape(fshape(X, m, False), m, True) = X. Default False.

    Returns
        (`ndarray` of float):
            Reshaped data is spatial*n_measure dimension by periodic time dimension.
            The order of the first column of the reshape is, e.g., [x1(t1), x2(t1), x1(t2) x2(t2),...]^T
    '''
    if invert:
        return data.T.reshape(n_measure*data.shape[1], -1).T
    else:
        return data.T.reshape(-1, n_measure*data.shape[0]).T
```

This experiment involves driving the qubit at an off-resononace frequency with a coupling strength where there is some discrepancy betweeen the RWA and the true dynamics. One feature we will observe with Floquet-based DMD methods is the ability to extract these differences.

```python
# Frequencies
w0 = 1
w1 = 1.1
Omega = 2*np.pi*w1
Delta = 2*np.pi*w0

# Time series parameters
sim_discretization = 32
T = 2*np.pi/Omega
sim_dt = 2*np.pi/Omega/sim_discretization
t0 = 0
tf = 13*T
ts = np.linspace(t0, tf, (tf-t0)//sim_dt, endpoint=False)

# Hamiltonians
# -- Lab frame
fc = lambda t, args:  args['V']*np.cos(Omega*t)
H = [Delta/2*qt.sigmaz(), [qt.sigmax(), fc]]

# Choose measurements
measure_list = [qt.sigmax(), qt.sigmay(), qt.sigmaz()]
measure_name = ['<{}>'.format(s) for s in ['x','y','z']]
n_measure = len(measure_list)

# Init.
psi0 = qt.basis(2,1)

# Simulate the dynamics
c_args = {'Omega': Omega, 'V': 1}
result = qt.mesolve(H, psi0, ts, [], measure_list, args=c_args)
```

The key parameter for the plots I will make in this section is the 'mm', which denotes the number of measurements per Floquet period. This number can be 1, but additional measurements let you resolve some of the fast dynamics that the rotating wave approximation does not capture.

```python
# Fit the model
mm = 4
ttsplit = 4 # number of periods for training/test split
measurement_skip = sim_discretization//mm

res = np.array(result.expect)
X = fshape(res[:,::measurement_skip],  mm)

X_train = X[:, :ttsplit]
ts_train = ts[::sim_discretization][:ttsplit]
model = dmd.DMD.from_full(X_train, ts_train, threshold_type='percent', threshold=1e-3)
msort = np.argsort(model.eigs)[::-1]

# Fit a full model -- no split (for later use)
full_X = fshape(res,  sim_discretization)
full_model = dmd.DMD.from_full(full_X[:, :ttsplit], ts_train, threshold_type='percent', threshold=1e-3)
fmsort = np.argsort(model.eigs)[::-1]
```

```python
# Look at the Floquet Hamiltonian
fig, ax = plt.subplots(1, figsize=[5,3])
qt.hinton(model.A, ax=ax)
```

In the next figures, we will compare the (substrobe-measured) DMD eigenmodes to the Floquet eigenmodes and the RWA eigenmodes.

I will use a shortcut to compute the exact Floquet eigenmodes. I will compute DMD for the fully measured system. All I need to show is that the substrobed-DMD modes match the full modes.

We also need to compute the rotating wave approximation eigenmodes and eigenvalues.

Take the adjoint representation of the Hamiltonian in the rotating wave approximation:
\begin{equation}
    \frac{\Omega}{2} \mathbb{1}+\frac{\left(\Delta-\Omega\right)}{2} \sigma_z+\frac{V}{2} \sigma_x
    \mapsto \frac{\left(\Delta-\Omega\right)}{2} L_z + \frac{V}{2} L_x \,.
\end{equation}
so the dynamics are $\dot{\vec{\rho}} = L_H \vec{\rho}$ where
\begin{equation}
    L_H = \frac{1}{2} \left(\begin{array}{ccc}
        0 & -(\Delta - \Omega) & 0 \\
        \Delta - \Omega & 0 & -V \\
        0 & V & 0 \end{array}\right)
\end{equation}

```python
Lh = np.array([[0, -(Delta - Omega), 0],[Delta-Omega, 0, -c_args['V']], [0, c_args['V'], 0]])
heigs, hmodes = np.linalg.eig(sp.linalg.expm(Lh*T))
hsort = np.argsort(heigs)[::-1]
```

```python
mpl.rcParams["font.family"] = "serif"
mpl.rcParams['xtick.labelsize'] = 14
mpl.rcParams['ytick.labelsize'] = 14

fig, ax = plt.subplots(1,figsize=[4,4])
ax.set_aspect('equal'), ax.set_xlim([-1.2, 1.2]), ax.set_ylim([-1.2, 1.2])

for i in range(3):
    color = cmap(6) if i == 0 else cmap(5)
    ax.plot(model.eigs[msort].real[i], model.eigs[msort].imag[i], ms=18, marker='x', markeredgewidth=2, ls='', color=color)
    ax.plot(full_model.eigs[fmsort].real[i], full_model.eigs[fmsort].imag[i], ms=20, marker='o', markeredgewidth=3, markerfacecolor='None', ls='', color=color)
    ax.plot(heigs[hsort].real[i], heigs[hsort].imag[i], ms=18, marker='+', markeredgewidth=2, markerfacecolor='None', ls='', color=color)

ax.add_artist(plt.Circle((0, 0), 1, color='k', linestyle='--', fill=False))

custom_lines = [mpl.lines.Line2D([0], [0], color='k', ls='', ms=10, markerfacecolor='None', marker='o', markeredgewidth=2),
                mpl.lines.Line2D([0], [0], color='k', ls='', ms=10, marker='x', markeredgewidth=2),
                mpl.lines.Line2D([0], [0], color='k', ls='', ms=10, marker='+', markeredgewidth=2),]
ax.legend(custom_lines, ['Exact', 'DMD', 'RWA'], fontsize=14, ncol=1, loc='center',frameon=False)
```

```python
mpl.rcParams['xtick.labelsize'] = 12
mpl.rcParams['ytick.labelsize'] = 12
mpl.rcParams["font.family"] = "serif"

figs = []
axes = []
for i in range(3):
    f, a = plt.subplots(2, 3, figsize=[6,3])
    figs.append(f)
    axes.append(a)

full_tsmode = ts[:sim_discretization]
tsmode = ts[::measurement_skip][:mm]

for imode, amode in enumerate(hmodes.T):
    Mmode = np.vstack([sp.linalg.expm(Lh*t)@amode for t in full_tsmode]).T
    Mmode = Mmode/Mmode[0,0]
    for j in range(n_measure):
        axes[imode][0,j].plot(full_tsmode, Mmode[j].real, ls='--', lw=1, alpha=1, color=cmap(j))
        axes[imode][1, j].plot(full_tsmode, Mmode[j].imag, ls='--', lw=1, alpha=1, color=cmap(j))

for imode, amode in enumerate(model.modes.T):
    norm = amode[0]
    for j in range(n_measure):
        jmode = amode[j::n_measure]
        jmode = jmode/norm
        axes[imode][0, j].plot(tsmode, jmode.real, color=cmap(j), ls='', markeredgewidth=1.5, ms=10, marker='x', markerfacecolor='None')
        axes[imode][1, j].plot(tsmode, jmode.imag, color=cmap(j), ls='', markeredgewidth=1.5, ms=10, marker='x', markerfacecolor='None')

for imode, amode in enumerate(full_model.modes.T):
    norm = amode[0]
    for j in range(n_measure):
        jmode = amode[j::n_measure]
        jmode = jmode/norm
        axes[imode][0,j].plot(full_tsmode, jmode.real, lw=1, color=cmap(j))
        axes[imode][1,j].plot(full_tsmode, jmode.imag, lw=1, color=cmap(j))

ytick = [-1.5,0,1.5]
for ia, a in enumerate(axes):
    for i in range(2):
        for j in range(3):
            a[i,j].set_ylim([-2.1, 2.1])
            a[i,j].set_xlim([-0.05,.9])
            a[i,j].set_yticks(ytick)
            if j == 0:
                a[i,j].set_yticklabels(ytick)
            else:
                a[i,j].set_yticklabels([])
            if i == 0 or ia == 0:
                a[i,j].set_xticklabels([])
for i, f in enumerate(figs):
    f.tight_layout()
    f.subplots_adjust(wspace=0, hspace=0)
```

```python
mpl.rcParams['xtick.labelsize'] = 14
mpl.rcParams['ytick.labelsize'] = 14

# Make new prediction beyond training range
ts_test = ts[::sim_discretization][ttsplit:] - ts[ttsplit*sim_discretization] # need to predict from 0
X_predict = model.predict_dst(ts_test, X[:,ttsplit])

ts_predict = ts[::measurement_skip][ttsplit*mm:]
res_predict = fshape(X_predict, mm, True).real

fig, ax = plt.subplots(1,figsize=[10,2.25])
y_lim = [-1.2,1.2]
for m in range(n_measure):
    ax.plot(ts, result.expect[m], color=cmap(m), label=measure_name[m], alpha=.5)

    ax.plot(ts[::measurement_skip][ttsplit*mm:], res[m,::measurement_skip][ttsplit*mm:],
            ls='', ms=10, color=cmap(m), marker='o', markerfacecolor='None', markeredgewidth=1)

    ax.plot(ts_predict[mm:], res_predict[m][mm:], color=cmap(m), marker='x', ls='', markeredgewidth=2)
    ax.plot(ts_predict[:mm], res_predict[m][:mm], color=cmap(m), marker='o', ls='', markeredgewidth=2)

ax.vlines(ts[::sim_discretization], y_lim[0], y_lim[1], ls='--', color='k', alpha=0.33)
ax.vlines(ts_predict[::mm][0], y_lim[0], y_lim[1], ls='-', color='k', alpha=1)

ax.set_ylim(y_lim)
ax.set_xlim([0,11.8])
fig.tight_layout()
```

```python

```

# Example 3

This experiment will explore stroboscopic dynamics for the case where we are limited in the measurements we can take. This is like working in the interaction picture with the rotating wave approximation. Indeed, if you only consider the amplitudes of the resonance frequency we are studying a discretized envelope with a discretization length related to the period of the bare oscillation (or, more accurately, the carrier that we have tuned to the bare oscillation). What is more, if we recall what we saw from the Floquet DMD results in Figure 2, the use of the Floquet theory makes these models accurate independent of the RWA.

Based on these ideas, the Floquet theory stuff we talk about here really only makes sense if you're dealing with strong coupling where the RWA breaks down. If you are not, you can take the RWA and work in an interaction frame with envelopes like we talked about in the variation on Example 1.

Also, as an aside, if you are in the interaction picture then the amplitudes attached to the resonance frequency are like the value of the constant time step on the envelope. However, this framework lets you worry about other frequencies.

## Construct a model from Fourier coefficients

```python
# Hamiltonian
w0 = 1
V0 = 1
Delta = 2*np.pi*w0
H0 = Delta/2*qt.sigmaz()
H1 = V0*1/2*qt.sigmax()

# Measurements
measure_list = [qt.sigmax(), qt.sigmay(), qt.sigmaz()]
measure_name = ['<{}>'.format(s) for s in ['x','y','z']]
n_measure = len(measure_list)
y_lim = [-1.5,1.5]

# Fourier/Floquet parameters
Omega = Delta # -- decide how often you want to measure
T = (2*np.pi/Omega)
n_fourier = 7

# Time parameters
n_discretization = 16
n_periods = 8
n_strobes = 2

# Count of indicies between various lengths of time
i_total = n_periods*n_discretization*n_fourier
i_period = n_fourier*n_discretization
i_strobe = i_period//n_strobes

# Time series for simulation
ts = np.linspace(0, n_periods*T, i_total, endpoint=True)
dt_fourier = ts[n_discretization] - ts[0]

# Generate monotone control signals
control_list = []
repetition = 1
ampl_list = np.linspace(-1,1,41)[::-1].tolist()
for ampl in ampl_list*repetition:
    for u_hat in hf.dft_coef_span(n_fourier):
        control_list.append({'u_hat': ampl*u_hat, 'dt': dt_fourier})
n_experiments = len(control_list)
```

```python
def run_experiment(args, order):
    # Make a data containers
    QbData = QubitExperiment()

    # Resample the control
    resampled_args = hf.resample_u_dft(args, n_discretization)
    control = hf.u_idft(ts, resampled_args)

    # Run the experiment
    psi0 = qt.basis(2,1)
    H = [H0, [H1, control]]
    result = qt.mesolve(H, psi0, ts, [], measure_list)
    X = np.array(result.expect)

    # Collect strobe data
    Y = fshape(X[:,::i_strobe],  n_strobes)
    Ty = ts[::i_period]

    # We only have monotone controls so we just repeat the pulse
    control_library = [hf.make_control_library(args['u_hat'], order)]*n_periods
    Uy = np.hstack(control_library)

    # Save data
    QbData.save(X, Y, Uy, Ty, resampled_args)

    return QbData
```

```python
# Library order
order = 3

# Run multi-thread experiments
def run_wrapper(arg):
    return run_experiment(arg, order=order)
with pool.Pool(6) as p:
    r = list(tqdm(p.imap(run_wrapper, control_list), total=n_experiments))
QbData = join_qubit_experiments(r)
```

    HBox(children=(IntProgress(value=0, max=287), HTML(value='')))

```python
# Fit model
model = dmd.biDMD(QbData.Y2_flat, QbData.Y1_flat, QbData.U1y_flat, QbData.T1y_flat, threshold_type='percent', threshold=1e-5)
fig, ax = dmd.plot_eigs(model.eigs, figsize=[3,3])
ax.set_xlim([-1.1,1.5])
model.eigs
```

These eigenvalues only tell part of the story because there's a lot of dynamics in the control operator. But we can get some sense of a bare oscillation from this, maybe.

Let's look at how our model handles reproducing the training data. Try going into the training and increasing the order of the multinomial powers if you do not see a good enough fit in $\texttt{run_experiment(arg, order)}$. Also, try lowering the order to see the fit break down.

```python
# Plot only one ampltiude set of the training pulses
plot_ampl = 0
plot_reps = len(ampl_list)*repetition
for i in range(plot_ampl*n_experiments//plot_reps, (plot_ampl+1)*n_experiments//plot_reps):
    fig,axes = plt.subplots(1,2, figsize=[12,3])
    ax=axes[0]
    Y_prediction = model.predict_dst(control=QbData.Uy[i], x0=QbData.Y[i][:,0])
    prediction = fshape(Y_prediction, n_strobes, True)
    for m in range(3):
        ax.plot(ts, QbData.X[i][m], color=cmap2(2*m+1))
    for m in range(3):
        ax.plot(ts[::i_strobe], QbData.X[i][m,::i_strobe], ls='', marker='o', markeredgewidth=2, markerfacecolor='None', markeredgecolor=cmap2(2*m), markersize=10)
        ax.plot(ts[::i_strobe], prediction[m], ls='', markeredgewidth=1, marker='x', color=cmap2(2*m), markersize=8)
    ax.set_ylim([-1.2,1.2])

    ax=axes[1]
    control = hf.u_idft(ts[:i_period], QbData.Args_List[i])
    ax.plot(ts[:i_period], control)
    ax.set_ylim([-1.5,1.5])
```

At this point I want to comment that the data-drive approach means that when you measure something like $\sigma_x$ with a fast oscillation in the lab frame, you do not need to be exactly correct. That is, if you are slightly rotated off the correct value by a fixed phase when you apply the pulse to generate the $\sigma_x$ measurement, if you are off by the same phase, the data-drive model will account for this.

Said another way, we don't actually need to measure the Bloch vector. We only need to measure a set of linearly independent observables sufficient to reconstruct the dynamics.

## Apply the model to new controls

```python
# In this version you supply a list of control arguments, one for each Floquet period.
# You might need a control function for the true signal if it's outside the Fourier span.
# We'll also mess with the initial condition.
def run_experiment_v2(args_list, order, control_fn=None):
    # Make data container
    QbData = QubitExperiment()

    control = []
    control_library = []
    saved_args = []
    for i, args in enumerate(args_list):
        # Get the appropriate time window
        ts_window = ts[i*i_period:(i+1)*i_period]
        # Resample the control
        saved_args.append(hf.resample_u_dft(args, n_discretization))
        # Store the control and control library
        if control_fn:
            control.append(control_fn(ts_window))
        else:
            control.append(hf.u_idft(ts_window, saved_args[-1]))
        control_library.append(hf.make_control_library(args['u_hat'], order))
    control = np.hstack(control)
    Uy = np.hstack(control_library)

    # Run the experiment
    psi0 = qt.basis(2,1)
    psi0 = psi0/psi0.norm()
    H = [H0, [H1, control]]
    result = qt.mesolve(H, psi0, ts, [], measure_list)
    X = np.array(result.expect)

    # Collect strobe data
    Y = fshape(X[:,::i_strobe],  n_strobes)
    Ty = ts[::i_period]

    # Save data
    QbData.save(X, Y, Uy, Ty, saved_args)

    return QbData
```

Now we run a three test examples. One at resonance, one random but in the training span, and one outside the training span.

```python
n_tests = 3
QbList_Test = []
ts_fourier = ts[:i_period][::n_discretization]

# 1. Resonance
u_hat = np.fft.fft(np.cos(Delta*ts_fourier), n_fourier)
QbList_Test.append(run_experiment_v2([{'u_hat': u_hat, 'dt': dt_fourier}]*n_periods, order, lambda t: np.cos(Delta*t)))

# 2. Random control -- Generate a random pulse with all the Frequency content with only positive Fourier coefficients (so something happens)
test_list = [{'u_hat': 1.5*hf.dft_coef_rand(n_fourier, 3, True), 'dt': dt_fourier} for i in range(n_periods)]
QbList_Test.append(run_experiment_v2(test_list, order))

# 3. Sawtooth wave at resonance
u_hat = np.fft.fft(sp.signal.sawtooth(Delta*ts_fourier), n_fourier)
QbList_Test.append(run_experiment_v2([{'u_hat': u_hat, 'dt': dt_fourier}]*n_periods, order, lambda t: sp.signal.sawtooth(Delta*t)))

# Plot the experiments
QbTest = join_qubit_experiments(QbList_Test)
for i in range(n_tests):
    fig,axes = plt.subplots(1,2, figsize=[12,3])
    ax=axes[0]
    Y_prediction = model.predict_dst(control=QbTest.Uy[i], x0=QbTest.Y[i][:,0])
    prediction = fshape(Y_prediction, n_strobes, True)
    for m in range(3):
        ax.plot(ts, QbTest.X[i][m], color=cmap2(2*m+1))
    for m in range(3):
        ax.plot(ts[::i_strobe], QbTest.X[i][m,::i_strobe], ls='', marker='o', markeredgewidth=2, markerfacecolor='None', markeredgecolor=cmap2(2*m), markersize=10)
        ax.plot(ts[::i_strobe], prediction[m], ls='', markeredgewidth=1, marker='x', color=cmap2(2*m), markersize=8)
        ax.set_ylim([-1.4,1.4])

    ax=axes[1]
    for window, args in enumerate(QbTest.Args_List[i]):
        ts_window = ts[window*i_period:(window+1)*i_period]
        ax.plot(ts_window, hf.u_idft(ts_window, args), color='k')
        ax.set_ylim([-1.4,1.4])
```

```python

```

```python

```

```python

```