Fix demos to use jax that require grads (#1226)

Doesn't build due to the fact that [this PR](PennyLaneAI/pennylane#6096) is not yet merged. JSON files need to be updated as well. **Summary:** Fixes remaining demos that checked for requires_grad. Won't build for now due to missing a fix in [this PR](PennyLaneAI/pennylane#6324) [[sc-69776](https://app.shortcut.com/xanaduai/story/69776)] [[sc-69778](https://app.shortcut.com/xanaduai/story/69778)] --------- Co-authored-by: GitHub Nightly Merge Action <[email protected]> Co-authored-by: Mudit Pandey <[email protected]> Co-authored-by: Mikhail Andrenkov <[email protected]> Co-authored-by: David Wierichs <[email protected]> Co-authored-by: Korbinian Kottmann <[email protected]> Co-authored-by: Ivana Kurečić <[email protected]> Co-authored-by: Paul Finlay <[email protected]> Co-authored-by: Jack Brown <[email protected]> Co-authored-by: bellekaplan <[email protected]> Co-authored-by: obliviateandsurrender <[email protected]> Co-authored-by: soranjh <[email protected]> Co-authored-by: Soran Jahangiri <[email protected]>
PennyLaneAI · Dec 13, 2024 · 95ef8d3 · 95ef8d3
1 parent 07dbec8
commit 95ef8d3
Show file tree

Hide file tree

Showing 10 changed files with 103 additions and 103 deletions.
diff --git a/demonstrations/tutorial_differentiable_HF.metadata.json b/demonstrations/tutorial_differentiable_HF.metadata.json
@@ -6,7 +6,7 @@
         }
     ],
     "dateOfPublication": "2022-05-09T00:00:00+00:00",
-    "dateOfLastModification": "2024-10-07T00:00:00+00:00",
+    "dateOfLastModification": "2024-12-13T00:00:00+00:00",
     "categories": [
         "Quantum Chemistry"
     ],

diff --git a/demonstrations/tutorial_differentiable_HF.py b/demonstrations/tutorial_differentiable_HF.py
@@ -101,20 +101,21 @@
 For the hydrogen molecule we have
 """
 
-from autograd import grad
 import pennylane as qml
-from pennylane import numpy as np
+import numpy as np
+import jax
+import jax.numpy as jnp
 import matplotlib.pyplot as plt
 np.set_printoptions(precision=5)
+jax.config.update("jax_enable_x64", True)
 
 symbols = ["H", "H"]
 # optimized geometry at the Hartree-Fock level
-geometry = np.array([[-0.672943567415407, 0.0, 0.0],
-                     [ 0.672943567415407, 0.0, 0.0]], requires_grad=True)
+geometry = jnp.array([[-0.672943567415407, 0.0, 0.0],
+                     [ 0.672943567415407, 0.0, 0.0]])
 
 ##############################################################################
-# The use of ``requires_grad=True`` specifies that the nuclear coordinates are differentiable
-# parameters. We can now compute the Hartree-Fock energy and its gradient with respect to the
+# We can now compute the Hartree-Fock energy and its gradient with respect to the
 # nuclear coordinates. To do that, we create a molecule object that stores all the molecular
 # parameters needed to perform a Hartree-Fock calculation.
 
@@ -124,12 +125,12 @@
 # The Hartree-Fock energy can now be computed with the
 # :func:`~.pennylane.qchem.hf_energy` function which is a function transform
 
-qml.qchem.hf_energy(mol)(geometry)
+qml.qchem.hf_energy(mol)()
 
 ##############################################################################
 # We now compute the gradient of the energy with respect to the nuclear coordinates
 
-grad(qml.qchem.hf_energy(mol))(geometry)
+jax.grad(qml.qchem.hf_energy(mol), argnums=0)(geometry, mol.coeff, mol.alpha)
 
 ##############################################################################
 # The obtained gradients are equal or very close to zero because the geometry we used here has been
@@ -181,13 +182,13 @@
 # are those of the hydrogen atoms by default and are therefore treated as differentiable parameters
 # by PennyLane.
 
-qml.qchem.overlap_integral(S1, S2)([geometry[0], geometry[1]])
+qml.qchem.overlap_integral(S1, S2)()
 
 ##############################################################################
 # You can verify that the overlap integral between two identical atomic orbitals is equal to one.
 # We can now compute the gradient of the overlap integral with respect to the orbital centres
 
-grad(qml.qchem.overlap_integral(S1, S2))([geometry[0], geometry[1]])
+jax.grad(qml.qchem.overlap_integral(S1, S2))(geometry, mol.coeff, mol.alpha)
 
 ##############################################################################
 # Can you explain why some of the computed gradients are zero?
@@ -228,13 +229,13 @@
 # plane.
 
 n = 30 # number of grid points along each axis
-
+qml.qchem.hf_energy(mol)()
 mol.mo_coefficients = mol.mo_coefficients.T
 mo = mol.molecular_orbital(0)
 x, y = np.meshgrid(np.linspace(-2, 2, n),
                    np.linspace(-2, 2, n))
 val = np.vectorize(mo)(x, y, 0)
-val = np.array([val[i][j]._value for i in range(n) for j in range(n)]).reshape(n, n)
+val = np.array([val[i][j] for i in range(n) for j in range(n)]).reshape(n, n)
 
 fig, ax = plt.subplots()
 co = ax.contour(x, y, val, 10, cmap='summer_r', zorder=0)
@@ -252,7 +253,7 @@
 # over molecular orbitals that can be used to construct the molecular Hamiltonian with the
 # :func:`~.pennylane.qchem.molecular_hamiltonian` function.
 
-hamiltonian, qubits = qml.qchem.molecular_hamiltonian(mol, args=[mol.coordinates])
+hamiltonian, qubits = qml.qchem.molecular_hamiltonian(mol)
 print(hamiltonian)
 
 ##############################################################################
@@ -265,11 +266,12 @@
 # hydrogen atoms.
 
 dev = qml.device("default.qubit", wires=4)
-def energy(mol):
-    @qml.qnode(dev, interface="autograd")
+def energy():
+    @qml.qnode(dev, interface="jax")
     def circuit(*args):
         qml.BasisState(np.array([1, 1, 0, 0]), wires=range(4))
-        qml.DoubleExcitation(*args[0][0], wires=[0, 1, 2, 3])
+        qml.DoubleExcitation(*args[0], wires=[0, 1, 2, 3])
+        mol = qml.qchem.Molecule(symbols, geometry, alpha=args[3], coeff=args[2])
         H = qml.qchem.molecular_hamiltonian(mol, args=args[1:])[0]
         return qml.expval(H)
     return circuit
@@ -282,23 +284,24 @@ def circuit(*args):
 # coordinate gradients are simply the forces on the atomic nuclei.
 
 # initial value of the circuit parameter
-circuit_param = [np.array([0.0], requires_grad=True)]
+circuit_param = jnp.array([0.0])
 
-for n in range(36):
+geometry = jnp.array([[0.0, 0.02, -0.672943567415407],
+                     [0.1, 0.0, 0.672943567415407]])
 
-    args = [circuit_param, geometry]
+for n in range(36):
     mol = qml.qchem.Molecule(symbols, geometry)
-
+    args = [circuit_param, geometry, mol.coeff, mol.alpha]
     # gradient for circuit parameters
-    g_param = grad(energy(mol), argnum = 0)(*args)
+    g_param = jax.grad(energy(), argnums = 0)(*args)
     circuit_param = circuit_param - 0.25 * g_param[0]
 
     # gradient for nuclear coordinates
-    forces = -grad(energy(mol), argnum = 1)(*args)
-    geometry = geometry + 0.5 * forces
+    forces = jax.grad(energy(), argnums = 1)(*args)
+    geometry = geometry - 0.5 * forces
 
     if n % 5 == 0:
-        print(f'n: {n}, E: {energy(mol)(*args):.8f}, Force-max: {abs(forces).max():.8f}')
+        print(f'n: {n}, E: {energy()(*args):.8f}, Force-max: {abs(forces).max():.8f}')
 
 ##############################################################################
 # After 35 steps of optimization, the forces on the atomic nuclei and the gradient of the
@@ -313,44 +316,41 @@ def circuit(*args):
 # coordinates and the basis set parameters are all differentiable parameters that can be optimized
 # simultaneously.
 
+symbols = ["H", "H"]
 # initial values of the nuclear coordinates
-geometry = np.array([[0.0, 0.0, -0.672943567415407],
-                     [0.0, 0.0, 0.672943567415407]], requires_grad=True)
-
-# initial values of the basis set exponents
-alpha = np.array([[3.42525091, 0.62391373, 0.1688554],
-                  [3.42525091, 0.62391373, 0.1688554]], requires_grad=True)
+geometry = jnp.array([[0.0, 0.0, -0.672943567415407],
+                     [0.0, 0.0, 0.672943567415407]])
 
 # initial values of the basis set contraction coefficients
-coeff = np.array([[0.1543289673, 0.5353281423, 0.4446345422],
-                  [0.1543289673, 0.5353281423, 0.4446345422]], requires_grad=True)
+coeff = jnp.array([[0.1543289673, 0.5353281423, 0.4446345422],
+                  [0.1543289673, 0.5353281423, 0.4446345422]])
+
+# initial values of the basis set exponents
+alpha = jnp.array([[3.42525091, 0.62391373, 0.1688554],
+                  [3.42525091, 0.62391373, 0.1688554]])
 
 # initial value of the circuit parameter
-circuit_param = [np.array([0.0], requires_grad=True)]
+circuit_param = jnp.array([0.0])
 
-for n in range(36):
+mol = qml.qchem.Molecule(symbols, geometry, coeff=coeff, alpha=alpha)
+args = [circuit_param, geometry, coeff, alpha]
 
-    args = [circuit_param, geometry, alpha, coeff]
+for n in range(36):
+    args = [circuit_param, geometry, coeff, alpha]
     mol = qml.qchem.Molecule(symbols, geometry, alpha=alpha, coeff=coeff)
 
     # gradient for circuit parameters
-    g_param = grad(energy(mol), argnum=0)(*args)
+    g_param = jax.grad(energy(), argnums=[0, 1, 2, 3])(*args)[0]
     circuit_param = circuit_param - 0.25 * g_param[0]
 
     # gradient for nuclear coordinates
-    forces = -grad(energy(mol), argnum=1)(*args)
-    geometry = geometry + 0.5 * forces
-
-    # gradient for basis set exponents
-    g_alpha = grad(energy(mol), argnum=2)(*args)
-    alpha = alpha - 0.25 * g_alpha
-
-    # gradient for basis set contraction coefficients
-    g_coeff = grad(energy(mol), argnum=3)(*args)
-    coeff = coeff - 0.25 * g_coeff
+    value, gradients = jax.value_and_grad(energy(), argnums=[1, 2, 3])(*args)
+    geometry = geometry - 0.5 * gradients[0]
+    alpha = alpha - 0.25 * gradients[2]
+    coeff = coeff - 0.25 * gradients[1]
 
     if n % 5 == 0:
-        print(f'n: {n}, E: {energy(mol)(*args):.8f}, Force-max: {abs(forces).max():.8f}')
+        print(f'n: {n}, E: {value:.8f}, Force-max: {abs(gradients[0]).max():.8f}')
 
 ##############################################################################
 # You can also print the gradients of the circuit and basis set parameters and confirm that they are
@@ -390,4 +390,4 @@ def circuit(*args):
 #
 # About the author
 # ----------------
-# .. include:: ../_static/authors/soran_jahangiri.txt
+# .. include:: ../_static/authors/soran_jahangiri.txt
diff --git a/demonstrations/tutorial_fermionic_operators.metadata.json b/demonstrations/tutorial_fermionic_operators.metadata.json
@@ -6,7 +6,7 @@
         }
     ],
     "dateOfPublication": "2023-06-27T00:00:00+00:00",
-    "dateOfLastModification": "2024-10-30T00:00:00+00:00",
+    "dateOfLastModification": "2024-12-13T00:00:00+00:00",
     "categories": [
         "Quantum Chemistry"
     ],

diff --git a/demonstrations/tutorial_fermionic_operators.py b/demonstrations/tutorial_fermionic_operators.py
@@ -108,7 +108,7 @@
 # The matrix representation of the qubit Hamiltonian in the computational basis can be diagonalized
 # to get its eigenpairs.
 
-from pennylane import numpy as np
+import numpy as np
 
 val, vec = np.linalg.eigh(h.sparse_matrix().toarray())
 print(f"eigenvalues:\n{val}")
@@ -136,9 +136,10 @@
 # the hydrogen molecule as an example. We first define the atom types and the atomic coordinates.
 
 import pennylane as qml
+from jax import numpy as jnp
 
 symbols = ["H", "H"]
-geometry = np.array([[-0.67294, 0.0, 0.0], [0.67294, 0.0, 0.0]], requires_grad=False)
+geometry = jnp.array([[-0.67294, 0.0, 0.0], [0.67294, 0.0, 0.0]])
 
 ##############################################################################
 # Then we compute the one- and two-electron integrals, which are the coefficients :math:`c` in the

diff --git a/demonstrations/tutorial_mapping.metadata.json b/demonstrations/tutorial_mapping.metadata.json
@@ -6,7 +6,7 @@
         }
     ],
     "dateOfPublication": "2024-05-06T00:00:00+00:00",
-    "dateOfLastModification": "2024-11-06T00:00:00+00:00",
+    "dateOfLastModification": "2024-12-13T00:00:00+00:00",
     "categories": [
         "Algorithms",
         "Quantum Computing",

diff --git a/demonstrations/tutorial_mapping.py b/demonstrations/tutorial_mapping.py
@@ -202,11 +202,14 @@
 # coordinates.
 
 from pennylane import qchem
-from pennylane import numpy as np
+from jax import numpy as jnp
+import jax
+
+jax.config.update("jax_enable_x64", True)
 
 symbols  = ['H', 'H']
-geometry = np.array([[0.0, 0.0, -0.69434785],
-                     [0.0, 0.0,  0.69434785]], requires_grad = False)
+geometry = jnp.array([[0.0, 0.0, -0.69434785],
+                     [0.0, 0.0,  0.69434785]])
 
 mol = qchem.Molecule(symbols, geometry)
 
@@ -289,7 +292,7 @@
 # Note that we need to exponentiate these operators to be able to use them in the circuit
 # [#Yordanov]_. We also use a set of pre-defined parameters to construct the excitation gates.
 
-params = np.array([0.22347661, 0.0, 0.0])
+params = jnp.array([0.22347661, 0.0, 0.0])
 
 dev = qml.device("default.qubit", wires=qubits)
 

diff --git a/demonstrations/tutorial_mol_geo_opt.metadata.json b/demonstrations/tutorial_mol_geo_opt.metadata.json
@@ -6,7 +6,7 @@
         }
     ],
     "dateOfPublication": "2021-06-30T00:00:00+00:00",
-    "dateOfLastModification": "2024-10-07T00:00:00+00:00",
+    "dateOfLastModification": "2024-12-13T00:00:00+00:00",
     "categories": [
         "Quantum Chemistry"
     ],