diff --git a/python/ffsim/_lib.pyi b/python/ffsim/_lib.pyi index 65dbf50d8..582882fcc 100644 --- a/python/ffsim/_lib.pyi +++ b/python/ffsim/_lib.pyi @@ -97,6 +97,27 @@ def contract_num_op_sum_spin_into_buffer( occupations: np.ndarray, out: np.ndarray, ) -> None: ... +def contract_fermion_operator_into_buffer( + vec: np.ndarray, + scalar_coeffs: np.ndarray, + alpha_only_src: np.ndarray, + alpha_only_dst: np.ndarray, + alpha_only_scale: np.ndarray, + beta_only_src: np.ndarray, + beta_only_dst: np.ndarray, + beta_only_scale: np.ndarray, + mixed_coeffs: np.ndarray, + mixed_alpha_indptr: np.ndarray, + mixed_alpha_src: np.ndarray, + mixed_alpha_dst: np.ndarray, + mixed_alpha_phase: np.ndarray, + mixed_beta_indptr: np.ndarray, + mixed_beta_src: np.ndarray, + mixed_beta_dst: np.ndarray, + mixed_beta_phase: np.ndarray, + out: np.ndarray, + reverse: bool, +) -> None: ... def jordan_wigner_qiskit( op: FermionOperator, norb: int, tol: float = ... ) -> tuple[list[tuple[str, list[int], complex]], int]: ... diff --git a/python/ffsim/protocols/linear_operator_protocol.py b/python/ffsim/protocols/linear_operator_protocol.py index 495fc3b6b..9fdcbf913 100644 --- a/python/ffsim/protocols/linear_operator_protocol.py +++ b/python/ffsim/protocols/linear_operator_protocol.py @@ -12,12 +12,15 @@ from __future__ import annotations +from dataclasses import dataclass from typing import Any, Protocol import numpy as np -import pyscf.fci +from pyscf.fci import cistring from scipy.sparse.linalg import LinearOperator +from ffsim._cistring import make_strings +from ffsim._lib import contract_fermion_operator_into_buffer from ffsim.operators import FermionOperator from ffsim.states.dimensions import dim, dims @@ -39,6 +42,41 @@ def _linear_operator_( """ +@dataclass(frozen=True) +class _SpinTermData: + src_addrs: np.ndarray + dst_addrs: np.ndarray + phases: np.ndarray + is_identity: bool = False + + +@dataclass(frozen=True) +class _FermionTermData: + coeff: complex + alpha: _SpinTermData + beta: _SpinTermData + + +@dataclass(frozen=True) +class _ContractionData: + scalar_coeffs: np.ndarray + alpha_only_src: np.ndarray + alpha_only_dst: np.ndarray + alpha_only_scale: np.ndarray + beta_only_src: np.ndarray + beta_only_dst: np.ndarray + beta_only_scale: np.ndarray + mixed_coeffs: np.ndarray + mixed_alpha_indptr: np.ndarray + mixed_alpha_src: np.ndarray + mixed_alpha_dst: np.ndarray + mixed_alpha_phase: np.ndarray + mixed_beta_indptr: np.ndarray + mixed_beta_src: np.ndarray + mixed_beta_dst: np.ndarray + mixed_beta_phase: np.ndarray + + def linear_operator( obj: Any, norb: int, nelec: int | tuple[int, int] ) -> LinearOperator: @@ -65,6 +103,19 @@ def linear_operator( def _fermion_operator_to_linear_operator( operator: FermionOperator, norb: int, nelec: int | tuple[int, int] ): + """Return a SciPy LinearOperator representing a FermionOperator. + + The operator is normal ordered and packed into transition data for fast + contraction in a fixed particle-number and spin-Z sector. + + Args: + operator: The FermionOperator to convert to a LinearOperator. + norb: The number of spatial orbitals. + nelec: The number of alpha and beta electrons. + + Returns: + A Scipy LinearOperator representing the FermionOperator. + """ if not (operator.conserves_particle_number() and operator.conserves_spin_z()): raise ValueError( "The given FermionOperator could not be converted to a LinearOperator " @@ -77,81 +128,323 @@ def _fermion_operator_to_linear_operator( nelec = (nelec, 0) dim_ = dim(norb, nelec) - adjoint = operator.adjoint() + dims_ = dims(norb, nelec) + # Drop terms outside the active orbital space. + bounded_operator = FermionOperator( + { + term: coeff + for term, coeff in operator.items() + if all(0 <= orb < norb for _, _, orb in term) + } + ) + fermion_term_data: list[_FermionTermData] = [] + # Split each term into alpha and beta transition data. + for term, coeff in bounded_operator.normal_ordered().items(): + if not coeff: + continue + + alpha_term = [] + beta_term = [] + beta_count = 0 + inversions = 0 + for action, spin, orb in term: + if spin: + beta_term.append((action, spin, orb)) + beta_count += 1 + else: + alpha_term.append((action, spin, orb)) + inversions += beta_count + if inversions & 1: + coeff = -coeff + fermion_term_data.append( + _FermionTermData( + coeff=coeff, + alpha=_make_spin_term_data(tuple(alpha_term), norb=norb, nocc=nelec[0]), + beta=_make_spin_term_data(tuple(beta_term), norb=norb, nocc=nelec[1]), + ) + ) + scalar_coeff = 0j + alpha_only_src_parts = [] + alpha_only_dst_parts = [] + alpha_only_scale_parts = [] + beta_only_src_parts = [] + beta_only_dst_parts = [] + beta_only_scale_parts = [] + mixed_terms = [] + + # Group terms by the spin sectors on which they act. + for term_data in fermion_term_data: + alpha_is_identity = term_data.alpha.is_identity + beta_is_identity = term_data.beta.is_identity + alpha_len = len(term_data.alpha.src_addrs) + beta_len = len(term_data.beta.src_addrs) + + if alpha_is_identity and beta_is_identity: + scalar_coeff += term_data.coeff + continue + if not alpha_is_identity and not alpha_len: + continue + if not beta_is_identity and not beta_len: + continue + + if beta_is_identity: + alpha_only_src_parts.append(term_data.alpha.src_addrs) + alpha_only_dst_parts.append(term_data.alpha.dst_addrs) + alpha_only_scale_parts.append(term_data.coeff * term_data.alpha.phases) + elif alpha_is_identity: + beta_only_src_parts.append(term_data.beta.src_addrs) + beta_only_dst_parts.append(term_data.beta.dst_addrs) + beta_only_scale_parts.append(term_data.coeff * term_data.beta.phases) + else: + mixed_terms.append(term_data) + + alpha_only_src, alpha_only_dst, alpha_only_scale = _combine_scaled_transitions( + alpha_only_src_parts, + alpha_only_dst_parts, + alpha_only_scale_parts, + ) + beta_only_src, beta_only_dst, beta_only_scale = _combine_scaled_transitions( + beta_only_src_parts, + beta_only_dst_parts, + beta_only_scale_parts, + ) + + # Pack mixed terms into flat arrays with per-term offsets. + n_mixed_terms = len(mixed_terms) + mixed_coeffs = np.empty(n_mixed_terms, dtype=np.complex128) + mixed_alpha_indptr = np.zeros(n_mixed_terms + 1, dtype=np.uintp) + mixed_beta_indptr = np.zeros(n_mixed_terms + 1, dtype=np.uintp) + + for index, term_data in enumerate(mixed_terms): + mixed_coeffs[index] = term_data.coeff + mixed_alpha_indptr[index + 1] = mixed_alpha_indptr[index] + len( + term_data.alpha.src_addrs + ) + mixed_beta_indptr[index + 1] = mixed_beta_indptr[index] + len( + term_data.beta.src_addrs + ) + + mixed_alpha_src = np.empty(mixed_alpha_indptr[-1], dtype=np.uintp) + mixed_alpha_dst = np.empty(mixed_alpha_indptr[-1], dtype=np.uintp) + mixed_alpha_phase = np.empty(mixed_alpha_indptr[-1], dtype=np.int8) + mixed_beta_src = np.empty(mixed_beta_indptr[-1], dtype=np.uintp) + mixed_beta_dst = np.empty(mixed_beta_indptr[-1], dtype=np.uintp) + mixed_beta_phase = np.empty(mixed_beta_indptr[-1], dtype=np.int8) + + for index, term_data in enumerate(mixed_terms): + alpha_start = mixed_alpha_indptr[index] + alpha_end = mixed_alpha_indptr[index + 1] + beta_start = mixed_beta_indptr[index] + beta_end = mixed_beta_indptr[index + 1] + + mixed_alpha_src[alpha_start:alpha_end] = term_data.alpha.src_addrs + mixed_alpha_dst[alpha_start:alpha_end] = term_data.alpha.dst_addrs + mixed_alpha_phase[alpha_start:alpha_end] = term_data.alpha.phases + mixed_beta_src[beta_start:beta_end] = term_data.beta.src_addrs + mixed_beta_dst[beta_start:beta_end] = term_data.beta.dst_addrs + mixed_beta_phase[beta_start:beta_end] = term_data.beta.phases + + contraction_data = _ContractionData( + scalar_coeffs=( + np.array([scalar_coeff], dtype=np.complex128) + if scalar_coeff + else np.empty(0, dtype=np.complex128) + ), + alpha_only_src=alpha_only_src, + alpha_only_dst=alpha_only_dst, + alpha_only_scale=alpha_only_scale, + beta_only_src=beta_only_src, + beta_only_dst=beta_only_dst, + beta_only_scale=beta_only_scale, + mixed_coeffs=mixed_coeffs, + mixed_alpha_indptr=mixed_alpha_indptr, + mixed_alpha_src=mixed_alpha_src, + mixed_alpha_dst=mixed_alpha_dst, + mixed_alpha_phase=mixed_alpha_phase, + mixed_beta_indptr=mixed_beta_indptr, + mixed_beta_src=mixed_beta_src, + mixed_beta_dst=mixed_beta_dst, + mixed_beta_phase=mixed_beta_phase, + ) def matvec(vec: np.ndarray): - result = np.zeros(dim_, dtype=complex) - for term, coeff in operator.items(): - result += coeff * _apply_fermion_term(vec, term, norb, nelec) - return result + return _apply_fermion_terms( + vec, contraction_data=contraction_data, dims_=dims_, reverse=False + ) def rmatvec(vec: np.ndarray): - result = np.zeros(dim_, dtype=complex) - for term, coeff in adjoint.items(): - result += coeff * _apply_fermion_term(vec, term, norb, nelec) - return result + return _apply_fermion_terms( + vec, contraction_data=contraction_data, dims_=dims_, reverse=True + ) return LinearOperator( shape=(dim_, dim_), matvec=matvec, rmatvec=rmatvec, dtype=complex ) -def _apply_fermion_term( - vec: np.ndarray, - term: tuple[tuple[bool, bool, int], ...], - norb: int, - nelec: tuple[int, int], -) -> np.ndarray: - """Apply a product of ladder operators to a state vector. +def _make_spin_term_data( + term: tuple[tuple[bool, bool, int], ...], norb: int, nocc: int +) -> _SpinTermData: + """Return transition data for a spin block of a FermionOperator term. + + The empty term is treated as the identity. + + Args: + term: The spin block of the FermionOperator term. + norb: The number of spatial orbitals. + nocc: The number of occupied orbitals in the spin sector. + + Returns: + The source addresses, destination addresses, and fermionic phases of the + spin block. + """ + if not term: + empty = np.array([], dtype=np.intp) + return _SpinTermData( + src_addrs=empty.astype(np.uintp, copy=False), + dst_addrs=empty.astype(np.uintp, copy=False), + phases=np.array([], dtype=np.int8), + is_identity=True, + ) + + creation_mask = 0 + annihilation_mask = 0 + for action, _, orb in term: + if action: + creation_mask |= 1 << orb + else: + annihilation_mask |= 1 << orb + creation_only_mask = creation_mask & ~annihilation_mask + + src_addrs = [] + dst_strings = [] + phases = [] + # prefilter source strings + for src_addr, string in enumerate(make_strings(range(norb), nocc)): + string = int(string) + if string & annihilation_mask != annihilation_mask: + continue + if string & creation_only_mask: + continue + + phase = 1 + transformed = string + for action, _, orb in reversed(term): + mask = 1 << orb + # accumulate the fermionic phase + if (transformed & (mask - 1)).bit_count() & 1: + phase = -phase + if action: + if transformed & mask: + break + transformed |= mask + else: + if not transformed & mask: + break + transformed ^= mask + else: + src_addrs.append(src_addr) + dst_strings.append(transformed) + phases.append(phase) + + if not src_addrs: + empty = np.array([], dtype=np.uintp) + return _SpinTermData( + src_addrs=empty, + dst_addrs=empty, + phases=np.array([], dtype=np.int8), + ) + + dst_addrs = cistring.strs2addr( + norb=norb, nelec=nocc, strings=np.asarray(dst_strings, dtype=np.int64) + ).astype(np.uintp, copy=False) + return _SpinTermData( + src_addrs=np.asarray(src_addrs, dtype=np.uintp), + dst_addrs=dst_addrs, + phases=np.asarray(phases, dtype=np.int8), + ) - Given a state vector and a string of ladder operators that conserves particle number - and the z component of spin, return the state vector that results from applying the - ladder operators to the given state vector. The string of ladder operators is - represented as a sequence of (``action``, ``spin``, ``orbital``) tuples, where: - - ``action`` is a bool. False indicates a destruction operator and True indicates - a creation operator. - - ``spin`` is a bool. False indicates spin alpha and True indicates spin beta. - - ``orbital`` is an integer giving the index of the spatial orbital to act on. +def _combine_scaled_transitions( + src_parts: list[np.ndarray], + dst_parts: list[np.ndarray], + scale_parts: list[np.ndarray], +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Combine scaled transition data into canonical flat arrays. - The string of ladder operators acts on a state vector by left multiplication, - so the resulting state vector is obtained by applying the ladder operators to the - initial vector in the reverse order in which they are given. + Each entry represents a contribution from a source address to a destination + address, scaled by the corresponding coefficient. Repeated source-destination + pairs are summed, and pairs with zero total scale are removed. Args: - vec: The state vector. - term: The product of ladder operators to apply to the state vector. - norb: The number of spatial orbitals. - nelec: The number of alpha and beta electrons. + src_parts: The source address arrays to combine. + dst_parts: The destination address arrays to combine. + scale_parts: The scale arrays to combine. Returns: - The result of applying the ladder operators to the state vector. + The combined source addresses, destination addresses, and scales. """ - result = _apply_fermion_term_real(vec.real, term, norb, nelec) - result += 1j * _apply_fermion_term_real(vec.imag, term, norb, nelec) - return result + if not src_parts: + return ( + np.empty(0, dtype=np.uintp), + np.empty(0, dtype=np.uintp), + np.empty(0, dtype=np.complex128), + ) + + src = np.concatenate(src_parts).astype(np.uintp, copy=False) + dst = np.concatenate(dst_parts).astype(np.uintp, copy=False) + scale = np.concatenate(scale_parts).astype(np.complex128, copy=False) + order = np.lexsort((dst, src)) + src = src[order] + dst = dst[order] + scale = scale[order] -def _apply_fermion_term_real( + change = np.empty(len(src), dtype=bool) + change[0] = True + change[1:] = (src[1:] != src[:-1]) | (dst[1:] != dst[:-1]) + starts = np.flatnonzero(change) + src = src[starts] + dst = dst[starts] + scale = np.add.reduceat(scale, starts) + + nonzero = scale != 0 + src = src[nonzero] + dst = dst[nonzero] + scale = scale[nonzero] + + order = np.lexsort((src, dst)) + return src[order], dst[order], scale[order] + + +def _apply_fermion_terms( vec: np.ndarray, - term: tuple[tuple[bool, bool, int], ...], - norb: int, - nelec: tuple[int, int], + contraction_data: _ContractionData, + dims_: tuple[int, int], + *, + reverse: bool, ) -> np.ndarray: - action_funcs = { - # key: (action, spin) - (False, False): pyscf.fci.addons.des_a, - (False, True): pyscf.fci.addons.des_b, - (True, False): pyscf.fci.addons.cre_a, - (True, True): pyscf.fci.addons.cre_b, - } - (dim_a, dim_b) = dims(norb, nelec) - transformed = vec.reshape((dim_a, dim_b)) - this_nelec = list(nelec) - for action, spin, orb in reversed(term): - action_func = action_funcs[(action, spin)] - transformed = action_func(transformed, norb, this_nelec, orb) - this_nelec[spin] += 1 if action else -1 - if this_nelec[spin] < 0 or this_nelec[spin] > norb: - return np.zeros(dim_a * dim_b, dtype=complex) - return transformed.reshape(-1).astype(complex, copy=False) + transformed = np.ascontiguousarray(vec, dtype=complex).reshape(dims_) + result = np.zeros(dims_, dtype=complex) + contract_fermion_operator_into_buffer( + transformed, + contraction_data.scalar_coeffs, + contraction_data.alpha_only_src, + contraction_data.alpha_only_dst, + contraction_data.alpha_only_scale, + contraction_data.beta_only_src, + contraction_data.beta_only_dst, + contraction_data.beta_only_scale, + contraction_data.mixed_coeffs, + contraction_data.mixed_alpha_indptr, + contraction_data.mixed_alpha_src, + contraction_data.mixed_alpha_dst, + contraction_data.mixed_alpha_phase, + contraction_data.mixed_beta_indptr, + contraction_data.mixed_beta_src, + contraction_data.mixed_beta_dst, + contraction_data.mixed_beta_phase, + result, + reverse=reverse, + ) + return result.reshape(-1) diff --git a/src/contract/fermion_operator.rs b/src/contract/fermion_operator.rs new file mode 100644 index 000000000..700c7619b --- /dev/null +++ b/src/contract/fermion_operator.rs @@ -0,0 +1,232 @@ +// (C) Copyright IBM 2026 +// +// 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. + +use numpy::Complex64; +use numpy::PyReadonlyArray1; +use numpy::PyReadonlyArray2; +use numpy::PyReadwriteArray2; +use pyo3::prelude::*; + +#[inline] +fn src_dst(src: &[usize], dst: &[usize], index: usize, reverse: bool) -> (usize, usize) { + if reverse { + (dst[index], src[index]) + } else { + (src[index], dst[index]) + } +} + +struct ScaledTransitions<'a> { + src: &'a [usize], + dst: &'a [usize], + scale: &'a [Complex64], +} + +struct PhasedTransitions<'a> { + indptr: &'a [usize], + src: &'a [usize], + dst: &'a [usize], + phase: &'a [i8], +} + +#[inline] +fn directed_scale(scale: Complex64, reverse: bool) -> Complex64 { + if reverse { + scale.conj() + } else { + scale + } +} + +fn contract_scalar_terms( + vec: &[Complex64], + scalar_coeffs: &[Complex64], + out: &mut [Complex64], + reverse: bool, +) { + for &scale in scalar_coeffs { + let scale = directed_scale(scale, reverse); + for index in 0..out.len() { + out[index] += scale * vec[index]; + } + } +} + +fn contract_alpha_only_terms( + vec: &[Complex64], + out: &mut [Complex64], + dim_b: usize, + transitions: &ScaledTransitions, + reverse: bool, +) { + for index in 0..transitions.scale.len() { + let (source_alpha, target_alpha) = + src_dst(transitions.src, transitions.dst, index, reverse); + let scale = directed_scale(transitions.scale[index], reverse); + let source_offset = source_alpha * dim_b; + let target_offset = target_alpha * dim_b; + for beta_index in 0..dim_b { + out[target_offset + beta_index] += scale * vec[source_offset + beta_index]; + } + } +} + +fn contract_beta_only_terms( + vec: &[Complex64], + out: &mut [Complex64], + dim_a: usize, + dim_b: usize, + transitions: &ScaledTransitions, + reverse: bool, +) { + for alpha_index in 0..dim_a { + let row_offset = alpha_index * dim_b; + for index in 0..transitions.scale.len() { + let (source_beta, target_beta) = + src_dst(transitions.src, transitions.dst, index, reverse); + let scale = directed_scale(transitions.scale[index], reverse); + out[row_offset + target_beta] += scale * vec[row_offset + source_beta]; + } + } +} + +// Mixed terms apply: +// out[t_alpha, t_beta] += coeff * p_alpha * p_beta * vec[s_alpha, s_beta]. +fn contract_mixed_terms( + vec: &[Complex64], + out: &mut [Complex64], + dim_b: usize, + coeffs: &[Complex64], + alpha: &PhasedTransitions, + beta: &PhasedTransitions, + reverse: bool, +) { + for (term_index, &coeff) in coeffs.iter().enumerate() { + let coeff = directed_scale(coeff, reverse); + + let alpha_start = alpha.indptr[term_index]; + let alpha_end = alpha.indptr[term_index + 1]; + let beta_start = beta.indptr[term_index]; + let beta_end = beta.indptr[term_index + 1]; + + if alpha_end - alpha_start <= beta_end - beta_start { + for (alpha_offset, &alpha_phase_value) in + alpha.phase[alpha_start..alpha_end].iter().enumerate() + { + let alpha_index = alpha_start + alpha_offset; + let (source_alpha, target_alpha) = + src_dst(alpha.src, alpha.dst, alpha_index, reverse); + let alpha_scale = coeff * Complex64::new(alpha_phase_value as f64, 0.0); + let source_alpha_offset = source_alpha * dim_b; + let target_alpha_offset = target_alpha * dim_b; + for (beta_offset, &beta_phase_value) in + beta.phase[beta_start..beta_end].iter().enumerate() + { + let beta_index = beta_start + beta_offset; + let (source_beta, target_beta) = + src_dst(beta.src, beta.dst, beta_index, reverse); + let scale = alpha_scale * Complex64::new(beta_phase_value as f64, 0.0); + out[target_alpha_offset + target_beta] += + scale * vec[source_alpha_offset + source_beta]; + } + } + continue; + } + + for (beta_offset, &beta_phase_value) in beta.phase[beta_start..beta_end].iter().enumerate() + { + let beta_index = beta_start + beta_offset; + let (source_beta, target_beta) = src_dst(beta.src, beta.dst, beta_index, reverse); + let beta_scale = coeff * Complex64::new(beta_phase_value as f64, 0.0); + for (alpha_offset, &alpha_phase_value) in + alpha.phase[alpha_start..alpha_end].iter().enumerate() + { + let alpha_index = alpha_start + alpha_offset; + let (source_alpha, target_alpha) = + src_dst(alpha.src, alpha.dst, alpha_index, reverse); + let scale = beta_scale * Complex64::new(alpha_phase_value as f64, 0.0); + out[target_alpha * dim_b + target_beta] += + scale * vec[source_alpha * dim_b + source_beta]; + } + } + } +} + +/// Contract a packed FermionOperator into a buffer. +#[allow(clippy::too_many_arguments)] +#[pyfunction] +pub fn contract_fermion_operator_into_buffer( + vec: PyReadonlyArray2, + scalar_coeffs: PyReadonlyArray1, + alpha_only_src: PyReadonlyArray1, + alpha_only_dst: PyReadonlyArray1, + alpha_only_scale: PyReadonlyArray1, + beta_only_src: PyReadonlyArray1, + beta_only_dst: PyReadonlyArray1, + beta_only_scale: PyReadonlyArray1, + mixed_coeffs: PyReadonlyArray1, + mixed_alpha_indptr: PyReadonlyArray1, + mixed_alpha_src: PyReadonlyArray1, + mixed_alpha_dst: PyReadonlyArray1, + mixed_alpha_phase: PyReadonlyArray1, + mixed_beta_indptr: PyReadonlyArray1, + mixed_beta_src: PyReadonlyArray1, + mixed_beta_dst: PyReadonlyArray1, + mixed_beta_phase: PyReadonlyArray1, + mut out: PyReadwriteArray2, + reverse: bool, +) { + let vec = vec.as_array(); + let shape = vec.shape(); + let dim_a = shape[0]; + let dim_b = shape[1]; + let vec = vec.as_slice().unwrap(); + + let mut out = out.as_array_mut(); + let out = out.as_slice_mut().unwrap(); + + let scalar_coeffs = scalar_coeffs.as_slice().unwrap(); + let alpha_only = ScaledTransitions { + src: alpha_only_src.as_slice().unwrap(), + dst: alpha_only_dst.as_slice().unwrap(), + scale: alpha_only_scale.as_slice().unwrap(), + }; + let beta_only = ScaledTransitions { + src: beta_only_src.as_slice().unwrap(), + dst: beta_only_dst.as_slice().unwrap(), + scale: beta_only_scale.as_slice().unwrap(), + }; + let mixed_coeffs = mixed_coeffs.as_slice().unwrap(); + let mixed_alpha = PhasedTransitions { + indptr: mixed_alpha_indptr.as_slice().unwrap(), + src: mixed_alpha_src.as_slice().unwrap(), + dst: mixed_alpha_dst.as_slice().unwrap(), + phase: mixed_alpha_phase.as_slice().unwrap(), + }; + let mixed_beta = PhasedTransitions { + indptr: mixed_beta_indptr.as_slice().unwrap(), + src: mixed_beta_src.as_slice().unwrap(), + dst: mixed_beta_dst.as_slice().unwrap(), + phase: mixed_beta_phase.as_slice().unwrap(), + }; + + contract_scalar_terms(vec, scalar_coeffs, out, reverse); + contract_alpha_only_terms(vec, out, dim_b, &alpha_only, reverse); + contract_beta_only_terms(vec, out, dim_a, dim_b, &beta_only, reverse); + contract_mixed_terms( + vec, + out, + dim_b, + mixed_coeffs, + &mixed_alpha, + &mixed_beta, + reverse, + ); +} diff --git a/src/contract/mod.rs b/src/contract/mod.rs index ec57b57d0..2a75fdb6e 100644 --- a/src/contract/mod.rs +++ b/src/contract/mod.rs @@ -9,4 +9,5 @@ // that they have been altered from the originals. pub mod diag_coulomb; +pub mod fermion_operator; pub mod num_op_sum; diff --git a/src/lib.rs b/src/lib.rs index 6b167e531..58c4434da 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -56,6 +56,10 @@ fn _lib(m: &Bound<'_, PyModule>) -> PyResult<()> { contract::num_op_sum::contract_num_op_sum_spin_into_buffer, m )?)?; + m.add_function(wrap_pyfunction!( + contract::fermion_operator::contract_fermion_operator_into_buffer, + m + )?)?; m.add_function(pyo3::wrap_pyfunction!( jordan_wigner::jordan_wigner_qiskit, m diff --git a/tests/python/operators/fermion_operator_test.py b/tests/python/operators/fermion_operator_test.py index 41e16494e..6c763a00b 100644 --- a/tests/python/operators/fermion_operator_test.py +++ b/tests/python/operators/fermion_operator_test.py @@ -12,11 +12,15 @@ from __future__ import annotations +from typing import cast + import numpy as np import pytest import ffsim from ffsim import FermionOperator +from ffsim._slow.fermion_operator import FermionOperator as SlowFermionOperator +from ffsim.operators.fermion_action import FermionAction RNG = np.random.default_rng(308444818162923832831691142041679886023) @@ -620,6 +624,66 @@ def test_linear_operator_one_body(): np.testing.assert_allclose(actual, expected) +def test_linear_operator(): + """Test linear operator.""" + norb = 5 + nelec = (2, 2) + op = FermionOperator( + { + (): 0.125 - 0.25j, + (ffsim.cre_a(3), ffsim.des_a(0)): 0.75 + 0.125j, + (ffsim.cre_b(4), ffsim.des_b(1)): -0.5 + 0.375j, + ( + ffsim.cre_a(2), + ffsim.des_a(1), + ffsim.cre_b(3), + ffsim.des_b(0), + ): -0.125 + 0.5j, + (ffsim.des_a(2), ffsim.cre_a(4)): 0.25j, + } + ) + vec = ffsim.random.random_state_vector(ffsim.dim(norb, nelec), seed=2468) + original = vec.copy() + linop = ffsim.linear_operator(op, norb=norb, nelec=nelec) + slow_linop = SlowFermionOperator( + cast(dict[tuple[FermionAction, ...], complex], dict(op.items())) + )._linear_operator_(norb, nelec) + slow_adjoint_linop = SlowFermionOperator( + cast(dict[tuple[FermionAction, ...], complex], dict(op.adjoint().items())) + )._linear_operator_( + norb, + nelec, + ) + + np.testing.assert_allclose(linop @ vec, slow_linop @ vec, atol=1e-14) + np.testing.assert_allclose( + linop.adjoint() @ vec, slow_adjoint_linop @ vec, atol=1e-14 + ) + np.testing.assert_allclose(original, vec) + + +def test_linear_operator_out_of_range(): + """Test linear operator with out-of-range terms.""" + norb = 4 + nelec = (2, 1) + valid_term = (ffsim.cre_a(2), ffsim.des_a(0)) + op = FermionOperator( + { + valid_term: 0.75 - 0.125j, + (ffsim.des_b(norb), ffsim.cre_b(norb)): -1.5 + 0.25j, + (ffsim.des_a(-1), ffsim.cre_a(-1)): -0.75 + 0.5j, + } + ) + expected_op = FermionOperator({valid_term: 0.75 - 0.125j}) + vec = ffsim.random.random_state_vector(ffsim.dim(norb, nelec), seed=24680) + + linop = ffsim.linear_operator(op, norb=norb, nelec=nelec) + expected_linop = ffsim.linear_operator(expected_op, norb=norb, nelec=nelec) + + np.testing.assert_allclose(linop @ vec, expected_linop @ vec) + np.testing.assert_allclose(linop.adjoint() @ vec, expected_linop.adjoint() @ vec) + + def test_approx_eq(): """Test approximate equality.""" op1 = FermionOperator(