# Source code for pennylane.pulse.convenience_functions

# Copyright 2018-2023 Xanadu Quantum Technologies Inc.

# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

# Unless required by applicable law or agreed to in writing, software
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
"""This file contains convenience functions for pulse programming."""
from typing import Callable, List, Tuple, Union

import numpy as np

has_jax = True
try:
import jax.numpy as jnp
except ImportError:
has_jax = False

# pylint: disable=unused-argument
[docs]def constant(scalar, time):
"""Returns the given scalar, for use in defining a :class:~.ParametrizedHamiltonian with a
trainable coefficient.

Args:
scalar (float): the scalar to be returned
time (float): Time. This argument is not used, but is required to match the call
signature of :class:~.ParametrizedHamiltonian.
Returns:
float: The input scalar.

This function is mainly used to build a :class:~.ParametrizedHamiltonian that can be differentiated
with respect to its time-independent term. It is an alias for lambda scalar, t: scalar.

**Example**

The constant function can be used to create a parametrized Hamiltonian

>>> H = qml.pulse.constant * qml.PauliX(0)

When calling the parametrized Hamiltonian, constant will always return the input parameter

>>> params = 
>>> H(params, t=8)
5*(PauliX(wires=))

>>> H(params, t=5)
5*(PauliX(wires=))

We can differentiate the parametrized Hamiltonian with respect to the constant parameter:

.. code-block:: python

dev = qml.device("default.qubit.jax", wires=1)
@qml.qnode(dev, interface="jax")
def circuit(params):
qml.evolve(H)(params, t=2)
return qml.expval(qml.PauliZ(0))

>>> params = jnp.array([5.0])
>>> circuit(params)
Array(0.40808904, dtype=float32)

Array([-3.6517754], dtype=float32)
"""
return scalar

[docs]def rect(x: Union[float, Callable], windows: Union[Tuple[float], List[Tuple[float]]] = None):
"""Takes a scalar or a scalar-valued function, x, and applies a rectangular window to it, such that the
returned function is x inside the window and 0 outside it.

Creates a callable for defining a :class:~.ParametrizedHamiltonian.

Args:
x (Union[float, Callable]): either a scalar, or a function that accepts two arguments: the trainable
parameters and time
windows (Union[Tuple[float], List[Tuple[float]]]): List of tuples containing time windows where x is
evaluated. If None it is always evaluated. Defaults to None.

Returns:
callable: A callable f(p, t) which evaluates the given function/scalar x inside the time windows defined in
windows, and otherwise returns 0.

.. note::
If x is a function, it must accept two arguments: the trainable parameters and time. The primary use
of rect is for numerical simulations via :class:ParametrizedEvolution, which assumes t to be a single scalar
argument. If you need to efficiently compute multiple times, you need to broadcast over t via
jax.vmap <https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html>_ (see examples below).

**Example**

Here we use :func:~.rect to create a parametrized coefficient that has a value of 0 outside the time interval
t=(1, 7), and is defined by jnp.polyval(p, t) within the interval:

.. code-block:: python3

def f(p, t):
return jnp.polyval(p, t)

p = jnp.array([1, 2, 3])
time = jnp.linspace(0, 10, 1000)
windows = [(1, 7)]

windowed_f = qml.pulse.rect(f, windows=windows)

y1 = f(p, time)
y2 = jax.vmap(windowed_f, (None, 0))(p, time)

plt.plot(time, y1, label=f"polyval(p={p}, t)")
plt.plot(time, y2, label=f"rect(polyval, windows={windows})(p={p}, t)")
plt.legend()
plt.xlabel("t")
plt.show()

.. figure:: ../../_static/pulse/rect_example.png
:align: center
:width: 60%
:target: javascript:void(0);

Note that in order to efficiently create y2, we broadcasted windowed_f over the
time argument using jax.vmap <https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html>_.

rect can be used to create a :class:~.ParametrizedHamiltonian in the following way:

>>> H = qml.pulse.rect(jnp.polyval, windows=[(1, 7)]) * qml.PauliX(0)

The resulting Hamiltonian will be non-zero only inside the window.

>>> H([[1, 3]], t=2)  # inside the window
5.0*(PauliX(wires=))

>>> H([[1, 3]], t=0.5 )  # outside the window
0.0*(PauliX(wires=))

It is also possible to define multiple windows for the same function:

.. code-block:: python

windows = [(1, 7), (9, 14)]
H = qml.pulse.rect(jnp.polyval, windows) * qml.PauliX(0)

When calling the :class:.ParametrizedHamiltonian, rect will evaluate the given function only
inside the time windows, and otherwise return 0.

One can also pass a scalar to the rect function

>>> H = qml.pulse.rect(10, (1, 7)) * qml.PauliX(0)

In this case, rect will return the given scalar only when the time is inside the provided
time windows

>>> params = [None]  # the parameter value won't be used!
>>> H(params, t=8)
0.0*(PauliX(wires=))

>>> H(params, t=5)
10.0*(PauliX(wires=))
"""
if not has_jax:
raise ImportError(
"Module jax is required for any pulse-related convenience function. "
"You can install jax via: pip install jax==0.4.10 jaxlib==0.4.10"
)
if windows is not None:
is_nested = any(hasattr(w, "__len__") for w in windows)
single_window = len(windows) == 2 and not is_nested
if single_window:
windows = [windows]
elif not all(hasattr(w, "__len__") and len(w) == 2 for w in windows):
raise ValueError("At least one provided window is not a two-element sequence.")

if not callable(x):

def _f(_, __):
return jnp.array(x, dtype=float)

else:
_f = x

def f(p, t):
p = jnp.array(p, dtype=float)  # if p is an integer, f(p, t) will be cast to an integer
if windows is not None:
ti, tf = zip(*windows)
ti, tf = jnp.array(ti), jnp.array(tf)
return jnp.where(jnp.any((t >= ti) & (t <= tf)), _f(p, t), 0)
return _f(p, t)

return f

[docs]def pwc(timespan):
"""Takes a time span and returns a callable for creating a function that is piece-wise constant in time. The returned
function takes arguments (p, t), where p is an array that defines the bin values for the function.

Creates a callable for defining a :class:~.ParametrizedHamiltonian.

Args:
timespan(Union[float, tuple(float, float)]): The time span defining the region where the function is non-zero.
If an integer is provided, the time span is defined as (0, timespan).

Returns:
callable: a function that takes two arguments: an array of trainable parameters, and a float defining the
time at which the function is evaluated.

The convenience function pwc essentially implements

.. code-block:: python3

def pwc(timespan):
def wrapped(p, t):
return p[int(t/len(p))]
return wrapped

This function can be used to create a parametrized coefficient function that is piece-wise constant
within the interval t, and 0 outside it.

When creating the callable, only the time span is passed. The number
of bins and values for the parameters are set when params is passed to the callable. Each bin value is set by
an element of the params array. The variable t is used to select the value of the parameter array
corresponding to the specified time, based on the assigned binning.

.. code-block:: python3

params = jnp.array([1, 2, 3, 4, 5])
time = jnp.linspace(0, 10, 1000)
timespan=(2, 7)
y = qml.pulse.pwc(timespan)(params, time)
plt.plot(time, y, label=f"params={params}, timespan={timespan}")
plt.legend()
plt.show()

.. figure:: ../../_static/pulse/pwc_example.png
:align: center
:width: 60%
:target: javascript:void(0);

.. warning::
The final time in the time span indicates the time at which the function output switches from params[-1] to 0.
As such, the above function returns 5 for a time slightly smaller than the final time in timespan,
but it returns 0 for the final time itself:

>>> qml.pulse.pwc(timespan)(params, 6.999999)
Array(5., dtype=float32)

>>> qml.pulse.pwc(timespan)(params, 7.)
Array(0., dtype=float32)

**Example**

>>> timespan = (2, 7)
>>> f1 = qml.pulse.pwc(timespan)
>>> H = f1 * qml.PauliX(0)

The resulting function f1 has the call signature f1(params, t). If passed an array of parameters and
a time, it will assign the array as the constants in the piece-wise function, and select the constant corresponding
to the specified time, based on the time interval defined by timespan.

In the following example, passing an array to pwc((2, 7)) evenly distributes the array values in the
interval t=2 to t=7. The time t is then used to select one of the array values based on this distribution.

>>> H(params=[[11, 12, 13, 14, 15]], t=2.3)
11.0*(PauliX(wires=))

>>> H(params=[[11, 12, 13, 14, 15]], t=2.5) # different time, same bin, same result
11.0*(PauliX(wires=))

>>> H(params=[[11, 12, 13, 14, 15]], t=3.1) # next bin
12.0*(PauliX(wires=))

>>> H(params=[[11, 12, 13, 14, 15]], t=8) # outside the window returns 0
0.0*(PauliX(wires=))

"""
if not has_jax:
raise ImportError(
"Module jax is required for any pulse-related convenience function. "
"You can install jax via: pip install jax==0.4.3 jaxlib==0.4.3"
)

if isinstance(timespan, (tuple, list)):
t0, t1 = timespan
else:
t0 = 0
t1 = timespan

def func(params, t):
num_bins = len(params)
params = jnp.concatenate([jnp.array(params), jnp.zeros(1)])
# get idx from timestamp, then set idx=0 if idx is out of bounds for the array
idx = num_bins / (t1 - t0) * (t - t0)
idx = jnp.where((idx >= 0) & (idx <= num_bins), jnp.array(idx, dtype=int), -1)

return params[idx]

return func

[docs]def pwc_from_function(timespan, num_bins):
"""
Decorates a smooth function, creating a piece-wise constant function that approximates it.

Creates a callable for defining a :class:~.ParametrizedHamiltonian.

Args:
timespan(Union[float, tuple(float)]): The time span defining the region where the function is non-zero.
If a float is provided, the time span is defined as (0, timespan).
num_bins(int): number of bins for time-binning the function

Returns:
callable: a function that takes some smooth function f(params, t) and converts it to a
piece-wise constant function spanning time t in num_bins bins.

**Example**

.. code-block:: python3

def smooth_function(params, t):
return params * t + params

timespan = 10
num_bins = 10

binned_function = qml.pulse.pwc_from_function(timespan, num_bins)(smooth_function)

>>> binned_function([2, 4], 3), smooth_function([2, 4], 3)  # t = 3
(Array(10.666667, dtype=float32), 10)

>>> binned_function([2, 4], 3.2), smooth_function([2, 4], 3.2)  # t = 3.2
(Array(10.666667, dtype=float32), 10.4)

>>> binned_function([2, 4], 4.5), smooth_function([2, 4], 4.5)  # t = 4.5
(Array(12.888889, dtype=float32), 13.0)

The same effect can be achieved by decorating the smooth function:

.. code-block:: python

from pennylane.pulse.convenience_functions import pwc_from_function

@pwc_from_function(timespan, num_bins)
def fn(params, t):
return params * t + params

>>> fn([2, 4], 3)
Array(10.666667, dtype=float32)

"""
if not has_jax:
raise ImportError(
"Module jax is required for any pulse-related convenience function. "
"You can install jax via: pip install jax==0.4.3 jaxlib==0.4.3"
)

if isinstance(timespan, tuple):
t0, t1 = timespan
else:
t0 = 0
t1 = timespan

def inner(fn):
time_bins = np.linspace(t0, t1, num_bins)

def wrapper(params, t):
constants = jnp.array(list(fn(params, time_bins)) + )

idx = num_bins / (t1 - t0) * (t - t0)
# check interval is within 0 to num_bins, then cast to int, to avoid casting outcomes between -1 and 0 as 0
idx = jnp.where((idx >= 0) & (idx <= num_bins), jnp.array(idx, dtype=int), -1)

return constants[idx]

return wrapper

return inner


Using PennyLane

Development

API

Internals