I am trying to make a sort of Rabi experiment with Pennylane pulse module, using as hamiltonian the qml.pulse.transmon_interaction and driving qml.pulse.transmon_drive with constant amplitude. I have seen some weird behaviour. Let me first show you my code:
import pennylane as qml
import jax
from jax import numpy as jnp
import matplotlib.pyplot as plt
w0 = 1 / (2 * jnp.pi)
w1 = 3 / (2 * jnp.pi)
g = 0.1 / (2 * jnp.pi)
amp = lambda A, t: A
phase = lambda phi0, t: phi0
ham = qml.pulse.transmon_interaction(qubit_freq=[w0, w1], connections=[[0,1]], coupling=g, wires=[0,1])
ham += qml.pulse.transmon_drive(amplitude=amp, phase=phase, freq=w0, wires=0)
dev = qml.device("default.qubit", wires=range(2))
@jax.jit
@qml.qnode(dev, interface="jax")
def qnode(params, duration):
qml.evolve(ham)(params, duration)
return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))
def rabi_oscillations(time_duration: float, max_amplitude: float):
phi0 = 0
func = jax.vmap(qnode, in_axes=(0,0))
As = jnp.linspace(1e-2, max_amplitude, 100)
durations = time_duration * jnp.ones_like(As)
params = jnp.stack([As, jnp.full_like(As, phi0)], axis=-1) # [[As[0], phi0], [As(1), phi0], ...]
z0s, z1s = func(params, durations)
plt.plot(As, z0s, label='z0s')
plt.plot(As, z1s, label='z1s')
plt.legend()
The idea was just to act on qubit 0 while qubit 1 is at rest, and see how the Rabi oscillations are. As suggested by some bibliography I have seen, I took |\omega_1 - \omega_2| \gg g to avoid interaction between qubits. I should then see how expected value of qubit 1 is constant (|0\rangle) while qubit 0 suffers a X rotation, whose angle is basically A \cdot T being T the pulse duration.
This is the case when I do: rabi_oscillations(time_duration=0.5, max_amplitude=10)
:
But things mess up if I increase time duration, for instance rabi_oscillations(time_duration=10, max_amplitude=2)
:
or decreasing the amplitude range:
I have been searching in the qml.evolve
documentation and source code and I see that the differential equation is solve with scipy.integrate.odeint
function, so qml.evolve
admits some kwargs. Among them , atol
is supposed to control the absolute error tolerance. I think this non-sense plot has to do with numerical approximation precision, since if I do qml.evolve(ham)(params, duration, atol=atol)
inside the qnode
function, I get different results. However, I do not suceed at controlling it, and do not understand its behaviour.
According to ParametrizedEvolution
Pennylane class (which is in charge for solving the Schrödinger evolution when calling qml.evolve
), the default value is 1.4e-8
. If I decrease this tolerance, I do not get a different plot than the ones I displayed above (I did up to atol=1e-20
). I think that very small values of atol
mess things up even more, since atol=1-30
returns nan
for z0s
and z1s
arrays, and atol=1e-25
returns some nan values as well. However, I am sure changing atol
has some effects on solutions because setting atol=1e-4
returns awful solutions (probabilities beyond 1):
I also played a bit with rtol
parameter but it seems to have no effect at all. hmax
and mxstep
arguments default to ´np.inf´, which I think that gives the best accuracy. I think there is not much I can do with numerical integration from my side.
As conclussion, I think that or I manually change the numerical integration function or that the problem has nothing to do with it. But in the latter, I cannot explain the results. Maybe some of you could help me since I am stuck