Enable piece-wise-constant functions

[lillian542]

Context:
We want to add a feature that users can use to do something like

H = pwc * H1 + pwc * H2
..
qml.Evolve(H)(params, ..)

I.e. provide a pwc that is a callable and gets thus recognized by the dunder method of operators to create ParametrizedHamiltonian.

One of the problems is that it needs information about the duration and/or how many samples there are such that it can assign based on the float parameter t in the ODE integrator the corresponding function value.

Description of the Change:
Two functions added in pennylane.math.utils:

• pwc: takes t and returns a pwc function with call signature f(params, t) that returns a value from params[index] based on t and the interval dt.
• pwc_from_function : takes t and num_bins, and decorates a smooth function to return a piecewise constant function with call signature f(params, t).

[AlbertMitjans]

Happy to approve once these comments have been addressed!

[Qottmann]

Awesome, great work @lillian542 !!! I have one nitpicky performance concern:

I noticed a regression between the self-written wrapper and qml.pulse.pwc, I used the following benchmark (comparing changing pwc to pwc2 in the envelope). While the execution is the same, the compilation time increases significantly (up to factor 2). I think I already found the cause of the problem. Can you change it and re-run the benchmark? If jit-time and exec time are both more or less the same I am happy to approve.

I am not entirely sure the behavior of outputting 0 for times out of bound is good, because this can lead to the function "working" by outputting zero instead of raising an error when it is incorrectly used. Additionally the jnp.where business for the index might be another source for the regression. If a user wants pwc to only work in a certain window they should use rect. Happy to hear a third opinion @trbromley . Not a hill I would die on, but I'd like to open the discussion one more time.

import pennylane as qml
import pennylane.numpy as np
import jax.numpy as jnp
import jax

from jax.experimental.ode import odeint as jaxodeint
from functools import partial

import matplotlib.pyplot as plt

from datetime import datetime
data = qml.data.load("qchem", molname="H2O", basis="STO-3G", bondlength=1.9)[0]
H_obj = data.hamiltonian
n_wires = len(H_obj.wires)
omega = jnp.ones(n_wires)
g = jnp.ones(n_wires)

H_D = qml.ops.dot(omega, [qml.PauliZ(i) @ qml.PauliZ(i) for i in range(n_wires)])
H_D += qml.ops.dot(g, [qml.PauliY(i) @ qml.PauliY((i+1)%n_wires) + qml.PauliY((i+1)%n_wires) @ qml.PauliY(i) for i in range(n_wires)])
# TODO use official convenience functions once merged
def pwc2(duration):
    def wrapped(params, t):
        N = len(params)
        idx = jnp.array(N/(duration) * t, dtype=int) # corresponding sample
        return params[idx]

    return wrapped

from pennylane.pulse import pwc

def envelope(duration):
    # assuming p = (len(t_bins) + 1) for the frequency nu
    def wrapped(p, t):
        return 0.02*pwc(duration)(p[:-1], t) * jnp.cos(p[-1]*t)
    return wrapped

duration = 20.

fs = [envelope(duration) for i in range(n_wires)]
ops = [qml.PauliX(i) for i in range(n_wires)]

H_C = qml.ops.dot(fs, ops)

H_pulse = H_D + H_C
t_bins = 40 # number of time bins
theta = jnp.array([jnp.ones(t_bins + 1, dtype=float) for _ in range(n_wires)])

dev = qml.device("default.qubit", wires=range(n_wires))

ts = jnp.linspace(0., duration, t_bins)

@jax.jit
@qml.qnode(dev, interface="jax")
def qnode(theta, t=ts):
    qml.BasisState(data.hf_state, wires=H_obj.wires)
    qml.evolve(H_pulse)(params=theta, t=t)
    return qml.expval(H_obj)
value_and_grad = jax.jit(jax.value_and_grad(qnode, argnums=0))
time0 = datetime.now()
val, grad = jax.block_until_ready(value_and_grad(theta))
time1 = datetime.now()
print(f"grad and val compilation time: {time1 - time0}")
# run in notebook
%timeit jax.block_until_ready(value_and_grad(theta))

[Qottmann]

All concerns resolved, thanks for the quick fix @lillian542 🚅

(For the record: Lillian made a good point about negative indexing, and I also realized that combining it with rect actually wouldnt work in the way I thought lambda p, t: pwc(duration)(p, t) * rect(duration2)(p, t))

[AlbertMitjans]

Good job!!