From a51163622e4e65393d967aa2bc4655ed523aabb0 Mon Sep 17 00:00:00 2001 From: "Kevin J. Sung" Date: Wed, 30 Oct 2024 06:59:22 -0400 Subject: [PATCH] enable check-untyped-defs mypy option (#337) Co-authored-by: Bart Andrews <38833761+bartandrews@users.noreply.github.com> --- pyproject.toml | 1 + python/ffsim/_lib.pyi | 7 +++-- python/ffsim/qiskit/gates/diag_coulomb.py | 4 +-- python/ffsim/variational/givens.py | 2 +- python/ffsim/variational/hopgate.py | 23 ++++++++++++++ python/ffsim/variational/num_num.py | 4 ++- tests/python/molecular_data_test.py | 5 ++++ .../python/operators/fermion_operator_test.py | 2 +- tests/python/states/bitstring_test.py | 10 +++---- tests/python/states/slater_test.py | 4 +-- tests/python/variational/hopgate_test.py | 6 +--- tests/python/variational/num_num_test.py | 4 ++- tests/python/variational/uccsd_test.py | 7 +---- .../variational/ucj_spin_balanced_test.py | 22 ++------------ .../variational/ucj_spin_unbalanced_test.py | 30 +++++-------------- tests/python/variational/ucj_spinless_test.py | 22 ++------------ 16 files changed, 63 insertions(+), 90 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 0bd14acd7..e78607d3c 100644 --- a/pyproject.toml +++ b/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"] diff --git a/python/ffsim/_lib.pyi b/python/ffsim/_lib.pyi index 7b54fb579..de736c70d 100644 --- a/python/ffsim/_lib.pyi +++ b/python/ffsim/_lib.pyi @@ -1,8 +1,8 @@ -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: ... @@ -10,7 +10,10 @@ class FermionOperator(Mapping[tuple[tuple[bool, bool, int], ...], complex]): 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: ... diff --git a/python/ffsim/qiskit/gates/diag_coulomb.py b/python/ffsim/qiskit/gates/diag_coulomb.py index 1cf1982b2..7df0b63a7 100644 --- a/python/ffsim/qiskit/gates/diag_coulomb.py +++ b/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( diff --git a/python/ffsim/variational/givens.py b/python/ffsim/variational/givens.py index 9ad5a4a33..e99b5543e 100644 --- a/python/ffsim/variational/givens.py +++ b/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( diff --git a/python/ffsim/variational/hopgate.py b/python/ffsim/variational/hopgate.py index 0b4cded26..1e07c9fda 100644 --- a/python/ffsim/variational/hopgate.py +++ b/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 diff --git a/python/ffsim/variational/num_num.py b/python/ffsim/variational/num_num.py index a548e132a..2aa7df6e5 100644 --- a/python/ffsim/variational/num_num.py +++ b/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: diff --git a/tests/python/molecular_data_test.py b/tests/python/molecular_data_test.py index f7aab426e..731095c03 100644 --- a/tests/python/molecular_data_test.py +++ b/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) diff --git a/tests/python/operators/fermion_operator_test.py b/tests/python/operators/fermion_operator_test.py index f9b9f6bfc..0e36c3939 100644 --- a/tests/python/operators/fermion_operator_test.py +++ b/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(): diff --git a/tests/python/states/bitstring_test.py b/tests/python/states/bitstring_test.py index 6edfcb74e..47d75e695 100644 --- a/tests/python/states/bitstring_test.py +++ b/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( [ diff --git a/tests/python/states/slater_test.py b/tests/python/states/slater_test.py index e7f120705..ef0c13bb7 100644 --- a/tests/python/states/slater_test.py +++ b/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 diff --git a/tests/python/variational/hopgate_test.py b/tests/python/variational/hopgate_test.py index 25b81e157..5cf846e90 100644 --- a/tests/python/variational/hopgate_test.py +++ b/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) diff --git a/tests/python/variational/num_num_test.py b/tests/python/variational/num_num_test.py index 7ecbed1c7..7548f9921 100644 --- a/tests/python/variational/num_num_test.py +++ b/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 ) diff --git a/tests/python/variational/uccsd_test.py b/tests/python/variational/uccsd_test.py index 1c6a58ba1..ef04df901 100644 --- a/tests/python/variational/uccsd_test.py +++ b/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(): diff --git a/tests/python/variational/ucj_spin_balanced_test.py b/tests/python/variational/ucj_spin_balanced_test.py index cbd839fd3..1b88adf56 100644 --- a/tests/python/variational/ucj_spin_balanced_test.py +++ b/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(): diff --git a/tests/python/variational/ucj_spin_unbalanced_test.py b/tests/python/variational/ucj_spin_unbalanced_test.py index 5045411ce..39b7cdba8 100644 --- a/tests/python/variational/ucj_spin_unbalanced_test.py +++ b/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)) diff --git a/tests/python/variational/ucj_spinless_test.py b/tests/python/variational/ucj_spinless_test.py index 0ed27c546..2aaee174e 100644 --- a/tests/python/variational/ucj_spinless_test.py +++ 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():