Problem in running the geometry optimization tutorial for Trihydrogen cation (H3+)

Hello Everyone,

I am trying to run the Geometry Optimization tutorial for the Trihydrogen cation (H3+) available on the Pennylane tutorial website.

I have downloaded the following code directly from the website:

import jax
from jax import numpy as jnp

jax.config.update("jax_enable_x64", True)

symbols = ["H", "H", "H"]
x = jnp.array([[0.028, 0.054, 0.0], [0.986, 1.610, 0.0], [1.855, 0.002, 0.0]])

import pennylane as qml


def H(x):
    molecule = qml.qchem.Molecule(symbols, x, charge=1)
    return qml.qchem.molecular_hamiltonian(molecule)[0]

hf = qml.qchem.hf_state(electrons=2, orbitals=6)
print(hf)

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


@qml.qnode(dev, interface="jax")
def circuit(params, obs, wires):
    qml.BasisState(hf, wires=wires)
    qml.DoubleExcitation(params[0], wires=[0, 1, 2, 3])
    qml.DoubleExcitation(params[1], wires=[0, 1, 4, 5])

    return qml.expval(obs)


def cost(params, x):
    hamiltonian = H(x)
    return circuit(params, obs=hamiltonian, wires=range(num_wires))


def finite_diff(f, x, delta=0.01):
    gradient = []
    x = jnp.ravel(x)

    for i in range(len(x)):
        shift = jnp.zeros_like(x)
        shift = shift.at[i].set(0.5*delta)
        res = (f(x + shift) - f(x - shift)) * delta**-1
        gradient.append(res)

    return gradient


def grad_x(params, x):
    grad_h = finite_diff(H, x)
    grad = [circuit(params, obs=obs, wires=range(num_wires)) for obs in grad_h]
    return jnp.array(grad).reshape(x.shape)


theta = jnp.array([0.0, 0.0])

# store the values of the cost function
energies = []

# store the values of the bond length
bond_length = []

# Factor to convert from Bohrs to Angstroms
bohr_angs = 0.529177210903

for n in range(36):
    # gradient for params
    g_param = jax.grad(cost, argnums=[0])(theta, x)[0]
    theta = theta - 0.8 * g_param

    # gradient for coordinates
    value, _ = jax.value_and_grad(cost, argnums=1)(theta, x)
    grad = grad_x(theta, x)
    x = x - 0.8 * grad
    energies.append(value)
    bond_length.append(jnp.linalg.norm(x[0] - x[1]) * bohr_angs) 

    if n % 4 == 0:
        print(f"Step = {n},  E = {energies[-1]:.8f} Ha,  bond length = {bond_length[-1]:.5f} A")

    # Check maximum component of the nuclear gradient
    if jnp.max(grad_x(theta, x)) <= 1e-04:
        break

print("\n" f"Final value of the ground-state energy = {energies[-1]:.8f} Ha")
print("\n" "Ground-state equilibrium geometry")
print("%s %4s %8s %8s" % ("symbol", "x", "y", "z"))
for i, atom in enumerate(symbols):
    print(f"  {atom}    {x[0][i]:.4f}   {x[1][i]:.4f}   {x[2][i]:.4f}")

import matplotlib.pyplot as plt

fig = plt.figure()
fig.set_figheight(5)
fig.set_figwidth(12)

# Add energy plot on column 1
E_fci = -1.27443765658
E_vqe = jnp.array(energies)
ax1 = fig.add_subplot(121)
ax1.plot(range(n + 1), E_vqe - E_fci, "go", ls="dashed")
ax1.plot(range(n + 1), jnp.full(n + 1, 0.001), color="red")
ax1.set_xlabel("Optimization step", fontsize=13)
ax1.set_ylabel("$E_{VQE} - E_{FCI}$ (Hartree)", fontsize=13)
ax1.text(5, 0.0013, r"Chemical accuracy", fontsize=13)
plt.yscale("log")
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)

# Add bond length plot on column 2
d_fci = 0.986
ax2 = fig.add_subplot(122)
ax2.plot(range(n + 1), bond_length, "go", ls="dashed")
ax2.plot(range(n + 1), jnp.full(n + 1, d_fci), color="red")
ax2.set_ylim([0.965, 0.99])
ax2.set_xlabel("Optimization step", fontsize=13)
ax2.set_ylabel("bond length ($\AA$)", fontsize=13)
ax2.text(5, 0.9865, r"Equilibrium bond length", fontsize=13)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)

plt.subplots_adjust(wspace=0.3)
plt.show()

Please find the full error message below:

Traceback (most recent call last):
  File "/lustre04/scratch/sanjeets/QUANTUM-COMPUTING/H3OPLUS-OPT.py", line 74, in <module>
    value, _ = jax.value_and_grad(cost, argnums=1)(theta, x)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/lustre04/scratch/sanjeets/QUANTUM-COMPUTING/H3OPLUS-OPT.py", line 33, in cost
    hamiltonian = H(x)
                  ^^^^
  File "/lustre04/scratch/sanjeets/QUANTUM-COMPUTING/H3OPLUS-OPT.py", line 14, in H
    return qml.qchem.molecular_hamiltonian(molecule)[0]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sanjeets/.local/lib/python3.11/site-packages/pennylane/qchem/hamiltonian.py", line 376, in molecular_hamiltonian
    return _molecular_hamiltonian_dispatch(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/cvmfs/soft.computecanada.ca/easybuild/software/2023/x86-64-v4/Compiler/gcccore/python/3.11.5/lib/python3.11/functools.py", line 909, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sanjeets/.local/lib/python3.11/site-packages/pennylane/qchem/hamiltonian.py", line 409, in _
    return _molecular_hamiltonian(
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sanjeets/.local/lib/python3.11/site-packages/pennylane/qchem/hamiltonian.py", line 516, in _molecular_hamiltonian
    geometry_dhf = qml.numpy.array(coordinates)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sanjeets/.local/lib/python3.11/site-packages/pennylane/numpy/wrapper.py", line 117, in _wrapped
    res = obj(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^
  File "/home/sanjeets/.local/lib/python3.11/site-packages/autograd/numpy/numpy_wrapper.py", line 75, in array
    return _array_from_scalar_or_array(args, kwargs, A)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sanjeets/.local/lib/python3.11/site-packages/autograd/tracer.py", line 54, in f_wrapped
    return f_raw(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sanjeets/.local/lib/python3.11/site-packages/autograd/numpy/numpy_wrapper.py", line 89, in _array_from_scalar_or_array
    return _np.array(scalar, *array_args, **array_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on traced array with shape float64[3,3].
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.```

Can you please help me in resolving the issue?

Cheers,
Sanjeet

Hi @sanjeets , welcome to the Forum!

This looks like a package version issue. Can you please post the output of qml.about() and your version of JAX?

The latest PennyLane version (v0.42) supports up to JAX==0.6.0
You can upgrade your PennyLane version with:

pip install pennylane --upgrade

I hope this helps!

Hi @CatalinaAlbornoz,

Thank you for your reply!!

Here is the output for qml.about():

Name: PennyLane
Version: 0.39.0+computecanada
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author:
Author-email:
License:
Location: /home/sanjeets/.local/lib/python3.11/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, toml, typing-extensions
Required-by: PennyLane_Lightning

Platform info:           Linux-4.18.0-553.56.1.el8_10.x86_64-x86_64-Intel-R-_Xeon-R-_Gold_6148_CPU_@_2.40GHz-with-glibc2.37
Python version:          3.11.5
Numpy version:           1.25.2
Scipy version:           1.11.2
Installed devices:
- lightning.qubit (PennyLane_Lightning-0.39.0+computecanada)
- default.clifford (PennyLane-0.39.0+computecanada)
- default.gaussian (PennyLane-0.39.0+computecanada)
- default.mixed (PennyLane-0.39.0+computecanada)
- default.qubit (PennyLane-0.39.0+computecanada)
- default.qutrit (PennyLane-0.39.0+computecanada)
- default.qutrit.mixed (PennyLane-0.39.0+computecanada)
- default.tensor (PennyLane-0.39.0+computecanada)
- null.qubit (PennyLane-0.39.0+computecanada)
- reference.qubit (PennyLane-0.39.0+computecanada)

None

I have tried with the latest version of jax (0.6.0) also by using the following command:

pip install --upgrade jax

I also tried with the older version of jax: 0.4.28.

Can you please point out the issue?

Hi @CatalinaAlbornoz,

Now, I am able to run the tutorial using the following versions:

Pannylane: 0.41.1
jax: 0.4.28
jaxlib: 0.4.28

Thank you for your help!!

Cheers,

Sanjeet

I’m glad this worked @sanjeets !

I would recommend always trying to stay up to date with the latest version of PennyLane if possible (v0.42 at the moment) since this will contain the latest features and bug fixes. So if you encounter issues along the way this is one of the first things you can try.

Enjoy using PennyLane!

1 Like