Running FineDrag experiment from Qiskit Experiments with a dynamic backend with JAX results in a faulty datapoint (that's almost zero)

[jaeunkim]

Informations

• Qiskit Dynamics version: 0.4.0
• Python version: 3.11.3
• Operating system: MacOS Ventura, version 13.3.1(a)
• JAX version: 0.4.10

What is the current behavior?

Running the FineDrag experiment provided by Qiskit Experiments with a dynamic backend with JAX results in a faulty datapoint that's almost zero. The faulty datapoint is always the first datapoint. This ruins the subsequent curve fitting, because the sinusoidal fitting function assumes that its offset (called base) is the average between the maximum and the minimum datapoint. Although all the datapoint should be around 0.5, the average becomes 0.25 due to this outlying faulty datapoint that's almost zero. The problem doesn't occur if you configure the Dynamics solver without JAX.

• FineDrag class from Qiskit Experiments: https://qiskit.org/ecosystem/experiments/stubs/qiskit_experiments.library.characterization.FineDrag.html

• ErrorAmplificationAnalysis. Variable base is set to a wrong number in the method _generate_fit_guesses() https://qiskit.org/ecosystem/experiments/_modules/qiskit_experiments/curve_analysis/standard_analysis/error_amplification_analysis.html#ErrorAmplificationAnalysis

• Result with JAX

• Result without JAX

Steps to reproduce the problem

• With JAX

  
  import jax
  import numpy as np
  from qiskit import pulse
  from qiskit.circuit import Parameter, Gate
  from qiskit.circuit.library import RZGate, SXGate
  from qiskit.transpiler import InstructionProperties
  from qiskit_dynamics import Solver, DynamicsBackend
  from qiskit_experiments.library.characterization import FineDrag

setup JAX

jax.config.update("jax_enable_x64", True)
jax.config.update('jax_platform_name', 'cpu')

# use Array class provided by qiskit_dynamics for jax
from qiskit_dynamics.array import Array
Array.set_default_backend('jax')

# Setup Hamiltonian and Dynamics Solver
r = 168761338.45929572
v = 5073462814.921423
anharm = -328548437.7601424
dt = 0.000000000222
dim = 4

a = np.diag(np.sqrt(np.arange(1, dim)), 1)
adag = np.diag(np.sqrt(np.arange(1, dim)), -1)
N = np.diag(np.arange(dim))

static_hamiltonian = 2 * np.pi * v * N + np.pi * anharm * N * (N - np.eye(dim))
drive_hamiltonian = 2 * np.pi * r * (a + adag)

ham_solver = Solver(
    hamiltonian_operators=[drive_hamiltonian],
    static_hamiltonian=static_hamiltonian,
    rotating_frame=static_hamiltonian,
    hamiltonian_channels=["d0"],
    channel_carrier_freqs={"d0": v},
    dt=dt,
)
solver_options = {"method": "jax_odeint", "atol": 1e-6, "rtol": 1e-8}

backend = DynamicsBackend(
    solver=ham_solver,
    solver_options=solver_options,  # to be used every time run is called
)
target = backend.target

# Add RZ instruction as phase shift, because it's required for FineDrag sequence
phi = Parameter("phi")
with pulse.build() as rz0:
    pulse.shift_phase(phi, pulse.DriveChannel(0))

target.add_instruction(
    RZGate(phi),
    {(0,): InstructionProperties(calibration=rz0)}
)

# Add SX gate because it's required for FineDrag sequence
with pulse.build() as SX:
    pulse.play(pulse.Gaussian(duration=160, amp=0.07824882, sigma=40), pulse.DriveChannel(0))
# note: we checked that the drag component miscalibration for this SX gate does not affect the outcome

target.add_instruction(
        SXGate(),
        {(0,): InstructionProperties(calibration=SX)}
)

# This is the schedule to be calibrated
# Let's start with the known optimal beta. This is equivalent to beginning with the result from RoughDrag. 
# Thus, FineDrag result is expected to be almost flat.
with pulse.build() as half_pi_reproduce:
    pulse.play(pulse.Drag(duration=160, amp=0.07824882, sigma=40, beta=-1.681), pulse.DriveChannel(0))

# Setting up the FineDrag experiment
my_gate = Gate("FineDragTest", 1, [])
exp = FineDrag((0,), gate=my_gate, backend=backend)
exp.set_experiment_options(schedule=half_pi_reproduce)
exp.set_experiment_options(repetitions=range(0, 20))

exp.analysis.plotter.set_figure_options(figure_title=f"rotation angle=half pi")
exp_data = exp.run().block_for_results()

exp_data.figure(0)

- Without JAX
```Python
import numpy as np
from qiskit import pulse
from qiskit.circuit import Parameter, Gate
from qiskit.circuit.library import RZGate, SXGate
from qiskit.transpiler import InstructionProperties
from qiskit_dynamics import Solver, DynamicsBackend
from qiskit_experiments.library.characterization import FineDrag

# use Array class provided by qiskit_dynamics, but with numpy
    from qiskit_dynamics.array import Array
    Array.set_default_backend('numpy')

    # Setup Hamiltonian and Dynamics Solver
    r = 168761338.45929572
    v = 5073462814.921423
    anharm = -328548437.7601424
    dt = 0.000000000222
    dim = 4

    a = np.diag(np.sqrt(np.arange(1, dim)), 1)
    adag = np.diag(np.sqrt(np.arange(1, dim)), -1)
    N = np.diag(np.arange(dim))

    static_hamiltonian = 2 * np.pi * v * N + np.pi * anharm * N * (N - np.eye(dim))
    drive_hamiltonian = 2 * np.pi * r * (a + adag)

    ham_solver = Solver(
        hamiltonian_operators=[drive_hamiltonian],
        static_hamiltonian=static_hamiltonian,
        rotating_frame=static_hamiltonian,
        hamiltonian_channels=["d0"],
        channel_carrier_freqs={"d0": v},
        dt=dt,
    )
    # note: using the default "method", instead of jax_odeint.
    solver_options = {"atol": 1e-6, "rtol": 1e-8}

    backend = DynamicsBackend(
        solver=ham_solver,
        solver_options=solver_options,  # to be used every time run is called
    )
    target = backend.target

    # Add RZ instruction as phase shift, because it's required for FineDrag sequence
    phi = Parameter("phi")
    with pulse.build() as rz0:
        pulse.shift_phase(phi, pulse.DriveChannel(0))

    target.add_instruction(
        RZGate(phi),
        {(0,): InstructionProperties(calibration=rz0)}
    )

    # Add SX gate because it's required for FineDrag sequence
    with pulse.build() as SX:
        pulse.play(pulse.Gaussian(duration=160, amp=0.07824882, sigma=40), pulse.DriveChannel(0))
    # note: we checked that the drag component miscalibration for this SX gate does not affect the outcome

    target.add_instruction(
            SXGate(),
            {(0,): InstructionProperties(calibration=SX)}
    )

    # This is the schedule to be calibrated
    # Let's start with the known optimal beta. This is equivalent to beginning with the result from RoughDrag. 
    # Thus, FineDrag result is expected to be almost flat.
    with pulse.build() as half_pi_reproduce:
        pulse.play(pulse.Drag(duration=160, amp=0.07824882, sigma=40, beta=-1.681), pulse.DriveChannel(0))

    # Setting up the FineDrag experiment
    my_gate = Gate("FineDragTest", 1, [])
    exp = FineDrag((0,), gate=my_gate, backend=backend)
    exp.set_experiment_options(schedule=half_pi_reproduce)
    exp.set_experiment_options(repetitions=range(0, 20))

    exp.analysis.plotter.set_figure_options(figure_title=f"rotation angle=half pi")
    exp_data = exp.run().block_for_results()

    exp_data.figure(0)

What is the expected behavior?

The result with JAX should be the same with the result without JAX.

Suggested solutions

[DanPuzzuoli]

Thank you for this detailed report. This is actually another realization of a recurring issue. If you add: "hmax": dt to solver_options in your JAX code, it will work properly.

See my response to https://github.com/Qiskit-Extensions/qiskit-dynamics/issues/238 for an explanation.

[jaeunkim]

Thank you @DanPuzzuoli for pointing me to the explanation. Now I understand what was going on. Moreover, adding "hmax": dt solved the problem 😊