enable check-untyped-defs mypy option (#337)

Co-authored-by: Bart Andrews <[email protected]>

pyproject.toml

@@ -62,6 +62,7 @@ module-name = "ffsim._lib"
 testpaths = ["tests"]
 
 [tool.mypy]
+check_untyped_defs = true
 ignore_missing_imports = true
 files = ["python/**/*.py", "python/**/*.pyi", "tests/**/*.py", "docs/**/*.py"]
 

python/ffsim/_lib.pyi

@@ -1,16 +1,19 @@
-from collections.abc import Iterator, Mapping
+from collections.abc import Iterator, MutableMapping
 
 import numpy as np
 
-class FermionOperator(Mapping[tuple[tuple[bool, bool, int], ...], complex]):
+class FermionOperator(MutableMapping[tuple[tuple[bool, bool, int], ...], complex]):
     def __init__(
         self, coeffs: dict[tuple[tuple[bool, bool, int], ...], complex]
     ) -> None: ...
     def normal_ordered(self) -> "FermionOperator": ...
     def conserves_particle_number(self) -> bool: ...
     def conserves_spin_z(self) -> bool: ...
     def many_body_order(self) -> int: ...
+    def copy(self) -> FermionOperator: ...
     def __getitem__(self, key: tuple[tuple[bool, bool, int], ...]) -> complex: ...
+    def __setitem__(self, key: tuple[tuple[bool, bool, int], ...], complex) -> None: ...
+    def __delitem__(self, key: tuple[tuple[bool, bool, int], ...]) -> None: ...
     def __contains__(self, key: object) -> bool: ...
     def __iter__(self) -> Iterator[tuple[tuple[bool, bool, int], ...]]: ...
     def __len__(self) -> int: ...

python/ffsim/qiskit/gates/diag_coulomb.py

@@ -145,9 +145,7 @@ def _define(self):
 
     def inverse(self):
         """Inverse gate."""
-        return DiagCoulombEvolutionSpinlessJW(
-            self.norb, self.mat, -self.time, z_representation=self.z_representation
-        )
+        return DiagCoulombEvolutionSpinlessJW(self.norb, self.mat, -self.time)
 
 
 def _diag_coulomb_evo_num_rep_spinless_jw(

python/ffsim/variational/givens.py

@@ -57,7 +57,7 @@ def __post_init__(self):
         if len(self.thetas) != len(self.interaction_pairs):
             raise ValueError(
                 "The number of thetas must equal the number of interaction pairs. "
-                f"Got {len(self.phis)} and {len(self.interaction_pairs)}."
+                f"Got {len(self.thetas)} and {len(self.interaction_pairs)}."
             )
         if self.phis is not None and len(self.phis) != len(self.interaction_pairs):
             raise ValueError(

python/ffsim/variational/hopgate.py

@@ -14,6 +14,7 @@
 
 import itertools
 from dataclasses import dataclass
+from typing import cast
 
 import numpy as np
 
@@ -113,3 +114,25 @@ def from_parameters(
             thetas=params,
             final_orbital_rotation=final_orbital_rotation,
         )
+
+    def _approx_eq_(self, other, rtol: float, atol: float) -> bool:
+        if isinstance(other, HopGateAnsatzOperator):
+            if self.norb != other.norb:
+                return False
+            if self.interaction_pairs != other.interaction_pairs:
+                return False
+            if not np.allclose(self.thetas, other.thetas, rtol=rtol, atol=atol):
+                return False
+            if (self.final_orbital_rotation is None) != (
+                other.final_orbital_rotation is None
+            ):
+                return False
+            if self.final_orbital_rotation is not None:
+                return np.allclose(
+                    cast(np.ndarray, self.final_orbital_rotation),
+                    cast(np.ndarray, other.final_orbital_rotation),
+                    rtol=rtol,
+                    atol=atol,
+                )
+            return True
+        return NotImplemented

python/ffsim/variational/num_num.py

@@ -85,7 +85,9 @@ def _apply_unitary_(
         )
 
     @staticmethod
-    def n_params(interaction_pairs: list[tuple[int, int]]) -> int:
+    def n_params(
+        interaction_pairs: tuple[list[tuple[int, int]], list[tuple[int, int]]],
+    ) -> int:
         """Return the number of parameters of an ansatz with given settings.
 
         Args:

tests/python/molecular_data_test.py

@@ -89,10 +89,15 @@ def test_molecular_data_run_methods():
     mol_data.run_sci()
     mol_data.run_fci()
 
+    assert isinstance(mol_data.mp2_energy, float)
     np.testing.assert_allclose(mol_data.mp2_energy, -108.58852784026)
+    assert isinstance(mol_data.ccsd_energy, float)
     np.testing.assert_allclose(mol_data.ccsd_energy, -108.5933309085008)
+    assert isinstance(mol_data.cisd_energy, float)
     np.testing.assert_allclose(mol_data.cisd_energy, -108.5878344909782)
+    assert isinstance(mol_data.sci_energy, float)
     np.testing.assert_allclose(mol_data.sci_energy, -108.59598682615388)
+    assert isinstance(mol_data.fci_energy, float)
     np.testing.assert_allclose(mol_data.fci_energy, -108.595987350986)
 
 

tests/python/operators/fermion_operator_test.py

@@ -215,7 +215,7 @@ def test_pow():
     assert op**3 == op * op * op
     assert pow(op, 2) == op * op
     with pytest.raises(ValueError, match="mod argument"):
-        _ = pow(op, 2, 2)
+        _ = pow(op, 2, 2)  # type: ignore
 
 
 def test_normal_ordered():

tests/python/states/bitstring_test.py

@@ -23,7 +23,7 @@
 def test_indices_to_strings_string():
     """Test converting statevector indices to strings, output type string."""
     norb = 3
-    nelec = 2
+    nelec: int | tuple[int, int] = 2
     dim = ffsim.dim(norb, nelec)
     strings = ffsim.indices_to_strings(range(dim), norb, nelec)
     assert strings == [
@@ -84,7 +84,7 @@ def test_indices_to_strings_string():
 def test_indices_to_strings_int():
     """Test converting statevector indices to strings, output type int."""
     norb = 3
-    nelec = 2
+    nelec: int | tuple[int, int] = 2
     dim = ffsim.dim(norb, nelec)
     strings = ffsim.indices_to_strings(
         range(dim), norb, nelec, bitstring_type=ffsim.BitstringType.INT
@@ -159,7 +159,7 @@ def test_indices_to_strings_int():
 def test_indices_to_strings_bit_array():
     """Test converting statevector indices to strings, output type bit array."""
     norb = 3
-    nelec = 2
+    nelec: int | tuple[int, int] = 2
     dim = ffsim.dim(norb, nelec)
     strings = ffsim.indices_to_strings(
         range(dim), norb, nelec, bitstring_type=ffsim.BitstringType.BIT_ARRAY
@@ -536,7 +536,7 @@ def test_addresses_to_strings_large_address():
 def test_strings_to_addresses_int():
     """Test converting statevector strings to addresses, input type int."""
     norb = 3
-    nelec = 2
+    nelec: int | tuple[int, int] = 2
     dim = ffsim.dim(norb, nelec)
     indices = ffsim.strings_to_addresses(
         [0b011, 0b101, 0b110],
@@ -571,7 +571,7 @@ def test_strings_to_addresses_int():
 def test_strings_to_addresses_string():
     """Test converting statevector strings to indices, input type string."""
     norb = 3
-    nelec = 2
+    nelec: int | tuple[int, int] = 2
     dim = ffsim.dim(norb, nelec)
     indices = ffsim.strings_to_addresses(
         [

tests/python/states/slater_test.py

@@ -155,7 +155,7 @@ def test_sample_slater_determinant_large():
     shots = 5000
     rotation_a = ffsim.random.random_unitary(norb, seed=rng)
     rotation_b = ffsim.random.random_unitary(norb, seed=rng)
-    occupied_orbitals = [(0, 2, 3), (2, 4)]
+    occupied_orbitals = ((0, 2, 3), (2, 4))
     rdm_a, rdm_b = ffsim.slater_determinant_rdms(
         norb, occupied_orbitals, (rotation_a, rotation_b)
     )
@@ -178,7 +178,7 @@ def test_sample_slater_determinant_restrict():
     nelec = (4, 3)
 
     shots = 10
-    occupied_orbitals = [(0, 2, 3, 5), (2, 3, 4)]
+    occupied_orbitals = ((0, 2, 3, 5), (2, 3, 4))
     rdm_a, rdm_b = ffsim.slater_determinant_rdms(norb, occupied_orbitals)
     samples = ffsim.sample_slater_determinant(
         (rdm_a, rdm_b), norb, nelec, orbs=([1, 2, 5], [3, 4, 5]), shots=shots

tests/python/variational/hopgate_test.py

@@ -46,8 +46,4 @@ def ncycles(iterable, n):
         interaction_pairs=interaction_pairs,
         with_final_orbital_rotation=True,
     )
-
-    np.testing.assert_allclose(roundtripped.thetas, operator.thetas)
-    np.testing.assert_allclose(
-        roundtripped.final_orbital_rotation, operator.final_orbital_rotation
-    )
+    assert ffsim.approx_eq(roundtripped, operator)

tests/python/variational/num_num_test.py

@@ -36,7 +36,9 @@ def test_parameters_roundtrip():
         thetas=(thetas_aa, thetas_ab),
     )
     assert (
-        operator.n_params(interaction_pairs=(pairs_aa, pairs_ab))
+        ffsim.NumNumAnsatzOpSpinBalanced.n_params(
+            interaction_pairs=(pairs_aa, pairs_ab)
+        )
         == len(operator.to_parameters())
         == 2 * norb
     )

tests/python/variational/uccsd_test.py

@@ -65,12 +65,7 @@ def test_parameters_roundtrip():
             nocc=nocc,
             with_final_orbital_rotation=with_final_orbital_rotation,
         )
-        np.testing.assert_allclose(roundtripped.t1, operator.t1)
-        np.testing.assert_allclose(roundtripped.t2, operator.t2)
-        if with_final_orbital_rotation:
-            np.testing.assert_allclose(
-                roundtripped.final_orbital_rotation, operator.final_orbital_rotation
-            )
+        assert ffsim.approx_eq(roundtripped, operator)
 
 
 def test_approx_eq():

tests/python/variational/ucj_spin_balanced_test.py

@@ -83,16 +83,7 @@ def test_parameters_roundtrip_all_to_all():
             n_reps=n_reps,
             with_final_orbital_rotation=with_final_orbital_rotation,
         )
-        np.testing.assert_allclose(
-            roundtripped.diag_coulomb_mats, operator.diag_coulomb_mats
-        )
-        np.testing.assert_allclose(
-            roundtripped.orbital_rotations, operator.orbital_rotations
-        )
-        if with_final_orbital_rotation:
-            np.testing.assert_allclose(
-                roundtripped.final_orbital_rotation, operator.final_orbital_rotation
-            )
+        assert ffsim.approx_eq(roundtripped, operator)
 
 
 def test_parameters_roundtrip_interaction_pairs():
@@ -116,16 +107,7 @@ def test_parameters_roundtrip_interaction_pairs():
             interaction_pairs=interaction_pairs,
             with_final_orbital_rotation=with_final_orbital_rotation,
         )
-        np.testing.assert_allclose(
-            roundtripped.diag_coulomb_mats, operator.diag_coulomb_mats
-        )
-        np.testing.assert_allclose(
-            roundtripped.orbital_rotations, operator.orbital_rotations
-        )
-        if with_final_orbital_rotation:
-            np.testing.assert_allclose(
-                roundtripped.final_orbital_rotation, operator.final_orbital_rotation
-            )
+        assert ffsim.approx_eq(roundtripped, operator)
 
 
 def test_t_amplitudes_energy():

tests/python/variational/ucj_spin_unbalanced_test.py

@@ -86,16 +86,7 @@ def test_parameters_roundtrip_all_to_all():
             n_reps=n_reps,
             with_final_orbital_rotation=with_final_orbital_rotation,
         )
-        np.testing.assert_allclose(
-            roundtripped.diag_coulomb_mats, operator.diag_coulomb_mats
-        )
-        np.testing.assert_allclose(
-            roundtripped.orbital_rotations, operator.orbital_rotations
-        )
-        if with_final_orbital_rotation:
-            np.testing.assert_allclose(
-                roundtripped.final_orbital_rotation, operator.final_orbital_rotation
-            )
+        assert ffsim.approx_eq(roundtripped, operator)
 
 
 def test_parameters_roundtrip_interaction_pairs():
@@ -119,16 +110,7 @@ def test_parameters_roundtrip_interaction_pairs():
             interaction_pairs=interaction_pairs,
             with_final_orbital_rotation=with_final_orbital_rotation,
         )
-        np.testing.assert_allclose(
-            roundtripped.diag_coulomb_mats, operator.diag_coulomb_mats
-        )
-        np.testing.assert_allclose(
-            roundtripped.orbital_rotations, operator.orbital_rotations
-        )
-        if with_final_orbital_rotation:
-            np.testing.assert_allclose(
-                roundtripped.final_orbital_rotation, operator.final_orbital_rotation
-            )
+        assert ffsim.approx_eq(roundtripped, operator)
 
 
 def test_t_amplitudes_energy():
@@ -149,7 +131,7 @@ def test_t_amplitudes_energy():
     mol_hamiltonian = mol_data.hamiltonian
 
     # Construct UCJ operator
-    n_reps = 4
+    n_reps: int | tuple[int, int] = 4
     operator = ffsim.UCJOpSpinUnbalanced.from_t_amplitudes(
         ccsd.t2, t1=ccsd.t1, n_reps=n_reps
     )
@@ -198,7 +180,8 @@ def test_t_amplitudes_random_n_reps():
     nvrt_a = norb - nocc_a
     nvrt_b = norb - nocc_b
 
-    for n_reps in [5, 50, (10, 5)]:
+    n_reps: int | tuple[int, int]
+    for n_reps in [5, 50, (10, 5)]:  # type: ignore
         t2aa = ffsim.random.random_t2_amplitudes(norb, nocc_a, seed=rng, dtype=float)
         t2ab = rng.standard_normal((nocc_a, nocc_b, nvrt_a, nvrt_b))
         t2bb = ffsim.random.random_t2_amplitudes(norb, nocc_b, seed=rng, dtype=float)
@@ -223,7 +206,8 @@ def test_t_amplitudes_zero_n_reps():
     nvrt_a = norb - nocc_a
     nvrt_b = norb - nocc_b
 
-    for n_reps in [5, 50, (10, 5)]:
+    n_reps: int | tuple[int, int]
+    for n_reps in [5, 50, (10, 5)]:  # type: ignore
         t2aa = np.zeros((nocc_a, nocc_a, nvrt_a, nvrt_a))
         t2ab = np.zeros((nocc_a, nocc_b, nvrt_a, nvrt_b))
         t2bb = np.zeros((nocc_b, nocc_b, nvrt_b, nvrt_b))

tests/python/variational/ucj_spinless_test.py

@@ -78,16 +78,7 @@ def test_parameters_roundtrip_all_to_all():
             n_reps=n_reps,
             with_final_orbital_rotation=with_final_orbital_rotation,
         )
-        np.testing.assert_allclose(
-            roundtripped.diag_coulomb_mats, operator.diag_coulomb_mats
-        )
-        np.testing.assert_allclose(
-            roundtripped.orbital_rotations, operator.orbital_rotations
-        )
-        if with_final_orbital_rotation:
-            np.testing.assert_allclose(
-                roundtripped.final_orbital_rotation, operator.final_orbital_rotation
-            )
+        assert ffsim.approx_eq(roundtripped, operator)
 
 
 def test_parameters_roundtrip_interaction_pairs():
@@ -111,16 +102,7 @@ def test_parameters_roundtrip_interaction_pairs():
             interaction_pairs=interaction_pairs,
             with_final_orbital_rotation=with_final_orbital_rotation,
         )
-        np.testing.assert_allclose(
-            roundtripped.diag_coulomb_mats, operator.diag_coulomb_mats
-        )
-        np.testing.assert_allclose(
-            roundtripped.orbital_rotations, operator.orbital_rotations
-        )
-        if with_final_orbital_rotation:
-            np.testing.assert_allclose(
-                roundtripped.final_orbital_rotation, operator.final_orbital_rotation
-            )
+        assert ffsim.approx_eq(roundtripped, operator)
 
 
 def test_t_amplitudes_energy():