Add wrapper to generate molecular Hamiltonian as InteractionOperator (#37)

* add generate_molecular_hamiltonian

* lint

Author: kevinsung

openfermionpyscf/__init__.py

@@ -14,10 +14,11 @@
 OpenFermion plugin to interface with PySCF.
 """
 
-try:
-    from ._pyscf_molecular_data import PyscfMolecularData
-    from ._run_pyscf import prepare_pyscf_molecule, run_pyscf
-except ImportError:
-    raise Exception("Please install PySCF.")
+from ._pyscf_molecular_data import PyscfMolecularData
+
+from ._run_pyscf import (
+        generate_molecular_hamiltonian,
+        prepare_pyscf_molecule,
+        run_pyscf)
 
 from ._version import __version__

openfermionpyscf/_run_pyscf.py

@@ -14,12 +14,12 @@
 
 from __future__ import absolute_import
 
+from functools import reduce
+
 import numpy
 import pyscf
-from functools import reduce
-from pyscf import ci, cc, fci, mp
 
-from openfermion.config import *
+from openfermion import MolecularData
 from openfermionpyscf import PyscfMolecularData
 
 
@@ -204,3 +204,54 @@ def run_pyscf(molecule,
     pyscf_molecular_data.__dict__.update(molecule.__dict__)
     pyscf_molecular_data.save()
     return pyscf_molecular_data
+
+
+def generate_molecular_hamiltonian(
+        geometry,
+        basis,
+        multiplicity,
+        charge=0,
+        n_active_electrons=None,
+        n_active_orbitals=None):
+    """Generate a molecular Hamiltonian with the given properties.
+
+    Args:
+        geometry: A list of tuples giving the coordinates of each atom.
+            An example is [('H', (0, 0, 0)), ('H', (0, 0, 0.7414))].
+            Distances in angstrom. Use atomic symbols to
+            specify atoms.
+        basis: A string giving the basis set. An example is 'cc-pvtz'.
+            Only optional if loading from file.
+        multiplicity: An integer giving the spin multiplicity.
+        charge: An integer giving the charge.
+        n_active_electrons: An optional integer specifying the number of
+            electrons desired in the active space.
+        n_active_orbitals: An optional integer specifying the number of
+            spatial orbitals desired in the active space.
+
+    Returns:
+        The Hamiltonian as an InteractionOperator.
+    """
+
+    # Run electronic structure calculations
+    molecule = run_pyscf(
+            MolecularData(geometry, basis, multiplicity, charge)
+    )
+
+    # Freeze core orbitals and truncate to active space
+    if n_active_electrons is None:
+        n_core_orbitals = 0
+        occupied_indices = None
+    else:
+        n_core_orbitals = (molecule.n_electrons - n_active_electrons) // 2
+        occupied_indices = list(range(n_core_orbitals))
+
+    if n_active_orbitals is None:
+        active_indices = None
+    else:
+        active_indices = list(range(n_core_orbitals,
+                                    n_core_orbitals + n_active_orbitals))
+
+    return molecule.get_molecular_hamiltonian(
+            occupied_indices=occupied_indices,
+            active_indices=active_indices)

openfermionpyscf/_run_pyscf_test.py

@@ -0,0 +1,30 @@
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+import openfermion
+import openfermionpyscf
+
+
+def test_load_molecular_hamiltonian():
+    geometry = [('Li', (0., 0., 0.)), ('H', (0., 0., 1.4))]
+
+    lih_hamiltonian = openfermionpyscf.generate_molecular_hamiltonian(
+            geometry, 'sto-3g', 1, 0, 2, 2)
+    assert openfermion.count_qubits(lih_hamiltonian) == 4
+
+    lih_hamiltonian = openfermionpyscf.generate_molecular_hamiltonian(
+            geometry, 'sto-3g', 1, 0, 2, 3)
+    assert openfermion.count_qubits(lih_hamiltonian) == 6
+
+    lih_hamiltonian = openfermionpyscf.generate_molecular_hamiltonian(
+            geometry, 'sto-3g', 1, 0, None, None)
+    assert openfermion.count_qubits(lih_hamiltonian) == 12