Gatchas

• At the moment I can access GPUs in a single node. cuQuantum seems to have multi-node capability in works (see here).
• The multi-node support status of the Pennylane  as of now is here. 

  • Currently cuQuantum multi-node support for a single state vector is not yet implemented (but they plan to implement it in the future)

  • Distrubuted computing examples mentioned in reply: 1 2 3

• pennylane._device.DeviceError: Observable Hermitian not supported on device lightning.gpu

• pennylane.QuantumFunctionError: Adjoint differentiation method does not support expectation return type mixed with other return types

•  

Prerequisites

• The preparation of the cuQuantum+Pennylane singularity container is explained elsewhere.
• I am using PopOs which is Ubuntu based. I followed the "setuid installation" instructions here for my own laptop. TRUBA already has apptainer installed.

TRUBA

• Details on how to use TRUBA are here.
• In order to use A100 GPUs, you need to use palamut. These machines are seperate from TRUBA, and exclusive to AI research. Details on how to reach palamut-ui are here.

Python code

• This is slightly modified from "Differentiable Hartree-Fock" Pennylane example.
• How to use multiple GPUs for the task is detailed under Parallel adjoint differentiation support of lightning.gpu doc.

from autograd import grad
import pennylane as qml
from pennylane import numpy as np

np.set_printoptions(precision=5)

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)

##############################################################################
# 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
# nuclear coordinates. To do that, we create a molecule object that stores all the molecular
# parameters needed to perform a Hartree-Fock calculation.

mol = qml.qchem.Molecule(symbols, geometry)

##############################################################################
# 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)

##############################################################################
# We now compute the gradient of the energy with respect to the nuclear coordinates

grad(qml.qchem.hf_energy(mol))(geometry)

##############################################################################
# The obtained gradients are equal or very close to zero because the geometry we used here has been
# previously optimized at the Hartree-Fock level. You can use a different geometry and verify that
# the newly computed gradients are not all zero.
#
# We can also compute the values and gradients of several other quantities that are obtained during
# the Hartree-Fock procedure. These include all integrals over basis functions, matrices formed from
# these integrals and the one- and two-body integrals over molecular orbitals [#arrazola2021]_.
# Let's look at a few examples.
#
# We first compute the overlap integral between the two S-type atomic orbitals of the hydrogen
# atoms. Here we are using the `STO-3G <https://en.wikipedia.org/wiki/STO-nG_basis_sets>`_
# basis set in which each of the atomic orbitals is represented by one basis function composed of
# three primitive Gaussian functions. These basis functions can be accessed from the molecule
# object as

S1 = mol.basis_set[0]
S2 = mol.basis_set[1]

##############################################################################
# We can check the parameters of the basis functions as

for param in S1.params:
    print(param)

##############################################################################
# This returns the exponents, contraction coefficients and the centres of the three Gaussian
# functions of the STO-3G basis set. These parameters can be also obtained individually by using
# ``S1.alpha``, ``S1.coeff`` and ``S1.r``, respectively. You can verify that both of the ``S1`` an
# ``S2`` atomic orbitals have the same exponents and contraction coefficients but are centred on
# different hydrogen atoms. You can also verify that the orbitals are S-type by printing the angular
# momentum quantum numbers with

S1.l

##############################################################################
# This gives us a tuple of three integers, representing the exponents of the :math:`x`, :math:`y`
# and :math:`z` components in the Gaussian functions [#arrazola2021]_.
#
# We can now compute the overlap integral,
#
# .. math::
#
#     S_{\mu \nu} = \int \chi_\mu^* (r) \chi_\nu (r) dr
#
# between the atomic orbitals :math:`\chi`, by passing the orbitals and the initial values of their
# centres to the :func:`~.pennylane.qchem.overlap_integral` function. The centres of the orbitals
# 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]])

##############################################################################
# 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]])

##############################################################################
# Can you explain why some of the computed gradients are zero?
#
# Let's now plot the atomic orbitals and their overlap. We can do it by using
# the :py:meth:`~.pennylane.qchem.Molecule.atomic_orbital` function, which evaluates the
# atomic orbital at a given coordinate. For instance, the value of the S orbital on the first
# hydrogen atom can be computed at the origin as

V1 = mol.atomic_orbital(0)
V1(0.0, 0.0, 0.0)

##############################################################################
# We can evaluate this orbital at different points along the :math:`x` axis and plot it.

x = np.linspace(-5, 5, 1000)

##############################################################################
# We can also plot the second S orbital and visualize the overlap between them

V2 = mol.atomic_orbital(1)

##############################################################################
# By looking at the orbitals, can you guess at what distance the value of the overlap becomes
# negligible? Can you verify your guess by computing the overlap at that distance?
#
# Similarly, we can plot the molecular orbitals of the hydrogen molecule obtained from the
# Hartree-Fock calculations. We plot the cross-section of the bonding orbital on the :math:`x-y`
# plane.

n = 30 # number of grid points along each axis

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)


##############################################################################
# VQE simulations
# ---------------
#
# By performing the Hartree-Fock calculations, we obtain a set of one- and two-body integrals
# 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.symbols,
                                                      mol.coordinates,
                                                      args=[mol.coordinates])
print(hamiltonian)

##############################################################################
# The Hamiltonian contains 15 terms and, importantly, the coefficients of the Hamiltonian are all
# differentiable. We can construct a circuit and perform a VQE simulation in which both of the
# circuit parameters and the nuclear coordinates are optimized simultaneously by using the computed
# gradients. We will have two sets of differentiable parameters: the first set is the
# rotation angles of the excitation gates which are applied to the reference Hartree-Fock state
# to construct the exact ground state. The second set contains the nuclear coordinates of the
# hydrogen atoms.

dev = qml.device("lightning.gpu", wires=4,batch_obs=False)
def energy(mol):
    @qml.qnode(dev)
    def circuit(*args):
        qml.BasisState(np.array([1, 1, 0, 0]), wires=range(4))
        qml.DoubleExcitation(*args[0][0], wires=[0, 1, 2, 3])
        H = qml.qchem.molecular_hamiltonian(mol.symbols, mol.coordinates, alpha=mol.alpha,
                                            coeff=mol.coeff, args=args[1:])[0]
        return qml.expval(H)
    return circuit

##############################################################################
# Note that we only use the :class:`~.pennylane.DoubleExcitation` gate as the
# :class:`~.pennylane.SingleExcitation` ones can be neglected in this particular example
# [#szabo1996]_. We now compute the gradients of the energy with respect to the circuit parameter
# and the nuclear coordinates and update the parameters iteratively. Note that the nuclear
# 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)]

for n in range(36):

    args = [circuit_param, geometry]
    mol = qml.qchem.Molecule(symbols, geometry)

    # gradient for circuit parameters
    g_param = grad(energy(mol), argnum = 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

    if n % 5 == 0:
        print(f'n: {n}, E: {energy(mol)(*args):.8f}, Force-max: {abs(forces).max():.8f}')

##############################################################################
# After 35 steps of optimization, the forces on the atomic nuclei and the gradient of the
# circuit parameter are both approaching zero, and the energy of the molecule is that of the
# optimized geometry at the
# `full-CI <https://en.wikipedia.org/wiki/Full_configuration_interaction>`_ level:
# :math:`-1.1373060483` Ha. You can print the optimized geometry and verify that the final bond
# length of hydrogen is identical to the one computed with full-CI which is :math:`1.389`
# `Bohr <https://en.wikipedia.org/wiki/Bohr_radius>`_.
#
# We are now ready to perform a full optimization where the circuit parameters, the nuclear
# coordinates and the basis set parameters are all differentiable parameters that can be optimized
# simultaneously.

# 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)

# 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)

# initial value of the circuit parameter
circuit_param = [np.array([0.0], requires_grad=True)]

for n in range(36):

    args = [circuit_param, geometry, alpha, coeff]
    mol = qml.qchem.Molecule(symbols, geometry, alpha=alpha, coeff=coeff)

    # gradient for circuit parameters
    g_param = grad(energy(mol), argnum=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

    # 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

    if n % 5 == 0:
        print(f'n: {n}, E: {energy(mol)(*args):.8f}, Force-max: {abs(forces).max():.8f}')

##############################################################################
# You can also print the gradients of the circuit and basis set parameters and confirm that they are
# approaching zero. The computed energy of :math:`-1.14040160` Ha is
# lower than the full-CI energy, :math:`-1.1373060483` Ha (obtained with the STO-3G basis set for
# the hydrogen molecule) because we have optimized the basis set parameters in our example. This
# means that we can reach a lower energy for hydrogen without increasing the basis set size, which
# would otherwise lead to a larger number of qubits.
#
# Conclusions
# -----------
# This tutorial introduces an important feature of PennyLane that allows performing
# fully-differentiable Hartree-Fock and subsequently VQE simulations. This feature provides two
# major benefits: i) All gradient computations needed for parameter optimization can be carried out
# with the elegant methods of automatic differentiation which facilitates simultaneous optimizations
# of circuit and Hamiltonian parameters in applications such as VQE molecular geometry
# optimizations. ii) By optimizing the molecular parameters such as the exponent and contraction
# coefficients of Gaussian functions of the basis set, one can reach a lower energy without
# increasing the number of basis functions. Can you think of other interesting molecular parameters
# that can be optimized along with the nuclear coordinates and basis set parameters that we
# optimized in this tutorial?
#
# References
# ----------
#
# .. [#arrazola2021]
#
#     Juan Miguel Arrazola, Soran Jahangiri, Alain Delgado, Jack Ceroni *et al.*, "Differentiable
#     quantum computational chemistry with PennyLane". `arXiv:2111.09967
#     <https://arxiv.org/abs/2111.09967>`__
#
# .. [#szabo1996]
#
#     Attila Szabo, Neil S. Ostlund, "Modern Quantum Chemistry: Introduction to Advanced Electronic
#     Structure Theory". Dover Publications, 1996.
#
#

HOWTO

• ssh to levrek1
• ssh to palamut-ui
• Create a directory, put the script, download and extract the dataset
  • mkdir cuQuantum-test
  • cd cuQuantum-test
• Get an interactive acces to a palamut node with 2 gpus
  • srun -A --partition palamut-cuda --job-name "cuQuantum" --ntasks-per-node=32 --gres=gpu:2 --time 01:00:00 -N 1 --pty bash -i
• Run the example
  • singularity run --nv ~/QC-progs.sif python3 tutorial_differentiable_HF.py
• You can check the utilization of GPUs by nvidia-smi -l

Conclusions

It works, but there are significant limitations due to Observable Hermitian and Adjoint differentiation gatchas. Also pyTorch hybrid neural networks use only a single gpu.