Quantum Tape for Spatial Gradients in Physic Informed NNs

Hello, I am trying to use qml to do physics informed quantum machine learning within Tensorflow. I know with TF, to get spatial derivatives of features in the network, you can use with tf.GradientTape() as tape and define a function representing a partial differential equation as:
@tf.function
def physic_loss(t, x):
u0 = u(t, x)
u_t = tf.gradients(u0, t)[0]
u_x = tf.gradients(u0, x)[0]
u_xx = tf.gradients(u_x, x)[0]
F = u_t + u0*u_x - (0.01/np.pi)*u_xx
return tf.reduce_mean(tf.square(F))

It doesn’t seem like the qnode components are contributing to this loss. My guess is that I need to use quantum grad tape somehow to get derivatives w.r.t x and t. Any guidance would be great!!!

Hey @Corey, welcome to the forum! :rocket:

I think you should check out our demo on how to transform QNodes to talk nicely with Tensorflow / Keras: Turning quantum nodes into Keras Layers | PennyLane Demos. Essentially, you can make a QNode look, smell, and taste just like a Tensorflow layer / function with qml.qnn.KerasLayer :sweat_smile:. Let me know if this helps!

Hello! Your link is to a very basic purely data-driven loss. I need gradients of the embedded features to include equation residuals into my loss. For example,
Loss = a * d2u/dx^2 + b * du/dx + c.

This requires taking gradients of the circuit w.r.t. the inputs - not parameters. I’m getting a little close with this:


import pennylane as qml
from pennylane import numpy as np

import autograd as ag

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

params = np.random.random([3], requires_grad=True)

x = np.array([[0.1, 0.2], [0.3, 0.4], [1.1, 1.2], [1.3, 1.4]], requires_grad=False)
opt = qml.GradientDescentOptimizer(0.9)


@qml.qnode(dev)
def circuit(params, x):
    qml.RX(x[:, 0], wires=0)
    qml.RY(x[:, 1], wires=1)
    qml.RZ(x[:, 0], wires=2)

    qml.RX(params[0], wires=0)
    qml.RY(params[1], wires=1)
    qml.RZ(params[2], wires=2)

    return qml.expval(qml.PauliZ(0))


def loss(params, x):
    grad_circ = ag.elementwise_grad(circuit)
    u_x = grad_circ(params, x)[0]

    return u_x

(params,  _), l = opt.step_and_cost(loss, params, x)

Hi @Corey,

Thank you for your question and sorry for the delay in our response.

We have forwarded this question to members of our technical team who will be getting back to you within a week. Feel free to post any updates to your question here in the thread in the meantime!

Hi @Corey,

If I’m following correctly, your end-goal is to compute the loss of a second-order function representing a physics problem. PennyLane can definitely get that done for you! Firstly, I see you used TensorFlow in one example, then Autograd in another. While both ML interfaces work with PennyLane (as well as Jax and PyTorch if those suit you better), it is strongly recommended to stick to a single interface throughout an experiment.

All data - parameters, inputs, TensorFlow, Autograd, whatever you may pass to a QNode - is trainable! The demo linked above regarding KerasLayer highlights a specific use-case where the inputs variable is not typically used to train a model, but there are many ways to get this done. Our demo on quantum analytic descent highlights an experiment that computes second-order partial derivatives (labelled hessian), hopefully that is closer to what you’re looking for. As an example, your first comment might look something like this (with some made-up u function):

dev = qml.device("default.qubit")

@qml.qnode(dev, max_diff=2)
def u(params):
    t, x = params
    qml.RX(x, wires=0)
    qml.RY(t, wires=0)
    return qml.expval(qml.PauliZ(0))

def physic_loss(params):
    with tf.GradientTape() as tape1:
        with tf.GradientTape() as tape2:
            u0 = u(params)
        grad = tape2.gradient(u0, params)

    u_t, u_x = grad
    u_xx = tape1.jacobian(grad, params)
    F = u_t + u0*u_x - (0.01/np.pi)*u_xx
    return tf.reduce_mean(tf.square(F))

params = tf.Variable([1.0, 2.0], dtype=tf.float64)
physic_loss(params)

Once you have that working, you can hook it in with your favourite TensorFlow optimizer! For more information on training with TensorFlow, check out this doc that features some comprehensive examples. Let me know if there’s anything else I can help you with.

Thank you so much Timmy! In your code, it looks like you are taking derivatives w.r.t. to the circuit parameters to me, not feature inputs. I’m not trying to optimize the features, just calculated their first and second order derivatives to add the residual to the loss function. I am still only optimizing the parameters.

For example, the circuit would more generally look like one of the following:

def u(features,params):
    qml.RX(features[0], wires=0)
    qml.RY(features[1], wires=1)
    qml.CNOT(wires=[0, 1])
    qml.PhaseShift(params[0], wires=0)
    return qml.expval(qml.PauliZ(0))

or as a qnode …

@qml.qnode(dev)
def qnode(inputs=features, weights=params):
    qml.AngleEmbedding(inputs, wires=range(n_qubits))
    qml.templates.StronglyEntanglingLayers(weights, wires=range(n_qubits))
    return [qml.expval(qml.PauliZ(wires=i)) for i in range(n_qubits)]

I appreciate you!

Happy to help, let’s get to the bottom of this :male_detective: I think I see what you mean. I’ve created a dummy example that computes the first and second order gradients of the features around the QNode, and uses those results to update the loss value which we train the weights with. Let me know if this is what you’re looking for:

import pennylane as qml
import tensorflow as tf
import numpy as np

n_qubits = 2

dev = qml.device("default.qubit")

@qml.qnode(dev)
def qnode(inputs, weights):
    qml.AngleEmbedding(inputs, wires=range(n_qubits))
    qml.templates.StronglyEntanglingLayers(weights, wires=range(n_qubits))
    return [qml.expval(qml.PauliZ(wires=i)) for i in range(n_qubits)]

weight_shape = qml.StronglyEntanglingLayers.shape(3, 2)

features = tf.Variable([1.1, 2.2], dtype=tf.float64)
weights = tf.Variable(np.random.uniform(size=weight_shape), dtype=tf.float64)

def cost_fn(features, weights):
    with tf.GradientTape() as qnode_tape1:
        with tf.GradientTape() as qnode_tape2:
            res = qnode(features, weights)
        df = qnode_tape2.gradient(res, features)
    df2 = qnode_tape1.gradient(df, features)
    # arbitrary function that uses df and df2
    return tf.reduce_mean(tf.square(res + df + (df2 ** 2 / 2)))

with tf.GradientTape() as tape:
    loss = cost_fn(features, weights)

weights_grad = tape.gradient(loss, weights)

That’s it! I see it’s actually much simpler than I thought. Thank you so much. FYI, this is me looking into quantum physics informed machine learning - which incorporates both data and a mass/momentum conservation residual. After implementing what you have, I seem to be able to run things now, but the loss curve is wildly fluctuating. It could be the nature of the problem, but it also could be a bug.

import pennylane as qml
import tensorflow as tf
import numpy as np
from silence_tensorflow import silence_tensorflow
silence_tensorflow() # to remove typcast warnings ....
tf.keras.backend.set_floatx("float64")
tf.keras.backend.clear_session()
def flatten(xss):
    return [x for xs in xss for x in xs]

n_qubits = 5
n_layers = 10
max_iterations = 1000

dev = qml.device("default.qubit")

@qml.qnode(dev)
def qnode(inputs, weights):
    qml.AngleEmbedding(inputs, wires=range(n_qubits))
    qml.templates.StronglyEntanglingLayers(weights, wires=range(n_qubits))
    #return [qml.expval(qml.PauliZ(wires=i)) for i in range(n_qubits)]
    return qml.expval(qml.PauliZ(0))

weight_shape = qml.StronglyEntanglingLayers.shape(n_layers, n_qubits)

####################################################
####################################################
# Prep input data for embedding and finite difference derivatives
# testcase 3: 
def fun(x):
    u = (x*(1 - np.e) + np.exp(x) - 1)/(1 - np.e)
    return u
def fun_resid(dudx,ddudx):
    return (-ddudx + dudx - 1) 
xL = 0.0
xR = 1.0
x = np.linspace(xL, xR, n_qubits) # number of point must be <= nqbits
ua = fun(x)
x = [[i] for i in x]
#print("x:", *(f"{x:.5f}" for x in flatten(x)))
####################################################
####################################################

np.random.seed(0) 
features = tf.Variable(x, dtype=tf.float64)
weights = tf.Variable(np.random.uniform(low=0, high=np.pi, size=weight_shape), dtype=tf.float64)

print("features: ",features)
print("weights: ",weights)

def cost_fn(features, weights):
    with tf.GradientTape() as qnode_tape1:
        with tf.GradientTape() as qnode_tape2:
            u = qnode(features, weights)
        print("     QML u: ",u)
        print("analytic u: ",ua)
        # Boundary condition data loss
        boundary_data_loss = 0.5 * (tf.abs(u[0] - ua[0]) + tf.abs(u[-1] - ua[-1]))
        df = qnode_tape2.gradient(u, features)
    df2 = qnode_tape1.gradient(df, features)
    
    # Physics-informed Loss
    piml_loss = tf.reduce_mean(tf.square(fun_resid(df,df2)))

    # Total Loss
    loss = boundary_data_loss + piml_loss

    print("boundary_data_loss: ",boundary_data_loss);
    print("piml_loss ",piml_loss);
    print("total_loss ",loss);
    
    # arbitrary function that uses df and df2
    return loss

opt = tf.optimizers.Adam(learning_rate = 0.01)
cost = [cost_fn(features, weights)] # Store the values of the cost function
for n in range(max_iterations):
    with tf.GradientTape() as tape:
         loss = cost_fn(features, weights)
    weights_grad = tape.gradient(loss, weights)
    opt.apply_gradients(zip(weights_grad, [weights]))
    #print("new_weights: ",weights)
    cost.append(loss,)
    print("*****************************************")
    print(f"Step = {n},  Cost function = {cost[-1]:.8f} ")
    print("*****************************************")

I’m still trying to convince myself that qnode should return
return qml.expval(qml.PauliZ(0))
and not
return [qml.expval(qml.PauliZ(wires=i)) for i in range(n_qubits)]

Hi again @Corey, and sorry for the delay!

Loss functions using quantum circuits (esp. things like StronglyEntanglingLayers :sweat_smile: ) can often fluctuate very wildly, so it’s hard to say. From what I’m seeing, there are no obvious signs that it is a bug. If you do end up finding a bug, please let us know on our bug tracker! As for what to return, it’s totally up to you as the chief designer of your program :smile: I’d suggest tinkering with both and seeing which one gives you the info you need. Tweaking the operators and parameters can also help you get the optimal loss curve, but it all depends on the experiment you’re working on. Good luck!!

1 Like