# 1.3 Performing AM in Code Space

## Introduction

**This section corresponds to Section 3.3 of the original paper.**

Learning density models $p(x)$ that directly approximate the data distribution can be difficult, or even nearly impossible for complex datasets. Generative models do not explicitly provide the density function but are able to sample from it through the following steps:

1. Sample from a simple distribution $q(z) \sim \mathcal{N}(0,I)$ defined in some abstract code space $\mathcal{Z}$.
2. Apply to the sample a decoding function $g : \mathcal{Z} \rightarrow \mathcal{X}$ that maps it back to the original input domain.

One such model that have gained popularity over the recent years is the Generative Adversarial Network (GAN). It learns a decoding function $g$ such that the generated images are theoretically impossible to distinguish from real images. The decoding function (generator) and the discriminant (discriminator) are typically neural networks. Here are some [great](https://wiseodd.github.io/techblog/2016/09/17/gan-tensorflow/) [blogs](http://blog.aylien.com/introduction-generative-adversarial-networks-code-tensorflow/) on GANs, if you are not familiar with them.

[Nguyen et al.](https://arxiv.org/abs/1605.09304) proposed a method of building prototype images $x^{*}$ for each labels $\omega_c$ by incorporating a pretrained generative model into the activation maximization framework. The optimization objective is redefined as (check tutorial 1.1 if you don't remember the original objective):

\begin{equation}
\max_{z \, \in \, \mathcal{Z}} \log p(\omega_c \, | \, g(z)) - \lambda \lVert z\rVert^2
\end{equation}

Now, instead of optimizing an image, we are optimizing the code $z$ such that the generated image $g(z)$ will maximize the activation for a particular class $\omega_c$. Once the solution $z^{*}$ to the optimization problem is found, the prototype for $\omega_c$ is generated by passing $z^{*}$ through the generator. That is, $x^{*} = g(z^{*})$. Hopefully, the prototype will turn out to be more realistic than vanilla AM, as the generator knows how to generate natural-looking images.

## Training Details

Remember how I said we incorporate a *pretrained* generative model into the AM framework? Here I will explain how we train the generator. The paper for this particular GAN framework by Nguyen et al. can be found [here](https://arxiv.org/abs/1602.02644). First take a look at the overall schematic for the GAN framework.

The training process involves four networks:

1. A pretrained encoder network $E$ to be inverted.
2. A generator network $G$.
3. A pretrained classifier $C$.
4. A discriminator $D$.

Note that we only have three networks in the above schematic while we require four. This is because the pretrained classifier serves as both $E$ and $C$. The activation of the hidden layer of the pretrained network acts as an encoding of the input as well as a means of comparing prominent features of two images. Now that we have figured out the components of this framework, let's see how the loss functions for $G$ and $D$ are defined. There are three parts.

#### Loss in Feature Space

Given a classifier $C: \mathbb{R}^{W \times H \times C} \rightarrow \mathbb{R}^{F}$, we define

\begin{equation}
L_{feat} = \sum_i \lVert C(G(x_i)) - C(y_i) \rVert^2
\end{equation}

where $y_i$ is the real image and $x_i$ is the encoding of $y_i$. With this loss, we are trying to encourage the generator to produce images whose features are similar to those of real images.

#### Adversarial Loss

The discriminator loss $L_{discr}$ and the generator loss $L_{adv}$ is given as follows:

\begin{align}
L_{discr} = - \sum_i \log (D(y_i)) + \log (1 - D(G(x_i)))
L_{adv} = - \sum_i \log (D(G(x_i))
\end{align}

This is just the GAN objective that constrains the generator to produce natural-looking images.

#### Loss in Image Space

Adding a pixel-wise loss stabilizes training.

\begin{equation}
L_{img} = \sum_i \lVert G(x_i) -  y_i \rVert^2
\end{equation}

Now that we have a pretrained generator, we can simply plug $G$ into the AM framework. As I mentioned above, we are **not** training the generator nor the classifier. We are optimizing the code $z$ that goes into the generator such that the generated image maximizes the activation for the class $\omega_c$.

## Tensorflow Walkthrough

### 1. Import Dependencies

We import the classifier (DNN), generator and the discriminator-the three components necessary for this AM framework.

```python
import os

from tensorflow.examples.tutorials.mnist import input_data
import matplotlib.pyplot as plt
import tensorflow as tf
import numpy as np

from models.models_1_3 import MNIST_DNN, MNIST_G, MNIST_D
from utils import plot

%matplotlib inline

mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)

logdir = './tf_logs/1_3_AM_Code/'
ckptdir = logdir + 'model'

if not os.path.exists(logdir):
    os.mkdir(logdir)
```

    Extracting MNIST_data/train-images-idx3-ubyte.gz
    Extracting MNIST_data/train-labels-idx1-ubyte.gz
    Extracting MNIST_data/t10k-images-idx3-ubyte.gz
    Extracting MNIST_data/t10k-labels-idx1-ubyte.gz

### 2. Building DNN Graph

In this step, we initialize a DNN classifier and attach necessary nodes for model training onto the computation graph.

```python
with tf.name_scope('Classifier'):

    # Initialize neural network
    DNN = MNIST_DNN('DNN')

    # Setup training process
    lmda = tf.placeholder_with_default(0.01, shape=[], name='lambda')
    X = tf.placeholder(tf.float32, [None, 784], name='X')
    Y = tf.placeholder(tf.float32, [None, 10], name='Y')

    tf.add_to_collection('placeholders', lmda)
    tf.add_to_collection('placeholders', X)
    tf.add_to_collection('placeholders', Y)

    code, logits = DNN(X)

    cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=Y))

    optimizer = tf.train.AdamOptimizer().minimize(cost, var_list=DNN.vars)

    correct_prediction = tf.equal(tf.argmax(logits, 1), tf.argmax(Y, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

cost_summary = tf.summary.scalar('Cost', cost)
accuray_summary = tf.summary.scalar('Accuracy', accuracy)
summary = tf.summary.merge_all()
```

### 3. Building GAN Subgraph

We now build the GAN part of the computation graph. If you see the right hand side of the $G$ cost, you can see that it is comprised of three terms. They are $L_{adv}$, $L_{img}$ and $L_{feat}$ from left to right.

```python
with tf.name_scope('GAN'):

    G = MNIST_G(z_dim=100, name='Generator')
    D = MNIST_D(name='Discriminator')

    X_fake = G(code)
    D_real = D(X)
    D_fake = D(X_fake, reuse=True)
    code_fake, logits_fake = DNN(X_fake, reuse=True)

    D_cost = -tf.reduce_mean(tf.log(D_real + 1e-7) + tf.log(1 - D_fake + 1e-7))
    G_cost = -tf.reduce_mean(tf.log(D_fake + 1e-7)) + tf.nn.l2_loss(X_fake - X) + tf.nn.l2_loss(code_fake - code)

    D_optimizer = tf.train.AdamOptimizer().minimize(D_cost, var_list=D.vars)
    G_optimizer = tf.train.AdamOptimizer().minimize(G_cost, var_list=G.vars)
```

### 4. Building Subgraph for Generating Prototypes

Before training the network, a subgraph for generating prototypes is added onto the graph. This subgraph will be used after training the model.

```python
with tf.name_scope('Prototype'):

    code_mean = tf.placeholder(tf.float32, [10, 100], name='code_mean')
    code_prototype = tf.get_variable('code_prototype', shape=[10, 100], initializer=tf.random_normal_initializer())

    X_prototype = G(code_prototype, reuse=True)
    Y_prototype = tf.one_hot(tf.cast(tf.lin_space(0., 9., 10), tf.int32), depth=10)
    _, logits_prototype = DNN(X_prototype, reuse=True)

    # Objective function definition
    cost_prototype = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits_prototype, labels=Y_prototype)) \
                     + lmda * tf.nn.l2_loss(code_prototype - code_mean)

    optimizer_prototype = tf.train.AdamOptimizer().minimize(cost_prototype, var_list=[code_prototype])

# Add the subgraph nodes to a collection so that they can be used after training of the network
tf.add_to_collection('prototype', code)
tf.add_to_collection('prototype', code_mean)
tf.add_to_collection('prototype', code_prototype)
tf.add_to_collection('prototype', X_prototype)
tf.add_to_collection('prototype', Y_prototype)
tf.add_to_collection('prototype', logits_prototype)
tf.add_to_collection('prototype', cost_prototype)
tf.add_to_collection('prototype', optimizer_prototype)
```

This is the general structure of the computation graph visualized using tensorboard.

### 5. Training Network

This is the step where the DNN is trained to classify the 10 digits of the MNIST images. Summaries are written into the logdir and you can visualize the statistics using tensorboard by typing this command: `tensorboard --lodir=./tf_logs`

```python
sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())

saver = tf.train.Saver()
file_writer = tf.summary.FileWriter(logdir, tf.get_default_graph())

# Hyper parameters
training_epochs = 15
batch_size = 100

for epoch in range(training_epochs):
    total_batch = int(mnist.train.num_examples / batch_size)
    avg_cost = 0
    avg_acc = 0

    for i in range(total_batch):
        batch_xs, batch_ys = mnist.train.next_batch(batch_size)
        _, c, a, summary_str = sess.run([optimizer, cost, accuracy, summary], feed_dict={X: batch_xs, Y: batch_ys})
        avg_cost += c / total_batch
        avg_acc += a / total_batch

        file_writer.add_summary(summary_str, epoch * total_batch + i)

    print('Epoch: {:04d} cost = {:.9f} accuracy = {:.9f}'.format(epoch + 1, avg_cost, avg_acc))

    saver.save(sess, ckptdir)

print('Accuracy:', sess.run(accuracy, feed_dict={X: mnist.test.images, Y: mnist.test.labels}))
```

    Epoch: 0001 cost = 0.237656591 accuracy = 0.930690911
    Epoch: 0002 cost = 0.088291109 accuracy = 0.972981828
    Epoch: 0003 cost = 0.055177446 accuracy = 0.982690919
    Epoch: 0004 cost = 0.037826146 accuracy = 0.988145463
    Epoch: 0005 cost = 0.029179757 accuracy = 0.990963644
    Epoch: 0006 cost = 0.023033684 accuracy = 0.992272734
    Epoch: 0007 cost = 0.017406837 accuracy = 0.994400005
    Epoch: 0008 cost = 0.015977912 accuracy = 0.994672732
    Epoch: 0009 cost = 0.013727595 accuracy = 0.995818186
    Epoch: 0010 cost = 0.011904787 accuracy = 0.996363640
    Epoch: 0011 cost = 0.015322508 accuracy = 0.994763641
    Epoch: 0012 cost = 0.010558971 accuracy = 0.996327276
    Epoch: 0013 cost = 0.009135132 accuracy = 0.997018185
    Epoch: 0014 cost = 0.008601392 accuracy = 0.997127275
    Epoch: 0015 cost = 0.008166363 accuracy = 0.997127275
    Accuracy: 0.98

### 6. Training GAN

We now train $G$ (and $D$) according to the description above.

```python
# Hyper parameters
training_epochs = 25
batch_size = 100
img_epoch = 1

for epoch in range(training_epochs):
    total_batch = int(mnist.train.num_examples / batch_size)
    avg_D_cost = 0
    avg_G_cost = 0

    for i in range(total_batch):
        batch_xs, _ = mnist.train.next_batch(batch_size)
        feed_dict = {X: batch_xs}

        _, D_c = sess.run([D_optimizer, D_cost], feed_dict=feed_dict)
        _, G_c = sess.run([G_optimizer, G_cost], feed_dict=feed_dict)

        avg_D_cost += D_c / total_batch
        avg_G_cost += G_c / total_batch

    print('Epoch: {:04d} G cost = {:.9f} D cost = {:.9f}'.format(epoch + 1, avg_G_cost, avg_D_cost))

# Uncomment this code if you want to see the generated images.
#
#     if (epoch + 1) % img_epoch == 0:
#         samples = sess.run(X_fake, feed_dict={X: mnist.test.images[:16, :]})
#         fig = plot(samples, 784, 1)
#         plt.savefig('./assets/1_3_AM_Code/G_{:04d}.png'.format(epoch), bbox_inches='tight')
#         plt.close(fig)

    saver.save(sess, ckptdir)

sess.close()
```

    Epoch: 0001 G cost = 5281.714079368 D cost = 0.017211454
    Epoch: 0002 G cost = 2569.694126864 D cost = 0.006985294
    Epoch: 0003 G cost = 2222.643753551 D cost = 0.007533360
    Epoch: 0004 G cost = 2040.322469150 D cost = 0.007319166
    Epoch: 0005 G cost = 1915.783002042 D cost = 0.007434430
    Epoch: 0006 G cost = 1825.571832386 D cost = 0.004931086
    Epoch: 0007 G cost = 1751.774513494 D cost = 0.005131674
    Epoch: 0008 G cost = 1695.891795987 D cost = 0.005583258
    Epoch: 0009 G cost = 1648.915287420 D cost = 0.004081625
    Epoch: 0010 G cost = 1606.690468084 D cost = 0.004811222
    Epoch: 0011 G cost = 1574.847379927 D cost = 0.002461862
    Epoch: 0012 G cost = 1548.405506925 D cost = 0.004549469
    Epoch: 0013 G cost = 1523.959982688 D cost = 0.004773995
    Epoch: 0014 G cost = 1502.840756392 D cost = 0.001094513
    Epoch: 0015 G cost = 1484.560579057 D cost = 0.003348438
    Epoch: 0016 G cost = 1470.541124157 D cost = 0.001267561
    Epoch: 0017 G cost = 1452.866942472 D cost = 0.006172946
    Epoch: 0018 G cost = 1441.055687811 D cost = 0.001345314
    Epoch: 0019 G cost = 1430.559412731 D cost = 0.000703439
    Epoch: 0020 G cost = 1417.898278809 D cost = 0.000637727
    Epoch: 0021 G cost = 1408.476033159 D cost = 0.005673418
    Epoch: 0022 G cost = 1400.744808683 D cost = 0.000871855
    Epoch: 0023 G cost = 1390.985546875 D cost = 0.000661015
    Epoch: 0024 G cost = 1384.504009011 D cost = 0.002004977
    Epoch: 0025 G cost = 1374.515778143 D cost = 0.000687619

This is a visualization of the GAN training process. The image quality initially improves although it somewhat plateaus out in the later epochs.

### 7. Restoring Subgraph

Here we first rebuild the DNN graph from metagraph, restore DNN parameters from the checkpoint and then gather the necessary nodes for prototype generation using the `tf.get_collection()` function (recall prototype subgraph nodes were added onto the 'prototype' collection at step 4).

```python
tf.reset_default_graph()

sess = tf.InteractiveSession()

new_saver = tf.train.import_meta_graph(ckptdir + '.meta')
new_saver.restore(sess, tf.train.latest_checkpoint(logdir))

# Get necessary placeholders
placeholders = tf.get_collection('placeholders')
lmda = placeholders[0]
X = placeholders[1]

# Get prototype nodes
prototype = tf.get_collection('prototype')
code = prototype[0]
code_mean = prototype[1]
X_prototype = prototype[3]
cost_prototype = prototype[6]
optimizer_prototype = prototype[7]
```

    INFO:tensorflow:Restoring parameters from ./tf_logs/1_3_AM_Code/model

### 8. Generating Prototype Images

Before performing gradient ascent, we calculate the image means $\overline{z}$ that will be used to regularize the prototype images. Then, we generate prototype images that maximize $\log p(\omega_c \, | \, g(z)) - \lambda \lVert z - \overline{z}\rVert^2$. I used 0.1 for lambda (lmda), but fine tuning may produce better prototype images.

```python
images = mnist.train.images
labels = mnist.train.labels

code_means = []
for i in range(10):
    imgs = images[np.argmax(labels, axis=1) == i]
    img_codes = sess.run(code, feed_dict={X: imgs})
    code_means.append(np.mean(img_codes, axis=0))

for epoch in range(15000):
    _, c = sess.run([optimizer_prototype, cost_prototype], feed_dict={lmda: 0.1, code_mean: code_means})

    if epoch % 500 == 0:
        print('Epoch: {:05d} Cost = {:.9f}'.format(epoch, c))

X_prototypes = sess.run(X_prototype)

sess.close()
```

    Epoch: 00000 Cost = 358.779754639
    Epoch: 00500 Cost = 275.063323975
    Epoch: 01000 Cost = 214.894088745
    Epoch: 01500 Cost = 168.886154175
    Epoch: 02000 Cost = 132.707305908
    Epoch: 02500 Cost = 103.784980774
    Epoch: 03000 Cost = 80.557037354
    Epoch: 03500 Cost = 61.938266754
    Epoch: 04000 Cost = 47.098171234
    Epoch: 04500 Cost = 35.386249542
    Epoch: 05000 Cost = 26.258031845
    Epoch: 05500 Cost = 19.242551804
    Epoch: 06000 Cost = 13.927100182
    Epoch: 06500 Cost = 9.950549126
    Epoch: 07000 Cost = 7.006939888
    Epoch: 07500 Cost = 4.849003792
    Epoch: 08000 Cost = 3.284199476
    Epoch: 08500 Cost = 2.165551901
    Epoch: 09000 Cost = 1.380895495
    Epoch: 09500 Cost = 0.844008684
    Epoch: 10000 Cost = 0.488746434
    Epoch: 10500 Cost = 0.264367282
    Epoch: 11000 Cost = 0.131287798
    Epoch: 11500 Cost = 0.058500197
    Epoch: 12000 Cost = 0.022639306
    Epoch: 12500 Cost = 0.007277159
    Epoch: 13000 Cost = 0.001836700
    Epoch: 13500 Cost = 0.000343904
    Epoch: 14000 Cost = 0.000050971
    Epoch: 14500 Cost = 0.000013767

### 9. Displaying Images

By incorporating $G$ into the AM framework, we are now able to produce realistic images. Recall that just using AM resulted in blurry prototype images.

```python
plt.figure(figsize=(15,15))
for i in range(5):
    plt.subplot(5, 2, 2 * i + 1)
    plt.imshow(np.reshape(X_prototypes[2 * i], [28, 28]), cmap='gray', interpolation='none')
    plt.title('Digit: {}'.format(2 * i))
    plt.colorbar()

    plt.subplot(5, 2, 2 * i + 2)
    plt.imshow(np.reshape(X_prototypes[2 * i + 1], [28, 28]), cmap='gray', interpolation='none')
    plt.title('Digit: {}'.format(2 * i + 1))
    plt.colorbar()

plt.tight_layout()
```

```python

```


# Optimization of a State-to-State Transfer in a Two-Level-System

```python
# NBVAL_IGNORE_OUTPUT
%load_ext watermark
#%load_ext autoreload
#%autoreload 2
import qutip
import numpy as np
import scipy
import matplotlib
import matplotlib.pylab as plt
import krotov
%watermark -v --iversions
π = np.pi
sqrt = np.sqrt
basis = qutip.basis
tensor = qutip.tensor
coherent = qutip.coherent

```

    The watermark extension is already loaded. To reload it, use:
      %reload_ext watermark
    matplotlib.pylab 1.15.4
    numpy            1.15.4
    scipy            1.2.0
    qutip            4.3.1
    matplotlib       3.0.2
    krotov           0.3.0
    CPython 3.7.2
    IPython 7.2.0

$\newcommand{tr}[0]{\operatorname{tr}}
\newcommand{diag}[0]{\operatorname{diag}}
\newcommand{abs}[0]{\operatorname{abs}}
\newcommand{pop}[0]{\operatorname{pop}}
\newcommand{aux}[0]{\text{aux}}
\newcommand{opt}[0]{\text{opt}}
\newcommand{tgt}[0]{\text{tgt}}
\newcommand{init}[0]{\text{init}}
\newcommand{lab}[0]{\text{lab}}
\newcommand{rwa}[0]{\text{rwa}}
\newcommand{bra}[1]{\langle#1\vert}
\newcommand{ket}[1]{\vert#1\rangle}
\newcommand{Bra}[1]{\left\langle#1\right\vert}
\newcommand{Ket}[1]{\left\vert#1\right\rangle}
\newcommand{Braket}[2]{\left\langle #1\vphantom{#2} \mid
#2\vphantom{#1}\right\rangle}
\newcommand{op}[1]{\hat{#1}}
\newcommand{Op}[1]{\hat{#1}}
\newcommand{dd}[0]{\,\text{d}}
\newcommand{Liouville}[0]{\mathcal{L}}
\newcommand{DynMap}[0]{\mathcal{E}}
\newcommand{identity}[0]{\mathbf{1}}
\newcommand{Norm}[1]{\lVert#1\rVert}
\newcommand{Abs}[1]{\left\vert#1\right\vert}
\newcommand{avg}[1]{\langle#1\rangle}
\newcommand{Avg}[1]{\left\langle#1\right\rangle}
\newcommand{AbsSq}[1]{\left\vert#1\right\vert^2}
\newcommand{Re}[0]{\operatorname{Re}}
\newcommand{Im}[0]{\operatorname{Im}}$
The purpose of this example is not to solve an especially interesting physical
problem but to give a rather simple example of how the package can be used in
order to solve an optimization problem.

## Define the Hamiltonian

In the
following the Hamiltonian, guess field and
states are defined.

The Hamiltonian
$\op{H}_{0} = - \omega \op{\sigma}_{z}$
represents a
simple qubit with energy
level splitting $\omega$ in the basis
$\{\ket{0},\ket{1}\}$. The control
field
$\epsilon(t)$ is assumed to couple via
the
Hamiltonian $\op{H}_{1}(t) =
\epsilon(t) \op{\sigma}_{x}$ to the qubit,
i.e., the control
field effectively
drives
transitions between both qubit
states. For now, we initialize the control
field as constant.

```python
N = 1
L = 3
α = 2
```

```python
qutip.visualization.plot_wigner(coherent(100,α)+coherent(100,-α))
```

# Plotting functions

```python
def to_two_level(state):
    if state.type is 'oper':
        return qutip.Qobj(state[0:2,0:2])
    else:
        return qutip.Qobj(state[0:2])

def plot_population(n, tlist):
    fig, ax = plt.subplots()
    leg = []
    for i in range(len(n)):
        ax.plot(tlist, n[i], label=str(i))
        leg.append(str(i))
    ax.legend()
    ax.set_xlabel('Time (ns)')
    ax.set_ylabel('Occupation')
    ax.legend(leg)
    plt.show(fig)

def plot_pulse(pulse, tlist):
    fig, ax = plt.subplots(figsize=(15,4))
    if callable(pulse):
        pulse = np.array([pulse(t, args=None) for t in tlist])
    if np.any(np.iscomplex(pulse)):
        ax.plot(tlist, np.real(pulse))
        ax.plot(tlist, np.imag(pulse))
        ax.legend(['Re', 'Im'])
    else:
        ax.plot(tlist, pulse)
    ax.set_xlabel('Time (ns)')
    ax.set_ylabel('Pulse amplitude')
    plt.show(fig)

def plot_system(ψ):
    bl = qutip.Bloch()
    bl.add_states(to_two_level(ψ.ptrace(0)))
    bl.show()
    qutip.visualization.plot_wigner_fock_distribution(to_two_level(ψ.ptrace(1)))
def plot_resonator(ψ):
    fig, ax = plt.subplots(1,len(ψ), figsize=(3*len(ψ),3))
    if len(ψ)<=1:
        ψ = ψ[0]
        axis = ax
        qutip.visualization.plot_wigner(ψ.ptrace(1), fig=fig, ax=axis, alpha_max = 2*α)
        axis.axis_equal = True
    else:
        for (ϕ, axis) in zip(ψ, ax):
            qutip.visualization.plot_wigner(ϕ.ptrace(1), fig=fig, ax=axis, alpha_max = 2*α)
            axis.axis_equal = True

def plot_cardinal(ψ):
    bl = qutip.Bloch()
    bl.vector_color = ['r','g','b','g','b','r']
    [bl.add_states(to_two_level(ϕ.ptrace(0)), 'vector') for ϕ in ψ]
    bl.show()

def plot_all(dyn, ψ):
    ψ_i = [g.states[0] for g in dyn]
    ψ_f = [g.states[-1] for g in dyn]
    ψ_t = [ϕ[1] for ϕ in ψ]
    plot_cardinal(ψ_i)
    plot_resonator(ψ_i)
    plot_cardinal(ψ_t)
    plot_resonator(ψ_t)
    plot_cardinal(ψ_f)
    plot_resonator(ψ_f)
def plot_evolution(dyn, steps=1):
    for d in dyn:
        points = [to_two_level(s.ptrace(0)) for s in d.states[0:-1:steps]]
        bl = qutip.Bloch()
        bl.vector_color = 'r'
        bl.point_color = 'r'
        bl.point_marker = 'o'
        bl.add_states(points, 'point')
        bl.show()
        bl = qutip.Bloch()
        bl.vector_color = 'r'
        bl.point_color = 'r'
        bl.point_marker = 'o'
        bl.view = [bl.view[0], 80]
        bl.add_states(points, 'point')
        bl.show()
```

```python
def cat(N,α,P = 1):
    return (coherent(N,α) + P*coherent(N,-α)).unit()
def fid(result, target):
    return (np.abs((result.states[-1].dag()*target).full())**2)[0][0]
def proj(ψ, ϕ=None):
    if ϕ is None:
        return ψ * ψ.dag()
    else:
        return ψ * ϕ.dag()
```

```python
Ri = qutip.operators.identity(N)
Si = qutip.operators.identity(L)
ZERO = tensor(qutip.operators.qzero(L), qutip.operators.qzero(N))

σ_z = proj(qutip.basis(L, 0)) - proj(qutip.basis(L, 1))
σ_y = 1j*(proj(qutip.basis(L, 1),qutip.basis(L, 0)) - proj(qutip.basis(L, 0), qutip.basis(L, 1)))
σ_x = proj(qutip.basis(L, 0),qutip.basis(L, 1)) - proj(qutip.basis(L, 1), qutip.basis(L, 0))
a = qutip.operators.destroy(N)

σ_z = qutip.tensor(σ_z,Ri)
σ_y = qutip.tensor(σ_y,Ri)
σ_x = qutip.tensor(σ_x,Ri)
b = tensor(qutip.operators.destroy(L), Ri)

#a = qutip.tensor(Si,a)
I = qutip.tensor(Si,Ri)

cat_0 = cat(N,α)
cat_1 = cat(N,α*1j)
ω_r = 8.3056 * 2 * π      # resonator frequency
ω_q = 6.2815 * 2 * π      # qubit frequency
use_rotating = True
def hamiltonian(ω=1.0, ampl0=1, use_rotating=True, pulses=None, tlist=None):
    """Two-level-system Hamiltonian

    Args:
        ω (float): energy separation of the qubit levels
        ampl0 (float): constant amplitude of the driving field
    """

    K_r   = 2*π*0.45e-3   # Kerr res
    K_q   = -2*π*297e-3    # Kerr qubit 200-300 MHz
    ω_r = 8.3056 * 2 * π      # resonator frequency
    ω_q = 6.2815 * 2 * π      # qubit frequency
    χ = 0.025 * 2 * π   # parameter in the dispersive hamiltonian

    Δ = abs(ω_r - ω_q)    # detuning
    g = sqrt(Δ * χ)  # coupling strength that is consistent with chi
    γ = 1e1

    #H_occ = w_r*a.dag()*a + w_q*b.dag()*b
    if L==3:
        H_occ_q = tensor(qutip.Qobj(np.diag([0, ω_q, 2*ω_q])), Ri)
    else:
        H_occ_q = tensor(qutip.Qobj(np.diag([0, ω_q])), Ri)
    H_occ_r = ω_r * tensor(Si, a.dag()*a)
    if N == 1:
        H_occ =  H_occ_q
    else:
        H_occ =  H_occ_q  + H_occ_r

    use_dispersive = True
    use_kerr = False
    if use_dispersive:
        #H_coup = - chi_qr * a.dag()*a * b.dag()*b
        H_coup =  χ * (tensor(Si, a.dag()*a) + I/2) * σ_z
    else:
        #H_coup = g * (a.dag() * b + a * b.dag())
        H_coup = g * σ_x *tensor(Si, a.dag() + a)
    if use_kerr:
        H_kerr = - K_r/2 * tensor(Si, a.dag()**2 * a**2) + K_q/2 * b.dag()**2 * b**2
    else:
        H_kerr = tensor(qutip.operators.qzero(L), qutip.operators.qzero(N))

    H_d = ZERO

    if use_rotating:
        H_d += H_kerr

        H_qr = (b.dag() + b)
        H_qi = 1j*(b.dag() - b)
        H_rr = (a + a.dag())
        H_ri = 1j*(a.dag() - a)

        ϵ_qr = lambda t, args: ampl0
        ϵ_qi = lambda t, args: ampl0
        ϵ_rr = lambda t, args: ampl0
        ϵ_ri = lambda t, args: ampl0

        # Random pulses (doesn't really work)
        #ϵ = lambda t, tlist, R: R[np.where(tlist<=t)[0][-1]]
        #O = np.random.rand(len(tlist))
        #ϵ_qr = lambda t, args: ϵ(t, tlist, O)
        #O = np.random.rand(len(tlist))
        #ϵ_qi = lambda t, args: ϵ(t, tlist, O)

        if pulses:
            ϵ_qr = pulses[0]
            ϵ_qi = pulses[1]
            ϵ_rr = np.zeros(len(pulses[0]))
            ϵ_ri = np.zeros(len(pulses[0]))

        return [H_d, [H_qr, ϵ_qr], [H_qi, ϵ_qi], [H_rr, ϵ_rr], [H_ri, ϵ_ri]]
    else:
        H_d += H_occ + H_kerr#+ H_coup

        H_q = b
        H_qc = b.dag()
        H_rr = ZERO
        H_ri = ZERO

        ϵ_q = lambda t, args: 1j*ampl0*np.exp(1j*ω_q*t)
        ϵ_qc = lambda t, args: -1j*ampl0*np.exp(-1j*ω_q*t)
        ϵ_rr = lambda t, args: ampl0
        ϵ_ri = lambda t, args: ampl0

        if pulses:
            ϵ_q = pulses[0]
            ϵ_qc = pulses[1]
            ϵ_rr = np.zeros(len(pulses[0]))
            ϵ_ri = np.zeros(len(pulses[0]))

        return [H_d, [H_q, ϵ_q], [H_qc, ϵ_qc], [H_rr, ϵ_rr], [H_ri, ϵ_ri]]

def coeffs_to_state(c,init = True):
    if init:
        ψ = tensor((c[0]*basis(L,0) + c[1]*basis(L,1)).unit() , (basis(N,0)))
    else:
        ψ = tensor((basis(L,0)) , (c[0]*cat_0 + c[1]*cat_1).unit())
    return ψ

def states(coeffs):
    return [[coeffs_to_state(c,True),coeffs_to_state(c,False)] for c in coeffs]
```

## Define the optimization target

First we define the time grid of the
dynamics, i.e., by taking the following
values as an example, we define the
initial state to be at time $t=0$ and
consider a total propagation time of
$T=5$. The entire time grid is divided into
$n_{t}=500$ equidistant time steps.

```python
T_r = (2*π)/ω_r
T_q = (2*π)/ω_q
T = 0.5 * T_q
T = 40
steps = 2000
tlist = np.linspace(0, T, steps)

H = hamiltonian(ampl0=1, use_rotating=True)
#coeffs = [(1,0), (1,-1), (1,1j), (1,1), (1,-1j), (0,1)]
#st = states(coeffs)
#ϕ = [[ tensor(basis(L,0), coherent(N,α)), tensor((basis(L,0)-basis(L,1)).unit(), coherent(N,α)) ]]
ϕ = [[ tensor(basis(L,0), qutip.Qobj(1)), tensor((basis(L,0)-basis(L,1)).unit(), qutip.Qobj(1)) ]]
F_oc_tar = 1-1e-4
```

Next, we define the optimization targets, which is technically a list of
objectives, but here it has just one entry defining a simple state-to-state
transfer
from initial state $\ket{\Psi_{\init}} = \ket{0}$ to the target state
$\ket{\Psi_{\tgt}} = \ket{1}$, which we want to reach at final time $T$. Note
that we also have to pass the Hamiltonian $\op{H}(t)$ that determines the
dynamics of
the system to the optimization objective.

```python
def state_rot(ϕ):
    if L == 3:
        rot_evo = tensor(qutip.Qobj([[1, 0, 0],[0, np.exp(-1j * ω_q * T), 0],[0, 0, 0]]), Ri)
    else:
        rot_evo = tensor(qutip.Qobj([[1, 0],[0, np.exp(-1j * ω_q * T)]]), Ri)
    ϕ[0][1] = rot_evo * ϕ[0][1]
    return ϕ
if use_rotating:
    objectives = [krotov.Objective(initial_state=ψ[0], target=ψ[1], H=H) for ψ in state_rot(ϕ)]
else:
    objectives = [krotov.Objective(initial_state=ψ[0], target=ψ[1], H=H) for ψ in ϕ]
```

In addition, we have to define and assign a shape function $S(t)$ for the update
in each control iteration to each
control field that will be updated. This shape
usually takes care of
experimental limits such as the necessity of finite ramps
at the beginning and
end of the control field or other conceivable limitations
for field shapes: wherever $S(t)$ is zero, the optimization will not change the
value of the control from the original guess.

```python
def S(t):
    """Shape function for the field update"""
    return 1
    return krotov.shapes.flattop(t, t_start=0, t_stop=T, t_rise=2, t_fall=2, func='sinsq')
def S2(t):
    return 0
def start_pulse(t):
    return 1
    return krotov.shapes.blackman(t, t_start=0, t_stop=π/2+4)
```

At this point, we also change the initial control field $\epsilon_{0}(t)$ from a
constant to a shaped pulse that switches on smoothly from zero and again
switches off at the final time $T$. We re-use the shape function $S(t)$ that we
defined for the updates for this purpose (although generally, $S(t)$ for the
updates has nothing to with the shape of the control field).

```python
def shape_field(ϵ, sf):
    """Applies the shape function S(t) to the guess field"""
    ϵ_shaped = lambda t, args: ϵ(t, args)*sf(t)
    return ϵ_shaped

S_funs = [start_pulse,start_pulse,S2,S2]
for i, H_i in enumerate(H[1:]):
    H_i[1] = shape_field(H_i[1], S_funs[i])

```

Having defined the shape function $S(t)$ and having shaped the guess field, we
now tell the optimization to also use $S(t)$ as the update-shape for
$\epsilon_0(t)$. In addition, we have to choose `lambda_a` for each control
field. It controls the update magnitude of the respective field in each
iteration.

```python

```

It is convenient to introduce the function `print_fidelity`, which can be passed
to the optimization procedure and will be called after each iteration and thus
provides additional feedback about the optimization progress.

## Simulate dynamics of the guess field

Before heading towards the optimization
procedure, we first simulate the
dynamics under the guess field
$\epsilon_{0}(t)$.

The following plot shows the guess field $\epsilon_{0}(t)$, which is, as chosen
above, just a constant field (with a smooth switch-on and switch-off)

```python
for H_i in H[1:]:
    plot_pulse(H_i[1], tlist)
```

The next line solves the equation of motion for the defined objective, which
contains the initial state $\ket{\Psi_{\init}}$ and the Hamiltonian $\op{H}(t)$
defining its evolution.

```python
guess_dynamics = [ob.mesolve(tlist, progress_bar=True, options=qutip.Options(nsteps=50000)) for ob in objectives]
```

```python
def qubit_occupation(dyn):
    occ = [tensor(basis(L,i)*basis(L,i).dag(),Ri) for i in range(0,L)]
    n = qutip.expect(occ, dyn.states)
    plot_population(n, dyn.times)

def plot_norm(result):
    state_norm = lambda i: result.states[i].norm()
    states_norm=np.vectorize(state_norm)

    fig, ax = plt.subplots()
    ax.plot(result.times, states_norm(np.arange(len(result.states))))
    ax.set_title('Norm loss', fontsize = 15)
    ax.set_xlabel('Time (ns)')
    ax.set_ylabel('State norm')
    plt.show(fig)
```

```python
qubit_occupation(guess_dynamics[0])
#plot_norm(guess_dynamics[0])
```

The plot of the population dynamics shows that the guess field does not transfer
the initial state $\ket{\Psi_{\init}} = \ket{0}$ to the desired target state
$\ket{\Psi_{\tgt}} = \ket{1}$.

```python
#plot_all(guess_dynamics, ϕ)
plot_evolution(guess_dynamics, steps=1)
```

```python
#plot_all(guess_dynamics, ϕ)
```

## Optimize

In the following we optimize the guess field $\epsilon_{0}(t)$ such
that the intended state-to-state transfer $\ket{\Psi_{\init}} \rightarrow
\ket{\Psi_{\tgt}}$ is solved.

The cell below carries out the optimization. It
requires, besides the
previously
defined optimization `objectives`, information
about the
optimization functional
$F$ (via `chi_constructor`) and the
propagation method that should be used. In
addition, the number of total
iterations is required and, as an option, we pass
an info-hook that after each
iteration combines a complete printout of the state
of the optimization with the
`print_fidelity` function defined above.

Here, we
choose $F = F_{re}$ with
\begin{equation}
F_{re}
=
\Re\Braket{\Psi(T)}{\Psi_{\tgt}}
\end{equation}

with
$\ket{\Psi(T)}$ the
forward propagated state of $\ket{\Psi_{\init}}$.

```python
c = 1e2
c_list = np.ones(len(H)-1)*c
pulse_options = {H_i[1]: dict(lambda_a=0.5, shape=S_funs[i]) for i, H_i in enumerate(H[1:])}
```

```python
def F_oc(fw_states_T, objectives, tau_vals=None, **kwargs):
    return np.abs(krotov.functionals.f_tau(fw_states_T, objectives, tau_vals, **kwargs))**2

def calc_fidelity(tau_vals):
    return np.abs(np.sum(tau_vals)/len(tau_vals))**2

def print_fidelity(**args):
    fid = calc_fidelity(np.array(args['tau_vals']))
    print("          F_t = {} | F = {} | F_t - F = {}".format(F_oc_tar, fid, F_oc_tar-fid))
    #plot_fid_convergence(args['info_vals']+[fid])
    #return fid
#def lambda_fid(fid, c, F=F_oc_tar):
#    return 1/(F-fid+1/c)-1/(F+1/c)
def modify_params(**kwargs):
    # Update λ
    #steps = 4
    λₐ = kwargs['lambda_vals'][0]
    #if kwargs['iteration'] % steps == 1:
    #    fid = calc_fidelity(kwargs['tau_vals'])
    #    fids = kwargs['info_vals'] + [fid]
    #    if dec_lambda(fids, steps=steps):
    #        λₐ /= 2
    #else if inc_lambda(fids):
    #    λₐ *= 2
    for i in range(len(kwargs['lambda_vals'])):
        kwargs['lambda_vals'][i] = λₐ
    print("λₐ = {}".format(kwargs['lambda_vals']))

    #tlist = kwargs['tlist']  # for numpy.fft.fftfreq
    #for pulse in kwargs['optimized_pulses']:
        # pulse is a numpy array that can be modified in place:
        # - call numpy.fft
        # - multiply values of spectrum by some filter function in [0, 1]
        # - call numpy.ifft
    #    pass
def plot_fid_convergence(info_vals):
    fig, ax = plt.subplots(1,1)
    ax.plot(info_vals)
    ax.set_xticks(np.arange(0, len(info_vals), step=1))
    ax.set_xlabel('Iteration')
    ax.set_ylabel('Fidelity')
    #ax.set_ylim((-0.2,.2))
    plt.show()
```

```python
# Reset results
oct_result = None
```

```python
iters = 25
if oct_result is not None:
    iters = oct_result.iters[-1] + iters

oct_result = krotov.optimize_pulses(
    objectives,
    pulse_options=pulse_options,
    tlist=tlist,
    propagator=krotov.propagators.expm,
    chi_constructor=krotov.functionals.chis_ss,
    info_hook=krotov.info_hooks.chain(
        krotov.info_hooks.print_table(J_T=F_oc),
        print_fidelity
    ),
    check_convergence=krotov.convergence.Or(
        krotov.convergence.value_above(F_oc_tar, name='F_oc'),
        #krotov.convergence.check_monotonic_fidelity,
    ),
    modify_params_after_iter = modify_params,
    iter_stop=iters,
    continue_from = oct_result,
)
```

```python
def inc_lambda(fids):
    fids = np.array(fids)
    fid_last = fids[-1]
    return np.any(fid_last<fids)
def dec_lambda(fids, steps=5):
    fids = np.array(fids)
    return is_alternating(np.diff(fids), steps=steps)
def is_alternating(data, steps = 5):
    if len(data)<steps:
        return False
    signs = np.sign(data[-steps:])
    first = signs[0]
    return np.all(signs[0::2]==first) and np.all(signs[1::2]==-first)
plot_fid_convergence(fids)
print(dec_lambda(fids[0:9]))
#plot_fid_convergence(dec_lambda(fids))
plot_fid_convergence(fids2)
print(dec_lambda(fids2[0:9]))
#plot_fid_convergence(dec_lambda(fids2))
plot_fid_convergence(fids3)
print(dec_lambda(fids3[0:9]))
#plot_fid_convergence(dec_lambda(fids3))

```

```python
plot_fid_convergence(oct_result.info_vals)
oct_result
```

## Simulate dynamics of the optimized field

Having obtained the optimized
control field, we can now
plot it and calculate the
population dynamics under
this field.

```python
[plot_pulse(c, tlist) for c in oct_result.optimized_controls];
```

```python
plot_pulse(np.abs(oct_result.optimized_controls[0]+1j*oct_result.optimized_controls[1]), tlist)
plot_pulse(np.angle(oct_result.optimized_controls[0]+1j*oct_result.optimized_controls[1]), tlist)
```

In contrast to the dynamics under the guess field, the optimized field indeed
drives the initial state $\ket{\Psi_{\init}} = \ket{0}$ to the desired target
state $\ket{\Psi_{\tgt}} = \ket{1}$.

```python
opt_dynamics = [ob.mesolve(tlist, progress_bar=True) for ob in oct_result.optimized_objectives]
```

    /home/wello/miniconda3/envs/qutip-env/lib/python3.7/site-packages/Cython/Compiler/Main.py:367: FutureWarning: Cython directive 'language_level' not set, using 2 for now (Py2). This will change in a later release! File: /mnt/c/Users/johan/Development/cat-state-encoding/Krotov/rhs320.pyx
      tree = Parsing.p_module(s, pxd, full_module_name)

    10.0%. Run time:   0.38s. Est. time left: 00:00:00:03
    20.0%. Run time:   0.63s. Est. time left: 00:00:00:02
    30.0%. Run time:   0.82s. Est. time left: 00:00:00:01
    40.0%. Run time:   1.00s. Est. time left: 00:00:00:01
    50.0%. Run time:   1.24s. Est. time left: 00:00:00:01
    60.0%. Run time:   1.47s. Est. time left: 00:00:00:00
    70.0%. Run time:   1.69s. Est. time left: 00:00:00:00
    80.0%. Run time:   1.95s. Est. time left: 00:00:00:00
    90.0%. Run time:   2.24s. Est. time left: 00:00:00:00
    Total run time:   2.53s

```python
qubit_occupation(opt_dynamics[0])
```

```python
plot_evolution(opt_dynamics)
#plot_all(opt_dynamics, ϕ)
```

# Simulate dynamics in lab frame

```python
#rotating_pulses = [np.vectorize(H_i[1])(tlist, None) for H_i in H[1:]]
rotating_pulses = oct_result.optimized_controls

tlist2 = np.linspace(0, T, int(np.ceil(200*T/T_q)))
Ω = rotating_pulses[0]+1j*rotating_pulses[1]
Ω = np.interp(tlist2, tlist, Ω)
pulses_lab = [Ω*np.exp(1j*ω_q*tlist2), np.conj(Ω)*np.exp(-1j*ω_q*tlist2)]
H_lab = hamiltonian(ω=1.0, ampl0=1, use_rotating=False, pulses=pulses_lab)
objectives_lab = [krotov.Objective(initial_state=ψ[0], target=ψ[1], H=H_lab) for ψ in ϕ]
opt_dynamics_lab = [ob.mesolve(tlist2, progress_bar=True, options=qutip.Options(nsteps=50000)) for ob in objectives_lab]
```

    /home/wello/miniconda3/envs/qutip-env/lib/python3.7/site-packages/Cython/Compiler/Main.py:367: FutureWarning: Cython directive 'language_level' not set, using 2 for now (Py2). This will change in a later release! File: /mnt/c/Users/johan/Development/cat-state-encoding/Krotov/rhs321.pyx
      tree = Parsing.p_module(s, pxd, full_module_name)

    10.0%. Run time:   7.14s. Est. time left: 00:00:01:04
    20.0%. Run time:  14.30s. Est. time left: 00:00:00:57
    30.0%. Run time:  21.60s. Est. time left: 00:00:00:50
    40.0%. Run time:  29.11s. Est. time left: 00:00:00:43
    50.0%. Run time:  36.20s. Est. time left: 00:00:00:36
    60.0%. Run time:  43.47s. Est. time left: 00:00:00:28
    70.0%. Run time:  50.85s. Est. time left: 00:00:00:21
    80.0%. Run time:  57.96s. Est. time left: 00:00:00:14
    90.0%. Run time:  65.35s. Est. time left: 00:00:00:07
    Total run time:  72.25s

```python
plot_pulse(Ω, tlist2)
[plot_pulse(p, tlist2) for p in pulses_lab];
```

```python
qubit_occupation(opt_dynamics_lab[0])
#plot_evolution(opt_dynamics_lab)
```

```python
plot_cardinal([opt_dynamics_lab[0].states[0]])
plot_cardinal([opt_dynamics_lab[0].states[-1]])
```

# Frequency analysis

```python
qubit_pulses = pulses_lab
time_list = tlist2
```

```python
qubit_pulses = oct_result.optimized_controls
time_list = tlist
```

```python
def plot_spectrum(pulse, tlist, mark_freq=None, pos=1):
    samples = len(tlist)
    sample_interval = tlist[-1]/samples
    time = np.linspace(0, samples*sample_interval, samples)

    signal_qubit = pulse
    signal_spectrum = np.fft.fftshift(np.fft.fft(signal_qubit))
    freqs = np.fft.fftshift(np.fft.fftfreq(samples, d=sample_interval))

    plt.figure(figsize=(10,5))
    plt.plot(freqs, np.abs(signal_spectrum))  # in GHz
    if mark_freq is not None:
        if not isinstance(mark_freq, list):
            mark_freq = [mark_freq]
        mf = np.array(mark_freq)/(2*π)
        if pos==1:
            plt.xlim(0, 2*mf[0])
        elif pos==-1:
            plt.xlim(-2*mf[0], 0)
        [plt.axvline(x=pos*m_f, ymin=0, ymax=1, color='k', linestyle='--', linewidth=1) for m_f in mf]
    plt.title('Qubit pulse spectrum')
    plt.xlabel('f (GHz)');

plot_spectrum(qubit_pulses[0], time_list, mark_freq=[ω_q, ω_q + K_q, ω_q - K_q], pos=1)
plot_spectrum(qubit_pulses[1], time_list, mark_freq=[ω_q, ω_q + K_q, ω_q - K_q], pos=-1)
```

```python

```

    12.563