diff --git a/python/ffsim/slow.py b/python/ffsim/slow.py deleted file mode 100644 index e77846bb2..000000000 --- a/python/ffsim/slow.py +++ /dev/null @@ -1,311 +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. - -"""Python versions of functions that were rewritten in Rust.""" - -from __future__ import annotations - -import itertools - -import numpy as np - -from ffsim.gates.gates import _apply_phase_shift - - -def apply_givens_rotation_in_place_slow( - vec: np.ndarray, - c: float, - s: float, - phase: complex, - slice1: np.ndarray, - slice2: np.ndarray, -) -> None: - """Apply a Givens rotation to slices of a state vector.""" - phase_conj = phase.conjugate() - for i, j in zip(slice1, slice2): - vec[i] *= phase_conj - tmp = c * vec[i] + s * vec[j] - vec[j] = c * vec[j] - s * vec[i] - vec[i] = tmp - vec[i] *= phase - - -def gen_orbital_rotation_index_in_place_slow( - norb: int, - nocc: int, - linkstr_index: np.ndarray, - diag_strings: np.ndarray, - off_diag_strings: np.ndarray, - off_diag_strings_index: np.ndarray, - off_diag_index: np.ndarray, -) -> None: - """Generate orbital rotation index.""" - diag_counter = np.zeros(norb, dtype=np.uint) - off_diag_counter = np.zeros(norb, dtype=np.uint) - for str0, tab in enumerate(linkstr_index[:, :, 0]): - for orb in tab[:nocc]: - count = diag_counter[orb] - diag_strings[orb, count] = str0 - diag_counter[orb] += 1 - for orb in tab[nocc:norb]: - count = off_diag_counter[orb] - off_diag_strings[orb, count] = str0 - off_diag_strings_index[orb, str0] = count - off_diag_counter[orb] += 1 - - index_counter = np.zeros_like(off_diag_strings) - for str0, tab in enumerate(linkstr_index): - for orb_c, orb_d, str1, sign in tab[nocc:]: - # str0 -> annihilate orb_d -> create orb_c -> str1 - index = off_diag_strings_index[orb_c, str0] - count = index_counter[orb_c, index] - off_diag_index[orb_c, index, count] = orb_d, str1, sign - index_counter[orb_c, index] += 1 - - -def apply_single_column_transformation_in_place_slow( - vec: np.ndarray, - column: np.ndarray, - diag_val: complex, - diag_strings: np.ndarray, - off_diag_strings: np.ndarray, - off_diag_index: np.ndarray, -) -> None: - """Apply a single-column orbital rotation.""" - for str0, tab in zip(off_diag_strings, off_diag_index): - for orb, str1, sign in tab: - vec[str0] += sign * column[orb] * vec[str1] - for str0 in diag_strings: - vec[str0] *= diag_val - - -def apply_num_op_sum_evolution_in_place_slow( - vec: np.ndarray, phases: np.ndarray, occupations: np.ndarray -) -> None: - """Apply time evolution by a sum of number operators in-place.""" - for row, orbs in zip(vec, occupations): - phase = 1 - for orb in orbs: - phase *= phases[orb] - row *= phase - - -def apply_diag_coulomb_evolution_in_place_num_rep_slow( - vec: np.ndarray, - mat_exp: np.ndarray, - norb: int, - mat_alpha_beta_exp: np.ndarray, - occupations_a: np.ndarray, - occupations_b: np.ndarray, -) -> None: - """Apply time evolution by a diagonal Coulomb operator in-place.""" - dim_a, dim_b = vec.shape - alpha_phases = np.empty((dim_a,), dtype=complex) - beta_phases = np.empty((dim_b,), dtype=complex) - phase_map = np.ones((dim_a, norb), dtype=complex) - - for i, orbs in enumerate(occupations_b): - phase = 1 - for orb_1, orb_2 in itertools.combinations_with_replacement(orbs, 2): - phase *= mat_exp[orb_1, orb_2] - beta_phases[i] = phase - - for i, (row, orbs) in enumerate(zip(phase_map, occupations_a)): - phase = 1 - for j in range(len(orbs)): - row *= mat_alpha_beta_exp[orbs[j]] - for k in range(j, len(orbs)): - phase *= mat_exp[orbs[j], orbs[k]] - alpha_phases[i] = phase - - for row, alpha_phase, phase_map in zip(vec, alpha_phases, phase_map): - for j, occ_b in enumerate(occupations_b): - phase = alpha_phase * beta_phases[j] - for orb_b in occ_b: - phase *= phase_map[orb_b] - row[j] *= phase - - -def apply_diag_coulomb_evolution_in_place_z_rep_slow( - vec: np.ndarray, - mat_exp: np.ndarray, - mat_exp_conj: np.ndarray, - norb: int, - mat_alpha_beta_exp: np.ndarray, - mat_alpha_beta_exp_conj: np.ndarray, - strings_a: np.ndarray, - strings_b: np.ndarray, -) -> None: - """Apply time evolution by a diagonal Coulomb operator in-place.""" - dim_a, dim_b = vec.shape - alpha_phases = np.empty((dim_a,), dtype=complex) - beta_phases = np.empty((dim_b,), dtype=complex) - phase_map = np.ones((dim_a, norb), dtype=complex) - - for i, str0 in enumerate(strings_b): - phase = 1 - for j in range(norb): - sign_j = str0 >> j & 1 - for k in range(j + 1, norb): - sign_k = str0 >> k & 1 - mat = mat_exp_conj if sign_j ^ sign_k else mat_exp - phase *= mat[j, k] - beta_phases[i] = phase - - for i, (row, str0) in enumerate(zip(phase_map, strings_a)): - phase = 1 - for j in range(norb): - sign_j = str0 >> j & 1 - mat = mat_alpha_beta_exp_conj if sign_j else mat_alpha_beta_exp - row *= mat[j] - for k in range(j + 1, norb): - sign_k = str0 >> k & 1 - mat = mat_exp_conj if sign_j ^ sign_k else mat_exp - phase *= mat[j, k] - alpha_phases[i] = phase - - for row, alpha_phase, phase_map in zip(vec, alpha_phases, phase_map): - for i, str0 in enumerate(strings_b): - phase = alpha_phase * beta_phases[i] - for j in range(norb): - phase_shift = phase_map[j] - if str0 >> j & 1: - phase_shift = phase_shift.conjugate() - phase *= phase_shift - row[i] *= phase - - -def apply_diag_coulomb_evolution_in_place_num_rep_numpy( - vec: np.ndarray, - mat_exp: np.ndarray, - norb: int, - nelec: tuple[int, int], - *, - mat_alpha_beta_exp: np.ndarray, -) -> None: - """Apply time evolution by a diagonal Coulomb operator in-place.""" - mat_alpha_beta_exp = mat_alpha_beta_exp.copy() - mat_alpha_beta_exp[np.diag_indices(norb)] **= 0.5 - for i, j in itertools.combinations_with_replacement(range(norb), 2): - for sigma in range(2): - orbitals: list[set[int]] = [set(), set()] - orbitals[sigma].add(i) - orbitals[sigma].add(j) - _apply_phase_shift( - vec, - mat_exp[i, j], - (tuple(orbitals[0]), tuple(orbitals[1])), - norb=norb, - nelec=nelec, - copy=False, - ) - orbitals = [set() for _ in range(2)] - orbitals[sigma].add(i) - orbitals[1 - sigma].add(j) - _apply_phase_shift( - vec, - mat_alpha_beta_exp[i, j], - (tuple(orbitals[0]), tuple(orbitals[1])), - norb=norb, - nelec=nelec, - copy=False, - ) - - -def contract_diag_coulomb_into_buffer_num_rep_slow( - vec: np.ndarray, - mat: np.ndarray, - norb: int, - mat_alpha_beta: np.ndarray, - occupations_a: np.ndarray, - occupations_b: np.ndarray, - out: np.ndarray, -) -> None: - dim_a, dim_b = vec.shape - alpha_coeffs = np.empty((dim_a,), dtype=complex) - beta_coeffs = np.empty((dim_b,), dtype=complex) - coeff_map = np.zeros((dim_a, norb), dtype=complex) - - for i, occ in enumerate(occupations_b): - coeff = 0 - for orb_1, orb_2 in itertools.combinations_with_replacement(occ, 2): - coeff += mat[orb_1, orb_2] - beta_coeffs[i] = coeff - - for i, (row, orbs) in enumerate(zip(coeff_map, occupations_a)): - coeff = 0 - for j in range(len(orbs)): - row += mat_alpha_beta[orbs[j]] - for k in range(j, len(orbs)): - coeff += mat[orbs[j], orbs[k]] - alpha_coeffs[i] = coeff - - for source, target, alpha_coeff, coeff_map in zip( - vec, out, alpha_coeffs, coeff_map - ): - for j, occ_b in enumerate(occupations_b): - coeff = alpha_coeff + beta_coeffs[j] - for orb_b in occ_b: - coeff += coeff_map[orb_b] - target[j] += coeff * source[j] - - -def contract_diag_coulomb_into_buffer_z_rep_slow( - vec: np.ndarray, - mat: np.ndarray, - norb: int, - mat_alpha_beta: np.ndarray, - strings_a: np.ndarray, - strings_b: np.ndarray, - out: np.ndarray, -) -> None: - dim_a, dim_b = vec.shape - alpha_coeffs = np.empty((dim_a,), dtype=complex) - beta_coeffs = np.empty((dim_b,), dtype=complex) - coeff_map = np.zeros((dim_a, norb), dtype=complex) - - for i, str0 in enumerate(strings_b): - coeff = 0 - for j in range(norb): - sign_j = -1 if str0 >> j & 1 else 1 - for k in range(j + 1, norb): - sign_k = -1 if str0 >> k & 1 else 1 - coeff += sign_j * sign_k * mat[j, k] - beta_coeffs[i] = coeff - - for i, (row, str0) in enumerate(zip(coeff_map, strings_a)): - coeff = 0 - for j in range(norb): - sign_j = -1 if str0 >> j & 1 else 1 - row += sign_j * mat_alpha_beta[j] - for k in range(j + 1, norb): - sign_k = -1 if str0 >> k & 1 else 1 - coeff += sign_j * sign_k * mat[j, k] - alpha_coeffs[i] = coeff - - for source, target, alpha_coeff, coeff_map in zip( - vec, out, alpha_coeffs, coeff_map - ): - for i, str0 in enumerate(strings_b): - coeff = alpha_coeff + beta_coeffs[i] - for j in range(norb): - sign_j = -1 if str0 >> j & 1 else 1 - coeff += sign_j * coeff_map[j] - target[i] += 0.25 * coeff * source[i] - - -def contract_num_op_sum_spin_into_buffer_slow( - vec: np.ndarray, coeffs: np.ndarray, occupations: np.ndarray, out: np.ndarray -) -> None: - for source_row, target_row, orbs in zip(vec, out, occupations): - coeff = 0 - for orb in orbs: - coeff += coeffs[orb] - target_row += coeff * source_row diff --git a/python/ffsim/slow/__init__.py b/python/ffsim/slow/__init__.py new file mode 100644 index 000000000..9739675a2 --- /dev/null +++ b/python/ffsim/slow/__init__.py @@ -0,0 +1,11 @@ +# (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. + +"""Python versions of functions that were rewritten in Rust.""" diff --git a/python/ffsim/slow/contract/__init__.py b/python/ffsim/slow/contract/__init__.py new file mode 100644 index 000000000..8b593392d --- /dev/null +++ b/python/ffsim/slow/contract/__init__.py @@ -0,0 +1,9 @@ +# (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. diff --git a/python/ffsim/slow/contract/diag_coulomb.py b/python/ffsim/slow/contract/diag_coulomb.py new file mode 100644 index 000000000..78b725540 --- /dev/null +++ b/python/ffsim/slow/contract/diag_coulomb.py @@ -0,0 +1,98 @@ +# (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. + + +from __future__ import annotations + +import itertools + +import numpy as np + + +def contract_diag_coulomb_into_buffer_num_rep_slow( + vec: np.ndarray, + mat: np.ndarray, + norb: int, + mat_alpha_beta: np.ndarray, + occupations_a: np.ndarray, + occupations_b: np.ndarray, + out: np.ndarray, +) -> None: + dim_a, dim_b = vec.shape + alpha_coeffs = np.empty((dim_a,), dtype=complex) + beta_coeffs = np.empty((dim_b,), dtype=complex) + coeff_map = np.zeros((dim_a, norb), dtype=complex) + + for i, occ in enumerate(occupations_b): + coeff = 0 + for orb_1, orb_2 in itertools.combinations_with_replacement(occ, 2): + coeff += mat[orb_1, orb_2] + beta_coeffs[i] = coeff + + for i, (row, orbs) in enumerate(zip(coeff_map, occupations_a)): + coeff = 0 + for j in range(len(orbs)): + row += mat_alpha_beta[orbs[j]] + for k in range(j, len(orbs)): + coeff += mat[orbs[j], orbs[k]] + alpha_coeffs[i] = coeff + + for source, target, alpha_coeff, coeff_map in zip( + vec, out, alpha_coeffs, coeff_map + ): + for j, occ_b in enumerate(occupations_b): + coeff = alpha_coeff + beta_coeffs[j] + for orb_b in occ_b: + coeff += coeff_map[orb_b] + target[j] += coeff * source[j] + + +def contract_diag_coulomb_into_buffer_z_rep_slow( + vec: np.ndarray, + mat: np.ndarray, + norb: int, + mat_alpha_beta: np.ndarray, + strings_a: np.ndarray, + strings_b: np.ndarray, + out: np.ndarray, +) -> None: + dim_a, dim_b = vec.shape + alpha_coeffs = np.empty((dim_a,), dtype=complex) + beta_coeffs = np.empty((dim_b,), dtype=complex) + coeff_map = np.zeros((dim_a, norb), dtype=complex) + + for i, str0 in enumerate(strings_b): + coeff = 0 + for j in range(norb): + sign_j = -1 if str0 >> j & 1 else 1 + for k in range(j + 1, norb): + sign_k = -1 if str0 >> k & 1 else 1 + coeff += sign_j * sign_k * mat[j, k] + beta_coeffs[i] = coeff + + for i, (row, str0) in enumerate(zip(coeff_map, strings_a)): + coeff = 0 + for j in range(norb): + sign_j = -1 if str0 >> j & 1 else 1 + row += sign_j * mat_alpha_beta[j] + for k in range(j + 1, norb): + sign_k = -1 if str0 >> k & 1 else 1 + coeff += sign_j * sign_k * mat[j, k] + alpha_coeffs[i] = coeff + + for source, target, alpha_coeff, coeff_map in zip( + vec, out, alpha_coeffs, coeff_map + ): + for i, str0 in enumerate(strings_b): + coeff = alpha_coeff + beta_coeffs[i] + for j in range(norb): + sign_j = -1 if str0 >> j & 1 else 1 + coeff += sign_j * coeff_map[j] + target[i] += 0.25 * coeff * source[i] diff --git a/python/ffsim/slow/contract/num_op_sum.py b/python/ffsim/slow/contract/num_op_sum.py new file mode 100644 index 000000000..9218795bc --- /dev/null +++ b/python/ffsim/slow/contract/num_op_sum.py @@ -0,0 +1,24 @@ +# (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. + + +from __future__ import annotations + +import numpy as np + + +def contract_num_op_sum_spin_into_buffer_slow( + vec: np.ndarray, coeffs: np.ndarray, occupations: np.ndarray, out: np.ndarray +) -> None: + for source_row, target_row, orbs in zip(vec, out, occupations): + coeff = 0 + for orb in orbs: + coeff += coeffs[orb] + target_row += coeff * source_row diff --git a/python/ffsim/slow/gates/__init__.py b/python/ffsim/slow/gates/__init__.py new file mode 100644 index 000000000..8b593392d --- /dev/null +++ b/python/ffsim/slow/gates/__init__.py @@ -0,0 +1,9 @@ +# (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. diff --git a/python/ffsim/slow/gates/diag_coulomb.py b/python/ffsim/slow/gates/diag_coulomb.py new file mode 100644 index 000000000..6f64e3a8c --- /dev/null +++ b/python/ffsim/slow/gates/diag_coulomb.py @@ -0,0 +1,139 @@ +# (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. + +from __future__ import annotations + +import itertools + +import numpy as np + +from ffsim.gates.gates import _apply_phase_shift + + +def apply_diag_coulomb_evolution_in_place_num_rep_slow( + vec: np.ndarray, + mat_exp: np.ndarray, + norb: int, + mat_alpha_beta_exp: np.ndarray, + occupations_a: np.ndarray, + occupations_b: np.ndarray, +) -> None: + """Apply time evolution by a diagonal Coulomb operator in-place.""" + dim_a, dim_b = vec.shape + alpha_phases = np.empty((dim_a,), dtype=complex) + beta_phases = np.empty((dim_b,), dtype=complex) + phase_map = np.ones((dim_a, norb), dtype=complex) + + for i, orbs in enumerate(occupations_b): + phase = 1 + for orb_1, orb_2 in itertools.combinations_with_replacement(orbs, 2): + phase *= mat_exp[orb_1, orb_2] + beta_phases[i] = phase + + for i, (row, orbs) in enumerate(zip(phase_map, occupations_a)): + phase = 1 + for j in range(len(orbs)): + row *= mat_alpha_beta_exp[orbs[j]] + for k in range(j, len(orbs)): + phase *= mat_exp[orbs[j], orbs[k]] + alpha_phases[i] = phase + + for row, alpha_phase, phase_map in zip(vec, alpha_phases, phase_map): + for j, occ_b in enumerate(occupations_b): + phase = alpha_phase * beta_phases[j] + for orb_b in occ_b: + phase *= phase_map[orb_b] + row[j] *= phase + + +def apply_diag_coulomb_evolution_in_place_z_rep_slow( + vec: np.ndarray, + mat_exp: np.ndarray, + mat_exp_conj: np.ndarray, + norb: int, + mat_alpha_beta_exp: np.ndarray, + mat_alpha_beta_exp_conj: np.ndarray, + strings_a: np.ndarray, + strings_b: np.ndarray, +) -> None: + """Apply time evolution by a diagonal Coulomb operator in-place.""" + dim_a, dim_b = vec.shape + alpha_phases = np.empty((dim_a,), dtype=complex) + beta_phases = np.empty((dim_b,), dtype=complex) + phase_map = np.ones((dim_a, norb), dtype=complex) + + for i, str0 in enumerate(strings_b): + phase = 1 + for j in range(norb): + sign_j = str0 >> j & 1 + for k in range(j + 1, norb): + sign_k = str0 >> k & 1 + mat = mat_exp_conj if sign_j ^ sign_k else mat_exp + phase *= mat[j, k] + beta_phases[i] = phase + + for i, (row, str0) in enumerate(zip(phase_map, strings_a)): + phase = 1 + for j in range(norb): + sign_j = str0 >> j & 1 + mat = mat_alpha_beta_exp_conj if sign_j else mat_alpha_beta_exp + row *= mat[j] + for k in range(j + 1, norb): + sign_k = str0 >> k & 1 + mat = mat_exp_conj if sign_j ^ sign_k else mat_exp + phase *= mat[j, k] + alpha_phases[i] = phase + + for row, alpha_phase, phase_map in zip(vec, alpha_phases, phase_map): + for i, str0 in enumerate(strings_b): + phase = alpha_phase * beta_phases[i] + for j in range(norb): + phase_shift = phase_map[j] + if str0 >> j & 1: + phase_shift = phase_shift.conjugate() + phase *= phase_shift + row[i] *= phase + + +def apply_diag_coulomb_evolution_in_place_num_rep_numpy( + vec: np.ndarray, + mat_exp: np.ndarray, + norb: int, + nelec: tuple[int, int], + *, + mat_alpha_beta_exp: np.ndarray, +) -> None: + """Apply time evolution by a diagonal Coulomb operator in-place.""" + mat_alpha_beta_exp = mat_alpha_beta_exp.copy() + mat_alpha_beta_exp[np.diag_indices(norb)] **= 0.5 + for i, j in itertools.combinations_with_replacement(range(norb), 2): + for sigma in range(2): + orbitals: list[set[int]] = [set(), set()] + orbitals[sigma].add(i) + orbitals[sigma].add(j) + _apply_phase_shift( + vec, + mat_exp[i, j], + (tuple(orbitals[0]), tuple(orbitals[1])), + norb=norb, + nelec=nelec, + copy=False, + ) + orbitals = [set() for _ in range(2)] + orbitals[sigma].add(i) + orbitals[1 - sigma].add(j) + _apply_phase_shift( + vec, + mat_alpha_beta_exp[i, j], + (tuple(orbitals[0]), tuple(orbitals[1])), + norb=norb, + nelec=nelec, + copy=False, + ) diff --git a/python/ffsim/slow/gates/num_op_sum.py b/python/ffsim/slow/gates/num_op_sum.py new file mode 100644 index 000000000..0b6875e87 --- /dev/null +++ b/python/ffsim/slow/gates/num_op_sum.py @@ -0,0 +1,24 @@ +# (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. + +from __future__ import annotations + +import numpy as np + + +def apply_num_op_sum_evolution_in_place_slow( + vec: np.ndarray, phases: np.ndarray, occupations: np.ndarray +) -> None: + """Apply time evolution by a sum of number operators in-place.""" + for row, orbs in zip(vec, occupations): + phase = 1 + for orb in orbs: + phase *= phases[orb] + row *= phase diff --git a/python/ffsim/slow/gates/orbital_rotation.py b/python/ffsim/slow/gates/orbital_rotation.py new file mode 100644 index 000000000..4c8df796a --- /dev/null +++ b/python/ffsim/slow/gates/orbital_rotation.py @@ -0,0 +1,80 @@ +# (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. + +from __future__ import annotations + +import numpy as np + + +def gen_orbital_rotation_index_in_place_slow( + norb: int, + nocc: int, + linkstr_index: np.ndarray, + diag_strings: np.ndarray, + off_diag_strings: np.ndarray, + off_diag_strings_index: np.ndarray, + off_diag_index: np.ndarray, +) -> None: + """Generate orbital rotation index.""" + diag_counter = np.zeros(norb, dtype=np.uint) + off_diag_counter = np.zeros(norb, dtype=np.uint) + for str0, tab in enumerate(linkstr_index[:, :, 0]): + for orb in tab[:nocc]: + count = diag_counter[orb] + diag_strings[orb, count] = str0 + diag_counter[orb] += 1 + for orb in tab[nocc:norb]: + count = off_diag_counter[orb] + off_diag_strings[orb, count] = str0 + off_diag_strings_index[orb, str0] = count + off_diag_counter[orb] += 1 + + index_counter = np.zeros_like(off_diag_strings) + for str0, tab in enumerate(linkstr_index): + for orb_c, orb_d, str1, sign in tab[nocc:]: + # str0 -> annihilate orb_d -> create orb_c -> str1 + index = off_diag_strings_index[orb_c, str0] + count = index_counter[orb_c, index] + off_diag_index[orb_c, index, count] = orb_d, str1, sign + index_counter[orb_c, index] += 1 + + +def apply_single_column_transformation_in_place_slow( + vec: np.ndarray, + column: np.ndarray, + diag_val: complex, + diag_strings: np.ndarray, + off_diag_strings: np.ndarray, + off_diag_index: np.ndarray, +) -> None: + """Apply a single-column orbital rotation.""" + for str0, tab in zip(off_diag_strings, off_diag_index): + for orb, str1, sign in tab: + vec[str0] += sign * column[orb] * vec[str1] + for str0 in diag_strings: + vec[str0] *= diag_val + + +def apply_givens_rotation_in_place_slow( + vec: np.ndarray, + c: float, + s: float, + phase: complex, + slice1: np.ndarray, + slice2: np.ndarray, +) -> None: + """Apply a Givens rotation to slices of a state vector.""" + phase_conj = phase.conjugate() + for i, j in zip(slice1, slice2): + vec[i] *= phase_conj + tmp = c * vec[i] + s * vec[j] + vec[j] = c * vec[j] - s * vec[i] + vec[i] = tmp + vec[i] *= phase diff --git a/tests/slow/__init__.py b/tests/slow/__init__.py new file mode 100644 index 000000000..8b593392d --- /dev/null +++ b/tests/slow/__init__.py @@ -0,0 +1,9 @@ +# (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. diff --git a/tests/slow/contract/__init__.py b/tests/slow/contract/__init__.py new file mode 100644 index 000000000..8b593392d --- /dev/null +++ b/tests/slow/contract/__init__.py @@ -0,0 +1,9 @@ +# (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. diff --git a/tests/slow/contract/diag_coulomb_test.py b/tests/slow/contract/diag_coulomb_test.py new file mode 100644 index 000000000..b9d886230 --- /dev/null +++ b/tests/slow/contract/diag_coulomb_test.py @@ -0,0 +1,109 @@ +# (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. + +"""Tests for gates.""" + +from __future__ import annotations + +import numpy as np +from pyscf.fci import cistring +from scipy.special import comb + +import ffsim +from ffsim._ffsim import ( + contract_diag_coulomb_into_buffer_num_rep, + contract_diag_coulomb_into_buffer_z_rep, +) +from ffsim.slow.contract.diag_coulomb import ( + contract_diag_coulomb_into_buffer_num_rep_slow, + contract_diag_coulomb_into_buffer_z_rep_slow, +) + + +def test_contract_diag_coulomb_into_buffer_num_rep_slow(): + """Test contracting diag Coulomb operator.""" + norb = 5 + rng = np.random.default_rng() + for _ in range(5): + n_alpha = rng.integers(1, norb + 1) + n_beta = rng.integers(1, norb + 1) + dim_a = comb(norb, n_alpha, exact=True) + dim_b = comb(norb, n_beta, exact=True) + occupations_a = cistring.gen_occslst(range(norb), n_alpha).astype( + np.uint, copy=False + ) + occupations_b = cistring.gen_occslst(range(norb), n_beta).astype( + np.uint, copy=False + ) + mat = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) + mat_alpha_beta = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) + vec = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( + (dim_a, dim_b) + ) + out_slow = np.zeros_like(vec) + out_fast = np.zeros_like(vec) + contract_diag_coulomb_into_buffer_num_rep_slow( + vec, + mat, + norb=norb, + mat_alpha_beta=mat_alpha_beta, + occupations_a=occupations_a, + occupations_b=occupations_b, + out=out_slow, + ) + contract_diag_coulomb_into_buffer_num_rep( + vec, + mat, + norb=norb, + mat_alpha_beta=mat_alpha_beta, + occupations_a=occupations_a, + occupations_b=occupations_b, + out=out_fast, + ) + np.testing.assert_allclose(out_slow, out_fast, atol=1e-8) + + +def test_contract_diag_coulomb_into_buffer_z_rep_slow(): + """Test contracting diag Coulomb operator.""" + norb = 5 + rng = np.random.default_rng() + for _ in range(5): + n_alpha = rng.integers(1, norb + 1) + n_beta = rng.integers(1, norb + 1) + dim_a = comb(norb, n_alpha, exact=True) + dim_b = comb(norb, n_beta, exact=True) + strings_a = cistring.make_strings(range(norb), n_alpha) + strings_b = cistring.make_strings(range(norb), n_beta) + mat = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) + mat_alpha_beta = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) + vec = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( + (dim_a, dim_b) + ) + out_slow = np.zeros_like(vec) + out_fast = np.zeros_like(vec) + contract_diag_coulomb_into_buffer_z_rep_slow( + vec, + mat, + norb=norb, + mat_alpha_beta=mat_alpha_beta, + strings_a=strings_a, + strings_b=strings_b, + out=out_slow, + ) + contract_diag_coulomb_into_buffer_z_rep( + vec, + mat, + norb=norb, + mat_alpha_beta=mat_alpha_beta, + strings_a=strings_a, + strings_b=strings_b, + out=out_fast, + ) + np.testing.assert_allclose(out_slow, out_fast, atol=1e-8) diff --git a/tests/slow/contract/num_op_sum_test.py b/tests/slow/contract/num_op_sum_test.py new file mode 100644 index 000000000..a1ff3f9bd --- /dev/null +++ b/tests/slow/contract/num_op_sum_test.py @@ -0,0 +1,48 @@ +# (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. + +"""Tests for gates.""" + +from __future__ import annotations + +import numpy as np +from pyscf.fci import cistring +from scipy.special import comb + +import ffsim +from ffsim._ffsim import contract_num_op_sum_spin_into_buffer +from ffsim.slow.contract.num_op_sum import contract_num_op_sum_spin_into_buffer_slow + + +def test_contract_num_op_sum_spin_into_buffer_slow(): + """Test applying num op sum evolution.""" + norb = 5 + rng = np.random.default_rng() + for _ in range(5): + n_alpha = rng.integers(1, norb + 1) + n_beta = rng.integers(1, norb + 1) + dim_a = comb(norb, n_alpha, exact=True) + dim_b = comb(norb, n_beta, exact=True) + occupations = cistring.gen_occslst(range(norb), n_alpha).astype( + np.uint, copy=False + ) + coeffs = np.random.uniform(size=norb) + vec = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( + (dim_a, dim_b) + ) + out_slow = np.zeros_like(vec) + out_fast = np.zeros_like(vec) + contract_num_op_sum_spin_into_buffer_slow( + vec, coeffs, occupations=occupations, out=out_slow + ) + contract_num_op_sum_spin_into_buffer( + vec, coeffs, occupations=occupations, out=out_fast + ) + np.testing.assert_allclose(out_slow, out_fast, atol=1e-8) diff --git a/tests/slow/gates/__init__.py b/tests/slow/gates/__init__.py new file mode 100644 index 000000000..8b593392d --- /dev/null +++ b/tests/slow/gates/__init__.py @@ -0,0 +1,9 @@ +# (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. diff --git a/tests/slow/gates/diag_coulomb_test.py b/tests/slow/gates/diag_coulomb_test.py new file mode 100644 index 000000000..d9430a932 --- /dev/null +++ b/tests/slow/gates/diag_coulomb_test.py @@ -0,0 +1,154 @@ +# (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. + +from __future__ import annotations + +import numpy as np +from pyscf.fci import cistring +from scipy.special import comb + +import ffsim +from ffsim._ffsim import ( + apply_diag_coulomb_evolution_in_place_num_rep, + apply_diag_coulomb_evolution_in_place_z_rep, +) +from ffsim.slow.gates.diag_coulomb import ( + apply_diag_coulomb_evolution_in_place_num_rep_numpy, + apply_diag_coulomb_evolution_in_place_num_rep_slow, + apply_diag_coulomb_evolution_in_place_z_rep_slow, +) + + +def test_apply_diag_coulomb_evolution_num_rep_slow(): + """Test applying time evolution of diagonal Coulomb operator.""" + norb = 5 + rng = np.random.default_rng() + for _ in range(5): + n_alpha = rng.integers(1, norb + 1) + n_beta = rng.integers(1, norb + 1) + dim_a = comb(norb, n_alpha, exact=True) + dim_b = comb(norb, n_beta, exact=True) + occupations_a = cistring.gen_occslst(range(norb), n_alpha).astype( + np.uint, copy=False + ) + occupations_b = cistring.gen_occslst(range(norb), n_beta).astype( + np.uint, copy=False + ) + time = 0.6 + mat = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) + mat_exp = np.exp(-1j * time * mat) + mat_alpha_beta = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) + mat_alpha_beta_exp = np.exp(-1j * time * mat_alpha_beta) + vec_slow = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( + (dim_a, dim_b) + ) + vec_fast = vec_slow.copy() + apply_diag_coulomb_evolution_in_place_num_rep_slow( + vec_slow, + mat_exp, + norb=norb, + mat_alpha_beta_exp=mat_alpha_beta_exp, + occupations_a=occupations_a, + occupations_b=occupations_b, + ) + apply_diag_coulomb_evolution_in_place_num_rep( + vec_fast, + mat_exp, + norb=norb, + mat_alpha_beta_exp=mat_alpha_beta_exp, + occupations_a=occupations_a, + occupations_b=occupations_b, + ) + np.testing.assert_allclose(vec_slow, vec_fast, atol=1e-8) + + +def test_apply_diag_coulomb_evolution_z_rep_slow(): + """Test applying time evolution of diagonal Coulomb operator.""" + norb = 5 + rng = np.random.default_rng() + for _ in range(5): + n_alpha = rng.integers(1, norb + 1) + n_beta = rng.integers(1, norb + 1) + dim_a = comb(norb, n_alpha, exact=True) + dim_b = comb(norb, n_beta, exact=True) + strings_a = cistring.make_strings(range(norb), n_alpha) + strings_b = cistring.make_strings(range(norb), n_beta) + time = 0.6 + mat = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) + mat_exp = np.exp(-1j * time * mat) + mat_alpha_beta = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) + mat_alpha_beta_exp = np.exp(-1j * time * mat_alpha_beta) + vec_slow = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( + (dim_a, dim_b) + ) + vec_fast = vec_slow.copy() + apply_diag_coulomb_evolution_in_place_z_rep_slow( + vec_slow, + mat_exp, + mat_exp.conj(), + norb=norb, + mat_alpha_beta_exp=mat_alpha_beta_exp, + mat_alpha_beta_exp_conj=mat_alpha_beta_exp.conj(), + strings_a=strings_a, + strings_b=strings_b, + ) + apply_diag_coulomb_evolution_in_place_z_rep( + vec_fast, + mat_exp, + mat_exp.conj(), + norb=norb, + mat_alpha_beta_exp=mat_alpha_beta_exp, + mat_alpha_beta_exp_conj=mat_alpha_beta_exp.conj(), + strings_a=strings_a, + strings_b=strings_b, + ) + np.testing.assert_allclose(vec_slow, vec_fast, atol=1e-8) + + +def test_apply_diag_coulomb_evolution_num_rep_numpy(): + """Test applying time evolution of diag Coulomb operator, numpy implementation.""" + norb = 5 + rng = np.random.default_rng() + for _ in range(5): + n_alpha = rng.integers(1, norb + 1) + n_beta = rng.integers(1, norb + 1) + dim_a = comb(norb, n_alpha, exact=True) + dim_b = comb(norb, n_beta, exact=True) + occupations_a = cistring.gen_occslst(range(norb), n_alpha).astype( + np.uint, copy=False + ) + occupations_b = cistring.gen_occslst(range(norb), n_beta).astype( + np.uint, copy=False + ) + time = 0.6 + mat = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) + mat_exp = np.exp(-1j * time * mat) + mat_alpha_beta = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) + mat_alpha_beta_exp = np.exp(-1j * time * mat_alpha_beta) + vec_slow = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( + (dim_a, dim_b) + ) + vec_fast = vec_slow.copy() + apply_diag_coulomb_evolution_in_place_num_rep_numpy( + vec_slow, + mat_exp, + norb=norb, + nelec=(n_alpha, n_beta), + mat_alpha_beta_exp=mat_alpha_beta_exp, + ) + apply_diag_coulomb_evolution_in_place_num_rep( + vec_fast, + mat_exp, + norb=norb, + mat_alpha_beta_exp=mat_alpha_beta_exp, + occupations_a=occupations_a, + occupations_b=occupations_b, + ) + np.testing.assert_allclose(vec_slow, vec_fast, atol=1e-8) diff --git a/tests/slow/gates/num_op_sum_test.py b/tests/slow/gates/num_op_sum_test.py new file mode 100644 index 000000000..99f122868 --- /dev/null +++ b/tests/slow/gates/num_op_sum_test.py @@ -0,0 +1,48 @@ +# (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. + +from __future__ import annotations + +import numpy as np +from pyscf.fci import cistring +from scipy.special import comb + +import ffsim +from ffsim._ffsim import apply_num_op_sum_evolution_in_place +from ffsim.slow.gates.num_op_sum import apply_num_op_sum_evolution_in_place_slow + + +def test_apply_num_op_sum_evolution_in_place_slow(): + """Test applying num op sum evolution.""" + norb = 5 + rng = np.random.default_rng() + for _ in range(5): + n_alpha = rng.integers(1, norb + 1) + n_beta = rng.integers(1, norb + 1) + dim_a = comb(norb, n_alpha, exact=True) + dim_b = comb(norb, n_beta, exact=True) + occupations = cistring.gen_occslst(range(norb), n_alpha).astype( + np.uint, copy=False + ) + exponents = np.random.uniform(0, 2 * np.pi, size=norb) + phases = np.exp(1j * exponents) + vec_slow = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( + (dim_a, dim_b) + ) + vec_fast = vec_slow.copy() + apply_num_op_sum_evolution_in_place_slow( + vec_slow, phases, occupations=occupations + ) + apply_num_op_sum_evolution_in_place( + vec_fast, + phases, + occupations=occupations, + ) + np.testing.assert_allclose(vec_slow, vec_fast, atol=1e-8) diff --git a/tests/slow/gates/orbital_rotation_test.py b/tests/slow/gates/orbital_rotation_test.py new file mode 100644 index 000000000..0d26caf99 --- /dev/null +++ b/tests/slow/gates/orbital_rotation_test.py @@ -0,0 +1,123 @@ +# (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. + +from __future__ import annotations + +import numpy as np +from pyscf.fci import cistring +from scipy.special import comb + +import ffsim +from ffsim._ffsim import ( + apply_givens_rotation_in_place, + apply_single_column_transformation_in_place, + gen_orbital_rotation_index_in_place, +) +from ffsim.gates.orbital_rotation import ( + _zero_one_subspace_indices, + gen_orbital_rotation_index, +) +from ffsim.slow.gates.orbital_rotation import ( + apply_givens_rotation_in_place_slow, + apply_single_column_transformation_in_place_slow, + gen_orbital_rotation_index_in_place_slow, +) + + +def test_apply_givens_rotation_in_place_slow(): + """Test applying Givens rotation.""" + norb = 5 + rng = np.random.default_rng() + for _ in range(5): + n_alpha = rng.integers(1, norb + 1) + n_beta = rng.integers(1, norb + 1) + dim_a = comb(norb, n_alpha, exact=True) + dim_b = comb(norb, n_beta, exact=True) + vec_slow = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( + (dim_a, dim_b) + ) + vec_fast = vec_slow.copy() + c = rng.uniform() + s = 1 - c**2 + phase = (1j) ** rng.uniform(0, 4) + indices = _zero_one_subspace_indices(norb, n_alpha, (1, 3)) + slice1 = indices[: len(indices) // 2] + slice2 = indices[len(indices) // 2 :] + apply_givens_rotation_in_place_slow(vec_slow, c, s, phase, slice1, slice2) + apply_givens_rotation_in_place(vec_fast, c, s, phase, slice1, slice2) + np.testing.assert_allclose(vec_slow, vec_fast, atol=1e-8) + + +def test_gen_orbital_rotation_index_in_place_slow(): + """Test generating orbital rotation index.""" + norb = 5 + rng = np.random.default_rng() + nocc = rng.integers(1, norb + 1) + linkstr_index = cistring.gen_linkstr_index(range(norb), nocc) + + dim_diag = comb(norb - 1, nocc - 1, exact=True) + dim_off_diag = comb(norb - 1, nocc, exact=True) + dim = dim_diag + dim_off_diag + + diag_strings_slow = np.empty((norb, dim_diag), dtype=np.uint) + off_diag_strings_slow = np.empty((norb, dim_off_diag), dtype=np.uint) + off_diag_index_slow = np.empty((norb, dim_off_diag, nocc, 3), dtype=np.int32) + off_diag_strings_index_slow = np.empty((norb, dim), dtype=np.uint) + + diag_strings_fast = np.empty((norb, dim_diag), dtype=np.uint) + off_diag_strings_fast = np.empty((norb, dim_off_diag), dtype=np.uint) + off_diag_index_fast = np.empty((norb, dim_off_diag, nocc, 3), dtype=np.int32) + off_diag_strings_index_fast = np.empty((norb, dim), dtype=np.uint) + + gen_orbital_rotation_index_in_place_slow( + norb=norb, + nocc=nocc, + linkstr_index=linkstr_index, + diag_strings=diag_strings_slow, + off_diag_strings=off_diag_strings_slow, + off_diag_strings_index=off_diag_strings_index_slow, + off_diag_index=off_diag_index_slow, + ) + gen_orbital_rotation_index_in_place( + norb=norb, + nocc=nocc, + linkstr_index=linkstr_index, + diag_strings=diag_strings_fast, + off_diag_strings=off_diag_strings_fast, + off_diag_strings_index=off_diag_strings_index_fast, + off_diag_index=off_diag_index_fast, + ) + np.testing.assert_array_equal(diag_strings_slow, diag_strings_fast) + np.testing.assert_array_equal(off_diag_strings_slow, off_diag_strings_fast) + np.testing.assert_array_equal(off_diag_index_slow, off_diag_index_fast) + + +def test_apply_single_column_transformation_in_place_slow(): + """Test applying single column transformation.""" + norb = 5 + rng = np.random.default_rng() + for _ in range(5): + n_alpha = rng.integers(1, norb + 1) + n_beta = rng.integers(1, norb + 1) + dim_a = comb(norb, n_alpha, exact=True) + dim_b = comb(norb, n_beta, exact=True) + orbital_rotation_index = gen_orbital_rotation_index(norb, n_alpha) + column = rng.uniform(size=norb) + 1j * rng.uniform(size=norb) + diag_val = rng.uniform() + 1j * rng.uniform() + vec_slow = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( + (dim_a, dim_b) + ) + vec_fast = vec_slow.copy() + index = [a[0] for a in orbital_rotation_index] + apply_single_column_transformation_in_place_slow( + vec_slow, column, diag_val, *index + ) + apply_single_column_transformation_in_place(vec_fast, column, diag_val, *index) + np.testing.assert_allclose(vec_slow, vec_fast, atol=1e-8) diff --git a/tests/slow_test.py b/tests/slow_test.py deleted file mode 100644 index 13cad8e53..000000000 --- a/tests/slow_test.py +++ /dev/null @@ -1,404 +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. - -"""Tests for gates.""" - -from __future__ import annotations - -import numpy as np -from pyscf.fci import cistring -from scipy.special import comb - -import ffsim -from ffsim._ffsim import ( - apply_diag_coulomb_evolution_in_place_num_rep, - apply_diag_coulomb_evolution_in_place_z_rep, - apply_givens_rotation_in_place, - apply_num_op_sum_evolution_in_place, - apply_single_column_transformation_in_place, - contract_diag_coulomb_into_buffer_num_rep, - contract_diag_coulomb_into_buffer_z_rep, - contract_num_op_sum_spin_into_buffer, - gen_orbital_rotation_index_in_place, -) -from ffsim.gates.orbital_rotation import ( - _zero_one_subspace_indices, - gen_orbital_rotation_index, -) -from ffsim.slow import ( - apply_diag_coulomb_evolution_in_place_num_rep_numpy, - apply_diag_coulomb_evolution_in_place_num_rep_slow, - apply_diag_coulomb_evolution_in_place_z_rep_slow, - apply_givens_rotation_in_place_slow, - apply_num_op_sum_evolution_in_place_slow, - apply_single_column_transformation_in_place_slow, - contract_diag_coulomb_into_buffer_num_rep_slow, - contract_diag_coulomb_into_buffer_z_rep_slow, - contract_num_op_sum_spin_into_buffer_slow, - gen_orbital_rotation_index_in_place_slow, -) - - -def test_apply_givens_rotation_in_place_slow(): - """Test applying Givens rotation.""" - norb = 5 - rng = np.random.default_rng() - for _ in range(5): - n_alpha = rng.integers(1, norb + 1) - n_beta = rng.integers(1, norb + 1) - dim_a = comb(norb, n_alpha, exact=True) - dim_b = comb(norb, n_beta, exact=True) - vec_slow = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( - (dim_a, dim_b) - ) - vec_fast = vec_slow.copy() - c = rng.uniform() - s = 1 - c**2 - phase = (1j) ** rng.uniform(0, 4) - indices = _zero_one_subspace_indices(norb, n_alpha, (1, 3)) - slice1 = indices[: len(indices) // 2] - slice2 = indices[len(indices) // 2 :] - apply_givens_rotation_in_place_slow(vec_slow, c, s, phase, slice1, slice2) - apply_givens_rotation_in_place(vec_fast, c, s, phase, slice1, slice2) - np.testing.assert_allclose(vec_slow, vec_fast, atol=1e-8) - - -def test_gen_orbital_rotation_index_in_place_slow(): - """Test generating orbital rotation index.""" - norb = 5 - rng = np.random.default_rng() - nocc = rng.integers(1, norb + 1) - linkstr_index = cistring.gen_linkstr_index(range(norb), nocc) - - dim_diag = comb(norb - 1, nocc - 1, exact=True) - dim_off_diag = comb(norb - 1, nocc, exact=True) - dim = dim_diag + dim_off_diag - - diag_strings_slow = np.empty((norb, dim_diag), dtype=np.uint) - off_diag_strings_slow = np.empty((norb, dim_off_diag), dtype=np.uint) - off_diag_index_slow = np.empty((norb, dim_off_diag, nocc, 3), dtype=np.int32) - off_diag_strings_index_slow = np.empty((norb, dim), dtype=np.uint) - - diag_strings_fast = np.empty((norb, dim_diag), dtype=np.uint) - off_diag_strings_fast = np.empty((norb, dim_off_diag), dtype=np.uint) - off_diag_index_fast = np.empty((norb, dim_off_diag, nocc, 3), dtype=np.int32) - off_diag_strings_index_fast = np.empty((norb, dim), dtype=np.uint) - - gen_orbital_rotation_index_in_place_slow( - norb=norb, - nocc=nocc, - linkstr_index=linkstr_index, - diag_strings=diag_strings_slow, - off_diag_strings=off_diag_strings_slow, - off_diag_strings_index=off_diag_strings_index_slow, - off_diag_index=off_diag_index_slow, - ) - gen_orbital_rotation_index_in_place( - norb=norb, - nocc=nocc, - linkstr_index=linkstr_index, - diag_strings=diag_strings_fast, - off_diag_strings=off_diag_strings_fast, - off_diag_strings_index=off_diag_strings_index_fast, - off_diag_index=off_diag_index_fast, - ) - np.testing.assert_array_equal(diag_strings_slow, diag_strings_fast) - np.testing.assert_array_equal(off_diag_strings_slow, off_diag_strings_fast) - np.testing.assert_array_equal(off_diag_index_slow, off_diag_index_fast) - - -def test_apply_single_column_transformation_in_place_slow(): - """Test applying single column transformation.""" - norb = 5 - rng = np.random.default_rng() - for _ in range(5): - n_alpha = rng.integers(1, norb + 1) - n_beta = rng.integers(1, norb + 1) - dim_a = comb(norb, n_alpha, exact=True) - dim_b = comb(norb, n_beta, exact=True) - orbital_rotation_index = gen_orbital_rotation_index(norb, n_alpha) - column = rng.uniform(size=norb) + 1j * rng.uniform(size=norb) - diag_val = rng.uniform() + 1j * rng.uniform() - vec_slow = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( - (dim_a, dim_b) - ) - vec_fast = vec_slow.copy() - index = [a[0] for a in orbital_rotation_index] - apply_single_column_transformation_in_place_slow( - vec_slow, column, diag_val, *index - ) - apply_single_column_transformation_in_place(vec_fast, column, diag_val, *index) - np.testing.assert_allclose(vec_slow, vec_fast, atol=1e-8) - - -def test_apply_num_op_sum_evolution_in_place_slow(): - """Test applying num op sum evolution.""" - norb = 5 - rng = np.random.default_rng() - for _ in range(5): - n_alpha = rng.integers(1, norb + 1) - n_beta = rng.integers(1, norb + 1) - dim_a = comb(norb, n_alpha, exact=True) - dim_b = comb(norb, n_beta, exact=True) - occupations = cistring.gen_occslst(range(norb), n_alpha).astype( - np.uint, copy=False - ) - exponents = np.random.uniform(0, 2 * np.pi, size=norb) - phases = np.exp(1j * exponents) - vec_slow = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( - (dim_a, dim_b) - ) - vec_fast = vec_slow.copy() - apply_num_op_sum_evolution_in_place_slow( - vec_slow, phases, occupations=occupations - ) - apply_num_op_sum_evolution_in_place( - vec_fast, - phases, - occupations=occupations, - ) - np.testing.assert_allclose(vec_slow, vec_fast, atol=1e-8) - - -def test_apply_diag_coulomb_evolution_num_rep_slow(): - """Test applying time evolution of diagonal Coulomb operator.""" - norb = 5 - rng = np.random.default_rng() - for _ in range(5): - n_alpha = rng.integers(1, norb + 1) - n_beta = rng.integers(1, norb + 1) - dim_a = comb(norb, n_alpha, exact=True) - dim_b = comb(norb, n_beta, exact=True) - occupations_a = cistring.gen_occslst(range(norb), n_alpha).astype( - np.uint, copy=False - ) - occupations_b = cistring.gen_occslst(range(norb), n_beta).astype( - np.uint, copy=False - ) - time = 0.6 - mat = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) - mat_exp = np.exp(-1j * time * mat) - mat_alpha_beta = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) - mat_alpha_beta_exp = np.exp(-1j * time * mat_alpha_beta) - vec_slow = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( - (dim_a, dim_b) - ) - vec_fast = vec_slow.copy() - apply_diag_coulomb_evolution_in_place_num_rep_slow( - vec_slow, - mat_exp, - norb=norb, - mat_alpha_beta_exp=mat_alpha_beta_exp, - occupations_a=occupations_a, - occupations_b=occupations_b, - ) - apply_diag_coulomb_evolution_in_place_num_rep( - vec_fast, - mat_exp, - norb=norb, - mat_alpha_beta_exp=mat_alpha_beta_exp, - occupations_a=occupations_a, - occupations_b=occupations_b, - ) - np.testing.assert_allclose(vec_slow, vec_fast, atol=1e-8) - - -def test_apply_diag_coulomb_evolution_z_rep_slow(): - """Test applying time evolution of diagonal Coulomb operator.""" - norb = 5 - rng = np.random.default_rng() - for _ in range(5): - n_alpha = rng.integers(1, norb + 1) - n_beta = rng.integers(1, norb + 1) - dim_a = comb(norb, n_alpha, exact=True) - dim_b = comb(norb, n_beta, exact=True) - strings_a = cistring.make_strings(range(norb), n_alpha) - strings_b = cistring.make_strings(range(norb), n_beta) - time = 0.6 - mat = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) - mat_exp = np.exp(-1j * time * mat) - mat_alpha_beta = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) - mat_alpha_beta_exp = np.exp(-1j * time * mat_alpha_beta) - vec_slow = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( - (dim_a, dim_b) - ) - vec_fast = vec_slow.copy() - apply_diag_coulomb_evolution_in_place_z_rep_slow( - vec_slow, - mat_exp, - mat_exp.conj(), - norb=norb, - mat_alpha_beta_exp=mat_alpha_beta_exp, - mat_alpha_beta_exp_conj=mat_alpha_beta_exp.conj(), - strings_a=strings_a, - strings_b=strings_b, - ) - apply_diag_coulomb_evolution_in_place_z_rep( - vec_fast, - mat_exp, - mat_exp.conj(), - norb=norb, - mat_alpha_beta_exp=mat_alpha_beta_exp, - mat_alpha_beta_exp_conj=mat_alpha_beta_exp.conj(), - strings_a=strings_a, - strings_b=strings_b, - ) - np.testing.assert_allclose(vec_slow, vec_fast, atol=1e-8) - - -def test_apply_diag_coulomb_evolution_num_rep_numpy(): - """Test applying time evolution of diag Coulomb operator, numpy implementation.""" - norb = 5 - rng = np.random.default_rng() - for _ in range(5): - n_alpha = rng.integers(1, norb + 1) - n_beta = rng.integers(1, norb + 1) - dim_a = comb(norb, n_alpha, exact=True) - dim_b = comb(norb, n_beta, exact=True) - occupations_a = cistring.gen_occslst(range(norb), n_alpha).astype( - np.uint, copy=False - ) - occupations_b = cistring.gen_occslst(range(norb), n_beta).astype( - np.uint, copy=False - ) - time = 0.6 - mat = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) - mat_exp = np.exp(-1j * time * mat) - mat_alpha_beta = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) - mat_alpha_beta_exp = np.exp(-1j * time * mat_alpha_beta) - vec_slow = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( - (dim_a, dim_b) - ) - vec_fast = vec_slow.copy() - apply_diag_coulomb_evolution_in_place_num_rep_numpy( - vec_slow, - mat_exp, - norb=norb, - nelec=(n_alpha, n_beta), - mat_alpha_beta_exp=mat_alpha_beta_exp, - ) - apply_diag_coulomb_evolution_in_place_num_rep( - vec_fast, - mat_exp, - norb=norb, - mat_alpha_beta_exp=mat_alpha_beta_exp, - occupations_a=occupations_a, - occupations_b=occupations_b, - ) - np.testing.assert_allclose(vec_slow, vec_fast, atol=1e-8) - - -def test_contract_diag_coulomb_into_buffer_num_rep_slow(): - """Test contracting diag Coulomb operator.""" - norb = 5 - rng = np.random.default_rng() - for _ in range(5): - n_alpha = rng.integers(1, norb + 1) - n_beta = rng.integers(1, norb + 1) - dim_a = comb(norb, n_alpha, exact=True) - dim_b = comb(norb, n_beta, exact=True) - occupations_a = cistring.gen_occslst(range(norb), n_alpha).astype( - np.uint, copy=False - ) - occupations_b = cistring.gen_occslst(range(norb), n_beta).astype( - np.uint, copy=False - ) - mat = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) - mat_alpha_beta = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) - vec = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( - (dim_a, dim_b) - ) - out_slow = np.zeros_like(vec) - out_fast = np.zeros_like(vec) - contract_diag_coulomb_into_buffer_num_rep_slow( - vec, - mat, - norb=norb, - mat_alpha_beta=mat_alpha_beta, - occupations_a=occupations_a, - occupations_b=occupations_b, - out=out_slow, - ) - contract_diag_coulomb_into_buffer_num_rep( - vec, - mat, - norb=norb, - mat_alpha_beta=mat_alpha_beta, - occupations_a=occupations_a, - occupations_b=occupations_b, - out=out_fast, - ) - np.testing.assert_allclose(out_slow, out_fast, atol=1e-8) - - -def test_contract_diag_coulomb_into_buffer_z_rep_slow(): - """Test contracting diag Coulomb operator.""" - norb = 5 - rng = np.random.default_rng() - for _ in range(5): - n_alpha = rng.integers(1, norb + 1) - n_beta = rng.integers(1, norb + 1) - dim_a = comb(norb, n_alpha, exact=True) - dim_b = comb(norb, n_beta, exact=True) - strings_a = cistring.make_strings(range(norb), n_alpha) - strings_b = cistring.make_strings(range(norb), n_beta) - mat = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) - mat_alpha_beta = ffsim.random.random_real_symmetric_matrix(norb, seed=rng) - vec = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( - (dim_a, dim_b) - ) - out_slow = np.zeros_like(vec) - out_fast = np.zeros_like(vec) - contract_diag_coulomb_into_buffer_z_rep_slow( - vec, - mat, - norb=norb, - mat_alpha_beta=mat_alpha_beta, - strings_a=strings_a, - strings_b=strings_b, - out=out_slow, - ) - contract_diag_coulomb_into_buffer_z_rep( - vec, - mat, - norb=norb, - mat_alpha_beta=mat_alpha_beta, - strings_a=strings_a, - strings_b=strings_b, - out=out_fast, - ) - np.testing.assert_allclose(out_slow, out_fast, atol=1e-8) - - -def test_contract_num_op_sum_spin_into_buffer_slow(): - """Test applying num op sum evolution.""" - norb = 5 - rng = np.random.default_rng() - for _ in range(5): - n_alpha = rng.integers(1, norb + 1) - n_beta = rng.integers(1, norb + 1) - dim_a = comb(norb, n_alpha, exact=True) - dim_b = comb(norb, n_beta, exact=True) - occupations = cistring.gen_occslst(range(norb), n_alpha).astype( - np.uint, copy=False - ) - coeffs = np.random.uniform(size=norb) - vec = ffsim.random.random_statevector(dim_a * dim_b, seed=rng).reshape( - (dim_a, dim_b) - ) - out_slow = np.zeros_like(vec) - out_fast = np.zeros_like(vec) - contract_num_op_sum_spin_into_buffer_slow( - vec, coeffs, occupations=occupations, out=out_slow - ) - contract_num_op_sum_spin_into_buffer( - vec, coeffs, occupations=occupations, out=out_fast - ) - np.testing.assert_allclose(out_slow, out_fast, atol=1e-8)