Repository: PennyLaneAI/pennylane
Pull request #3363 (2022-11-15 21:36:25)

Add function to load qchem basis sets
**Context:**
Adds functionality to load basis sets from the [Basis Set Exchange](https://www.basissetexchange.org/) repo.

**Description of the Change:**
The function `load_basisset` is added to `qchem/basis_data.py`.

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**


github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3363?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3363](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3363?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (53226f8) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/3144cced0e7ff5698b92411eb4f366ea786c542a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (3144cce) will **decrease** coverage by `0.00%`.
> The diff coverage is `89.47%`.

```diff
@@            Coverage Diff             @@
##           master    #3363      +/-   ##
==========================================
- Coverage   99.70%   99.69%   -0.01%     
==========================================
  Files         299      299              
  Lines       26146    26165      +19     
==========================================
+ Hits        26068    26085      +17     
- Misses         78       80       +2     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3363?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/qchem/basis\_data.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3363/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3FjaGVtL2Jhc2lzX2RhdGEucHk=) | `92.30% <89.47%> (-7.70%)` | :arrow_down: |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index f3ec6d421..08a693766 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -96,6 +96,10 @@
   >>> jax.grad(qml.qchem.hf_energy(mol))(*args)
   >>> DeviceArray([[0.0, 0.0, 0.3650435], [0.0, 0.0, -0.3650435]], dtype=float32)
   ```
+  
+* The function `load_basisset` is added to extract qchem basis set data from the Basis Set Exchange
+  library. 
+  [(#3363)](https://github.com/PennyLaneAI/pennylane/pull/3363)
 
 <h3>Improvements</h3>
 
@@ -142,6 +146,7 @@
 
 This release contains contributions from (in alphabetical order):
 
+Juan Miguel Arrazola
 Ikko Ashimine
 Utkarsh Azad
 Astral Cai
diff --git a/pennylane/qchem/__init__.py b/pennylane/qchem/__init__.py
index 59ef34d0d..84201eeeb 100644
--- a/pennylane/qchem/__init__.py
+++ b/pennylane/qchem/__init__.py
@@ -23,6 +23,7 @@ from .openfermion_obs import (
     decompose,
     molecular_hamiltonian,
 )
+from .basis_data import load_basisset
 from .basis_set import BasisFunction, atom_basis_data, mol_basis_data
 from .convert import import_operator
 from .dipole import dipole_integrals, fermionic_dipole, dipole_moment
diff --git a/pennylane/qchem/basis_data.py b/pennylane/qchem/basis_data.py
index 22c28868d..3c2de49e0 100644
--- a/pennylane/qchem/basis_data.py
+++ b/pennylane/qchem/basis_data.py
@@ -14,9 +14,12 @@
 """
 This module contains basis set parameters defining Gaussian-type orbitals for a selected number of
 atoms. The data are taken from the Basis Set Exchange `library <https://www.basissetexchange.org>`_.
-The current data include the STO-3G and 6-31G basis sets for atoms with atomic numbers 1-10.
+The current data includes the STO-3G, 6-31G, 6-311G and CC-PVDZ, basis sets for elements with atomic
+numbers 1-10.
 """
 
+import itertools
+
 atomic_numbers = {
     "H": 1,
     "He": 2,
@@ -754,3 +757,70 @@ basis_sets = {
     "6-311g": POPLE6311G,
     "cc-pvdz": CCPVDZ,
 }
+
+
+def load_basisset(basis, element):
+    r"""Extracts basis set data from the Basis Set Exchange library.
+
+    Args:
+        basis (str): name of the basis set
+        element (str): atomic symbol of the chemical element
+
+    Returns:
+        dict[str, list]: dictionary containing orbital names, and the exponents and contraction
+        coefficients of a basis function
+
+    **Example**
+
+    >>> basis = '6-31g'
+    >>> element = 'He'
+    >>> basis = load_basisset(basis, element)
+    >>> basis
+    {'orbitals': ['S', 'S'],
+     'exponents': [[38.421634, 5.77803, 1.241774], [0.297964]],
+     'coefficients': [[0.04013973935, 0.261246097, 0.7931846246], [1.0]]}
+    """
+    try:
+        # pylint: disable=import-outside-toplevel, unused-import, multiple-imports
+        import basis_set_exchange as bse
+    except ImportError as Error:
+        raise ImportError(
+            "This feature requires basis_set_exchange. "
+            "It can be installed with: pip install basis-set-exchange."
+        ) from Error
+
+    orbital_map = {
+        "[0]": "S",
+        "[0, 1]": "SP",
+        "[1]": "P",
+        "[2]": "D",
+        "[3]": "F",
+        "[4]": "G",
+        "[5]": "H",
+    }
+
+    element = str(atomic_numbers[element])
+
+    data = bse.get_basis(basis)["elements"][element]["electron_shells"]
+
+    orbitals = []
+    exponents = []
+    coefficients = []
+
+    for term in data:
+        if orbital_map[str(term["angular_momentum"])] == "SP":
+            orbitals.append(["S", "P"])
+        else:
+            orbitals.append(
+                [orbital_map[str(term["angular_momentum"])]] * len(term["coefficients"])
+            )
+        exponents.append(
+            [list(map(float, item)) for item in [term["exponents"]] * len(term["coefficients"])]
+        )
+        coefficients.append([list(map(float, item)) for item in term["coefficients"]])
+
+    return {
+        "orbitals": list(itertools.chain(*orbitals)),
+        "exponents": list(itertools.chain(*exponents)),
+        "coefficients": list(itertools.chain(*coefficients)),
+    }
diff --git a/pennylane/qchem/basis_set.py b/pennylane/qchem/basis_set.py
index 16b1204f6..6afac08f1 100644
--- a/pennylane/qchem/basis_set.py
+++ b/pennylane/qchem/basis_set.py
@@ -15,8 +15,8 @@
 This module contains functions and classes to create a
 :class:`~pennylane.qchem.basis_set.BasisFunction` object from standard basis sets such as STO-3G.
 """
-# pylint: disable=too-few-public-methods
-from .basis_data import basis_sets
+# pylint: disable=too-few-public-methods, too-many-branches
+from .basis_data import basis_sets, load_basisset
 
 
 class BasisFunction:
@@ -66,16 +66,18 @@ class BasisFunction:
         self.params = [self.alpha, self.coeff, self.r]
 
 
-def atom_basis_data(name, atom):
+def atom_basis_data(name, atom, load_data=False):
     r"""Generate default basis set parameters for an atom.
 
     This function extracts the angular momentum, exponents, and contraction coefficients of
-    Gaussian functions forming atomic orbitals for a given atom. These values are taken from the
-    basis set data provided in :mod:`~pennylane.qchem.basis_data`.
+    Gaussian functions forming atomic orbitals for a given atom. These values are taken, by default,
+    from the basis set data provided in :mod:`~pennylane.qchem.basis_data`. If `load_data = True`,
+    the basis set data is loaded from the basis-set-exchange library.
 
     Args:
         name (str): name of the basis set
         atom (str): atomic symbol of the chemical element
+        load_data (bool): flag to load data from the basis-set-exchange library
 
     Returns:
         list(tuple): tuple containing the angular momentum, the exponents and contraction
@@ -95,7 +97,11 @@ def atom_basis_data(name, atom):
     # for dxy, dxz, dyz, dxx, dyy, dzz, respectively:
     d = [(1, 1, 0), (1, 0, 1), (0, 1, 1), (2, 0, 0), (0, 2, 0), (0, 0, 2)]
 
-    basis = basis_sets[name][atom]
+    if load_data:
+        basis = load_basisset(name, atom)
+    else:
+        basis = basis_sets[name][atom]
+
     params = []
     sp_count = 0
     for i, j in enumerate(basis["orbitals"]):
@@ -122,7 +128,7 @@ def atom_basis_data(name, atom):
     return params
 
 
-def mol_basis_data(name, symbols):
+def mol_basis_data(name, symbols, load_data=False):
     r"""Generates default basis set parameters for a molecule.
 
     This function generates the default basis set parameters for a list of atomic symbols and
@@ -131,6 +137,7 @@ def mol_basis_data(name, symbols):
     Args:
         name (str): name of the basis set
         symbols (list[str]): symbols of the atomic species in the molecule
+        load_data (bool): flag to load data from the basis-set-exchange library
 
     Returns:
         tuple(list, tuple): the number of atomic basis functions and the basis set parameters for
@@ -148,7 +155,7 @@ def mol_basis_data(name, symbols):
     n_basis = []
     basis_set = []
     for s in symbols:
-        basis = atom_basis_data(name, s)
+        basis = atom_basis_data(name, s, load_data=load_data)
         n_basis += [len(basis)]
         basis_set += basis
     return n_basis, tuple(basis_set)
diff --git a/pennylane/qchem/molecule.py b/pennylane/qchem/molecule.py
index 402effb46..5cfc800b8 100644
--- a/pennylane/qchem/molecule.py
+++ b/pennylane/qchem/molecule.py
@@ -44,6 +44,7 @@ class Molecule:
             values of ``mult`` are :math:`1, 2, 3, \ldots`.
         basis_name (str): Atomic basis set used to represent the molecular orbitals. Currently, the
             only supported basis sets are 'STO-3G', '6-31G', '6-311G' and 'CC-PVDZ'.
+        load_data (bool): flag to load data from the basis-set-exchange library
         l (tuple[int]): angular momentum quantum numbers of the basis function
         alpha (array[float]): exponents of the primitive Gaussian functions
         coeff (array[float]): coefficients of the contracted Gaussian functions
@@ -67,6 +68,7 @@ class Molecule:
         charge=0,
         mult=1,
         basis_name="sto-3g",
+        load_data=False,
         l=None,
         alpha=None,
         coeff=None,
@@ -81,7 +83,8 @@ class Molecule:
         ]:
             raise ValueError(
                 "Currently, the only supported basis sets are 'sto-3g', '6-31g', '6-311g' and"
-                " 'cc-pvdz'."
+                " 'cc-pvdz'. Please consider using `load_data=True` to download the basis set from"
+                " an external library that can be installed with: pip install basis-set-exchange."
             )
 
         if set(symbols) - set(atomic_numbers):
@@ -93,7 +96,7 @@ class Molecule:
         self.mult = mult
         self.basis_name = basis_name.lower()
 
-        self.n_basis, self.basis_data = mol_basis_data(self.basis_name, self.symbols)
+        self.n_basis, self.basis_data = mol_basis_data(self.basis_name, self.symbols, load_data)
 
         if l is None:
             l = [i[0] for i in self.basis_data]
diff --git a/pennylane/qchem/openfermion_obs.py b/pennylane/qchem/openfermion_obs.py
index 87c36c30f..3e154b481 100644
--- a/pennylane/qchem/openfermion_obs.py
+++ b/pennylane/qchem/openfermion_obs.py
@@ -813,6 +813,7 @@ def molecular_hamiltonian(
     args=None,
     grouping_type=None,
     grouping_method="rlf",
+    load_data=False,
     convert_tol=1e012,
 ):  # pylint:disable=too-many-arguments
     r"""Generate the qubit Hamiltonian of a molecule.
@@ -853,7 +854,7 @@ def molecular_hamiltonian(
         method (str): Quantum chemistry method used to solve the
             mean field electronic structure problem. Available options are ``method="dhf"``
             to specify the built-in differentiable Hartree-Fock solver, or ``method="pyscf"``
-            to use the OpenFermion-PySCF plugin (this requires that ``openfermionpyscf`` be installed).
+            to use the OpenFermion-PySCF plugin (this requires ``openfermionpyscf`` to be installed).
         active_electrons (int): Number of active electrons. If not specified, all electrons
             are considered to be active.
         active_orbitals (int): Number of active orbitals. If not specified, all orbitals
@@ -871,6 +872,7 @@ def molecular_hamiltonian(
         grouping_type (str): method to group commuting observables
         grouping_method (str): the graph coloring heuristic to use in solving minimum clique cover
             for grouping
+        load_data (bool): flag to load data from the basis-set-exchange library
         convert_tol (float): Tolerance in `machine epsilon <https://numpy.org/doc/stable/reference/generated/numpy.real_if_close.html>`_
             for the imaginary part of the Hamiltonian coefficients created by openfermion.
             Coefficients with imaginary part less than 2.22e-16*tol are considered to be real.
@@ -926,6 +928,7 @@ def molecular_hamiltonian(
             charge=charge,
             mult=mult,
             basis_name=basis,
+            load_data=load_data,
             alpha=alpha,
             coeff=coeff,
         )


Test patch:
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index d0b9a7480..24bf18580 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -336,6 +336,9 @@ jobs:
       - name: Install openfermionpyscf
         run: pip3 install openfermionpyscf
 
+      - name: Install basis-set-exchange
+        run: pip3 install basis-set-exchange
+
       - name: Install PennyLane
         run: |
           python setup.py bdist_wheel
diff --git a/tests/qchem/test_basis_set.py b/tests/qchem/test_basis_set.py
index 21d09a503..eed82d54e 100644
--- a/tests/qchem/test_basis_set.py
+++ b/tests/qchem/test_basis_set.py
@@ -14,6 +14,8 @@
 """
 Unit tests for generating basis set default parameters.
 """
+import sys
+
 # pylint: disable=no-self-use
 import pytest
 
@@ -397,6 +399,33 @@ class TestBasis:
         assert alpha == alpha_ref
         assert coeff == coeff_ref
 
+    @pytest.mark.parametrize(
+        ("basis_name", "atom_name"),
+        [
+            (
+                "sto-3g",
+                "H",
+            ),
+            (
+                "6-311g",
+                "C",
+            ),
+            (
+                "cc-pvdz",
+                "O",
+            ),
+        ],
+    )
+    def test_data_basis_set_exchange(self, basis_name, atom_name):
+        """Test that correct basis set parameters are loaded from basis_set_exchange library."""
+
+        pytest.importorskip("basis_set_exchange")
+
+        data_load = qchem.atom_basis_data(basis_name, atom_name, load_data=True)
+        data_read = qchem.atom_basis_data(basis_name, atom_name, load_data=False)
+
+        assert data_load == data_read
+
     basis_data_HF = [
         (
             "sto-3g",
@@ -447,3 +476,321 @@ class TestBasis:
         assert n_basis == n_ref
 
         assert np.allclose(params, params_ref)
+
+
+class TestLoadBasis:
+    """Tests for loading data from external libraries."""
+
+    @pytest.mark.parametrize(
+        ("basis_name", "atom_name", "params_ref"),
+        [  # data manually copied from https://www.basissetexchange.org/
+            (
+                "sto-3g",
+                "H",
+                (
+                    [
+                        (
+                            "S",  # l
+                            [0.3425250914e01, 0.6239137298e00, 0.1688554040e00],  # alpha
+                            [0.1543289673e00, 0.5353281423e00, 0.4446345422e00],  # coeff
+                        )
+                    ]
+                ),
+            ),
+            (
+                "6-31g",
+                "H",
+                (
+                    (
+                        "S",  # l
+                        [0.1873113696e02, 0.2825394365e01, 0.6401216923e00],  # alpha
+                        [0.3349460434e-01, 0.2347269535e00, 0.8137573261e00],  # coeff
+                    ),
+                    (
+                        "S",  # l
+                        [0.1612777588e00],  # alpha
+                        [1.0000000],  # coeff
+                    ),
+                ),
+            ),
+            (
+                "6-31g",
+                "O",
+                (
+                    (
+                        "S",  # l
+                        [
+                            0.5484671660e04,
+                            0.8252349460e03,
+                            0.1880469580e03,
+                            0.5296450000e02,
+                            0.1689757040e02,
+                            0.5799635340e01,
+                        ],  # alpha
+                        [
+                            0.1831074430e-02,
+                            0.1395017220e-01,
+                            0.6844507810e-01,
+                            0.2327143360e00,
+                            0.4701928980e00,
+                            0.3585208530e00,
+                        ],  # coeff
+                    ),
+                    (
+                        "S",  # l
+                        [0.1553961625e02, 0.3599933586e01, 0.1013761750e01],  # alpha
+                        [-0.1107775495e00, -0.1480262627e00, 0.1130767015e01],
+                        # coeff
+                    ),
+                    (
+                        "P",  # l
+                        [0.1553961625e02, 0.3599933586e01, 0.1013761750e01],  # alpha
+                        [0.7087426823e-01, 0.3397528391e00, 0.7271585773e00],  # coeff
+                    ),
+                    (
+                        "S",  # l
+                        [0.2700058226e00],  # alpha
+                        [0.1000000000e01],  # coeff
+                    ),
+                    (
+                        "P",  # l
+                        [0.2700058226e00],  # alpha
+                        [0.1000000000e01],  # coeff
+                    ),
+                ),
+            ),
+            (
+                "6-311g",
+                "O",
+                (
+                    (
+                        "S",  # l
+                        [8588.5, 1297.23, 299.296, 87.3771, 25.6789, 3.74004],  # alpha
+                        [0.00189515, 0.0143859, 0.070732, 0.240001, 0.594797, 0.280802],
+                        # coeff
+                    ),
+                    (
+                        "S",  # l
+                        [42.11750, 9.628370, 2.853320],  # alpha
+                        [0.113889, 0.920811, -0.00327447],
+                        # coeff
+                    ),
+                    (
+                        "P",  # l
+                        [42.11750, 9.628370, 2.853320],  # alpha
+                        [0.0365114, 0.237153, 0.819702],  # coeff
+                    ),
+                    (
+                        "S",  # l
+                        [0.905661],  # alpha
+                        [1.000000],  # coeff
+                    ),
+                    (
+                        "P",  # l
+                        [0.905661],  # alpha
+                        [1.000000],  # coeff
+                    ),
+                    (
+                        "S",  # l
+                        [0.255611],  # alpha
+                        [1.000000],  # coeff
+                    ),
+                    (
+                        "P",  # l
+                        [0.255611],  # alpha
+                        [1.000000],  # coeff
+                    ),
+                ),
+            ),
+            (
+                "cc-pvdz",
+                "O",
+                (
+                    (
+                        "S",  # l
+                        [
+                            1.172000e04,
+                            1.759000e03,
+                            4.008000e02,
+                            1.137000e02,
+                            3.703000e01,
+                            1.327000e01,
+                            5.025000e00,
+                            1.013000e00,
+                            3.023000e-01,
+                        ],  # alpha
+                        [
+                            7.100000e-04,
+                            5.470000e-03,
+                            2.783700e-02,
+                            1.048000e-01,
+                            2.830620e-01,
+                            4.487190e-01,
+                            2.709520e-01,
+                            1.545800e-02,
+                            -2.585000e-03,
+                        ],
+                        # coeff
+                    ),
+                    (
+                        "S",  # l
+                        [
+                            1.172000e04,
+                            1.759000e03,
+                            4.008000e02,
+                            1.137000e02,
+                            3.703000e01,
+                            1.327000e01,
+                            5.025000e00,
+                            1.013000e00,
+                            3.023000e-01,
+                        ],  # alpha
+                        [
+                            -1.600000e-04,
+                            -1.263000e-03,
+                            -6.267000e-03,
+                            -2.571600e-02,
+                            -7.092400e-02,
+                            -1.654110e-01,
+                            -1.169550e-01,
+                            5.573680e-01,
+                            5.727590e-01,
+                        ],
+                        # coeff
+                    ),
+                    (
+                        "S",  # l
+                        [
+                            1.172000e04,
+                            1.759000e03,
+                            4.008000e02,
+                            1.137000e02,
+                            3.703000e01,
+                            1.327000e01,
+                            5.025000e00,
+                            1.013000e00,
+                            3.023000e-01,
+                        ],  # alpha
+                        [
+                            0.000000e00,
+                            0.000000e00,
+                            0.000000e00,
+                            0.000000e00,
+                            0.000000e00,
+                            0.000000e00,
+                            0.000000e00,
+                            0.000000e00,
+                            1.000000e00,
+                        ],
+                        # coeff
+                    ),
+                    (
+                        "P",  # l
+                        [1.770000e01, 3.854000e00, 1.046000e00, 2.753000e-01],
+                        # alpha
+                        [4.301800e-02, 2.289130e-01, 5.087280e-01, 4.605310e-01],
+                        # coeff
+                    ),
+                    (
+                        "P",  # l
+                        [1.770000e01, 3.854000e00, 1.046000e00, 2.753000e-01],
+                        # alpha
+                        [0.000000e00, 0.000000e00, 0.000000e00, 1.000000e00],
+                        # coeff
+                    ),
+                    (
+                        "D",  # l
+                        [1.185000e00],  # alpha
+                        [1.0000000],  # coeff
+                    ),
+                ),
+            ),
+            (
+                "6-31+G*",
+                "B",
+                (
+                    (
+                        "S",  # l
+                        [
+                            0.2068882250e04,
+                            0.3106495700e03,
+                            0.7068303300e02,
+                            0.1986108030e02,
+                            0.6299304840e01,
+                            0.2127026970e01,
+                        ],  # alpha
+                        [
+                            0.1866274590e-02,
+                            0.1425148170e-01,
+                            0.6955161850e-01,
+                            0.2325729330e00,
+                            0.4670787120e00,
+                            0.3634314400e00,
+                        ],  # coeff
+                    ),
+                    (
+                        "S",  # l
+                        [0.4727971071e01, 0.1190337736e01, 0.3594116829e00],  # alpha
+                        [-0.1303937974e00, -0.1307889514e00, 0.1130944484e01],
+                        # coeff
+                    ),
+                    (
+                        "P",  # l
+                        [0.4727971071e01, 0.1190337736e01, 0.3594116829e00],  # alpha
+                        [0.7459757992e-01, 0.3078466771e00, 0.7434568342e00],  # coeff
+                    ),
+                    (
+                        "S",  # l
+                        [0.1267512469e00],  # alpha
+                        [0.1000000000e01],  # coeff
+                    ),
+                    (
+                        "P",  # l
+                        [0.1267512469e00],  # alpha
+                        [0.1000000000e01],  # coeff
+                    ),
+                    (
+                        "D",  # l
+                        [0.6000000000e00],  # alpha
+                        [1.0000000],  # coeff
+                    ),
+                    (
+                        "S",  # l
+                        [0.3150000000e-01],  # alpha
+                        [0.1000000000e01],  # coeff
+                    ),
+                    (
+                        "P",  # l
+                        [0.3150000000e-01],  # alpha
+                        [0.1000000000e01],  # coeff
+                    ),
+                ),
+            ),
+        ],
+    )
+    def test_load_basis_data(self, basis_name, atom_name, params_ref):
+        """Test that correct basis set parameters are loaded for a given atom."""
+
+        pytest.importorskip("basis_set_exchange")
+
+        data = qchem.load_basisset(basis_name, atom_name)
+
+        l_ref = [p[0] for p in params_ref]
+        alpha_ref = [p[1] for p in params_ref]
+        coeff_ref = [p[2] for p in params_ref]
+
+        assert data["orbitals"] == l_ref
+        assert data["exponents"] == alpha_ref
+        assert data["coefficients"] == coeff_ref
+
+    def test_fail_import(self, monkeypatch):
+        """Test if an ImportError is raised when basis_set_exchange is requested but
+        not installed."""
+
+        pytest.importorskip("basis_set_exchange")
+
+        with monkeypatch.context() as m:
+            m.setitem(sys.modules, "basis_set_exchange", None)
+
+            with pytest.raises(ImportError, match="This feature requires basis_set_exchange"):
+                qchem.load_basisset("sto-3g", "H")
diff --git a/tests/qchem/test_molecule.py b/tests/qchem/test_molecule.py
index d9523fa41..8ea3db8f4 100644
--- a/tests/qchem/test_molecule.py
+++ b/tests/qchem/test_molecule.py
@@ -123,7 +123,7 @@ class TestMolecule:
         assert np.allclose(mol.basis_data, basis_data)
 
     @pytest.mark.parametrize(
-        ("symbols", "geometry", "l", "alpha", "coeff", "r"),
+        ("symbols", "geometry", "l", "alpha", "coeff", "r", "load_data"),
         [
             (
                 ["H", "H"],
@@ -132,6 +132,16 @@ class TestMolecule:
                 [[3.42525091, 0.62391373, 0.1688554], [3.42525091, 0.62391373, 0.1688554]],
                 [[0.15432897, 0.53532814, 0.44463454], [0.15432897, 0.53532814, 0.44463454]],
                 np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]),
+                False,
+            ),
+            (
+                ["H", "H"],
+                np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]),
+                [(0, 0, 0), (0, 0, 0)],
+                [[3.42525091, 0.62391373, 0.1688554], [3.42525091, 0.62391373, 0.1688554]],
+                [[0.15432897, 0.53532814, 0.44463454], [0.15432897, 0.53532814, 0.44463454]],
+                np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0]]),
+                True,
             ),
             (
                 ["H", "F"],
@@ -163,14 +173,15 @@ class TestMolecule:
                         [0.0, 0.0, 1.0],
                     ]
                 ),
+                False,
             ),
         ],
     )
-    def test_basisset(self, symbols, geometry, l, alpha, coeff, r):
+    def test_basisset(self, symbols, geometry, l, alpha, coeff, r, load_data):
         r"""Test that the molecule object contains the correct basis set and non-default basis data
         for a given molecule.
         """
-        mol = qchem.Molecule(symbols, geometry, normalize=False)
+        mol = qchem.Molecule(symbols, geometry, load_data=load_data, normalize=False)
 
         assert set(map(type, mol.basis_set)) == {qchem.BasisFunction}
         assert mol.l == l