Add mypy check to CI and fix mypy errors

.github/workflows/pythonpackage.yml

@@ -25,8 +25,21 @@ on:
     branches: [ master ]
 
 jobs:
-  build:
+  mypy:
+    name: Type check
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions/setup-python@v2
+        with:
+          python-version: '3.6'
+          architecture: 'x64'
+      - name: Install mypy
+        run: pip install mypy~=0.790.0
+      - name: Run mypy
+        run: dev/check/mypy
 
+  build:
     runs-on: ubuntu-latest
     strategy:
       matrix:

dev/.mypy/mypy.ini

@@ -1,4 +1,11 @@
 [mypy]
+
+[mypy-__main__]
 follow_imports = silent
 ignore_missing_imports = True
 strict_optional = False
+
+# 3rd-party libs for which we don't have stubs
+[mypy-numpy.*,scipy.*,pytest.*,joblib.*,openfermion.*,cirq.*]
+follow_imports = silent
+ignore_missing_imports = true

dev/check/mypy

@@ -0,0 +1,19 @@
+#!/usr/bin/env bash
+
+################################################################################
+# Runs mypy on the repository using a preconfigured mypy.ini file.
+#
+# Usage:
+#     check/mypy [--flags]
+################################################################################
+
+# Get the working directory to the repo root.
+cd "$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
+cd "$(git rev-parse --show-toplevel)"
+
+echo -e -n "\033[31m"
+mypy --config-file=dev/.mypy/mypy.ini "$@" src/fqe
+result=$?
+echo -e -n "\033[0m"
+
+exit ${result}

src/fqe/_fqe_control.py

@@ -254,9 +254,9 @@ def from_cirq(state: numpy.ndarray, thresh: float) -> 'wavefunction.Wavefunction
         for orb in occ:
             if numpy.absolute(state[orb[0]]) > thresh:
                 param.append([pnum, orb[1], norb])
-    param = set([tuple(p) for p in param])
-    param = [list(x) for x in param]
-    wfn = wavefunction.Wavefunction(param)
+    param_set = set([tuple(p) for p in param])
+    param_list = [list(x) for x in param_set]
+    wfn = wavefunction.Wavefunction(param_list)
     transform.from_cirq(wfn, state)
     return wfn
 
@@ -300,25 +300,25 @@ def expectationValue(wfn: 'wavefunction.Wavefunction',
     return wfn.expectationValue(ops, brawfn)
 
 
-def get_s2_operator() -> 's2_op.S2Operator':
+def get_s2_operator() -> 'S2Operator':
     """Return an S^2 operator.
     """
     return S2Operator()
 
 
-def get_sz_operator() -> 'sz_op.SzOperator':
+def get_sz_operator() -> 'SzOperator':
     """Return an S_zperator.
     """
     return SzOperator()
 
 
-def get_time_reversal_operator() -> 'tr_op.TimeReversalOp':
+def get_time_reversal_operator() -> 'TimeReversalOp':
     """Return a Time Reversal Operator
     """
     return TimeReversalOp()
 
 
-def get_number_operator() -> 'number_op.NumberOperator':
+def get_number_operator() -> 'NumberOperator':
     """Return the particle number operator
     """
     return NumberOperator()
@@ -435,7 +435,7 @@ def get_gso_hamiltonian(tensors: Tuple[numpy.ndarray, ...],
 
 
 def get_restricted_hamiltonian(tensors: Tuple[numpy.ndarray, ...],
-                               e_0: complex = 0. + 0.j) -> 'restricted_hamiltonian.Restricted':
+                               e_0: complex = 0. + 0.j) -> 'restricted_hamiltonian.RestrictedHamiltonian':
     """Initialize spin conserving spin restricted hamiltonian
 
     Args:

src/fqe/algorithm/adapt_vqe.py

@@ -1,5 +1,5 @@
 """Infrastructure for ADAPT VQE algorithm"""
-from typing import List, Union
+from typing import List, Tuple, Union
 import copy
 import openfermion as of
 import fqe
@@ -12,7 +12,6 @@
 from openfermion.chem.molecular_data import spinorb_from_spatial
 
 from fqe.hamiltonians.restricted_hamiltonian import RestrictedHamiltonian
-from fqe.hamiltonians.general_hamiltonian import General as GeneralFQEHamiltonian
 from fqe.hamiltonians.hamiltonian import Hamiltonian as ABCHamiltonian
 from fqe.fqe_decorators import build_hamiltonian
 from fqe.algorithm.brillouin_calculator import (
@@ -21,6 +20,7 @@
     one_rdo_commutator_symm,
 )
 from fqe.algorithm.generalized_doubles_factorization import doubles_factorization
+from fqe.wavefunction import Wavefunction
 
 
 class OperatorPool:
@@ -36,7 +36,7 @@ def __init__(self, norbs: int, occ: List[int], virt: List[int]):
         self.norbs = norbs
         self.occ = occ
         self.virt = virt
-        self.op_pool = []
+        self.op_pool: List[of.FermionOperator] = []
 
     def singlet_t2(self):
         """
@@ -153,7 +153,7 @@ def __init__(self, oei: np.ndarray, tei: np.ndarray, operator_pool,
         self.operator_pool = operator_pool
         self.stopping_eps = stopping_epsilon
 
-    def vbc(self, initial_wf: fqe.Wavefunction, update_rank=None,
+    def vbc(self, initial_wf: Wavefunction, update_rank=None,
             opt_method: str='L-BFGS-B',
             num_opt_var=None
             ):
@@ -174,8 +174,8 @@ def vbc(self, initial_wf: fqe.Wavefunction, update_rank=None,
         self.num_opt_var = num_opt_var
         nso = 2 * self.sdim
         operator_pool = []
-        operator_pool_fqe = []
-        existing_parameters = []
+        operator_pool_fqe: List[ABCHamiltonian] = []
+        existing_parameters: List[float] = []
         self.energies = []
         self.residuals = []
         iteration = 0
@@ -271,7 +271,7 @@ def vbc(self, initial_wf: fqe.Wavefunction, update_rank=None,
                 break
             iteration += 1
 
-    def adapt_vqe(self, initial_wf: fqe.Wavefunction,
+    def adapt_vqe(self, initial_wf: Wavefunction,
                   opt_method: str='L-BFGS-B'):
         """
         Run ADAPT-VQE using
@@ -281,8 +281,8 @@ def adapt_vqe(self, initial_wf: fqe.Wavefunction,
             opt_method: scipy optimizer to use
         """
         operator_pool = []
-        operator_pool_fqe = []
-        existing_parameters = []
+        operator_pool_fqe: List[ABCHamiltonian] = []
+        existing_parameters: List[float] = []
         self.gradients = []
         self.energies = []
         iteration = 0
@@ -336,10 +336,10 @@ def adapt_vqe(self, initial_wf: fqe.Wavefunction,
             iteration += 1
 
     def optimize_param(self, pool: Union[
-        List[of.FermionOperator], List[GeneralFQEHamiltonian]],
+        List[of.FermionOperator], List[ABCHamiltonian]],
                        existing_params: Union[List, np.ndarray],
-                       initial_wf: fqe.Wavefunction,
-                       opt_method: str) -> fqe.wavefunction:
+                       initial_wf: Wavefunction,
+                       opt_method: str) -> Tuple[np.ndarray, float]:
         """Optimize a wavefunction given a list of generators
 
         Args:

src/fqe/algorithm/brillouin_calculator.py

@@ -21,6 +21,7 @@
 
 import openfermion as of
 import fqe
+from fqe.wavefunction import Wavefunction
 from fqe.hamiltonians.restricted_hamiltonian import RestrictedHamiltonian
 
 try:
@@ -42,12 +43,6 @@ def get_fermion_op(coeff_tensor) -> of.FermionOperator:
     Returns:
         A FermionOperator object
     """
-    if len(coeff_tensor.shape) not in (2, 4):
-        raise ValueError(
-            "Arg `coeff_tensor` should have dimension 2 or 4 but has dimension"
-            f" {len(coeff_tensor.shape)}."
-        )
-
     if len(coeff_tensor.shape) == 4:
         nso = coeff_tensor.shape[0]
         fermion_op = of.FermionOperator()
@@ -57,17 +52,23 @@ def get_fermion_op(coeff_tensor) -> of.FermionOperator:
             fermion_op += fop
         return fermion_op
 
-    if len(coeff_tensor.shape) == 2:
+    elif len(coeff_tensor.shape) == 2:
         nso = coeff_tensor.shape[0]
         fermion_op = of.FermionOperator()
         for p, q in product(range(nso), repeat=2):
-            op = ((p, 1), (q, 0))
-            fop = of.FermionOperator(op, coefficient=coeff_tensor[p, q])
+            oper = ((p, 1), (q, 0))
+            fop = of.FermionOperator(oper, coefficient=coeff_tensor[p, q])
             fermion_op += fop
         return fermion_op
 
+    else:
+        raise ValueError(
+            "Arg `coeff_tensor` should have dimension 2 or 4 but has dimension"
+            f" {len(coeff_tensor.shape)}."
+        )
+
 
-def get_acse_residual_fqe(fqe_wf: fqe.Wavefunction,
+def get_acse_residual_fqe(fqe_wf: Wavefunction,
                           fqe_ham: RestrictedHamiltonian,
                           norbs: int) -> np.ndarray:
     """Get the ACSE block by using reduced density operators that are Sz spin

src/fqe/algorithm/low_rank.py

@@ -19,10 +19,11 @@
 from openfermion import givens_decomposition_square
 
 import fqe
+from fqe.wavefunction import Wavefunction
 from fqe.hamiltonians.diagonal_coulomb import DiagonalCoulomb
 
 
-def evolve_fqe_givens(wfn: fqe.Wavefunction, u: np.ndarray) -> np.ndarray:
+def evolve_fqe_givens(wfn: Wavefunction, u: np.ndarray) -> np.ndarray:
     """Evolve a wavefunction by u generated from a 1-body Hamiltonian.
 
     Args:
@@ -77,8 +78,8 @@ def evolve_fqe_givens(wfn: fqe.Wavefunction, u: np.ndarray) -> np.ndarray:
 
 
 def evolve_fqe_diagaonal_coulomb(
-    wfn: fqe.Wavefunction, vij_mat: np.ndarray, time=1
-) -> fqe.Wavefunction:
+    wfn: Wavefunction, vij_mat: np.ndarray, time=1
+) -> Wavefunction:
     r"""Utility for testing evolution of a full 2^{n} wavefunction via
 
     :math:`exp{-i time * \sum_{i,j, sigma, tau}v_{i, j}n_{i\sigma}n_{j\tau}}.`
@@ -96,8 +97,8 @@ def evolve_fqe_diagaonal_coulomb(
 
 
 def double_factor_trotter_evolution(
-    initial_wfn: fqe.Wavefunction, basis_change_unitaries, vij_mats, deltat
-) -> fqe.Wavefunction:
+    initial_wfn: Wavefunction, basis_change_unitaries, vij_mats, deltat
+) -> Wavefunction:
     r"""Doubled Factorized Trotter Evolution
 
     Evolves an initial according to the double factorized algorithm where each

src/fqe/algorithm/low_rank_api.py

@@ -75,7 +75,7 @@ def first_factorization(self, threshold: Optional[float] = None):
             threshold = self.icut
 
         # convert physics notation integrals into chemist notation
-        # and determine the first low-rank fractorization
+        # and determine the first low-rank factorization
         if self.spin_basis:
             (
                 eigenvalues,
@@ -95,7 +95,7 @@ def first_factorization(self, threshold: Optional[float] = None):
                 one_body_correction,
                 _,
             ) = low_rank_two_body_decomposition(
-                0.5 * self.tei,
+                0.5 * self.tei,  # type: ignore
                 truncation_threshold=threshold,
                 final_rank=self.lmax,
                 spin_basis=self.spin_basis,

src/fqe/bitstring.py

@@ -114,8 +114,7 @@ def lexicographic_bitstring_generator(str0: int, norb: int) -> List[int]:
             states
     """
     out = []
-    gs_bs = '{0:b}'.format(str0).zfill(norb)
-    gs_bs = [int(x) for x in gs_bs]
+    gs_bs = [int(x) for x in '{0:b}'.format(str0).zfill(norb)]
     n_elec = sum(gs_bs)
     n_orbs = len(gs_bs)
     for ones_positions in combinations(range(n_orbs), n_elec):

src/fqe/cirq_utils.py

@@ -76,7 +76,7 @@ def qubit_op_to_gate(operation: 'QubitOperator',
 def qubit_projection(ops: QubitOperator,
                      qubits: List[LineQubit],
                      state: numpy.ndarray,
-                     coeff: numpy.ndarray) -> numpy.ndarray:
+                     coeff: numpy.ndarray) -> None:
     """Find the projection of each set of qubit operators on a
     wavefunction.
 
@@ -89,12 +89,8 @@ def qubit_projection(ops: QubitOperator,
         state (numpy.array(dtype=numpy.complex64)) - a cirq wavefunction that \
             is being projected
 
-        coeff (numpy.array(dtype=numpy.complex64)) - a coeffcient array that \
+        coeff (numpy.array(dtype=numpy.complex64)) - a coefficient array that \
             will store the result of the projection.
-
-    Returns:
-        coeff (numpy.array(dtype=numpy.complex64)) - the coefficients modified \
-            in place.
     """
     qpu = Simulator(dtype=numpy.complex128)
     for indx, cluster in enumerate(ops.terms):

src/fqe/fqe_data.py

@@ -703,7 +703,7 @@ def _apply_array_spin123(self,
         if not from1234:
             (dveca, dvecb) = self.calculate_dvec_spin()
         else:
-            dveca, dvecb = dvec[0], dvec[1]
+            dveca, dvecb = dvec  # type: ignore
 
         if not from1234:
             evecaa = numpy.zeros((norb, norb, norb, norb, lena, lenb),
@@ -724,7 +724,7 @@ def _apply_array_spin123(self,
                     evecab[:, :, i, j, :, :] = tmp[0][:, :, :, :]
                     evecbb[:, :, i, j, :, :] = tmp[1][:, :, :, :]
         else:
-            evecaa, evecab, evecba, evecbb = evec[0], evec[1], evec[2], evec[3]
+            evecaa, evecab, evecba, evecbb = evec  # type: ignore
 
         symfac = 2.0 if not from1234 else 1.0
 

src/fqe/fqe_decorators.py

@@ -17,7 +17,7 @@
 #there are two places where access to protected members improves code quality
 #pylint: disable=protected-access
 
-from typing import Dict, Tuple, Union, Optional, List
+from typing import Any, Dict, Tuple, Union, Optional, List
 from functools import wraps
 
 import copy
@@ -75,6 +75,7 @@ def build_hamiltonian(ops: Union[FermionOperator, hamiltonian.Hamiltonian],
 
     assert is_hermitian(ops)
 
+    out: Any
     if len(ops.terms) <= 2:
         out = sparse_hamiltonian.SparseHamiltonian(ops, e_0=e_0)
 
@@ -84,7 +85,7 @@ def build_hamiltonian(ops: Union[FermionOperator, hamiltonian.Hamiltonian],
 
         ops = normal_ordered(ops)
 
-        ops_rank, e_0 = split_openfermion_tensor(ops)
+        ops_rank, e_0 = split_openfermion_tensor(ops)  # type: ignore
 
         if norb == 0:
             for term in ops_rank.values():

src/fqe/fqe_ops/fqe_operator.py

@@ -25,7 +25,7 @@ class FqeOperator(ABC):
 
     @abstractmethod
     def contract(
-        self, brastate: "Wavefunction", ketstate: Optional["Wavefunction"]
+        self, brastate: "Wavefunction", ketstate: "Wavefunction"
     ) -> complex:
         """Given two wavefunctions, generate the expectation value of the
         operator according to its representation.

src/fqe/hamiltonians/diagonal_hamiltonian.py

@@ -64,7 +64,7 @@ def diag_values(self) -> np.ndarray:
         """Returns the diagonal values packed into a single dimension."""
         return self._hdiag
 
-    def iht(self, time: float) -> Tuple[np.ndarray, ...]:
+    def iht(self, time: float) -> 'Diagonal':
         """Returns the matrices of the Hamiltonian prepared for time evolution.
 
         Args: