Source code for ffsim.states.states

# (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.

"""Fermionic quantum states."""

from __future__ import annotations

import math
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Optional, cast, overload

import numpy as np
import scipy.linalg
from pyscf.fci import cistring
from pyscf.fci.spin_op import contract_ss
from typing_extensions import deprecated

from ffsim import linalg
from ffsim.gates.orbital_rotation import apply_orbital_rotation
from ffsim.states.bitstring import (
    BitstringType,
    addresses_to_strings,
    concatenate_bitstrings,
    restrict_bitstrings,
)


[docs] @dataclass class StateVector: """A state vector in the FCI representation. Attributes: vec: Array of state vector coefficients. norb: The number of spatial orbitals. nelec: Either a single integer representing the number of fermions for a spinless system, or a pair of integers storing the numbers of spin alpha and spin beta fermions. """ vec: np.ndarray norb: int nelec: int | tuple[int, int] def __array__(self, dtype=None, copy=None): # TODO in Numpy 2.0 this can be simplified to # return np.array(self.vec, dtype=dtype, copy=copy) if copy: if dtype is None: return self.vec.copy() else: return self.vec.astype(dtype, copy=True) if dtype is None: return self.vec return self.vec.astype(dtype, copy=False)
[docs] def dims(norb: int, nelec: tuple[int, int]) -> tuple[int, int]: """Get the dimensions of the FCI space. Args: norb: The number of spatial orbitals. nelec: The number of alpha and beta electrons. Returns: A pair of integers (dim_a, dim_b) representing the dimensions of the alpha- and beta- FCI space. """ n_alpha, n_beta = nelec dim_a = math.comb(norb, n_alpha) dim_b = math.comb(norb, n_beta) return dim_a, dim_b
[docs] def dim(norb: int, nelec: int | tuple[int, int]) -> int: """Get the dimension of the FCI space. Args: norb: The number of spatial orbitals. nelec: Either a single integer representing the number of fermions for a spinless system, or a pair of integers storing the numbers of spin alpha and spin beta fermions. Returns: The dimension of the FCI space. """ if isinstance(nelec, int): return math.comb(norb, nelec) n_alpha, n_beta = nelec return math.comb(norb, n_alpha) * math.comb(norb, n_beta)
[docs] @deprecated( "Using one_hot from the ffsim namespace is deprecated. " "Instead, use ffsim.linalg.one_hot." ) def one_hot(shape: int | tuple[int, ...], index, *, dtype=complex): """Return an array of all zeros except for a one at a specified index. .. warning:: This function is deprecated. Use :func:`ffsim.linalg.one_hot` instead. Args: shape: The desired shape of the array. index: The index at which to place a one. Returns: The one-hot vector. """ vec = np.zeros(shape, dtype=dtype) vec[index] = 1 return vec
@overload def slater_determinant( norb: int, occupied_orbitals: Sequence[int], orbital_rotation: np.ndarray | None = None, ) -> np.ndarray: ... @overload def slater_determinant( norb: int, occupied_orbitals: tuple[Sequence[int], Sequence[int]], orbital_rotation: np.ndarray | tuple[np.ndarray | None, np.ndarray | None] | None = None, ) -> np.ndarray: ...
[docs] def slater_determinant( norb: int, occupied_orbitals: Sequence[int] | tuple[Sequence[int], Sequence[int]], orbital_rotation: np.ndarray | tuple[np.ndarray | None, np.ndarray | None] | None = None, ) -> np.ndarray: r"""Return a Slater determinant. A Slater determinant is a state of the form .. math:: \mathcal{U} \lvert x \rangle, where :math:`\mathcal{U}` is an :doc:`orbital rotation </explanations/orbital-rotation>` and :math:`\lvert x \rangle` is an electronic configuration. Args: norb: The number of spatial orbitals. occupied_orbitals: The occupied orbitals in the electronic configuration. This is either a list of integers specifying spinless orbitals, or a pair of lists, where the first list specifies the spin alpha orbitals and the second list specifies the spin beta orbitals. orbital_rotation: The optional orbital rotation. You can pass either a single Numpy array specifying the orbital rotation to apply to both spin sectors, or you can pass a pair of Numpy arrays specifying independent orbital rotations for spin alpha and spin beta. If passing a pair, you can use ``None`` for one of the values in the pair to indicate that no operation should be applied to that spin sector. Returns: The Slater determinant as a state vector. """ if norb == 0: return np.ones(1, dtype=complex) if not occupied_orbitals or isinstance(occupied_orbitals[0], (int, np.integer)): occupied_orbitals = (cast(Sequence[int], occupied_orbitals), []) alpha_orbitals, beta_orbitals = cast( tuple[Sequence[int], Sequence[int]], occupied_orbitals ) n_alpha = len(alpha_orbitals) n_beta = len(beta_orbitals) nelec = (n_alpha, n_beta) dim1, dim2 = dims(norb, nelec) alpha_bits = np.zeros(norb, dtype=bool) alpha_bits[list(alpha_orbitals)] = 1 alpha_string = int("".join("1" if b else "0" for b in alpha_bits[::-1]), base=2) alpha_index = cistring.str2addr(norb, n_alpha, alpha_string) beta_bits = np.zeros(norb, dtype=bool) beta_bits[list(beta_orbitals)] = 1 beta_string = int("".join("1" if b else "0" for b in beta_bits[::-1]), base=2) beta_index = cistring.str2addr(norb, n_beta, beta_string) vec = linalg.one_hot( (dim1, dim2), (alpha_index, beta_index), dtype=complex ).reshape(-1) if orbital_rotation is not None: vec = apply_orbital_rotation( vec, orbital_rotation, norb=norb, nelec=nelec, copy=False ) return vec
[docs] def hartree_fock_state(norb: int, nelec: int | tuple[int, int]) -> np.ndarray: """Return the Hartree-Fock state. Args: norb: The number of spatial orbitals. nelec: Either a single integer representing the number of fermions for a spinless system, or a pair of integers storing the numbers of spin alpha and spin beta fermions. Returns: The Hartree-Fock state as a state vector. """ if isinstance(nelec, int): return slater_determinant(norb, occupied_orbitals=range(nelec)) n_alpha, n_beta = nelec return slater_determinant(norb, occupied_orbitals=(range(n_alpha), range(n_beta)))
[docs] @deprecated( "ffsim.slater_determinant_rdm is deprecated. " "Instead, use ffsim.slater_determinant_rdms." ) def slater_determinant_rdm( norb: int, occupied_orbitals: tuple[Sequence[int], Sequence[int]], orbital_rotation: np.ndarray | tuple[np.ndarray | None, np.ndarray | None] | None = None, rank: int = 1, spin_summed: bool = True, ) -> np.ndarray: """Return the reduced density matrix of a `Slater determinant`_. .. warning:: This function is deprecated. Use :func:`ffsim.slater_determinant_rdms` instead. Note: Currently, only rank 1 is supported. Args: norb: The number of spatial orbitals. occupied_orbitals: The occupied orbitals in the electronic configuration. This is a pair of lists of integers, where the first list specifies the spin alpha orbitals and the second list specifies the spin beta orbitals. orbital_rotation: The optional orbital rotation. You can pass either a single Numpy array specifying the orbital rotation to apply to both spin sectors, or you can pass a pair of Numpy arrays specifying independent orbital rotations for spin alpha and spin beta. If passing a pair, you can use ``None`` for one of the values in the pair to indicate that no operation should be applied to that spin sector. rank: The rank of the reduced density matrix. I.e., rank 1 corresponds to the one-particle RDM, rank 2 corresponds to the 2-particle RDM, etc. spin_summed: Whether to sum over the spin index. Returns: The reduced density matrix of the Slater determinant. .. _Slater determinant: ffsim.html#ffsim.slater_determinant """ if rank == 1: rdm_a = np.zeros((norb, norb), dtype=complex) rdm_b = np.zeros((norb, norb), dtype=complex) alpha_orbitals = np.array(occupied_orbitals[0]) beta_orbitals = np.array(occupied_orbitals[1]) if len(alpha_orbitals): rdm_a[(alpha_orbitals, alpha_orbitals)] = 1 if len(beta_orbitals): rdm_b[(beta_orbitals, beta_orbitals)] = 1 if orbital_rotation is not None: if isinstance(orbital_rotation, np.ndarray): orbital_rotation_a: np.ndarray | None = orbital_rotation orbital_rotation_b: np.ndarray | None = orbital_rotation else: orbital_rotation_a, orbital_rotation_b = orbital_rotation if orbital_rotation_a is not None: rdm_a = orbital_rotation_a.conj() @ rdm_a @ orbital_rotation_a.T if orbital_rotation_b is not None: rdm_b = orbital_rotation_b.conj() @ rdm_b @ orbital_rotation_b.T if spin_summed: return rdm_a + rdm_b return scipy.linalg.block_diag(rdm_a, rdm_b) raise NotImplementedError( f"Returning the rank {rank} reduced density matrix is currently not supported." )
@overload def slater_determinant_rdms( norb: int, occupied_orbitals: Sequence[int], orbital_rotation: np.ndarray | None = None, *, rank: int = 1, ) -> np.ndarray: ... @overload def slater_determinant_rdms( norb: int, occupied_orbitals: tuple[Sequence[int], Sequence[int]], orbital_rotation: np.ndarray | tuple[np.ndarray | None, np.ndarray | None] | None = None, *, rank: int = 1, ) -> np.ndarray: ...
[docs] def slater_determinant_rdms( norb: int, occupied_orbitals: Sequence[int] | tuple[Sequence[int], Sequence[int]], orbital_rotation: np.ndarray | tuple[np.ndarray | None, np.ndarray | None] | None = None, *, rank: int = 1, ) -> np.ndarray: """Return the reduced density matrices of a `Slater determinant`_. Note: Currently, only rank 1 is supported. Args: norb: The number of spatial orbitals. occupied_orbitals: The occupied orbitals in the electronic configuration. This is either a list of integers specifying spinless orbitals, or a pair of lists, where the first list specifies the spin alpha orbitals and the second list specifies the spin beta orbitals. orbital_rotation: The optional orbital rotation. You can pass either a single Numpy array specifying the orbital rotation to apply to both spin sectors, or you can pass a pair of Numpy arrays specifying independent orbital rotations for spin alpha and spin beta. If passing a pair, you can use ``None`` for one of the values in the pair to indicate that no operation should be applied to that spin sector. rank: The rank of the reduced density matrix. I.e., rank 1 corresponds to the one-particle RDM, rank 2 corresponds to the 2-particle RDM, etc. Returns: The reduced density matrices of the Slater determinant. All RDMs up to and including the specified rank are returned, in increasing order of rank. For example, if `rank=2` then a tuple `(rdm1, rdm2)` is returned. The representation of an RDM depends on whether `occupied_orbitals` is a sequence of integers (spinless case), or a pair of such sequences (spinful case). In the spinless case, the full RDM is returned. In the spinful case, each RDM is represented as a stacked Numpy array of sub-RDMs. For example, the 1-RDMs are: (alpha-alpha, alpha-beta), and the 2-RDMs are: (alpha-alpha, alpha-beta, beta-beta). .. _Slater determinant: ffsim.html#ffsim.slater_determinant """ if not occupied_orbitals or isinstance(occupied_orbitals[0], (int, np.integer)): # Spinless case occupied_orbitals = list(cast(Sequence[int], occupied_orbitals)) if rank == 1: rdm = np.zeros((norb, norb), dtype=complex) if occupied_orbitals: rdm[(occupied_orbitals, occupied_orbitals)] = 1 if orbital_rotation is not None: orbital_rotation = cast(np.ndarray, orbital_rotation) rdm = orbital_rotation.conj() @ rdm @ orbital_rotation.T return rdm raise NotImplementedError( f"Returning the rank {rank} reduced density matrix is currently not " "supported." ) else: # Spinful case alpha_orbitals, beta_orbitals = cast( tuple[Sequence[int], Sequence[int]], occupied_orbitals ) alpha_orbitals = list(alpha_orbitals) beta_orbitals = list(beta_orbitals) if rank == 1: rdm_a = np.zeros((norb, norb)) rdm_b = np.zeros((norb, norb)) if alpha_orbitals: rdm_a[(alpha_orbitals, alpha_orbitals)] = 1 if beta_orbitals: rdm_b[(beta_orbitals, beta_orbitals)] = 1 if orbital_rotation is not None: if ( isinstance(orbital_rotation, np.ndarray) and orbital_rotation.ndim == 2 ): orbital_rotation_a: np.ndarray | None = orbital_rotation orbital_rotation_b: np.ndarray | None = orbital_rotation else: orbital_rotation_a, orbital_rotation_b = orbital_rotation if orbital_rotation_a is not None: rdm_a = orbital_rotation_a.conj() @ rdm_a @ orbital_rotation_a.T if orbital_rotation_b is not None: rdm_b = orbital_rotation_b.conj() @ rdm_b @ orbital_rotation_b.T return np.stack([rdm_a, rdm_b]) raise NotImplementedError( f"Returning the rank {rank} reduced density matrix is currently not " "supported." )
# source: pyscf.fci.spin_op.spin_square0 # modified to support complex wavefunction
[docs] def spin_square(fcivec: np.ndarray, norb: int, nelec: tuple[int, int]): """Expectation value of spin squared operator on a state vector.""" if np.iscomplexobj(fcivec): ci1 = contract_ss(fcivec.real, norb, nelec).astype(complex) ci1 += 1j * contract_ss(fcivec.imag, norb, nelec) else: ci1 = contract_ss(fcivec, norb, nelec) return np.einsum("ij,ij->", fcivec.reshape(ci1.shape), ci1.conj()).real
[docs] def sample_state_vector( vec: np.ndarray | StateVector, *, norb: int | None = None, nelec: int | tuple[int, int] | None = None, orbs: Sequence[int] | tuple[Sequence[int], Sequence[int]] | None = None, shots: int = 1, concatenate: bool = True, bitstring_type: BitstringType = BitstringType.STRING, seed: np.random.Generator | int | None = None, ): """Sample bitstrings from a state vector. Args: vec: The state vector to sample from. norb: The number of spatial orbitals. nelec: Either a single integer representing the number of fermions for a spinless system, or a pair of integers storing the numbers of spin alpha and spin beta fermions. orbs: The spin-orbitals to sample. In the spinless case (when `nelec` is an integer), this is a list of integers ranging from ``0`` to ``norb``. In the spinful case, this is a pair of lists of such integers, with the first list storing the spin-alpha orbitals and the second list storing the spin-beta orbitals. If not specified, then all spin-orbitals are sampled. shots: The number of bitstrings to sample. concatenate: Whether to concatenate the spin-alpha and spin-beta parts of the bitstrings. If True, then a single list of concatenated bitstrings is returned. The strings are concatenated in the order :math:`s_b s_a`, that is, the alpha string appears on the right. If False, then two lists are returned, ``(strings_a, strings_b)``. Note that the list of alpha strings appears first, that is, on the left. In the spinless case (when `nelec` is an integer), this argument is ignored. bitstring_type: The desired type of bitstring output. seed: A seed to initialize the pseudorandom number generator. Should be a valid input to ``np.random.default_rng``. Returns: The sampled bitstrings, as a list of strings of length `shots`. Raises: TypeError: When passing vec as a Numpy array, norb and nelec must be specified. TypeError: When passing vec as a StateVector, norb and nelec must both be None. """ vec, norb, nelec = canonicalize_vec_norb_nelec(vec, norb, nelec) if isinstance(nelec, int): return _sample_state_vector_spinless( vec, norb=norb, nelec=nelec, orbs=cast(Optional[Sequence[int]], orbs), shots=shots, bitstring_type=bitstring_type, seed=seed, ) return _sample_state_vector_spinful( vec, norb=norb, nelec=nelec, orbs=cast(Optional[tuple[Sequence[int], Sequence[int]]], orbs), shots=shots, concatenate=concatenate, bitstring_type=bitstring_type, seed=seed, )
def _sample_state_vector_spinless( vec: np.ndarray, *, norb: int, nelec: int, orbs: Sequence[int] | None, shots: int, bitstring_type: BitstringType, seed: np.random.Generator | int | None, ): if orbs is None: orbs = range(norb) rng = np.random.default_rng(seed) probabilities = np.abs(vec) ** 2 samples = rng.choice(len(vec), size=shots, p=probabilities) strings = addresses_to_strings(samples, norb, nelec, bitstring_type=bitstring_type) if list(orbs) == list(range(norb)): return strings return restrict_bitstrings(strings, orbs, bitstring_type=bitstring_type) def _sample_state_vector_spinful( vec: np.ndarray, *, norb: int, nelec: tuple[int, int], orbs: tuple[Sequence[int], Sequence[int]] | None, shots: int, concatenate: bool = True, bitstring_type: BitstringType, seed: np.random.Generator | int | None, ): if orbs is None: orbs = range(norb), range(norb) rng = np.random.default_rng(seed) probabilities = np.abs(vec) ** 2 samples = rng.choice(len(vec), size=shots, p=probabilities) orbs_a, orbs_b = orbs if list(orbs_a) == list(orbs_b) == list(range(norb)): # All orbitals are sampled, so we can simply call addresses_to_strings return addresses_to_strings( samples, norb, nelec, concatenate=concatenate, bitstring_type=bitstring_type ) strings_a, strings_b = addresses_to_strings( samples, norb, nelec, concatenate=False, bitstring_type=bitstring_type ) strings_a = restrict_bitstrings(strings_a, orbs_a, bitstring_type=bitstring_type) strings_b = restrict_bitstrings(strings_b, orbs_b, bitstring_type=bitstring_type) if concatenate: return concatenate_bitstrings( strings_a, strings_b, bitstring_type=bitstring_type, length=len(orbs_a) ) return strings_a, strings_b def canonicalize_vec_norb_nelec( vec: np.ndarray | StateVector, norb: int | None, nelec: int | tuple[int, int] | None ) -> tuple[np.ndarray, int, int | tuple[int, int]]: """Canonicalize a state vector, number of orbitals, and number(s) of electrons. Args: vec: The state vector. norb: The number of spatial orbitals, or None if `vec` is a StateVector. nelec: Either a single integer representing the number of fermions for a spinless system, or a pair of integers storing the numbers of spin alpha and spin beta fermions. If `vec` is a StateVector then this should be None. Returns: The state vector as a Numpy array, the number of spatial orbitals, and the number or numbers of electrons. Raises: TypeError: When passing vec as a Numpy array, norb and nelec must be specified. TypeError: When passing vec as a StateVector, norb and nelec must both be None. """ if isinstance(vec, np.ndarray): if norb is None or nelec is None: raise TypeError( "When passing vec as a Numpy array, norb and nelec must be specified." ) else: if norb is not None or nelec is not None: raise TypeError( "When passing vec as a StateVector, norb and nelec must both be None." ) norb = vec.norb nelec = vec.nelec vec = vec.vec return vec, norb, nelec