From d54d7c72edcd87520cdc3b6ad2a6fc2acf66b7b2 Mon Sep 17 00:00:00 2001 From: "Ilya V. Schurov" Date: Thu, 4 Jul 2024 00:25:32 +0200 Subject: [PATCH] added rdms --- python/lattice_symmetries/rdms.py | 264 ++++++++++++++++++++++++++++++ python/test/test_rdms.py | 46 ++++++ 2 files changed, 310 insertions(+) create mode 100644 python/lattice_symmetries/rdms.py create mode 100644 python/test/test_rdms.py diff --git a/python/lattice_symmetries/rdms.py b/python/lattice_symmetries/rdms.py new file mode 100644 index 0000000..27753f0 --- /dev/null +++ b/python/lattice_symmetries/rdms.py @@ -0,0 +1,264 @@ +from typing import Literal, overload + +import numba +import numpy as np +import numpy.typing as npt +import scipy.special +from numba import float32, float64, uint64 + +from . import Basis + + +def index_to_spin(index: npt.NDArray[np.uint64], number_spins: int) -> npt.NDArray[np.bool_]: + return ( + (index.reshape(-1, 1).astype(np.int64) & (1 << np.arange(number_spins).astype(np.int64))) + ) > 0 + + +def spin_to_index( + spin: npt.NDArray[np.bool_ | np.int_ | np.uint], number_spins: int +) -> npt.NDArray[np.uint64]: + a = 2 ** np.arange(number_spins, dtype=np.int64) + return spin.dot(a) + + +@numba.jit("uint64(uint64, uint64)", nogil=True, nopython=True) +def _binom(n, k): + r"""Compute the number of ways to choose k elements out of a pile of n. + + :param n: the size of the pile of elements + :param k: the number of elements to take from the pile + :return: the number of ways to choose k elements out of a pile of n + """ + assert 0 <= n and n < 40 + assert 0 <= k and k <= n + + if k == 0 or k == n: + return 1 + total_ways = 1 + for i in range(min(k, n - k)): + total_ways = total_ways * (n - i) // (i + 1) + return total_ways + + +@numba.jit("uint64[:, :](uint64, uint64)", nogil=True, nopython=True) +def make_binom_cache(max_n, max_k): + assert 0 < max_k and max_k <= max_n + assert max_n < 40 + + cache = np.zeros((max_n + 1, max_k + 2), dtype=np.uint64) + for i in range(max_n + 1): + for j in range(min(i, max_k) + 2): + cache[i, j] = _binom(i, j) if j <= i else 0 + return cache + + +@numba.jit("uint64(uint64)", nogil=True, nopython=True) +def _hamming_weight(n: int): + # See https://stackoverflow.com/a/9830282 + n = (n & 0x5555555555555555) + ((n & 0xAAAAAAAAAAAAAAAA) >> 1) + n = (n & 0x3333333333333333) + ((n & 0xCCCCCCCCCCCCCCCC) >> 2) + n = (n & 0x0F0F0F0F0F0F0F0F) + ((n & 0xF0F0F0F0F0F0F0F0) >> 4) + n = (n & 0x00FF00FF00FF00FF) + ((n & 0xFF00FF00FF00FF00) >> 8) + n = (n & 0x0000FFFF0000FFFF) + ((n & 0xFFFF0000FFFF0000) >> 16) + n = (n & 0x00000000FFFFFFFF) + ((n & 0xFFFFFFFF00000000) >> 32) + return n + + +@numba.jit("uint64[:](uint64, uint64)", nogil=True, nopython=True) +def generate_binaries(length: int, hamming: int) -> npt.NDArray[np.uint64]: + assert length > 0 + assert hamming >= 0 and hamming <= length + if hamming == 0: + return np.zeros(1, dtype=np.uint64) + + size = _binom(length, hamming) + set_of_vectors = np.empty(size, dtype=np.uint64) + val = (1 << hamming) - 1 + for i in range(size): + set_of_vectors[i] = val + c = val & -val + r = val + c + val = (((r ^ val) >> 2) // c) | r + return set_of_vectors + + +@numba.jit("uint64(uint64, uint64, uint64)", nogil=True, nopython=True) +def _merge(vec1: int, vec2: int, dim2: int) -> int: + """ + (vec1 << dim2) | vec2 + """ + return (vec1 << dim2) | vec2 + + +@numba.jit("uint64(uint64, uint64, uint64[:, :])", nogil=True, nopython=True) +def _get_index(x: int, dim: int, binom_cache: npt.NDArray[np.uint64]) -> int: + r"""Compute the lexicographic index of a binary number among those with + the same Hamming weight. + + :param x: the binary number to find the index for + :param dim: the number of bits in the binary number + :param binom_cache: a 2D array of precomputed binomial coefficients + :return: the lexicographic index of the input binary number + + This function uses a precomputed binomial coefficient cache for efficiency. + It is optimized with Numba's just-in-time compilation and assumes x is a + valid binary number with dim bits. + """ + + n = 0 + h = 0 + if x & 1 == 1: + h += 1 + x >>= 1 + for i in range(1, dim): + if x & 1 == 1: + h += 1 + n += binom_cache[i, h] + x >>= 1 + return n + + +@numba.jit( + "complex128(uint64, uint64, uint64, uint64, uint64, complex128[::1], uint64[:, :])", + nogil=True, + nopython=True, +) +def _density_matrix_element( + sector_dim: int, + dim: int, + hamming: int, + vec1: int, + vec2: int, + amplitudes: npt.NDArray[np.complex128], + binom_cache: npt.NDArray[np.uint64], +) -> complex: + sector_hamming = _hamming_weight(vec1) + assert sector_hamming == _hamming_weight(vec2) + + smallest_number = 2 ** (hamming - sector_hamming) - 1 + k = dim - sector_dim + index1 = _get_index(_merge(vec1, smallest_number, k), dim, binom_cache) + index2 = _get_index(_merge(vec2, smallest_number, k), dim, binom_cache) + size_of_traced = _binom(k, hamming - sector_hamming) + matrix_element = np.dot( + amplitudes[index1 : index1 + size_of_traced].conj(), + amplitudes[index2 : index2 + size_of_traced], + ) + + return matrix_element + + +@numba.jit( + "Tuple((complex128[:, :], uint64[:]))(uint64, uint64, uint64, uint64, complex128[::1])", + nogil=True, + nopython=True, + parallel=False, +) +def _sector_density_matrix( + sector_dim: int, + dim: int, + sector_hamming: int, + hamming: int, + amplitudes: npt.NDArray[np.complex128], +) -> tuple[npt.NDArray[np.complex128], npt.NDArray[np.uint64]]: + assert sector_hamming <= hamming and sector_dim <= dim + assert hamming - sector_hamming <= dim - sector_dim + + sector_basis = generate_binaries(sector_dim, sector_hamming) + binom_cache = make_binom_cache(dim, hamming) + n = len(sector_basis) + matrix = np.empty((n, n), dtype=np.complex128) + for i in range(len(sector_basis)): + for j in range(i, len(sector_basis)): + matrix[i, j] = _density_matrix_element( + sector_dim, + dim, + hamming, + sector_basis[i], + sector_basis[j], + amplitudes, + binom_cache, + ) + matrix[j, i] = np.conj(matrix[i, j]) + + return matrix, sector_basis + + +def _density_matrix_blocks( + sub_dim: int, dim: int, hamming: int, amplitudes: npt.NDArray[np.complex128] +) -> tuple[list[npt.NDArray[np.complex128]], list[npt.NDArray[np.uint64]]]: + matrices = [] + sector_basis_states = [] + for sub_hamming in range(max(0, hamming - (dim - sub_dim)), min(hamming, sub_dim) + 1): + matrix, sector_basis = _sector_density_matrix( + sub_dim, dim, sub_hamming, hamming, amplitudes + ) + matrices.append(matrix) + sector_basis_states.append(sector_basis) + + return matrices, sector_basis_states + + +def _move_sites_to_back( + wavefunction: npt.NDArray[np.complex128], basis: ls.Basis, sites: list[int] +) -> npt.NDArray[np.complex128]: + unpacked_states = index_to_spin(basis.states, basis.number_sites) + permuted_sites = sorted(set(range(basis.number_sites)) - set(sites)) + sites + assert set(permuted_sites) == set(range(basis.number_sites)) + assert len(permuted_sites) == basis.number_sites + permuted_unpacked_states = unpacked_states[:, np.argsort(permuted_sites)] + permuted_states = spin_to_index(permuted_unpacked_states, basis.number_sites) + return wavefunction[basis.index(permuted_states)] + + +def _reduced_density_matrix_blocks( + wavefunction: npt.NDArray[np.complex128], basis: ls.Basis, sites: list[int] +) -> tuple[list[npt.NDArray[np.complex128]], list[npt.NDArray[np.uint64]]]: + if basis.hamming_weight is None: + raise ValueError("Only bases with a fixed hamming weight are supported.") + if basis.symmetries != []: + raise ValueError("Symmetries are not supported (yet)") + if basis.spin_inversion is not None: + raise ValueError("Spin inversion is not supported (yet)") + + hamming_weight = basis.hamming_weight + n_sites = basis.number_sites + wavefunction = _move_sites_to_back(wavefunction, basis, sites) + return _density_matrix_blocks(len(sites), n_sites, hamming_weight, wavefunction) + + +@overload +def reduced_density_matrix( + wavefunction: npt.NDArray[np.complex128], + basis: Basis, + sites: list[int], + return_states: Literal[True] = True, +) -> tuple[npt.NDArray[np.complex128], npt.NDArray[np.uint64]]: ... + + +@overload +def reduced_density_matrix( + wavefunction: npt.NDArray[np.complex128], + basis: Basis, + sites: list[int], + return_states: Literal[False] = True, +) -> npt.NDArray[np.complex128]: ... + + +def reduced_density_matrix( + wavefunction: npt.NDArray[np.complex128], + basis: Basis, + sites: list[int], + return_states: bool = True, +) -> npt.NDArray[np.complex128] | tuple[npt.NDArray[np.complex128], npt.NDArray[np.uint64]]: + matrices, sector_basis_states = _reduced_density_matrix_blocks(wavefunction, basis, sites) + if return_states: + return scipy.linalg.block_diag(*matrices), np.concatenate(sector_basis_states) + + all_states = np.arange(2 ** len(sites), dtype=np.uint64) + full_matrix = np.zeros((2 ** len(sites), 2 ** len(sites)), dtype=np.complex128) + for matrix, states in zip(matrices, sector_basis_states): + state_idxs = np.searchsorted(all_states, states) + full_matrix[np.ix_(state_idxs, state_idxs)] = matrix + return full_matrix diff --git a/python/test/test_rdms.py b/python/test/test_rdms.py new file mode 100644 index 0000000..46ee6ea --- /dev/null +++ b/python/test/test_rdms.py @@ -0,0 +1,46 @@ +import itertools + +import lattice_symmetries as ls +import numpy as np +import numpy.typing as npt +import scipy + + +def test_reduced_density_matrix_hamming(): + def rdms_reference( + ground_state: npt.NDArray[np.float64], sites: list[int], n_spins: int + ) -> npt.NDArray[np.float64]: + ground_state_reshaped = ground_state.reshape((2,) * n_spins).transpose( + range(n_spins - 1, -1, -1) + ) + idxs = tuple(sorted(set(range(n_spins)) - set(sites))) + ground_state_transposed = ground_state_reshaped.transpose(tuple(reversed(sites)) + idxs) + ground_state_transposed = ground_state_transposed.reshape( + (2 ** len(sites), 2 ** (n_spins - len(sites))) + ) + return ground_state_transposed @ ground_state_transposed.T + + expr_str = "σᶻ₀ σᶻ₁ + σᶻ₀ σᶻ₃ + σᶻ₀ σᶻ₄ + σᶻ₀ σᶻ₁₂ + 2.0 σ⁺₀ σ⁻₁ + 2.0 σ⁺₀ σ⁻₃ + 2.0 σ⁺₀ σ⁻₄ + 2.0 σ⁺₀ σ⁻₁₂ + 2.0 σ⁻₀ σ⁺₁ + 2.0 σ⁻₀ σ⁺₃ + 2.0 σ⁻₀ σ⁺₄ + 2.0 σ⁻₀ σ⁺₁₂ + σᶻ₁ σᶻ₂ + σᶻ₁ σᶻ₅ + σᶻ₁ σᶻ₁₃ + 2.0 σ⁺₁ σ⁻₂ + 2.0 σ⁺₁ σ⁻₅ + 2.0 σ⁺₁ σ⁻₁₃ + 2.0 σ⁻₁ σ⁺₂ + 2.0 σ⁻₁ σ⁺₅ + 2.0 σ⁻₁ σ⁺₁₃ + σᶻ₂ σᶻ₃ + σᶻ₂ σᶻ₆ + σᶻ₂ σᶻ₁₄ + 2.0 σ⁺₂ σ⁻₃ + 2.0 σ⁺₂ σ⁻₆ + 2.0 σ⁺₂ σ⁻₁₄ + 2.0 σ⁻₂ σ⁺₃ + 2.0 σ⁻₂ σ⁺₆ + 2.0 σ⁻₂ σ⁺₁₄ + σᶻ₃ σᶻ₇ + σᶻ₃ σᶻ₁₅ + 2.0 σ⁺₃ σ⁻₇ + 2.0 σ⁺₃ σ⁻₁₅ + 2.0 σ⁻₃ σ⁺₇ + 2.0 σ⁻₃ σ⁺₁₅ + σᶻ₄ σᶻ₅ + σᶻ₄ σᶻ₇ + σᶻ₄ σᶻ₈ + 2.0 σ⁺₄ σ⁻₅ + 2.0 σ⁺₄ σ⁻₇ + 2.0 σ⁺₄ σ⁻₈ + 2.0 σ⁻₄ σ⁺₅ + 2.0 σ⁻₄ σ⁺₇ + 2.0 σ⁻₄ σ⁺₈ + σᶻ₅ σᶻ₆ + σᶻ₅ σᶻ₉ + 2.0 σ⁺₅ σ⁻₆ + 2.0 σ⁺₅ σ⁻₉ + 2.0 σ⁻₅ σ⁺₆ + 2.0 σ⁻₅ σ⁺₉ + σᶻ₆ σᶻ₇ + σᶻ₆ σᶻ₁₀ + 2.0 σ⁺₆ σ⁻₇ + 2.0 σ⁺₆ σ⁻₁₀ + 2.0 σ⁻₆ σ⁺₇ + 2.0 σ⁻₆ σ⁺₁₀ + σᶻ₇ σᶻ₁₁ + 2.0 σ⁺₇ σ⁻₁₁ + 2.0 σ⁻₇ σ⁺₁₁ + σᶻ₈ σᶻ₉ + σᶻ₈ σᶻ₁₁ + σᶻ₈ σᶻ₁₂ + 2.0 σ⁺₈ σ⁻₉ + 2.0 σ⁺₈ σ⁻₁₁ + 2.0 σ⁺₈ σ⁻₁₂ + 2.0 σ⁻₈ σ⁺₉ + 2.0 σ⁻₈ σ⁺₁₁ + 2.0 σ⁻₈ σ⁺₁₂ + σᶻ₉ σᶻ₁₀ + σᶻ₉ σᶻ₁₃ + 2.0 σ⁺₉ σ⁻₁₀ + 2.0 σ⁺₉ σ⁻₁₃ + 2.0 σ⁻₉ σ⁺₁₀ + 2.0 σ⁻₉ σ⁺₁₃ + σᶻ₁₀ σᶻ₁₁ + σᶻ₁₀ σᶻ₁₄ + 2.0 σ⁺₁₀ σ⁻₁₁ + 2.0 σ⁺₁₀ σ⁻₁₄ + 2.0 σ⁻₁₀ σ⁺₁₁ + 2.0 σ⁻₁₀ σ⁺₁₄ + σᶻ₁₁ σᶻ₁₅ + 2.0 σ⁺₁₁ σ⁻₁₅ + 2.0 σ⁻₁₁ σ⁺₁₅ + σᶻ₁₂ σᶻ₁₃ + σᶻ₁₂ σᶻ₁₅ + 2.0 σ⁺₁₂ σ⁻₁₃ + 2.0 σ⁺₁₂ σ⁻₁₅ + 2.0 σ⁻₁₂ σ⁺₁₃ + 2.0 σ⁻₁₂ σ⁺₁₅ + σᶻ₁₃ σᶻ₁₄ + 2.0 σ⁺₁₃ σ⁻₁₄ + 2.0 σ⁻₁₃ σ⁺₁₄ + σᶻ₁₄ σᶻ₁₅ + 2.0 σ⁺₁₄ σ⁻₁₅ + 2.0 σ⁻₁₄ σ⁺₁₅" + number_sites = 16 + basis = ls.SpinBasis( + number_spins=number_sites, + hamming_weight=number_sites // 2, + ) + basis_no_hamming = ls.SpinBasis(number_spins=number_sites, hamming_weight=None) + basis.build() + basis_no_hamming.build() + expression = ls.Expr(expr_str, particle="spin-1/2") + hamiltonian = ls.Operator(expression, basis) + hamiltonian_no_hamming = ls.Operator(expression, basis_no_hamming) + ground_state = scipy.sparse.linalg.eigsh(hamiltonian, k=1, which="SA")[1].ravel() + ground_state_no_hamming = scipy.sparse.linalg.eigsh(hamiltonian_no_hamming, k=1, which="SA")[ + 1 + ].ravel() + + for length in [1, 2, 3]: + for sites in itertools.combinations(range(basis.number_sites), length): + reference = rdms_reference(ground_state_no_hamming, list(sites), basis.number_sites) + rdm = ls.rdms.reduced_density_matrix( + ground_state.astype(np.complex128), basis, list(sites), return_states=False + ) + assert np.allclose(reference, rdm)