From a74f8e236cb033455e51980a69c7fedadb1e4a3e Mon Sep 17 00:00:00 2001 From: "Kevin J. Sung" Date: Thu, 19 Oct 2023 13:48:00 -0500 Subject: [PATCH] remove hamiltonian_linop --- docs/tutorials/02-orbital-rotation.ipynb | 6 +- python/ffsim/contract/__init__.py | 2 +- python/ffsim/contract/hamiltonian.py | 150 ------------------ python/ffsim/contract/one_body.py | 90 +++++++++++ .../hamiltonians/molecular_hamiltonian.py | 36 +++-- tests/fermion_operator_test.py | 4 +- tests/gates/basic_gates_test.py | 12 +- tests/gates/diag_coulomb_test.py | 8 +- tests/gates/num_op_sum_test.py | 4 +- tests/gates/orbital_rotation_test.py | 12 +- tests/states_test.py | 4 +- tests/trotter/qdrift_test.py | 13 +- tests/wick_test.py | 8 +- 13 files changed, 140 insertions(+), 209 deletions(-) delete mode 100644 python/ffsim/contract/hamiltonian.py create mode 100644 python/ffsim/contract/one_body.py diff --git a/docs/tutorials/02-orbital-rotation.ipynb b/docs/tutorials/02-orbital-rotation.ipynb index 255eba1da..f79e41f8e 100644 --- a/docs/tutorials/02-orbital-rotation.ipynb +++ b/docs/tutorials/02-orbital-rotation.ipynb @@ -118,9 +118,7 @@ "one_body_tensor = ffsim.random.random_hermitian(norb, seed=1234)\n", "\n", "# Convert the Hamiltonian to a LinearOperator\n", - "hamiltonian = ffsim.contract.hamiltonian_linop(\n", - " one_body_tensor=one_body_tensor, norb=norb, nelec=nelec\n", - ")\n", + "hamiltonian = ffsim.contract.one_body_linop(one_body_tensor, norb=norb, nelec=nelec)\n", "\n", "# Get the ground state of the Hamiltonian\n", "eigs, vecs = scipy.sparse.linalg.eigsh(hamiltonian, which=\"LA\", k=1)\n", @@ -206,7 +204,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.11" + "version": "3.10.13" }, "orig_nbformat": 4 }, diff --git a/python/ffsim/contract/__init__.py b/python/ffsim/contract/__init__.py index f3108bd4e..7aef7f04f 100644 --- a/python/ffsim/contract/__init__.py +++ b/python/ffsim/contract/__init__.py @@ -11,5 +11,5 @@ """Functions for contracting tensors and constructing linear operators.""" from ffsim.contract.diag_coulomb import contract_diag_coulomb, diag_coulomb_linop -from ffsim.contract.hamiltonian import hamiltonian_linop, hamiltonian_trace from ffsim.contract.num_op_sum import contract_num_op_sum, num_op_sum_linop +from ffsim.contract.one_body import contract_one_body, one_body_linop diff --git a/python/ffsim/contract/hamiltonian.py b/python/ffsim/contract/hamiltonian.py deleted file mode 100644 index c5e86bc84..000000000 --- a/python/ffsim/contract/hamiltonian.py +++ /dev/null @@ -1,150 +0,0 @@ -# (C) Copyright IBM 2023. -# -# This code is licensed under the Apache License, Version 2.0. You may -# obtain a copy of this license in the LICENSE.txt file in the root directory -# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0. -# -# Any modifications or derivative works of this code must retain this -# copyright notice, and modified files need to carry a notice indicating -# that they have been altered from the originals. - -"""Contract a molecular Hamiltonian.""" - -from __future__ import annotations - -import numpy as np -import scipy.sparse.linalg -from pyscf.fci import cistring -from pyscf.fci.direct_nosym import absorb_h1e, contract_1e, make_hdiag -from pyscf.fci.fci_slow import contract_2e - -from ffsim.states import dim - - -def hamiltonian_trace( - *, - norb: int, - nelec: tuple[int, int], - one_body_tensor: np.ndarray | None = None, - two_body_tensor: np.ndarray | None = None, - constant: float = 0.0, -) -> float: - """Get the trace of the Hamiltonian in the FCI basis.""" - return constant * dim(norb, nelec) + np.sum( - make_hdiag(one_body_tensor, two_body_tensor, norb, nelec) - ) - - -def hamiltonian_linop( - *, - norb: int, - nelec: tuple[int, int], - one_body_tensor: np.ndarray | None = None, - two_body_tensor: np.ndarray | None = None, - constant: float = 0.0, -) -> scipy.sparse.linalg.LinearOperator: - r"""Convert a molecular Hamiltonian to a linear operator. - - A molecular Hamiltonian has the form - - .. math:: - - H = \sum_{pq, \sigma} h_{pq} a^\dagger_{p, \sigma} a_{q, \sigma} - + \frac12 \sum_{pqrs, \sigma} h_{pqrs, \sigma\tau} - a^\dagger_{p, \sigma} a^\dagger_{r, \tau} a_{s, \tau} a_{q, \sigma}. - - Here :math:`h_{pq}` is called the one-body tensor and :math:`h_{pqrs}` is called - the two-body tensor. - - Args: - one_body_tensor: The one-body tensor. - two_body_tensor: The two-body tensor. - norb: The number of spatial orbitals. - nelec: The number of alpha and beta electrons. - constant: The constant. - - Returns: - A LinearOperator that implements the action of the Hamiltonian. - """ - if one_body_tensor is None: - one_body_tensor = np.zeros((norb, norb)) - if two_body_tensor is None: - return _one_body_tensor_linop( - one_body_tensor, norb=norb, nelec=nelec, constant=constant - ) - - n_alpha, n_beta = nelec - linkstr_index_a = cistring.gen_linkstr_index(range(norb), n_alpha) - linkstr_index_b = cistring.gen_linkstr_index(range(norb), n_beta) - link_index = (linkstr_index_a, linkstr_index_b) - two_body = absorb_h1e(one_body_tensor, two_body_tensor, norb, nelec, 0.5) - dim_ = dim(norb, nelec) - - def matvec(vec: np.ndarray): - return constant * vec + contract_2e( - two_body, vec, norb, nelec, link_index=link_index - ) - - return scipy.sparse.linalg.LinearOperator( - shape=(dim_, dim_), matvec=matvec, rmatvec=matvec, dtype=complex - ) - - -def _one_body_tensor_linop( - mat: np.ndarray, norb: int, nelec: tuple[int, int], constant: float = 0.0 -) -> scipy.sparse.linalg.LinearOperator: - r"""Convert a one-body tensor to a linear operator. - - A one-body tensor has the form - - .. math:: - - \sum_{ij} M_{ij} a^\dagger{i} a_j - - where :math:`M` is a complex-valued matrix. - - Args: - mat: The one-body tensor. - norb: The number of spatial orbitals. - nelec: The number of alpha and beta electrons. - - Returns: - A LinearOperator that implements the action of the one-body tensor. - """ - dim_ = dim(norb, nelec) - n_alpha, n_beta = nelec - linkstr_index_a = cistring.gen_linkstr_index(range(norb), n_alpha) - linkstr_index_b = cistring.gen_linkstr_index(range(norb), n_beta) - link_index = (linkstr_index_a, linkstr_index_b) - - def matvec(vec: np.ndarray): - result = constant * vec.astype(complex) - result += contract_1e(mat.real, vec.real, norb, nelec, link_index=link_index) - result += 1j * contract_1e( - mat.imag, vec.real, norb, nelec, link_index=link_index - ) - result += 1j * contract_1e( - mat.real, vec.imag, norb, nelec, link_index=link_index - ) - result -= contract_1e(mat.imag, vec.imag, norb, nelec, link_index=link_index) - return result - - def rmatvec(vec: np.ndarray): - one_body_tensor_H = mat.T.conj() - result = contract_1e( - one_body_tensor_H.real, vec.real, norb, nelec, link_index=link_index - ).astype(complex) - result += 1j * contract_1e( - one_body_tensor_H.imag, vec.real, norb, nelec, link_index=link_index - ) - result += 1j * contract_1e( - one_body_tensor_H.real, vec.imag, norb, nelec, link_index=link_index - ) - result -= contract_1e( - one_body_tensor_H.imag, vec.imag, norb, nelec, link_index=link_index - ) - return result - - return scipy.sparse.linalg.LinearOperator( - shape=(dim_, dim_), matvec=matvec, rmatvec=rmatvec, dtype=complex - ) diff --git a/python/ffsim/contract/one_body.py b/python/ffsim/contract/one_body.py new file mode 100644 index 000000000..7e17973c8 --- /dev/null +++ b/python/ffsim/contract/one_body.py @@ -0,0 +1,90 @@ +# (C) Copyright IBM 2023. +# +# This code is licensed under the Apache License, Version 2.0. You may +# obtain a copy of this license in the LICENSE.txt file in the root directory +# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0. +# +# Any modifications or derivative works of this code must retain this +# copyright notice, and modified files need to carry a notice indicating +# that they have been altered from the originals. + +"""Contract one-body operator.""" + +from __future__ import annotations + +import numpy as np +import scipy.sparse.linalg +from pyscf.fci import cistring +from pyscf.fci.direct_nosym import contract_1e + +from ffsim.states import dim + + +def contract_one_body( + vec: np.ndarray, + mat: np.ndarray, + norb: int, + nelec: tuple[int, int], + *, + linkstr_index_a: np.ndarray, + linkstr_index_b: np.ndarray, +) -> np.ndarray: + link_index = (linkstr_index_a, linkstr_index_b) + result = contract_1e(mat.real, vec.real, norb, nelec, link_index=link_index).astype( + complex + ) + result += 1j * contract_1e(mat.imag, vec.real, norb, nelec, link_index=link_index) + result += 1j * contract_1e(mat.real, vec.imag, norb, nelec, link_index=link_index) + result -= contract_1e(mat.imag, vec.imag, norb, nelec, link_index=link_index) + return result + + +def one_body_linop( + mat: np.ndarray, norb: int, nelec: tuple[int, int] +) -> scipy.sparse.linalg.LinearOperator: + r"""Convert a one-body tensor to a linear operator. + + A one-body tensor has the form + + .. math:: + + \sum_{ij} M_{ij} a^\dagger{i} a_j + + where :math:`M` is a complex-valued matrix. + + Args: + mat: The one-body tensor. + norb: The number of spatial orbitals. + nelec: The number of alpha and beta electrons. + + Returns: + A LinearOperator that implements the action of the one-body tensor. + """ + dim_ = dim(norb, nelec) + n_alpha, n_beta = nelec + linkstr_index_a = cistring.gen_linkstr_index(range(norb), n_alpha) + linkstr_index_b = cistring.gen_linkstr_index(range(norb), n_beta) + + def matvec(vec: np.ndarray): + return contract_one_body( + vec, + mat, + norb=norb, + nelec=nelec, + linkstr_index_a=linkstr_index_a, + linkstr_index_b=linkstr_index_b, + ) + + def rmatvec(vec: np.ndarray): + return contract_one_body( + vec, + mat.T.conj(), + norb=norb, + nelec=nelec, + linkstr_index_a=linkstr_index_a, + linkstr_index_b=linkstr_index_b, + ) + + return scipy.sparse.linalg.LinearOperator( + shape=(dim_, dim_), matvec=matvec, rmatvec=rmatvec, dtype=complex + ) diff --git a/python/ffsim/hamiltonians/molecular_hamiltonian.py b/python/ffsim/hamiltonians/molecular_hamiltonian.py index 266f83414..7cd5d2061 100644 --- a/python/ffsim/hamiltonians/molecular_hamiltonian.py +++ b/python/ffsim/hamiltonians/molecular_hamiltonian.py @@ -13,9 +13,13 @@ import dataclasses import numpy as np +import scipy.sparse.linalg +from pyscf.fci import cistring +from pyscf.fci.direct_nosym import absorb_h1e, make_hdiag +from pyscf.fci.fci_slow import contract_2e from scipy.sparse.linalg import LinearOperator -from ffsim.contract.hamiltonian import hamiltonian_linop, hamiltonian_trace +from ffsim.states import dim @dataclasses.dataclass @@ -51,20 +55,26 @@ def norb(self): def _linear_operator_(self, norb: int, nelec: tuple[int, int]) -> LinearOperator: """Return a SciPy LinearOperator representing the object.""" - return hamiltonian_linop( - norb=norb, - nelec=nelec, - one_body_tensor=self.one_body_tensor, - two_body_tensor=self.two_body_tensor, - constant=self.constant, + n_alpha, n_beta = nelec + linkstr_index_a = cistring.gen_linkstr_index(range(norb), n_alpha) + linkstr_index_b = cistring.gen_linkstr_index(range(norb), n_beta) + link_index = (linkstr_index_a, linkstr_index_b) + two_body = absorb_h1e( + self.one_body_tensor, self.two_body_tensor, norb, nelec, 0.5 + ) + dim_ = dim(norb, nelec) + + def matvec(vec: np.ndarray): + return self.constant * vec + contract_2e( + two_body, vec, norb, nelec, link_index=link_index + ) + + return scipy.sparse.linalg.LinearOperator( + shape=(dim_, dim_), matvec=matvec, rmatvec=matvec, dtype=complex ) def _trace_(self, norb: int, nelec: tuple[int, int]) -> float: """Return the trace of the object.""" - return hamiltonian_trace( - norb=norb, - nelec=nelec, - one_body_tensor=self.one_body_tensor, - two_body_tensor=self.two_body_tensor, - constant=self.constant, + return self.constant * dim(norb, nelec) + np.sum( + make_hdiag(self.one_body_tensor, self.two_body_tensor, norb, nelec) ) diff --git a/tests/fermion_operator_test.py b/tests/fermion_operator_test.py index 8add3a2ea..b5da9bde9 100644 --- a/tests/fermion_operator_test.py +++ b/tests/fermion_operator_test.py @@ -507,8 +507,8 @@ def test_linear_operator(): one_body_tensor = np.zeros((norb, norb)) one_body_tensor[1, 2] = 0.5 one_body_tensor[2, 1] = -0.5 - expected_linop = ffsim.contract.hamiltonian_linop( - one_body_tensor=one_body_tensor, norb=norb, nelec=nelec + expected_linop = ffsim.contract.one_body_linop( + one_body_tensor, norb=norb, nelec=nelec ) vec = ffsim.random.random_statevector(ffsim.dim(norb, nelec), seed=rng) diff --git a/tests/gates/basic_gates_test.py b/tests/gates/basic_gates_test.py index 12d15b36b..51bfdd19e 100644 --- a/tests/gates/basic_gates_test.py +++ b/tests/gates/basic_gates_test.py @@ -96,9 +96,7 @@ def test_apply_givens_rotation(norb: int, nelec: tuple[int, int]): a, b = target_orbs generator[a, b] = theta generator[b, a] = -theta - linop = ffsim.contract.hamiltonian_linop( - one_body_tensor=generator, norb=norb, nelec=nelec - ) + linop = ffsim.contract.one_body_linop(generator, norb=norb, nelec=nelec) expected = scipy.sparse.linalg.expm_multiply( linop, vec, traceA=np.sum(np.abs(generator)) ) @@ -158,9 +156,7 @@ def test_apply_tunneling_interaction(norb: int, nelec: tuple[int, int]): generator = np.zeros((norb, norb)) generator[i, j] = theta generator[j, i] = theta - linop = ffsim.contract.hamiltonian_linop( - one_body_tensor=generator, norb=norb, nelec=nelec - ) + linop = ffsim.contract.one_body_linop(generator, norb=norb, nelec=nelec) expected = scipy.sparse.linalg.expm_multiply( 1j * linop, vec, traceA=np.sum(np.abs(generator)) ) @@ -217,9 +213,7 @@ def test_apply_num_interaction(norb: int, nelec: tuple[int, int]): ) generator = np.zeros((norb, norb)) generator[target_orb, target_orb] = theta - linop = ffsim.contract.hamiltonian_linop( - one_body_tensor=generator, norb=norb, nelec=nelec - ) + linop = ffsim.contract.one_body_linop(generator, norb=norb, nelec=nelec) expected = scipy.sparse.linalg.expm_multiply( 1j * linop, vec, traceA=np.sum(np.abs(generator)) ) diff --git a/tests/gates/diag_coulomb_test.py b/tests/gates/diag_coulomb_test.py index ac745219f..00d746284 100644 --- a/tests/gates/diag_coulomb_test.py +++ b/tests/gates/diag_coulomb_test.py @@ -50,8 +50,8 @@ def test_apply_diag_coulomb_evolution(z_representation: bool): op = ffsim.contract.diag_coulomb_linop( mat, norb=norb, nelec=nelec, z_representation=z_representation ) - orbital_op = ffsim.contract.hamiltonian_linop( - one_body_tensor=scipy.linalg.logm(orbital_rotation), norb=norb, nelec=nelec + orbital_op = ffsim.contract.one_body_linop( + scipy.linalg.logm(orbital_rotation), norb=norb, nelec=nelec ) expected = scipy.sparse.linalg.expm_multiply( -orbital_op, vec, traceA=np.sum(np.abs(orbital_rotation)) @@ -102,8 +102,8 @@ def test_apply_diag_coulomb_evolution_alpha_beta(z_representation: bool): mat_alpha_beta=mat_alpha_beta, z_representation=z_representation, ) - orbital_op = ffsim.contract.hamiltonian_linop( - one_body_tensor=scipy.linalg.logm(orbital_rotation), norb=norb, nelec=nelec + orbital_op = ffsim.contract.one_body_linop( + scipy.linalg.logm(orbital_rotation), norb=norb, nelec=nelec ) expected = scipy.sparse.linalg.expm_multiply( -orbital_op, vec, traceA=np.sum(np.abs(orbital_rotation)) diff --git a/tests/gates/num_op_sum_test.py b/tests/gates/num_op_sum_test.py index 0962c1a1e..29ad28824 100644 --- a/tests/gates/num_op_sum_test.py +++ b/tests/gates/num_op_sum_test.py @@ -78,9 +78,7 @@ def test_apply_quadratic_hamiltonian_evolution(): result = ffsim.apply_num_op_sum_evolution( vec, eigs, time, norb, nelec, orbital_rotation=vecs ) - op = ffsim.contract.hamiltonian_linop( - one_body_tensor=mat, norb=norb, nelec=nelec - ) + op = ffsim.contract.one_body_linop(mat, norb=norb, nelec=nelec) expected = scipy.sparse.linalg.expm_multiply( -1j * time * op, vec, traceA=np.sum(np.abs(mat)) ) diff --git a/tests/gates/orbital_rotation_test.py b/tests/gates/orbital_rotation_test.py index 9cad0bf8e..2afe01ad2 100644 --- a/tests/gates/orbital_rotation_test.py +++ b/tests/gates/orbital_rotation_test.py @@ -41,8 +41,8 @@ def test_apply_orbital_rotation(dtype: type, atol: float): original_vec = vec.copy() result = ffsim.apply_orbital_rotation(vec, mat, norb, nelec) - op = ffsim.contract.hamiltonian_linop( - one_body_tensor=scipy.linalg.logm(mat), norb=norb, nelec=nelec + op = ffsim.contract.one_body_linop( + scipy.linalg.logm(mat), norb=norb, nelec=nelec ) expected = scipy.sparse.linalg.expm_multiply(op, original_vec, traceA=1) np.testing.assert_allclose(result, expected, atol=atol) @@ -86,8 +86,8 @@ def test_apply_orbital_rotation_permutation(dtype: type, atol: float): vec, mat, norb, nelec, allow_col_permutation=True, copy=True ) np.testing.assert_allclose(np.linalg.norm(result), 1, atol=atol) - op = ffsim.contract.hamiltonian_linop( - one_body_tensor=scipy.linalg.logm(mat @ perm), norb=norb, nelec=nelec + op = ffsim.contract.one_body_linop( + scipy.linalg.logm(mat @ perm), norb=norb, nelec=nelec ) expected = scipy.sparse.linalg.expm_multiply(op, original_vec, traceA=1) np.testing.assert_allclose(result, expected, atol=atol) @@ -95,8 +95,8 @@ def test_apply_orbital_rotation_permutation(dtype: type, atol: float): result, perm = ffsim.apply_orbital_rotation( vec, mat, norb, nelec, allow_row_permutation=True, copy=False ) - op = ffsim.contract.hamiltonian_linop( - one_body_tensor=scipy.linalg.logm(perm @ mat), norb=norb, nelec=nelec + op = ffsim.contract.one_body_linop( + scipy.linalg.logm(perm @ mat), norb=norb, nelec=nelec ) expected = scipy.sparse.linalg.expm_multiply(op, original_vec, traceA=1) np.testing.assert_allclose(result, expected, atol=atol) diff --git a/tests/states_test.py b/tests/states_test.py index 5f7e20a3d..3d371fcbf 100644 --- a/tests/states_test.py +++ b/tests/states_test.py @@ -29,7 +29,5 @@ def test_slater_determinant(): norb, occupied_orbitals, orbital_rotation=orbital_rotation ) - hamiltonian = ffsim.contract.hamiltonian_linop( - one_body_tensor=one_body_tensor, norb=norb, nelec=nelec - ) + hamiltonian = ffsim.contract.one_body_linop(one_body_tensor, norb=norb, nelec=nelec) np.testing.assert_allclose(hamiltonian @ state, eig * state) diff --git a/tests/trotter/qdrift_test.py b/tests/trotter/qdrift_test.py index 8f28fea05..dd21c4888 100644 --- a/tests/trotter/qdrift_test.py +++ b/tests/trotter/qdrift_test.py @@ -54,8 +54,8 @@ def variance( def test_spectral_norm_one_body_tensor(norb: int, nelec: tuple[int, int]): """Test spectral norm of one-body operator.""" one_body_tensor = ffsim.random.random_hermitian(norb, seed=8034) - one_body_linop = ffsim.contract.hamiltonian_linop( - one_body_tensor=one_body_tensor, norb=norb, nelec=nelec + one_body_linop = ffsim.contract.one_body_linop( + one_body_tensor, norb=norb, nelec=nelec ) actual = spectral_norm_one_body_tensor(one_body_tensor, nelec=nelec) singular_vals = scipy.sparse.linalg.svds( @@ -127,10 +127,7 @@ def test_one_body_squared_decomposition(norb: int, nelec: tuple[int, int]): ) actual = sum( [ - ffsim.contract.hamiltonian_linop( - one_body_tensor=tensor, norb=norb, nelec=nelec - ) - ** 2 + ffsim.contract.one_body_linop(tensor, norb=norb, nelec=nelec) ** 2 for tensor in one_body_tensors ], start=zero, @@ -158,8 +155,8 @@ def test_variance_one_body_tensor(norb: int, nelec: tuple[int, int]): rng = np.random.default_rng() one_body_tensor = ffsim.random.random_hermitian(norb, seed=rng) - one_body_linop = ffsim.contract.hamiltonian_linop( - one_body_tensor=one_body_tensor, norb=norb, nelec=nelec + one_body_linop = ffsim.contract.one_body_linop( + one_body_tensor, norb=norb, nelec=nelec ) # generate a random Slater determinant diff --git a/tests/wick_test.py b/tests/wick_test.py index 5ea01e6ea..b144345b3 100644 --- a/tests/wick_test.py +++ b/tests/wick_test.py @@ -33,9 +33,7 @@ def test_expectation_product(): for _ in range(n_tensors): one_body_tensor = rng.standard_normal((norb, norb)).astype(complex) one_body_tensor += 1j * rng.standard_normal((norb, norb)) - linop = ffsim.contract.hamiltonian_linop( - one_body_tensor=one_body_tensor, norb=norb, nelec=nelec - ) + linop = ffsim.contract.one_body_linop(one_body_tensor, norb=norb, nelec=nelec) one_body_tensors.append(one_body_tensor) linops.append(linop) @@ -75,9 +73,7 @@ def test_expectation_power(): # generate a random one-body tensor one_body_tensor = rng.standard_normal((norb, norb)).astype(complex) one_body_tensor += 1j * rng.standard_normal((norb, norb)) - linop = ffsim.contract.hamiltonian_linop( - one_body_tensor=one_body_tensor, norb=norb, nelec=nelec - ) + linop = ffsim.contract.one_body_linop(one_body_tensor, norb=norb, nelec=nelec) # generate a random Slater determinant vecs = ffsim.random.random_unitary(norb, seed=rng)