A1. State preparation with Jaynes-Cummings controls for | 1 ⟩ + | 3 ⟩ fock state

A1. State preparation with Jaynes-Cummings controls for \( | 1 ⟩ + | 3 ⟩ \) fock state#

from feedback_grape.fgrape import optimize_pulse
from feedback_grape.utils.operators import (
    sigmap,
    sigmam,
    create,
    destroy,
    identity,
)
from feedback_grape.utils.states import basis, fock
from feedback_grape.utils.tensor import tensor
import jax.numpy as jnp
from jax.scipy.linalg import expm

As a preliminary step, we consider state preparation of a target state starting from a pure state. In addition, we assume that any coupling to an external environment is negligible and that the parametrized controls can be implemented perfectly.

Here no feedback is required, we are just testing the parameterized gates setup.

As a first example, we consider the state preparation of a cavity resonantly coupled to an externally driven qubit

Here, we consider a particular sequence of parametrized unitary gates originally introduced by Law and Eberly

N_cav = 30

def qubit_unitary(alphas):
    alpha_re = alphas[0]
    alpha_im = alphas[1]
    alpha = alpha_re + 1j * alpha_im
    return tensor(
        identity(N_cav),
        expm(-1j * (alpha * sigmap() + alpha.conjugate() * sigmam()) / 2),
    )

def qubit_cavity_unitary(beta_re):
    beta = beta_re
    return expm(
        -1j
        * (
            beta * (tensor(destroy(N_cav), sigmap()))
            + beta.conjugate() * (tensor(create(N_cav), sigmam()))
        )
        / 2
    )

In their groundbreaking work, Law and Eberly have shown that any arbitrary superposition of Fock states with maximal excitation number N can be prepared out of the ground state in a sequence of N such interleaved gates, also providing an algorithm to find the correct angles and interaction durations

First target is the state \( | 1 ⟩ + | 3 ⟩ \)#

time_steps = 4  # corressponds to maximal excitation number of an arbitrary Fock State Superposition

psi0 = tensor(basis(N_cav), basis(2))
psi_target = tensor((fock(N_cav, 1) + fock(N_cav, 3)) / jnp.sqrt(2), basis(2))

psi0.shape

(60, 1)

psi_target.shape

(60, 1)

print(fock(N_cav, 1))

[[0.+0.j]
 [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]
 [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]]

Law and Eberly provided an algorithm to determine the correct parameters for state preparation. These include:

The rotation angle \( |\alpha| \),
The azimuthal angle \( \arg\left(\frac{\alpha}{|\alpha|}\right) \),
The interaction duration \( |\beta| \).

So Goal is to find the best control vector (rather than control amplitudes, this time) that leads to finding the optimal state-preparation strategies. Performing as well as the Law-Eberly algorithm.

Optimizing#

Currently l-bfgs with the same learning rate of 0.3 converges at a local minimum of 0.5, adam also converges at 0.5 but at smaller learning rates

from feedback_grape.fgrape import Gate
import jax

key = jax.random.PRNGKey(42)
# not provideing param_constraints just propagates the same initial_parameters for each time step
qub_unitary = Gate(
    gate=qubit_unitary,
    initial_params=jax.random.uniform(
        key,
        shape=(2,),  # 2 for gamma and delta
        minval=-5 * jnp.pi,
        maxval=5 * jnp.pi,
        dtype=jnp.float64,
    ),
    measurement_flag=False,
)

qub_cav = Gate(
    gate=qubit_cavity_unitary,
    initial_params=jax.random.uniform(
        key,
        shape=(1,),
        minval=-5 * jnp.pi,
        maxval=5 * jnp.pi,
        dtype=jnp.float64,
    ),
    measurement_flag=False,
)

system_params = [qub_unitary, qub_cav]


result = optimize_pulse(
    U_0=psi0,
    C_target=psi_target,
    system_params=system_params,
    num_time_steps=time_steps,
    max_iter=1000,
    convergence_threshold=1e-16,
    evo_type="state",
    mode="no-measurement",
    goal="fidelity",
    learning_rate=0.05,
    batch_size=10,
    eval_batch_size=2,
)

len(result.returned_params)

# here makes sense for each batch size we have a different set of parameters since there are no measurements and therefore no stochasticisty or randomness
result.returned_params

[[Array([[ -0.20469016, -12.56470343],
         [ -0.20469016, -12.56470343]], dtype=float64),
  Array([[-2.71356038],
         [-2.71356038]], dtype=float64)],
 [Array([[ 7.91772224e-07, -1.57079633e+01],
         [ 7.91772224e-07, -1.57079633e+01]], dtype=float64),
  Array([[-1.66925597],
         [-1.66925597]], dtype=float64)],
 [Array([[ 6.76720049e-08, -1.57079633e+01],
         [ 6.76720049e-08, -1.57079633e+01]], dtype=float64),
  Array([[-1.03771436],
         [-1.03771436]], dtype=float64)],
 [Array([[ 1.50110351e-07, -1.57079633e+01],
         [ 1.50110351e-07, -1.57079633e+01]], dtype=float64),
  Array([[-3.14159264],
         [-3.14159264]], dtype=float64)]]

print(result.final_state)

[[[ 1.10550300e-09+1.19607397e-11j]
  [-4.03327373e-10+1.70038772e-08j]
  [ 2.87728545e-08-7.07106781e-01j]
  [ 1.45929005e-11-1.35150716e-09j]
  [ 4.45239623e-10+9.84414454e-12j]
  [-9.03587760e-11+4.14909862e-09j]
  [ 4.54459230e-08-7.07106782e-01j]
  [ 4.40660838e-12+1.11246573e-09j]
  [ 0.00000000e+00+0.00000000e+00j]
  [-8.75992484e-09-5.63002478e-16j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]]

 [[ 1.10550300e-09+1.19607397e-11j]
  [-4.03327373e-10+1.70038772e-08j]
  [ 2.87728545e-08-7.07106781e-01j]
  [ 1.45929005e-11-1.35150716e-09j]
  [ 4.45239623e-10+9.84414454e-12j]
  [-9.03587760e-11+4.14909862e-09j]
  [ 4.54459230e-08-7.07106782e-01j]
  [ 4.40660838e-12+1.11246573e-09j]
  [ 0.00000000e+00+0.00000000e+00j]
  [-8.75992484e-09-5.63002478e-16j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]
  [ 0.00000000e+00+0.00000000e+00j]]]

print(result)

FgResult(optimized_trainable_parameters=[[Array([ -0.20469016, -12.56470343], dtype=float64), Array([-2.71356038], dtype=float64)], [Array([ 7.91772224e-07, -1.57079633e+01], dtype=float64), Array([-1.66925597], dtype=float64)], [Array([ 6.76720049e-08, -1.57079633e+01], dtype=float64), Array([-1.03771436], dtype=float64)], [Array([ 1.50110351e-07, -1.57079633e+01], dtype=float64), Array([-3.14159264], dtype=float64)]], iterations=372, final_state=Array([[[ 1.10550300e-09+1.19607397e-11j],
        [-4.03327373e-10+1.70038772e-08j],
        [ 2.87728545e-08-7.07106781e-01j],
        [ 1.45929005e-11-1.35150716e-09j],
        [ 4.45239623e-10+9.84414454e-12j],
        [-9.03587760e-11+4.14909862e-09j],
        [ 4.54459230e-08-7.07106782e-01j],
        [ 4.40660838e-12+1.11246573e-09j],
        [ 0.00000000e+00+0.00000000e+00j],
        [-8.75992484e-09-5.63002478e-16j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j]],

       [[ 1.10550300e-09+1.19607397e-11j],
        [-4.03327373e-10+1.70038772e-08j],
        [ 2.87728545e-08-7.07106781e-01j],
        [ 1.45929005e-11-1.35150716e-09j],
        [ 4.45239623e-10+9.84414454e-12j],
        [-9.03587760e-11+4.14909862e-09j],
        [ 4.54459230e-08-7.07106782e-01j],
        [ 4.40660838e-12+1.11246573e-09j],
        [ 0.00000000e+00+0.00000000e+00j],
        [-8.75992484e-09-5.63002478e-16j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j],
        [ 0.00000000e+00+0.00000000e+00j]]], dtype=complex128), returned_params=[[Array([[ -0.20469016, -12.56470343],
       [ -0.20469016, -12.56470343]], dtype=float64), Array([[-2.71356038],
       [-2.71356038]], dtype=float64)], [Array([[ 7.91772224e-07, -1.57079633e+01],
       [ 7.91772224e-07, -1.57079633e+01]], dtype=float64), Array([[-1.66925597],
       [-1.66925597]], dtype=float64)], [Array([[ 6.76720049e-08, -1.57079633e+01],
       [ 6.76720049e-08, -1.57079633e+01]], dtype=float64), Array([[-1.03771436],
       [-1.03771436]], dtype=float64)], [Array([[ 1.50110351e-07, -1.57079633e+01],
       [ 1.50110351e-07, -1.57079633e+01]], dtype=float64), Array([[-3.14159264],
       [-3.14159264]], dtype=float64)]], final_purity=None, final_fidelity=Array(1., dtype=float64), fidelity_each_timestep=[Array([0., 0.], dtype=float64), Array([0., 0.], dtype=float64), Array([0.06916415, 0.06916415], dtype=float64), Array([7.23876266e-20, 7.23876266e-20], dtype=float64), Array([1., 1.], dtype=float64)], purity_each_timestep=[])

result.final_fidelity

Array(1., dtype=float64)

A1. State preparation with Jaynes-Cummings controls for | 1 ⟩ + | 3 ⟩ fock state

Contents

A1. State preparation with Jaynes-Cummings controls for \( | 1 ⟩ + | 3 ⟩ \) fock state#

First target is the state \( | 1 ⟩ + | 3 ⟩ \)#

Optimizing#