Pulse module and chaotic Rabi Oscillations

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):
image

But things mess up if I increase time duration, for instance rabi_oscillations(time_duration=10, max_amplitude=2):
image

or decreasing the amplitude range:

image

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):

image

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 :frowning:

Hi @Kernel123 I’m sorry to hear that you’re stuck and thank you for your very detailed post! I’ll ask one of our experts to look at your question to try to identify the issue.

1 Like

Thanks Catalina!!
I also aked this question on a Quantum Computing Stack Exchange post, and @Korbinian_Kottmann answered me!. Maybe we can continue discussing directly through the Stack Exchange channel.
In case anyone else is interested in the solution, here is my question pos

@Kernel123 glad you found a solution :open_hands:

If you have further questions in the future, just note that we don’t actively monitor stack exchange, but we do keep tabs on this forum, so it provides the best avenue to get a response from our team