# Ion Trap Reinforcement Learning Environment

## How it works

This exercise is supposed to be interactive, but it is up to you how interactive it really is. I will explain some stuff, and then, I will pose some **TASK**. You can choose to work for a time on the task and then move on to the solutions at any time. The solutions are usually provided in blocks of related tasks. Note that my solutions are not necessarily the most efficient or best. They are suggestions that work fine for what we are going to do here. If your code looks different, this is not necessarily an indication that it is wrong.

At the end of the last two chapters **Qudit quantum mechanics** and **Laser pulses as unitaries** there are *Wrappinp up* sections which summarize the most important code created in the respective chapters. The full environment is provided at the end and builds on that code.

Moreover, the entire code of the ion trap environment is also provided in accompanying *python* files. If you are really impatient you can just understand the code starting from there. However, you loose the benefit of understanding how the process of building such an environment goes.

The reinforcement learning environmentclass `IonTrapEnv` can be found in `ion_trap.py` which inherits from the `QuditQM` class in `qudit_qm.py`. The environment also uses methods defined in the `LaserGates` class from `laser_gates.py`.

**Enjoy it your way!**

## Introduction

**In this exercise we are going to build a reinforcement learning environment in which an agent can apply (simulated) laser pulses to an array of (simulated) ions (e.g., ${}^{137}Ba^{+})$ in a Paul trap. The goal for the agent is to create a large amount of quantum entanglement between as many ions as possible.**

Maybe you don't know anything about ions, or Paul traps. Well, lucky for you, me neither. So, let us together ignore our ignorance for the time being and move on. Ideally, smarter people have then already figured it all out in such a way that a theorist like myself can readily neglect everything about ions and lasers, and instead, just consider Hilbert spaces and unitaries. If this is the case, we can rephrase the introductory sentence in a way which is much more adequate for what we are going to code:

**In this exercise we are going to build a reinforcement learning environment in which an agent can apply unitaries to $n$ $d$-level quantum systems. The goal is to create $d$-dimensional GHZ states.**

In other words, we are dealing with complex vectors $|\psi\rangle\in(\mathbb{C}^d)^{\otimes n}$ and unitary matrices $U$ acting on them. Of course, what exactly these unitaries are will be determined by the laser pulses that can actually be implemented in the lab. Being an ignorant quantum information theorist by heart, I do not know what is and what is not possible in an experiment. Still, I am not ready to confront my ignorance just yet and will instead keep heading forward heedlessly.

What I like to do when I am about to code something like a reinforcement learning environment is to build a skeleton. That is, code or pseudocode which is not working but helps me understand what classes or methods are required to complete it. I like to think of it as a cloze where the blanks are methods or attributes which are to be filled in. So let's do that.

## Skeleton Environment

What a reinforcement learning environment is supposed to do can be easily described within a generic picture.

As far as the associated code goes, we can imagine something as follows,

```python
# dummy code -- do not run

env = FooEnv()

for i in range(100):
    done = False
    observation = env.reset()
    while not done:
        action = np.random.choice(env.num_actions)
        observation, reward, done = env.step(action)
```

Here, the agent is random and chooses actions at random. The environment however, has all the important methods:

* `reset()`: Resets the environment to its initial state. *Returns* the initial observation.
* `step(action)`: Performs an action (given by an action index) on the environment. *Returns* the new observation, an associated reward and a bool value `done` which indicates whether a terminal state has been reached.

These methods should be provided in every reinforcement learning environment and conform with the [*openAI* `gym`](https://github.com/openai/gym) standard.

### Build an empty class with all required methods

**TASK** Now, we can start building our skeleton. That is, we want to build a class `IonTrapEnv()` which features the methods `reset()` and `step(action)`. In addition, we can already implement the `num_actions` attribute with a placeholder value. You get the chance to do it yourself below. The solution will always be provided afterwards.

```python
# build skeleton class `IonTrapEnv`

```

### Solution

In this solution I already added some documentations loosely following the [google docs](https://sphinxcontrib-napoleon.readthedocs.io/en/latest/example_google.html) style guide. As good practice you should do the same. It makes working in a team much easier if you have proper docs and code.

```python
# dummy code -- do not run

class IonTrapEnv(object):
    def __init__(self):
        """
        """
        #int: Number of actions.
        self.num_actions = 0 # placeholder

    def reset():
        """
        Returns:
            observation (array): The observation received by the agent.
        """
        pass

    def step(action):
        """
        Args:
            action (int): The action index.

        Returns:
            observation (array): The observation received by the agent.
            reward (float): The reward quantifying good or bad behavior.
            done (bool): Whether or not the episode has finished.
        """
        pass
```

```python

```

### Start adding (pseudo-)functionality

Now this is what I call a bare minimum of a skeleton, no flesh whatsoever. The natural next step is to add some pseudocode to define its functionality, or maybe put into writing (e.g., as docs) what the environment and the methods are supposed to do.

#### init method

I will make the start and describe what I want the environment to do in general. In other words, I will give a global view starting with the `init` function and you can do the `step` and `reset` methods. Here is my interpretation of the environment in the docs:

```python
# dummy code -- do not run

class IonTrapEnv(object):
    def __init__(self):
        """
        This is the environment for a d-level ion trap quantum computer.
        An agent has a certain gate set at its disposal made up from elementary
        laser pulses (whatever that means).

        At each time step an agent can choose one gate to apply to the current
        quantum state of the environment. The state, as well as observation,
        is a d^n dimensional quantum state where d is the local dimension of
        each ion and n is the number of ions. The initial state is |0...0>.

        The goal is to create high-dimensional, multipartite entanglement between
        ions. Once such a state has been created, the agent receives a reward.

        If a maximum number of time steps is reached, an episoded is forcefully
        ended and should be restarted.
        """
        #int: Number of actions.
        self.num_actions = 0 # placeholder
```

Sounds high-level. So what does that imply about any attributes and methods that we are going to need? Well, let's go through the docs line by line:

* **gate set**: So there has to be some gate set. These are the unitaries which we are going to implement. I would consider this a `list()` of matrices such that its length is the number of actions and a single action can be associated with an index. In order to define and generate these gates, I would even suggest a new class `LaserGates()`. This class may provide an attribute `gates` which is our gate set.

* **apply gates**: This is going to be the `action` in the `step` method. We want something like `self.actions[action]` *applied to* `self.state` to implement an action.

* **d^n dimensional quantum state**: This tells us that we are dealing with $d^n$ dimensional vectors and $d^n\times d^n$ dimensional matrices. Good. So, probably we want some attributes for this, too. Let's define `dim` for $d$ and `num_ions` for $n$.

* **initial state**: So the initial state will be some tensor product of $|0\rangle$ states. This will be an attribute even though we haven't discussed yet how to describe them mathematically.

* **high-dimensional, multipartite entanglement**: This will be used in the `step` function to determine whether a reward is given or not. Therefore, there should be some goal states. Even though, we don't know what they are yet, we can put an attribute with placeholder value.

* **reward**: This is tied to *high-dimensional, multipartite entanglement* and is implemented in and received from the `step` method.

* **maximum time steps**: This tells us that there is a maximum number of time steps which should make the `step` method return `done=True` if exceeded.

Alright, let us put in some attributes in line with our docs.

```python
# dummy code -- do not run

class IonTrapEnv(object):
    def __init__(self):
        """
        This is the environment for a d-level ion trap quantum computer.
        An agent has a certain gate set at its disposal made up from elementary
        laser pulses (whatever that means).

        At each time step an agent can choose one gate to apply to the current
        quantum state of the environment. The state, as well as observation,
        is a d^n dimensional quantum state where d is the local dimension of
        each ion and n is the number of ions. The initial state is |0...0>.

        The goal is to create high-dimensional, multipartite entanglement between
        ions. Once such a state has been created, the agent receives a reward.

        If a maximum number of time steps is reached, an episoded is forcefully
        ended and should be restarted.
        """
        #:class:`LaserGates`: Some class that lets us generate the unitaries used as actions.
        self.gates = LaserGates()
        #list: The action gate set obtained from the class.
        self.actions = self.gates.gates
        #int: Number of actions.
        self.num_actions = len(self.actions)

        #int: Local ion dimension.
        self.dim = 3
        #int: Number of ions.
        self.num_ions = 3

        self.init_state = 0 # placeholder: the initial state
        self.state = self.init_state # placeholder: the internal state of the environment

        self.goal = 0 # placeholder: something that characterizes entanglement

        #int: Current number of time steps.
        self.time_step = 0
        #int: Maximum number of time steps.
        self.max_steps = 10
```

From this quick sketch of an `init` function, it becomes obvious that we need to start thinking about implementation of states and matrices. Also, we don't know yet how to classify entanglement. So, let's get this discussion out of the way now!

#### numerical python for linear algebra

I am a big fan of `numpy` and I wholeheartedly recommend it. It is fast, reliable and the most common tool for numerical computations in python. In addition, it supports `scipy` as an additional module for scientific computing with python. In case you don't know, let me here and now ensure you that `numpy` in combination with `scipy` gives us all the tools we need to handle quantum information and computation on a small scale with few discrete quantum objects (i.e., linear algebra). For instance, matrix multiplication between $M$ and $N$ is performed as `M@N`.
Henceforth, our states and unitiaries will therefore be of type `np.ndarray`. You may not be convinced yet, but I will have you warm up to `numpy` and `scipy` before too long.

#### entanglement characterization

Now, we also need to make a decision on how to characterize entanglement. This is a hard one. As good physicists, we should start with a quick Google search. A natural first would be *multidimensional multipartite entanglement characterization* (heh, I cheated because I know what we are looking for). There is a bunch of papers popping up and what ensues is a mind-numbing search for an adequate entanglement measure or monotone. To save breath, I will tell you that one of the first papers on your google search already contains the information we are looking for:

*The structure of multidimensional entanglement in multipartite systems*, M. Huber and J.I. de Vicente, [arxiv:1210.6876](https://arxiv.org/pdf/1210.6876.pdf).

You may have heard the word *Schmidt rank* before. It is a measure of entanglement for quantum states shared between *two* parties. In the paper, the authors generalize the Schmidt rank to multipartite systems by introducing a *Schmidt rank vector* (SRV). In particular, it provides us with an easy-to-calculate tool which also classifies GHZ-type entanglement. Perfect! We still don't know what it is exactly, but let's say we are going to use it.

#### reset and step methods

**TASK** With these decisions from the table, I ask you now to complete the skeleton for the `reset` and `step` methods below. I have taken the freedom to fill in the docs for you. Lucky you!

```python
# dummy code -- do not run

    def reset(self):
        """
        Resets the environment to its initial state, i.e., |0...0>.
        Resets counter.

        Returns:
            state (np.ndarray): Initial state vector.
        """
        # TODO: fill in code

    def step(self, action):
        """
        Performs and evaluates an action on the environment.

        (1) A quantum gate from our laser gate set is applied to the
            current state of the environment.
        (2) The Schmidt rank vector (SRV) is calculate for the current state
            of the environment. I
        (3) If the SRV is among the set of goal SRVs, a reward +1 is provided and
            the trial is done. If not, no reward is given. The trial is ended if
            the time exceed the maximum number of steps.

        Args:
            action (int): Index of action in gate set.

        Returns:
            state (np.ndarray): Current state of the environment.
            reward (float): Current reward.
            done (bool): Whether or not the episode is done.
        """
        # TODO: fill in missing code
        done = False
        reward = 0.
        # increment counter
        self.time_step += 1
        # (1) Apply gate to state (FILL IN)

        # (2) Calculate SRV
        srv = self.srv(self.state) # some method that we are still missing
        # (3) Finish episode and give reward if appropriate (FILL IN)

        return self.state, reward, done
```

#### Solution

```python
# dummy code -- do not run

    def reset(self):
        """
        Resets the environment to its initial state, i.e., |0...0>.
        Resets counter.

        Returns:
            state (np.ndarray): Initial state vector.
        """
        self.state = self.init_state
        self.time_step = 0

        return self.state

    def step(self, action):
        """
        Performs and evaluates an action on the environment.

        (1) A quantum gate from our laser gate set is applied to the
            current state of the environment.
        (2) The Schmidt rank vector (SRV) is calculate for the current state
            of the environment. I
        (3) If the SRV is among the set of goal SRVs, a reward is provided and
            the trial is done. If not, no reward is given. The trial is ended if
            the time exceed the maximum number of steps.

        Args:
            action (int): Index of action in gate set.

        Returns:
            state (np.ndarray): Current state of the environment.
            reward (float): Current reward.
            done (bool): Whether or not the episode is done.
        """
        done = False
        reward = 0.
        # increment counter
        self.time_step += 1
        # (1) Apply gate to state
        self.state = self.gates.gates[action]@self.state
        # (2) Calculate SRV
        srv = self.srv(self.state) # some method that we are still missing
        # (3) Finish episode and give reward if appropriate
        if srv in self.goal:
            done = True
            reward = 1.
        elif self.time_step >= self.max_steps:
            done = True

        return self.state, reward, done
```

### What's next

At this point, we should have a look at our skeleton class for our ion trap environment and ask ourselves what is missing.

```python
# dummy code -- do not run

class IonTrapEnv(object):
    def __init__(self):
        """
        This is the environment for a d-level ion trap quantum computer.
        An agent has a certain gate set at its disposal made up from elementary
        laser pulses (whatever that means).

        At each time step an agent can choose one gate to apply to the current
        quantum state of the environment. The state, as well as observation,
        is a d^n dimensional quantum state where d is the local dimension of
        each ion and n is the number of ions. The initial state is |0...0>.

        The goal is to create high-dimensional, multipartite entanglement between
        ions. Once such a state has been created, the agent receives a reward.

        If a maximum number of time steps is reached, an episoded is forcefully
        ended and should be restarted.
        """
        #:class:`LaserGates`: Some class that lets us generate the unitaries used as actions.
        self.gates = LaserGates()
        #list: The action gate set obtained from the class.
        self.actions = self.gates.gates
        #int: Number of actions.
        self.num_actions = len(self.actions)

        #int: Local ion dimension.
        self.dim = 3
        #int: Number of ions.
        self.num_ions = 3

        #np.ndarray: Initial state.
        self.init_state = 0 # placeholder
        #np.ndarray: Current state of the environment and observation.
        self.state = self.init_state

        #list: List of SRVs that are rewarded.
        self.goal = 0 # placeholder

        #int: Current number of time steps.
        self.time_step = 0
        #int: Maximum number of time steps.
        self.max_steps = 10

    def reset(self):
        """
        Resets the environment to its initial state, i.e., |0...0>.
        Resets counter.

        Returns:
            state (np.ndarray): Initial state vector.
        """
        self.state = self.init_state
        self.time_step = 0

        return self.state

    def step(self, action):
        """
        Performs and evaluates an action on the environment.

        (1) A quantum gate from our laser gate set is applied to the
            current state of the environment.
        (2) The Schmidt rank vector (SRV) is calculate for the current state
            of the environment. I
        (3) If the SRV is among the set of goal SRVs, a reward is provided and
            the trial is done. If not, no reward is given. The trial is ended if
            the time exceed the maximum number of steps.

        Args:
            action (int): Index of action in gate set.

        Returns:
            state (np.ndarray): Current state of the environment.
            reward (float): Current reward.
            done (bool): Whether or not the episode is done.
        """
        done = False
        reward = 0.
        # increment counter
        self.time_step += 1
        # (1) Apply gate to state
        self.state = self.gates.gates[action]@self.state
        # (2) Calculate SRV
        srv = self.srv(self.state) # some method that we are still missing
        # (3) Finish episode and give reward if appropriate (FILL IN)
        if srv in self.goal:
            done = True
            reward = 1.
        elif self.time_step >= self.max_steps:
            done = True

        return self.state, reward, done
```

Alright! Looks already not too bad. But there are a few obvious *blanks* that needs to be filled in.

* `LaserGates()`: We are missing this class altogether.
* `self.init_state`: We have not defined $|0...0\rangle$ in terms of `np.ndarrays` yet.
* `self.goal`: We have not defined the goal SRVs yet.
* `self.srv()`: We have no method to calculate an SRV.

**TASK** Good. That's well-arranged. The important question at this point is what's the next item on our to-do list? Think about it!

### Solution

Well, I do not believe it is the `LaserGates` since they will require methods which define states and matrices in the first place. So, actually I think we want to do some linear algebra with d-level quantum systems, or as we like to call them in quantum information, *qudits*.

## Qudit Quantum Mechanics

Trusting student that you are, you have probably taken me by my word that `numpy` and `scipy` are sufficient and good choices for all the things we want to do. Indeed, matrix multiplication between $M$ and $N$ is performed easily as `M@N` and the conjugate transposition $M^\dagger$ of a matrix $M$ is simply `M.conj().T`. The `scipy.linalg` module then fills in the few gaps in the `numpy` library. For instance, it supports methods such as `expm()`, i.e., matrix exponentials $e^M$, which we need to define gates in quantum computation. However, don't take my word for it, start using it yourself!

### Foolin' with qu*bits*

**TASK** Let us fool around a bit. The most fundamental object that is commonly used to describe spin $1/2$ systems in quantum mechanics are Pauli matrices. In quantum computation, we use them to describe a Clifford algebra for 2-level quantum gates. They are as follows,

\begin{align}
\sigma_x =
\begin{pmatrix}
0 & 1\\
1 & 0
\end{pmatrix},\:
\sigma_y =
\begin{pmatrix}
0 & -i\\
i & 0
\end{pmatrix},\:
\sigma_z =
\begin{pmatrix}
1 & 0\\
0 & -1
\end{pmatrix}.\nonumber
\end{align}

Hoping that you have at least a little experience with `numpy` I ask you to implement these. Also don't forget to always do a sanity check:

```python
# run this so ipython notebook generates multiple outputs
from IPython.core.interactiveshell import InteractiveShell
InteractiveShell.ast_node_interactivity = "all"
```

```python
import numpy as np
from scipy import linalg

# define Pauli matrices (FILL IN)

# check that Pauli matrices are hermitian (FILL IN)

# check XZ = -iY (FILL IN)

# check XZ=-ZX (FILL IN)

# define eigenstates of Z (FILL IN)
# |0>
zero = np.array([[1],[0]], dtype=complex)
# |1>

# check that above states are indeed eigenstates (FILL IN)

```

Alright, let us define an arbitrary $Z$ rotation on a single qubit. This is defined as follows,

\begin{align}
R_Z(\theta)=e^{i\theta Z}.\nonumber
\end{align}

Go ahead, the stage is all yours:

```python
def z_rot(theta):
    """
    Arbitrary Z rotation.

    Args:
        theta (float): Rotation angle.

    Returns:
        z_rot (np.ndarray): exp(i*theta*Z)
    """
    # TODO: fill in

# test rotation with some angles

```

### Solutions

```python
import numpy as np
from scipy import linalg

# define Pauli matrices
X = np.array([[0,1.], [1.,0]], dtype=complex)
Y = np.array([[0,-1.j], [1.j,0]], dtype=complex)
Z = np.array([[1.,0], [0,-1.]], dtype=complex)

# check that Pauli matrices are hermitian
(X.conj().T==X).all()
(Y.conj().T==Y).all()
(Z.conj().T==Z).all()

# check XZ = -iY
(X@Z==-1j*Y).all()
# check XZ=-ZX
(X@Z==-1*Z@X).all()

# define eigenstates of Z
# |0>
zero = np.array([[1],[0]], dtype=complex)
# |1>
one = np.array([[0],[1]], dtype=complex)
# check eigenstates of Z
(Z@zero==zero).all()
(Z@one==-1*one).all()
```

    True

    True

    True

    True

    True

    True

    True

```python
def z_rot(theta):
    """
    Arbitrary Z rotation.

    Args:
        theta (float): Rotation angle.

    Returns:
        z_rot (np.ndarray): exp(i*theta*Z)
    """

    return linalg.expm(1.j*theta*Z)

# test rotation with some angles
# theta=pi
np.around(z_rot(np.pi), decimals=10) # getting rid of rounding errors
# theta=pi/2
z_rot(np.pi/2)
# theta=pi/4
z_rot(np.pi/4)
```

    array([[-1.+0.j,  0.+0.j],
           [ 0.+0.j, -1.-0.j]])

    array([[0.+1.j, 0.+0.j],
           [0.+0.j, 0.-1.j]])

    array([[0.70710678+0.70710678j, 0.        +0.j        ],
           [0.        +0.j        , 0.70710678-0.70710678j]])

### Foolin' with qu*dits*

This was good for a warm-up. However, qubits are too simple. We want arbitrary $d$-level systems, *qudits*! So, we must ask whether we can define a generalized $X$- and $Z$-gate? Indeed, we can and we will do so now because it illustrates some nice functionality! In the interest of time and brevity, let us neglect $Z$ and consider a generalized $X$-operator.

Note the qubit $X$ operator maps $|0\rangle$ to $|1\rangle$. In general, we may define an $X$ operator mapping $|i\rangle$ to |i+1\textrm{mod }d\rangle$ for all $i=0,...,d-1$. This can be put into a neat formular:

\begin{align}
X = \sum_{i=0}^{d-1} |i+1\textrm{mod }d\rangle\langle i|\nonumber
\end{align}

That's easy. In fact, in python pseudocode I can directly write down this formular.

```python
# dummy code -- do not run
dim = 3
Xd = np.zeros((dim, dim), dtype=complex)
for i in range(dim):
    Xd += ket((i+1)%dim, dim)@bra(i, dim)
```

**TASK** So just do me the favor and define the `ket` and `bra` functions so that the above code actually works:

```python
def ket(i, dim):
    """
    Basis vector |i> = (0,...,0,1,0,...,0).

    Args:
        i (int): Index of basis vector.
        dim (int): Local dimension.

    Returns:
        vec (np.ndarray): Vector form of |i> as (d,1)-dimensional array.
    """
    # TODO: fill in

def bra(i, dim):
    """
    Basis vector <i| = (0,...,0,1,0,...,0)^T.

    Args:
        i (int): Index of basis vector.
        dim (int): Local dimension.

    Returns:
        vec (np.ndarray): Vector form of <i| as (1,d)-dimensional array.
    """
    # TODO: fill in

# check that above code works (FILL IN)

# check action of X on |0>,|1> and |2> (FILL IN)

```

Fantastic! That's a big step already. You have learned to translate a matrix written as a sum of ket's and bra's into python code. And it wasn't even that hard. That will for sure come in handy once we start implementing laser gates.

**TASK** So put your fresh skills to the test and implement the famous *creation* and *annihilation* operators defined for $d$-level systems (where $d=2s+1$) as follows:

\begin{align}
S_{+} = \sum_{l=-s}^{s-1}\sqrt{s(s+1)-l(l+1)}|l+s+1\rangle\langle{l+s}|\nonumber\\
S_{-} = \sum_{l=-s}^{s-1}\sqrt{s(s+1)-l(l+1)}|l+s\rangle\langle{l+s+1}|.\nonumber
\end{align}

Hands free! Please also do some docs because that's something one should learn early on. It facilitates good teamwork.

```python
def creation(dim):
    """
    Creation operator S_{+} for odd dimensions.

    Args:

    Returns:

    """
    # TODO: fill in

def annihilation(dim):
    """
    Annihilation operator S_{-} for odd dimensions.

    Args:

    Returns:

    """
    # TODO: fill in

# look at example with d=3

```

### Solutions

```python
def ket(i, dim):
    """
    Basis vector |i> = (0,...,0,1,0,...,0).

    Args:
        i (int): Index of basis vector.
        dim (int): Local dimension.

    Returns:
        vec (np.ndarray): Vector form of |i> as (d,1)-dimensional array.
    """
    vec = np.zeros((dim, 1), dtype=complex)
    vec[i] = 1.

    return vec

def bra(i, dim):
    """
    Basis vector <i| = (0,...,0,1,0,...,0)^T.

    Args:
        i (int): Index of basis vector.
        dim (int): Local dimension.

    Returns:
        vec (np.ndarray): Vector form of <i| as (1,d)-dimensional array.
    """
    vec = ket(i, dim).conj().T

    return vec

# check that above code works
dim = 3
Xd = np.zeros((dim, dim), dtype=complex)
for i in range(dim):
    Xd += ket((i+1)%dim, dim)@bra(i, dim)
Xd

# check action of X on |0>,|1> and |2>
(Xd@ket(0,3)==ket(1,3)).all()
(Xd@ket(1,3)==ket(2,3)).all()
(Xd@ket(2,3)==ket(0,3)).all()

```

    array([[0.+0.j, 0.+0.j, 1.+0.j],
           [1.+0.j, 0.+0.j, 0.+0.j],
           [0.+0.j, 1.+0.j, 0.+0.j]])

    True

    True

    True

```python
def creation(dim):
    """
    Creation operator S_{+} for odd dimensions.

    Args:
        dim (int): Local dimension.

    Returns:
        creation (np.ndarray): Creation operator of given dimension.
    """
    # define empty creation operator
    creation = np.zeros((dim, dim), dtype=complex)

    # add summand |l+s+1><l+s| as defined in Eq. (5)
    s = int((dim - 1)/2)
    for l in range(-s, s):
        m = np.sqrt(s*(s+1) - l*(l+1))*(ket(l+s+1, dim)@bra(l+s, dim))
        creation += m

    return creation

def annihilation(dim):
    """
    Annihilation operator S_{-} for odd dimensions.

    Args:
        dim (int): Local dimension.

    Returns:
        annihilation (np.ndarray): Annihilation operator of given dimension.
    """
    # define empty annihilation operator
    annihilation = np.zeros((dim, dim), dtype=complex)
    s = int((dim - 1)/2)

    # add summand |l+s><l+s+1| as defined in Eq. (5)
    for l in range(-s, s):
        m = np.sqrt(s*(s+1) - l*(l+1))*(ket(l+s, dim)@bra(l+s+1, dim))
        annihilation += m

    return annihilation

# look at example with d=3
creation(3)
annihilation(3)
```

    array([[0.        +0.j, 0.        +0.j, 0.        +0.j],
           [1.41421356+0.j, 0.        +0.j, 0.        +0.j],
           [0.        +0.j, 1.41421356+0.j, 0.        +0.j]])

    array([[0.        +0.j, 1.41421356+0.j, 0.        +0.j],
           [0.        +0.j, 0.        +0.j, 1.41421356+0.j],
           [0.        +0.j, 0.        +0.j, 0.        +0.j]])

### Multiple qudits

Wow, we are deep into the matter. We have a way to define and manipulate single $d$-level quantum systems. But let's not forget our goal: Entanglement between *many* quantum systems. It is time to introduce *composite systems* to the code. For that, we need to define a *tensor product* for matrices and vectors. Luckily, `numpy` provides us with this method *for free*. It is called `np.kron` and can, for our purposes, be used interchangeably with $\otimes$.

**TASK** Let's try it out! Given a `num_ions`, try to create an initial state $|0\rangle\otimes...\otimes |0\rangle$ (**Hint:** Use `reduce` from `functools` to nest a function).
Now, you should also be able to create the infamous GHZ state for $n$ qudits:

\begin{align}
|\mathrm{GHZ}_{n,d}\rangle = \frac{1}{\sqrt{d}}\sum_{i=0}^{d-1}|i\rangle_1\otimes\cdots\otimes|i\rangle_n\nonumber
\end{align}

For example, for 2 qutrits $|\mathrm{GHZ}_{2,3}\rangle=\frac{1}{\sqrt{3}}(|00\rangle+|11\rangle+|22\rangle)$.

```python
from functools import reduce

dim = 3
num_ions = 2
# create initial state |0...0> for any number of ions (FILL IN)

# check dimension (FILL IN)

# create ghz state (FILL IN)

```

Okay. We have learned to define all the states we want by just using `reduce`, `np.kron` as well as our `ket` and `bra` methods. Easy! So how do we create gates on multiple qudits? Well, as long as we keep defining matrices in terms of `ket` and `bra`, we can just continue as always but throw some `reduce` and `np.kron` into the mix. There is however, one trick that might turn out to be useful.

**TASK** Say, we want to perform the single-qudit $X$-gate from before but have two qudits. How do we extend this such that it acts on the first qudit of our two qutrit initial state $|00\rangle$? **HINT**: You will need to define an additional identity $X\otimes I$.

```python
# define Xd \otimes I

# check dimension

# check gate functionality by applying to |00>

```

### Solutions

```python
from functools import reduce

dim = 3
num_ions = 2
# create initial state |0...0> for any number of ions (FILL IN)
init_state = reduce(np.kron, [ket(0, dim) for i in range(num_ions)])
init_state

# check dimension
init_state.shape[0]==dim**num_ions

# create ghz state
ghz = np.zeros((dim**num_ions, 1), dtype=complex)
for i in range(dim):
    ghz += 1/np.sqrt(dim)*reduce(np.kron, [ket(i, dim) for k in range(num_ions)])
ghz
```

    array([[1.+0.j],
           [0.+0.j],
           [0.+0.j],
           [0.+0.j],
           [0.+0.j],
           [0.+0.j],
           [0.+0.j],
           [0.+0.j],
           [0.+0.j]])

    True

    array([[0.57735027+0.j],
           [0.        +0.j],
           [0.        +0.j],
           [0.        +0.j],
           [0.57735027+0.j],
           [0.        +0.j],
           [0.        +0.j],
           [0.        +0.j],
           [0.57735027+0.j]])

```python
# define Xd \otimes I
# define list of identities with X at the first position
multi_gate = [np.eye(dim) for i in range(num_ions)]
multi_gate[0] = Xd
# apply tensor product
multi_gate = reduce(np.kron, multi_gate)

# check dimension
multi_gate.shape==(dim**num_ions,dim**num_ions)
# check gate functionality by applying to |00>
(multi_gate@init_state==reduce(np.kron, [ket(1,dim), ket(0,dim)])).all()
```

    True

    True

### Entanglement and Schmidt rank vector

Alright, we know how to define and manipulate composite systems. However, we still don't know how to classify them in terms of entanglement. We know that we are going to use some mathematical object called the *Schmidt rank vector* (SRV) but we don't know what it is or how to calculate it. So what follows is a brief introduction which you may skip at your own leisure. Practically speaking, all you need is the formular at the end.

The Schmidt rank of a pure, bipartite quantum state $|\psi\rangle\in\mathcal{H}_A\otimes\mathcal{H}_B$ is a measure of entanglement shared between two parties Alice $A$ and Bob $B$. It is defined as the rank of the reduced density matrix $\rho_A=\textrm{tr}_{B}(\rho)$ of $\rho=|\psi\rangle\langle\psi|$. The partial trace is defined as follows,

\begin{align}
\mathrm{tr}_B(\rho):=\sum_{i=0}^{d-1}I_A\otimes\langle i|_B \rho I_A\otimes |i\rangle_B\nonumber
\end{align}

The SRV then calculates the Schmidt rank of all single-particle marginal which answers the question of how many dimensions are at least required to faithfully represent the pure state and its correlation in any local basis. In other words, defining the reduced density matrix

\begin{align}
\rho_i:=\mathrm{tr}_{\bar{i}}(\rho)
\end{align}
where $\bar{i}$ is the complement of $\{i\}$ w.r.t. the set $\{1,...,n\}$, we can associate with it a Schmidt rank

\begin{align}
r^\rho_i=\mathrm{rank}(\rho_i).
\end{align}

Arranging these number in a vector for all reduced density matrices, we obtain the SRV:

\begin{align}
\vec{d}_\rho=(r_1^\rho,...,r_n^\rho).
\end{align}

These three equations tell us exactly how to calculate the SRV for any pure state $\rho=|\psi\rangle\langle\psi|$.

Last but not least, let me tell you one more property of the numbers $r_i^\rho$ that makes our life easier here: $\mathrm{rank}(\rho_i)=\mathrm{rank}(\rho_\bar{i})$. That is, it does not matter whether we calculate the rank for the subsystem $i$ or $\bar{i}$. So we only need to trace over $i$ in our calculations... Easy!

#### partial trace

Before, we go all physics again though, let us code a partial trace. Since we are rather versed in coding by now, let us do this as general as possible. That is, taking any number of systems that we want to trace over as arguments `*args`.

**TASK** Below, somebody implemented a method that produces the partial trace. Unfortunately, whoever implemented this method was rather sloppy or unwilling to document their code. So provide help for all the people who will follow and write the docs and some supporting comments about what is happening here. Alternatevely, you may also first try to understand the undocumented version for some time and then the documented version in the solution. I predict that you will notice a significant difference.

I hope that I can make you a believer in good code conduct. It is essential when working with people on the same code. Note that in-line comments can be helpful but are not strictly required by any style guide. And in fact, I did overcomment on purpose in this exercise. In general, make sure that your methods are small and sufficiently self-contained such that the code is readable even if uncommented.

```python
from itertools import product
from copy import deepcopy

def partial_tr(rho, dim, num_ions, *args):
    """
    """
    # TODO: fill in the docs and some comments at your own leisure
    dim_reduced = dim**(num_ions-len(args))
    reduced_rho = np.zeros((dim_reduced, dim_reduced), dtype=complex)
    basis = [np.eye(dim, dtype=complex) for i in range(num_ions)]
    for ion in args:
        basis[ion] = np.zeros((dim, 1), dtype=complex)
    for d_vec in product(range(dim), repeat=len(args)):
        contractor = deepcopy(basis)
        for n, ion in enumerate(args):
            contractor[ion][d_vec[n]] = 1.
        contractor = reduce(np.kron, contractor)
        reduced_rho += contractor.conj().T@rho@contractor
    return reduced_rho

# check rank of reduced density matrix for a 2-qutrit GHZ state
dim = 3
num_ions = 2
reduced_rho = partial_tr(ghz@ghz.conj().T, dim, num_ions, 1)
np.linalg.matrix_rank(reduced_rho)
```

    3

#### solution

```python
from itertools import product
from copy import deepcopy

def partial_tr(rho, dim, num_ions, *args):
    """
    Partial trace for a density matrix of multiple qudits (ions).
    Given a set of indices of qudits, this methods traces the respective
    qudits.

    Args:
        rho (np.ndarray): Density matrix which is to be traced.
        dim (int): Local dimension.
        num_ions (int): Number of ions.
        *args (int): Variable list of indices of qudits/ions which are to be
                     traced.

    Returns:
        reduced_rho (np.ndarray): Reduced density matrix.
    """
    # define empty reduced density matrix
    dim_reduced = dim**(num_ions-len(args))
    reduced_rho = np.zeros((dim_reduced, dim_reduced), dtype=complex)
    # define list of identities to be replaced by trace operator/contractor
    basis = [np.eye(dim, dtype=complex) for i in range(num_ions)]
    # replace id with all-zero ket vectors for ions which are to be traced
    for ion in args:
        basis[ion] = np.zeros((dim, 1), dtype=complex)
    # perform contraction <i_1,..,i_k|\rho|i_1,...,i_k> for all bases
    for d_vec in product(range(dim), repeat=len(args)):
        # fix basis state |i_1,...,i_k>
        contractor = deepcopy(basis)
        for n, ion in enumerate(args):
            contractor[ion][d_vec[n]] = 1.
        contractor = reduce(np.kron, contractor)
        # add contraction
        reduced_rho += contractor.conj().T@rho@contractor

    return reduced_rho

# check rank of reduced density matrix for a 2-qutrit GHZ state
dim = 3
num_ions = 2
reduced_rho = partial_tr(ghz@ghz.conj().T, dim, num_ions, 1)
np.linalg.matrix_rank(reduced_rho)
```

    3

#### calculating SRVs

Ok, now let's go physics and calculate SRVs!

**TASK** With the partial trace at hand you should be able to calculate the SRV for a pure state $|\psi\rangle$. Complete the method indicated below and test it on a three-qutrit GHZ state. This state is maximally entangled and should therefore be (3,3,3). *Hint*: This is the time to remember $\mathrm{rank}(\rho_i)=\mathrm{rank}(\rho_\bar{i})$.

```python
def srv(ket, dim, num_ions):
    """
    Calculates the SRV of an arbitrary pure state.

    Args:
        ket (np.ndarray): The input state.
        dim (int): Local dimension.
        num_ions (int): Number of ions.

    Returns:
        srv (list): List of integeres specifying the Schmidt ranks of all
                    single-qudit marginals.
    """
    # TODO: fill in blanks
    # get density matrix from pure state (FILL IN)

    # set of all ions
    ions = set(range(num_ions))
    # define empty list of SRs
    srv = []
    # calculate SR for all single-qudit marginals (FILL IN)

    return srv

# check SRV for three-qutrit GHZ state (FILL IN)

```

#### solution

```python
def srv(ket, dim, num_ions):
    """
    Calculates the SRV of an arbitrary pure state.

    Args:
        ket (np.ndarray): The input state.
        dim (int): Local dimension.
        num_ions (int): Number of ions.

    Returns:
        srv (list): List of integeres specifying the Schmidt ranks of all
                    single-qudit marginals.
    """
    # get density matrix from pure state
    rho = ket@ket.conj().T
    # set of all ions
    ions = set(range(num_ions))
    # define empty list of SRs
    srv = []
    # calculate SR for all single-qudit marginals
    for i in range(num_ions):
        # perform partial trace over the respective ion
        reduced_rho = partial_tr(rho, dim, num_ions, i)
        # calculate rank of reduce density matrix
        rank = np.linalg.matrix_rank(reduced_rho)
        # add to SRV
        srv.append(rank)

    return srv

# check SRV for three-qutrit GHZ state
dim = 3
num_ions = 3
# create ghz state
ghz = np.zeros((dim**num_ions, 1), dtype=complex)
for i in range(dim):
    ghz += 1/np.sqrt(dim)*reduce(np.kron, [ket(i, dim) for k in range(num_ions)])
# calculate SRV
srv(ghz, dim, num_ions)
```

    [3, 3, 3]

### Wrapping up

**Fantastic**. That should complete all the quantum mechanics methods that we need for our ion trap environment.
In summary, we have methods

* `ket()`, `bra()`: Allowing us to implement states and matrices of arbitray dimension. That is, we can define high-dimensional composite quantum states and represent multi-qudit unitaries.
* `srv()`: We can calculate the SRV for an arbitray pure quantum state.

Now, we have a way to fill in `self.goal`, `self.init_state` and `self.srv()` in our `IonTrapEnv` class. Give yourself a clap on the back, you earned it. Up next is the laser pulses which is likely going to borrow some of the functions from this chapter. Especially, `ket()` and `bra()` will find some applications there. Since the methods defined in this chapter are going to be used in both `IonTrapEnv` and `LaserGates` classes, we should make it another class and have the others inherit its methods. No worries, in anticipation of this, I already coded it.

**TASK** Make sure that you understand the following class which is built from the methods you have designed before. **Also, please run this code.**

```python
import numpy as np
from functools import reduce
from itertools import product
from copy import deepcopy

class QuditQM(object):
    def __init__(self, dim, num_ions):
        """
        Here we define some useful methods for n d-level quantum systems.
        Specifically, this includes basis states |k> and <k|, creation S_{+} and
        annihilation S_{-} operators. Further, we introduce some useful
        mathematical tools such as partial trace and Schmidt rank vector.

        Args:
            dim (int): The local (odd) dimension of an ion.
            num_ions (int): The number of ions.
        """
        if dim % 2 == 0:
            raise NotImplementedError('Our methods are only defined for odd ' +
                                      f'dimensions. Your dimension: {dim}.')
        self.dim = dim
        self.num_ions = num_ions

    def ket(self, i):
        """
        Basis vector |i> = (0,...,0,1,0,...,0).

        Args:
            i (int): Index of basis vector.

        Returns:
            vec (np.ndarray): Vector form of |i> as (d,1)-dimensional array.
        """
        vec = np.zeros((self.dim, 1), dtype=complex)
        vec[i] = 1.

        return vec

    def bra(self, i):
        """
        Basis vector <i| = (0,...,0,1,0,...,0)^T.

        Args:
            i (int): Index of basis vector.

        Returns:
            vec (np.ndarray): Vector form of <i| as (1,d)-dimensional array.
        """
        vec = self.ket(i).conj().T

        return vec

    def to_density(self, ket):
        """
        Creates density matrix from complex vector |psi> as |psi><psi|.

        Args:
            ket (np.ndarray): Complex (D,1)-dimensional array.

        Returns:
            rho (np.ndarray): Complex (D,D)-dimensional density matrix.
        """
        rho = ket@ket.conj().T

        return rho

    def creation(self):
        """
        Creation operator S_{+} for odd dimensions.

        Returns:
            creation (np.ndarray): Creation operator of given dimension.
        """
        # define empty creation operator
        creation = np.zeros((self.dim, self.dim), dtype=complex)

        # add summand |l+s+1><l+s| as defined in Eq. (5)
        s = int((self.dim - 1)/2)
        for l in range(-s, s):
            m = np.sqrt(s*(s+1) - l*(l+1))*(self.ket(l+s+1)@self.bra(l+s))
            creation += m

        return creation

    def annihilation(self):
        """
        Annihilation operator S_{-} for odd dimensions.

        Returns:
            annihilation (np.ndarray): Annihilation operator of given dimension.
        """
        # define empty annihilation operator
        annihilation = np.zeros((self.dim, self.dim), dtype=complex)
        s = int((self.dim - 1)/2)

        # add summand |l+s><l+s+1| as defined in Eq. (5)
        for l in range(-s, s):
            m = np.sqrt(s*(s+1) - l*(l+1))*(self.ket(l+s)@self.bra(l+s+1))
            annihilation += m

        return annihilation

    def partial_tr(self, rho, *args):
        """
        Partial trace for a density matrix of multiple qudits (ions).
        Given a set of indices of qudits, this methods traces the respective
        qudits.

        Args:
            rho (np.ndarray): Density matrix which is to be traced.
            *args (int): Variable list of indices of qudits/ions which are to be
                         traced.

        Returns:
            reduced_rho (np.ndarray): Reduced density matrix.
        """
        # define empty reduced density matrix
        dim_reduced = self.dim**(self.num_ions-len(args))
        reduced_rho = np.zeros((dim_reduced, dim_reduced), dtype=complex)
        # define list of identities to be replaced by trace operator/contractor
        basis = [np.eye(self.dim, dtype=complex) for i in range(self.num_ions)]
        # replace id with all-zero ket vectors for ions which are to be traced
        for ion in args:
            basis[ion] = np.zeros((self.dim, 1), dtype=complex)
        # perform contraction <i_1,..,i_k|\rho|i_1,...,i_k> for all bases
        for d_vec in product(range(self.dim), repeat=len(args)):
            # fix basis state |i_1,...,i_k>
            contractor = deepcopy(basis)
            for n, ion in enumerate(args):
                contractor[ion][d_vec[n]] = 1.
            contractor = reduce(np.kron, contractor)
            # add contraction
            reduced_rho += contractor.conj().T@rho@contractor

        return reduced_rho

    def srv(self, ket):
        """
        Calculates the SRV of an arbitrary pure state.

        Args:
            ket (np.ndarray): The input state.

        Returns:
            srv (list): List of integeres specifying the Schmidt ranks of all
                        single-qudit marginals.
        """
        # get density matrix from pure state
        rho = self.to_density(ket)
        # set of all ions
        ions = set(range(self.num_ions))
        # define empty list of SRs
        srv = []
        # calculate SR for all single-qudit marginals
        for i in range(self.num_ions):
            # perform partial trace over the respective ion
            reduced_rho = self.partial_tr(rho, i)
            # calculate rank of reduce density matrix
            rank = np.linalg.matrix_rank(reduced_rho)
            # add to SRV
            srv.append(rank)

        return srv

```

## Laser pulses as unitaries

It has finally come to this point: No longer can we ignore our ignorance. We need to become enlightened about lasers. Funnily, we can find both enlightment as well as ignorance at the same place: the internet. So let's venture there for help.

Before we do that however, let us briefly discuss what we are looking for. In order to create entanglement, we will need multi-qudit gates, i.e., unitaries which cannot be written as a tensor product of single-qudit unitaries. For $d=2$ the controlled-NOT gate is an example of such an entangling gate. It is defined as

\begin{align}
\mathrm{CNOT} = |0\rangle\langle 0|\otimes I + |1\rangle\langle 1|\otimes X\nonumber
\end{align}

and maps the product state $\frac{1}{\sqrt{2}}(|0\rangle + |1\rangle)\otimes|0\rangle$ onto an entangled Bell state $\frac{1}{\sqrt{2}}(|00\rangle + |11\rangle)$. We want something like this. However, to ensure that we can really create *arbitrary* states, we also want single-ion gates $U\otimes\otimes I$ which act as identity everywhere but on one qudit. Last but not least, we want these gates to be implementable in ion traps when more than two levels of the ions are under consideration.

**TASK** Literature search is a vital skill of any scientist. Unfortunately, it can also be difficult and not very rewarding. But keep your head up my fellow scientist, we live in an era of information where there is always too much rather than too little. Instead of going to the library, you just need to remain seated and watch the information scroll by.
So my quest to you is to find some helpful paper about *qudit quantum computation with ions*. Can you determine from a quick glance whether or not a paper is useful for what we are going to do and what we have already done? To answer this question you may require a lot of inuition especially when there is too much information. Nevertheless, this is a big part of research and should not be taken lightly.

### Solution

Luckily, it is 2021 and there are not too many papers around discussing ion trap quantum computers beyond qubits. My guess is that you may have stumbled over this paper:

*Practical trapped-ion protocols for universal qudit-based quantum computing*, P.J. Low, B.M. White, A.A. Cox, M.L. Day and C. Senko, [arXiv:1907.08569](https://arxiv.org/pdf/1907.08569.pdf).

If that is the case, we hopefully came to the same conclusion: That it is perfect for our purpose. I can spot two particularly relevant equations. Eqs. (2) and (6) both defining unitary operators which easily fit in our framework.

### Skeleton class

Skimming through the paper reveals that it is indeed possible to consider ions as $d$-level quantum systems. Good, so all our work is not for nothing. Furthermore, a second scan reveals the relevant equations: Eqs. (2) and (6) both define unitary operators that can be applied to the ions. In fact, Eq. (2) defines the action of a single-ion laser pulse and Eq. (6) defines a generalized Molmer-Sorensen (MS) gate. For those of you who have never heard of the MS gate, it is created by a *global* laser beam that interacts with all ions in a single trap. When considering qubits, the MS gate can entangle all ions at once. This seems to be its generalization to qudits, lucky us!

**TASK** Below is a skeleton class which inherits from the `QuditQM` class and contains methods `pulse` and `molmer_sorensen`. Make sure that this is in line with what we discussed and what you would have implemented.

```python
# dummy code -- do not run

class LaserGates(QuditQM):
    def __init__(self, dim, num_ions):
        """
        This class generates the required gate set for an ion trap quantum
        computer as described in https://arxiv.org/pdf/1907.08569.pdf.
        Specifically, it implements single-qudit gates of the form Eq. (2)and
        Molmer-Sorensen gates of the form Eq. (6).

        Args:
            dim (int): The local (odd) dimension of an ion.
            num_ions (int): The number of ions.
        """
        # add methods (such as `ket` or `bra`) from parent
        super().__init__(dim, num_ions)
        #`list` of `np.ndarray`: List of unitaries.
        self.gates = []

    def pulse(self, j, k, theta, phi):
        """
        Single-ion laser pulse of the form Eq. (2) in
        https://arxiv.org/pdf/1907.08569.pdf.

        NOTE: According to the paper only nearest neighbor transitions
        k = j+1 are practically implementable.

        Args:
            j (int): Basis vector |j> for transition |j><k|.
            k (int): Basis vector |k>for transition |j><k|.
            theta (float): Pulse angle (depending on the Rabi frequency).
            phi (float): Phase controlled by microwave or optical radiation.

        Returns:
            pulse (np.ndarray): The matrix representing the laser pulse's action
                                on an ion.
        """

    def molmer_sorensen(self, theta):
        """
        Molmer-Sorensen (MS) gate from Eq. (6) of
        https://arxiv.org/pdf/1907.08569.pdf.
        This is a global laser beam that can entangle all ions.

        Args:
            theta (float): The MS phase (depending on the Rabi frequency).

        Returns:
            ms (np.ndarray): MS gate unitary acting on all ions.
        """

```

### Laser quantum gates

**TASK** You have arrived at the last stretch, you can almost see the finish line now. You have learned a lot and done well to arrive here. So now, let us implement the `pulse` and `molmer_sorensen` methods using Eqs. (2) and (6). With your new skill set and methods built in the previous chapters, you should be able to translate Eqs. (2) and (6) directly into methods using `numpy` and `scipy.linalg.expm`.

For your convenience I already added a method `pad_gate(gate, ion)` which enables us to apply a single-ion `gate` generated by the `pulse` method to any `ion` of the `num_ions` ions. Practically, you have already done this in a previous exercise, so we skip it here.

```python
class LaserGates(QuditQM):
    def __init__(self, dim, num_ions):
        """
        This class generates the required gate set for an ion trap quantum
        computer as described in https://arxiv.org/pdf/1907.08569.pdf.
        Specifically, it implements single-qudit gates of the form Eq. (2)and
        Molmer-Sorensen gates of the form Eq. (6).

        Args:
            dim (int): The local (odd) dimension of an ion.
            num_ions (int): The number of ions.
        """
        # add methods (such as `ket` or `bra`) from parent
        super().__init__(dim, num_ions)
        #`list` of `np.ndarray`: List of unitaries.
        self.gates = []

    def pulse(self, j, k, theta, phi):
        """
        Single-ion laser pulse of the form Eq. (2) in
        https://arxiv.org/pdf/1907.08569.pdf.

        NOTE: According to the paper only nearest neighbor transitions
        k = j+1 are practically implementable.

        Args:
            j (int): Basis vector |j> for transition |j><k|.
            k (int): Basis vector |k>for transition |j><k|.
            theta (float): Pulse angle (depending on the Rabi frequency).
            phi (float): Phase controlled by microwave or optical radiation.

        Returns:
            pulse (np.ndarray): The matrix representing the laser pulse's action
                                on an ion.
        """
        # TODO: fill in

    def molmer_sorensen(self, theta):
        """
        Molmer-Sorensen (MS) gate from Eq. (6) of
        https://arxiv.org/pdf/1907.08569.pdf.
        This is a global laser beam that can entangle all ions.

        Args:
            theta (float): The MS phase (depending on the Rabi frequency).

        Returns:
            ms (np.ndarray): MS gate unitary acting on all ions.
        """
        # TODO: fill in

    def pad_gate(self, gate, ion):
        """
        Pads a single-ion gate with identities on other ions.

        Args:
            gate (np.ndarray): Single-ion gate to be padded.
            ion (int): The index of the ion on which the gate acts.

        Returns:
            multi_gate (np.ndarray): Tensor product of `gate` with identities on
                                     all other ions.
        """
        # define list of empty single-ion identities
        multi_gate = [np.eye(self.dim) for i in range(self.num_ions)]
        # add gate at the position of desired ion
        multi_gate[ion] = gate
        # apply tensor product between single-ion gates
        multi_gate = reduce(np.kron, multi_gate)

        return multi_gate

```

**TASK** Let us check that everything works as expected. Starting with the single-ion gates, Table VI in the paper gives us a good reference to check. Try reproducing the qutrit $X$-gate with the pulse sequences suggested in the table.

```python
dim = 3
num_ions = 3
laser = LaserGates(dim, num_ions)

# check that we can reproduce the Xd gate (FILL IN)

```

**TASK** What is the resulting state and its SRV when an MS gate with angle $-\frac{\pi}{2}$ is applied to some standard input states $|000\rangle$?

```python
# initial three-qutrit state
init_state = reduce(np.kron, [laser.ket(0) for i in range(num_ions)])

# resulting state (FILL IN)

# resulting state's SRV (FILL IN)

```

#### Solutions

```python
class LaserGates(QuditQM):
    def __init__(self, dim, num_ions):
        """
        This class generates the required gate set for an ion trap quantum
        computer as described in https://arxiv.org/pdf/1907.08569.pdf.
        Specifically, it implements single-qudit gates of the form Eq. (2)and
        Molmer-Sorensen gates of the form Eq. (6).

        Args:
            dim (int): The local (odd) dimension of an ion.
            num_ions (int): The number of ions.
        """
        # add methods (such as `ket` or `bra`) from parent
        super().__init__(dim, num_ions)
        #`list` of `np.ndarray`: List of unitaries.
        self.gates = []

    def pulse(self, j, k, theta, phi):
        """
        Single-ion laser pulse of the form Eq. (2) in
        https://arxiv.org/pdf/1907.08569.pdf.

        NOTE: According to the paper only nearest neighbor transitions
        k = j+1 are practically implementable.

        Args:
            j (int): Basis vector |j> for transition |j><k|.
            k (int): Basis vector |k>for transition |j><k|.
            theta (float): Pulse angle (depending on the Rabi frequency).
            phi (float): Phase controlled by microwave or optical radiation.

        Returns:
            pulse (np.ndarray): The matrix representing the laser pulse's action
                                on an ion.
        """
        # |j><k|
        jk = self.ket(j)@self.bra(k)
        # |k><j|
        kj = self.ket(k)@self.bra(j)
        # definition Eq. (2)
        pulse = linalg.expm(-1.j*theta*(np.exp(1.j*phi)*jk+np.exp(-1.j*phi)*kj))

        return pulse

    def molmer_sorensen(self, theta):
        """
        Molmer-Sorensen (MS) gate from Eq. (6) of
        https://arxiv.org/pdf/1907.08569.pdf.
        This is a global laser beam that can entangle all ions.

        Args:
            theta (float): The MS phase (depending on the Rabi frequency).

        Returns:
            ms (np.ndarray): MS gate unitary acting on all ions.
        """
        # define generalized X gate
        sx = (self.creation() + self.annihilation())/2.
        # define empty ms gate exponent
        ms = np.zeros((self.dim**self.num_ions, self.dim**self.num_ions),
                       dtype=complex
                     )
        # calculate matrix exponent, i.e. sum of Sx for all ions
        for n in range(self.num_ions):
            # define list of empty single-ion identities
            sx_n = [np.eye(self.dim) for i in range(self.num_ions)]
            # add Sx at position of the desired ion
            sx_n[n] = sx
            # apply tensor product between single-ion gates and add up
            ms += reduce(np.kron, sx_n)

        # definition Eq. (6)
        ms = linalg.expm(1.j*theta*(np.linalg.matrix_power(ms,2)))

        return ms

    def pad_gate(self, gate, ion):
        """
        Pads a single-ion gate with identities on other ions.

        Args:
            gate (np.ndarray): Single-ion gate to be padded.
            ion (int): The index of the ion on which the gate acts.

        Returns:
            multi_gate (np.ndarray): Tensor product of `gate` with identities on
                                     all other ions.
        """
        # define list of empty single-ion identities
        multi_gate = [np.eye(self.dim) for i in range(self.num_ions)]
        # add gate at the position of desired ion
        multi_gate[ion] = gate
        # apply tensor product between single-ion gates
        multi_gate = reduce(np.kron, multi_gate)

        return multi_gate

```

```python
dim = 3
num_ions = 3
laser = LaserGates(dim, num_ions)

# check that we can reproduce the Xd gate
(Xd==np.around(laser.pulse(0,1,np.pi/2,np.pi/2)@laser.pulse(1,2,np.pi/2,np.pi/2)@laser.pulse(0,1,np.pi,0),
               decimals=10
              )).all()
```

    True

```python
# initial three-qutrit state
init_state = reduce(np.kron, [laser.ket(0) for i in range(num_ions)])

# resulting state
np.around(laser.molmer_sorensen(-np.pi/2)@init_state, decimals=10)
# resulting state's SRV
laser.srv(laser.molmer_sorensen(-np.pi/2)@init_state)
```

    array([[ 0.5-0.5j],
           [ 0. +0.j ],
           [-0. -0.j ],
           [ 0. +0.j ],
           [-0. -0.j ],
           [ 0. +0.j ],
           [-0. -0.j ],
           [ 0. +0.j ],
           [-0. +0.j ],
           [ 0. +0.j ],
           [-0. +0.j ],
           [ 0. +0.j ],
           [-0. -0.j ],
           [ 0. +0.j ],
           [-0. +0.j ],
           [ 0. +0.j ],
           [-0. +0.j ],
           [ 0. +0.j ],
           [ 0. -0.j ],
           [ 0. +0.j ],
           [-0. +0.j ],
           [ 0. +0.j ],
           [-0. +0.j ],
           [ 0. +0.j ],
           [-0. +0.j ],
           [ 0. +0.j ],
           [-0.5-0.5j]])

    [2, 2, 2]

### Wrapping up

**Awesome**, we have a working class that implements all the quantum gates we could have wished for. However, it is already so far in the past that we have almost forgotten the purpose of this class: To generate an action set for our `IonTrapEnv`'s class. Well, that's what we intended the `gate` attribute to be. But it is still just an empty list in our class. We need to feed it with unitary quantum gates from our selection of gates! To this end, we need to select some angles for the different gates and add unitaries to the list.

**TASK** Because this is not very informative, I have completed the class in your stead. More or less at random, I have chose the following set of angles:

* $\theta$ in Eq. (2): ($\frac{\pi}{2}$) as `pulse_angles`
* $\phi$ in Eq. (2): ($0$, $\frac{\pi}{2}$, $\frac{\pi}{6}$) as `pulse_phases`
* $\theta_0$ in Eq. (6) ($-\frac{\pi}{2}$) as `ms_phases`

Make sure that you understand what happened in the code below. **Please, run the code below.**

```python
class LaserGates(QuditQM):
    def __init__(self, dim, num_ions):
        """
        This class generates the required gate set for an ion trap quantum
        computer as described in https://arxiv.org/pdf/1907.08569.pdf.
        Specifically, it implements single-qudit gates of the form Eq. (2)and
        Molmer-Sorensen gates of the form Eq. (6).

        Args:
            dim (int): The local (odd) dimension of an ion.
            num_ions (int): The number of ions.
        """
        # add methods (such as `ket` or `bra`) from parent
        super().__init__(dim, num_ions)
        #list: List of angles theta in Eq. (2)
        self.pulse_angles = [np.pi/2]
        #list: List of angles phi in Eq. (2)
        self.pulse_phases = [0, np.pi/2, np.pi/6]
        #list: List of angles theta in Eq. (6)
        self.ms_phases = [-np.pi/2]
        #`list` of `np.ndarray`: List of unitaries.
        self.gates = []
        # fills the `gates` with unitaries as specified by the angles.
        self._generate_gates()

    def pulse(self, j, k, theta, phi):
        """
        Single-ion laser pulse of the form Eq. (2) in
        https://arxiv.org/pdf/1907.08569.pdf.

        NOTE: According to the paper only nearest neighbor transitions
        k = j+1 are practically implementable.

        Args:
            j (int): Basis vector |j> for transition |j><k|.
            k (int): Basis vector |k>for transition |j><k|.
            theta (float): Pulse angle (depending on the Rabi frequency).
            phi (float): Phase controlled by microwave or optical radiation.

        Returns:
            pulse (np.ndarray): The matrix representing the laser pulse's action
                                on an ion.
        """
        # |j><k|
        jk = self.ket(j)@self.bra(k)
        # |k><j|
        kj = self.ket(k)@self.bra(j)
        # definition Eq. (2)
        pulse = linalg.expm(-1.j*theta*(np.exp(1.j*phi)*jk+np.exp(-1.j*phi)*kj))

        return pulse

    def molmer_sorensen(self, theta):
        """
        Molmer-Sorensen (MS) gate from Eq. (6) of
        https://arxiv.org/pdf/1907.08569.pdf.
        This is a global laser beam that can entangle all ions.

        Args:
            theta (float): The MS phase (depending on the Rabi frequency).

        Returns:
            ms (np.ndarray): MS gate unitary acting on all ions.
        """
        # define generalized X gate
        sx = (self.creation() + self.annihilation())/2.
        # define empty ms gate exponent
        ms = np.zeros((self.dim**self.num_ions, self.dim**self.num_ions),
                       dtype=complex
                     )
        # calculate matrix exponent, i.e. sum of Sx for all ions
        for n in range(self.num_ions):
            # define list of empty single-ion identities
            sx_n = [np.eye(self.dim) for i in range(self.num_ions)]
            # add Sx at position of the desired ion
            sx_n[n] = sx
            # apply tensor product between single-ion gates and add up
            ms += reduce(np.kron, sx_n)

        # definition Eq. (6)
        ms = linalg.expm(1.j*theta*(np.linalg.matrix_power(ms,2)))

        return ms

    def pad_gate(self, gate, ion):
        """
        Pads a single-ion gate with identities on other ions.

        Args:
            gate (np.ndarray): Single-ion gate to be padded.
            ion (int): The index of the ion on which the gate acts.

        Returns:
            multi_gate (np.ndarray): Tensor product of `gate` with identities on
                                     all other ions.
        """
        # define list of empty single-ion identities
        multi_gate = [np.eye(self.dim) for i in range(self.num_ions)]
        # add gate at the position of desired ion
        multi_gate[ion] = gate
        # apply tensor product between single-ion gates
        multi_gate = reduce(np.kron, multi_gate)

        return multi_gate

    # ---------------------- helper methods ------------------------------------

    def _generate_gates(self):
        """
        Adds single-ion and MS unitaries to the list of n-ion gates according to
        the provided list of phases.
        """
        # add ms gate for each angle theta
        for theta in self.ms_phases:
            self.gates.append(self.molmer_sorensen(theta))

        # adds all single-ion gates
        # for all ions
        for ion in range(self.num_ions):
            # for all transitions k <--> k+1
            for k in range(self.dim-1):
                # for all angles theta
                for theta in self.pulse_angles:
                    # for all angles phi
                    for phi in self.pulse_phases:
                        # create gate
                        gate = self.pulse(k, k+1, theta, phi)
                        # add n-ion gate
                        self.gates.append(self.pad_gate(gate, ion))

```

## Running our ion trap quantum computer environment

Now, we have everything ready to finally complete our `IonTrapEnv`. This should be easy since we already had everything there in the skeleton except for `self.init_state`, `self.goal`, `self.srv()` and the `QuditQM` methods. Below I implemented the last missing pieces with a few modifications, e.g., `dim`, `num_ions` and a set of `goal` SRVs are now args to the class.

**TASK** The `IonTrapEnv` is ready. Now, you need to give yourself a clap on the back. You are done. Congratulations, we can *play* now.

```python
class IonTrapEnv(QuditQM):
    def __init__(self, dim, num_ions, goal):
        """
        This is the environment for a d-level ion trap quantum computer as
        described in https://arxiv.org/pdf/1907.08569.pdf.
        An agent has a certain gate set at its disposal made up from elementary
        laser pulses and global molmer sorensen gates.

        At each time step an agent can choose one gate to apply to the current
        quantum state of the environment. The state, as well as observation,
        is a d^n dimensional quantum state where d is the local dimension of
        each ion and n is the number of ions. The initial state is |0...0>.

        The goal is to create multipartite high-dimensional entanglement between
        ions characterizsed by a Schmidt rank vector. Once such a state has
        been created, the agent receives a reward.

        If a maximum number of time steps is reached, an episoded is forcefully
        ended and should be restarted.

        Args:
            dim (int): Local ion dimension. Must be odd.
            num_ions (int): Number of ions.
            goal (list): List of SRVs that need to be obtained.
        """
        # add methods (such as `ket` or `bra`) from parent
        super().__init__(dim, num_ions)
        self.goal = goal

        #:class:`LaserGates`: Some class that lets us generate the unitaries used as actions.
        self.gates = LaserGates(self.dim, self.num_ions)
        #list: The action gate set obtained from the class.
        self.actions = self.gates.gates
        #int: Number of actions.
        self.num_actions = len(self.actions)

        #np.ndarray: Initial state.
        self.init_state = reduce(np.kron,
                                 [self.ket(0) for i in range(self.num_ions)]
                                )
        #np.ndarray: Current state of the environment and observation.
        self.state = self.init_state

        #int: Current number of time steps.
        self.time_step = 0
        #int: Maximum number of time steps.
        self.max_steps = 10

    def reset(self):
        """
        Resets the environment to its initial state, i.e., |0...0>.
        Resets counter.

        Returns:
            state (np.ndarray): Initial state vector.
        """
        self.state = self.init_state
        self.time_step = 0

        return self.state

    def step(self, action):
        """
        Performs and evaluates an action on the environment.

        (1) A quantum gate from our laser gate set is applied to the
            current state of the environment.
        (2) The Schmidt rank vector (SRV) is calculate for the current state
            of the environment. I
        (3) If the SRV is among the set of goal SRVs, a reward is provided and
            the trial is done. If not, no reward is given. The trial is ended if
            the time exceed the maximum number of steps.

        Args:
            action (int): Index of action in gate set.

        Returns:
            state (np.ndarray): Current state of the environment.
            reward (float): Current reward.
            done (bool): Whether or not the episode is done.
        """
        done = False
        reward = 0.
        # increment counter
        self.time_step += 1
        # (1) Apply gate to state
        self.state = self.gates.gates[action]@self.state
        # (2) Calculate SRV
        srv = self.srv(self.state) # some method that we are still missing
        # (3) Finish episode and give reward if appropriate
        if srv in self.goal:
            done = True
            reward = 1.
        elif self.time_step >= self.max_steps:
            done = True

        return self.state, reward, done

```

Remember the very first code we have written in this way too long notebook? The random agent interacting with the environment? Well, we can run it now to randomly search for $n$-qudit GHZ states,

\begin{align}
|\mathrm{GHZ}_{n,d}\rangle = \frac{1}{\sqrt{d}}\sum_{i=0}^{d-1}|i\rangle\otimes\cdots\otimes|i\rangle\nonumber
\end{align}

characterized by the maxmimum SRV.

\begin{align}
\vec{d}_{\mathrm{ghz}} = (d,\dots, d)\nonumber
\end{align}.

**TASK** Run the code below and fool around with the environment as you see fit. You earned it! **Have fun!**

```python
dim = 3
num_ions = 3
goal_srv = [[3,3,3]]
env = IonTrapEnv(dim, num_ions, goal_srv)

for i in range(1000):
    done = False
    observation = env.reset()
    action_seq = []
    while not done:
        action = np.random.choice(env.num_actions)
        action_seq.append(action)
        observation, reward, done = env.step(action)

        if reward > 0:
            print(f'The following action sequence obtained a reward: {action_seq}')
```

    The following action sequence obtained a reward: [0, 9, 10, 2, 0, 16, 10, 2, 6, 0]
    The following action sequence obtained a reward: [5, 0, 2, 0, 16, 9, 7, 7, 0]
    The following action sequence obtained a reward: [0, 3, 17, 0, 16, 8, 12, 11, 0]

```python

```