Merge pull request #100 from quantumlib/comment-#99

Docstrings update related to #99

src/fqe/_fqe_control.py

@@ -446,7 +446,9 @@ def get_restricted_hamiltonian(
     """Initialize spin conserving spin restricted hamiltonian
 
     Args:
-        tensors (Tuple[numpy.ndarray, ...]) - tensors for the Hamiltonian elements
+        tensors (Tuple[numpy.ndarray, ...]) - tensors for the Hamiltonian elements. \
+            Note that the tensor should be in a spin-free form; therefore, \
+            the size of each dimension is the number of spatial orbitals.
 
         e_0 (complex) - scalar part of the Hamiltonian
     """

src/fqe/fqe_decorators.py

@@ -67,8 +67,10 @@ def build_hamiltonian(ops: Union[FermionOperator, hamiltonian.Hamiltonian],
 
     if isinstance(ops, tuple):
         validate_tuple(ops)
-
-        return general_hamiltonian.General(ops, e_0=e_0)
+        if norb != 0 and ops[0].shape[0] == norb:
+            return restricted_hamiltonian.RestrictedHamiltonian(ops, e_0=e_0)
+        else:
+            return general_hamiltonian.General(ops, e_0=e_0)
 
     if not isinstance(ops, FermionOperator):
         raise TypeError('Expected FermionOperator' \
@@ -373,7 +375,9 @@ def convert(self, ops: Union['FermionOperator', 'hamiltonian.Hamiltonian']):
         Args:
             ops (FermionOperator or Hamiltonian) - input operator
         """
-        hamil = build_hamiltonian(ops, conserve_number=self.conserve_number())
+        hamil = build_hamiltonian(ops,
+                                  norb=self.norb(),
+                                  conserve_number=self.conserve_number())
         return apply(self, hamil)
 
     return convert
@@ -395,7 +399,9 @@ def convert(self,
 
             ops (FermionOperator or Hamiltonian) - input operator
         """
-        hamil = build_hamiltonian(ops, conserve_number=self.conserve_number())
+        hamil = build_hamiltonian(ops,
+                                  norb=self.norb(),
+                                  conserve_number=self.conserve_number())
         return time_evolve(self, time, hamil, inplace)
 
     return convert
@@ -433,7 +439,9 @@ def convert(self,
         Returns:
             newwfn (Wavefunction) - a new intialized wavefunction object
         """
-        hamil = build_hamiltonian(ops, conserve_number=self.conserve_number())
+        hamil = build_hamiltonian(ops,
+                                  norb=self.norb(),
+                                  conserve_number=self.conserve_number())
         return apply_generated_unitary(self,
                                        time,
                                        algo,

src/fqe/hamiltonians/restricted_hamiltonian.py

@@ -34,9 +34,10 @@ def __init__(self,
         Arguments:
             tensors: Variable length tuple containg between one and four
                      numpy.arrays of increasing rank. The tensors contain the
-                     n-body hamiltonian elements. Tensors up to the highest
-                     order must be included even if the lower terms are full of
-                     zeros.
+                     n-body hamiltonian elements in the spin-free form.
+                     Therefore, the size of each dimension is the number of
+                     spatial orbitals.  Tensors up to the highest order must be
+                     included even if the lower terms are full of zeros.
             e_0: Scalar potential associated with the Hamiltonian.
         """
         super().__init__(e_0=e_0)

src/fqe/property_test.py

@@ -22,7 +22,7 @@
 
 import fqe
 from fqe.wavefunction import Wavefunction
-from fqe.hamiltonians import general_hamiltonian
+from fqe.hamiltonians import restricted_hamiltonian
 from fqe.fqe_ops.fqe_ops import (
     NumberOperator,
     S2Operator,
@@ -47,7 +47,7 @@ def test_lih_energy(self):
         nele = nalpha + nbeta
         h1e, h2e, lih_ground = build_lih_data.build_lih_data('energy')
 
-        elec_hamil = general_hamiltonian.General((h1e, h2e))
+        elec_hamil = restricted_hamiltonian.RestrictedHamiltonian((h1e, h2e))
         wfn = Wavefunction([[nele, nalpha - nbeta, norb]])
         wfn.set_wfn(strategy='from_data',
                     raw_data={(nele, nalpha - nbeta): lih_ground})

src/fqe/wavefunction.py

@@ -38,7 +38,7 @@
 from fqe.util import sort_configuration_keys
 from fqe.util import vdot
 from fqe.hamiltonians import hamiltonian, sparse_hamiltonian
-from fqe.hamiltonians import diagonal_hamiltonian
+from fqe.hamiltonians import diagonal_hamiltonian, restricted_hamiltonian
 from fqe.bitstring import count_bits
 from fqe.fqe_ops import fqe_operator, fqe_ops_utils
 from fqe.wick import wick
@@ -369,6 +369,16 @@ def apply(self, hamil: 'hamiltonian.Hamiltonian') -> 'Wavefunction':
             if isinstance(hamil, diagonal_hamiltonian.Diagonal):
                 transformed = out._apply_diagonal(hamil)
             else:
+                if isinstance(hamil,
+                              restricted_hamiltonian.RestrictedHamiltonian):
+                    expected = self._norb
+                else:
+                    expected = self._norb * 2
+                if hamil.dim() != expected:
+                    raise ValueError('Hamiltonian has incorrect size:' \
+                                     + ' expected {}'.format(expected) \
+                                     + ' provided {}'.format(hamil.dim()))
+
                 transformed = out._apply_array(hamil.tensors(), hamil.e_0())
 
             if self._conserve_spin and not self._conserve_number:
@@ -388,6 +398,8 @@ def _apply_array(self, array: Tuple[numpy.ndarray, ...],
         Arg:
             array (numpy.array) - numpy array
 
+            e_0 (complex) - constant part of the Hamiltonian
+
         Returns:
             newwfn (Wavvefunction) - a new intialized wavefunction object
 
@@ -942,10 +954,12 @@ def time_evolve(self, time: float, hamil,
         either as raw operations or wrapped up in a Hamiltonian.
 
         Args:
-            ops (FermionOperators) - FermionOperators which are to be time evolved.
-
             time (float) - the duration by which to evolve the operators
 
+            hamil - Hamiltoninans or FermionOperators which are to be time evolved.
+
+            inplace (bool) - whether the result will be stored in place
+
         Returns:
             Wavefunction - a wavefunction object that has been time evolved.
         """
@@ -1291,7 +1305,6 @@ def rdm(self, string: str, brawfn: Optional['Wavefunction'] = None
 # TODO: Delete or make unit test?
 if __name__ == "__main__":
     from openfermion import FermionOperator
-    import numpy
     import fqe
     from fqe.unittest_data import build_lih_data, build_hamiltonian
 

src/fqe/wavefunction_test.py

@@ -34,6 +34,7 @@
 from fqe.hamiltonians import general_hamiltonian
 from fqe.hamiltonians import sparse_hamiltonian
 from fqe.hamiltonians import diagonal_hamiltonian
+from fqe.hamiltonians import restricted_hamiltonian
 from fqe import get_restricted_hamiltonian
 
 from fqe.unittest_data import build_wfn, build_hamiltonian
@@ -123,9 +124,18 @@ def test_apply_type_error(self):
         self.assertRaises(TypeError, wfn.apply, hamil)
         self.assertRaises(TypeError, wfn.time_evolve, 0.1, hamil)
 
+    def test_apply_value_error(self):
+        data = numpy.zeros((2, 2), dtype=numpy.complex128)
         wfn = Wavefunction([[2, 0, 2]])
-        hamil = get_restricted_hamiltonian((data,))
-        self.assertRaises(ValueError, wfn.time_evolve, 0.1, hamil, True)
+        hamil = general_hamiltonian.General((data,))
+        self.assertRaises(ValueError, wfn.apply, hamil)
+
+        data2 = numpy.zeros((3, 3), dtype=numpy.complex128)
+        hamil2 = restricted_hamiltonian.RestrictedHamiltonian((data2,))
+        self.assertRaises(ValueError, wfn.apply, hamil2)
+
+        hamil3 = restricted_hamiltonian.RestrictedHamiltonian((data,))
+        self.assertRaises(ValueError, wfn.time_evolve, 0.1, hamil3, True)
 
     def test_apply_individual_nbody_error(self):
         fop = FermionOperator('1^ 0')