Source code for ffsim.variational.givens

# (C) Copyright IBM 2024.
#
# 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.

"""Givens rotation ansatz."""

from __future__ import annotations

import cmath
import itertools
import math
from dataclasses import dataclass
from typing import cast

import numpy as np
from scipy.linalg.lapack import zrot

from ffsim import linalg
from ffsim.gates import apply_orbital_rotation


[docs] @dataclass(frozen=True) class GivensAnsatzOp: """A Givens rotation ansatz operator. The Givens rotation ansatz consists of a sequence of `Givens rotations`_ followed by a layer of single-orbital phase gates. Note that this ansatz does not implement any interactions between spin alpha and spin beta orbitals. Attributes: norb (int): The number of spatial orbitals. interaction_pairs (list[tuple[int, int]]): The orbital pairs to apply the Givens rotations to. thetas (np.ndarray): The angles for the Givens rotations. phis (np.ndarray | None): The optional phase angles for the Givens rotations. phase_angles (np.ndarray | None): The optional phase angles for the layer of single-orbital phase gates. .. _Givens rotations: ffsim.html#ffsim.apply_givens_rotation """ norb: int interaction_pairs: list[tuple[int, int]] thetas: np.ndarray phis: np.ndarray | None phase_angles: np.ndarray | None def __post_init__(self): if len(self.thetas) != len(self.interaction_pairs): raise ValueError( "The number of thetas must equal the number of interaction pairs. " f"Got {len(self.thetas)} and {len(self.interaction_pairs)}." ) if self.phis is not None and len(self.phis) != len(self.interaction_pairs): raise ValueError( "The number of phis must equal the number of interaction pairs. " f"Got {len(self.phis)} and {len(self.interaction_pairs)}." ) if self.phase_angles is not None and len(self.phase_angles) != self.norb: raise ValueError( "The number of phase angles must equal the number of orbitals. " f"Got {len(self.phase_angles)} and {self.norb}." ) def _apply_unitary_( self, vec: np.ndarray, norb: int, nelec: int | tuple[int, int], copy: bool ) -> np.ndarray: """Apply the operator to a vector.""" return apply_orbital_rotation( vec, self.to_orbital_rotation(), norb=norb, nelec=nelec, copy=copy )
[docs] @staticmethod def n_params( norb: int, interaction_pairs: list[tuple[int, int]], with_phis: bool = True, with_phase_angles: bool = True, ) -> int: """Return the number of parameters of an ansatz with given settings. Args: norb: The number of spatial orbitals. interaction_pairs: The orbital pairs to apply the Givens rotation gates to. with_phis: Whether to include complex phases for the Givens rotations. with_phase_angles: Whether to include a layer of single-orbital phase gates. """ return (1 + with_phis) * len(interaction_pairs) + with_phase_angles * norb
[docs] def to_parameters(self) -> np.ndarray: """Convert the operator to a real-valued parameter vector.""" if self.phis is not None and self.phase_angles is not None: return np.concatenate([self.thetas, self.phis, self.phase_angles]) if self.phis is not None: return np.concatenate([self.thetas, self.phis]) if self.phase_angles is not None: return np.concatenate([self.thetas, self.phase_angles]) return self.thetas
[docs] @staticmethod def from_parameters( params: np.ndarray, norb: int, interaction_pairs: list[tuple[int, int]], with_phis: bool = True, with_phase_angles: bool = True, ) -> GivensAnsatzOp: """Initialize the operator from a real-valued parameter vector. Args: params: The real-valued parameter vector. norb: The number of spatial orbitals. interaction_pairs: The orbital pairs to apply the Givens rotation gates to. with_phis: Whether to include complex phases for the Givens rotations. with_phase_angles: Whether to include a layer of single-orbital phase gates. """ n_params = (1 + with_phis) * len(interaction_pairs) + with_phase_angles * norb if len(params) != n_params: raise ValueError( "The number of parameters passed did not match the number expected " "based on the function inputs. " f"Expected {n_params} but got {len(params)}." ) thetas = params[: len(interaction_pairs)] phis = None phase_angles = None if with_phis and with_phase_angles: phis = params[len(interaction_pairs) : 2 * len(interaction_pairs)] phase_angles = params[2 * len(interaction_pairs) :] elif with_phis: phis = params[len(interaction_pairs) :] phase_angles = None elif with_phase_angles: phis = None phase_angles = params[len(interaction_pairs) :] return GivensAnsatzOp( norb=norb, interaction_pairs=interaction_pairs, thetas=thetas, phis=phis, phase_angles=phase_angles, )
[docs] @staticmethod def from_orbital_rotation(orbital_rotation: np.ndarray) -> GivensAnsatzOp: """Initialize the operator from an orbital rotation. Args: orbital_rotation: The orbital rotation. """ norb, _ = orbital_rotation.shape givens_rotations, phases = linalg.givens_decomposition(orbital_rotation) interaction_pairs = [] thetas = [] phis = [] for c, s, i, j in givens_rotations: interaction_pairs.append((i, j)) r, phi = cmath.polar(s) thetas.append(math.atan2(r, c)) phis.append(phi) interaction_pairs, thetas, phis = _brickwork_givens_rotations( interaction_pairs, thetas, phis, norb=norb ) return GivensAnsatzOp( norb=norb, interaction_pairs=interaction_pairs, thetas=np.array(thetas), phis=np.array(phis), phase_angles=np.angle(phases), )
[docs] def to_orbital_rotation(self) -> np.ndarray: """Convert the Givens ansatz operator to an orbital rotation.""" phis = self.phis phase_angles = self.phase_angles if phis is None: phis = np.zeros(len(self.interaction_pairs)) if phase_angles is None: phase_angles = np.zeros(self.norb) orbital_rotation = np.diag(np.exp(1j * phase_angles)) for (i, j), theta, phi in zip( self.interaction_pairs[::-1], self.thetas[::-1], phis[::-1] ): orbital_rotation[:, j], orbital_rotation[:, i] = zrot( orbital_rotation[:, j], orbital_rotation[:, i], math.cos(theta), cmath.rect(math.sin(theta), -phi), ) return orbital_rotation
def _approx_eq_(self, other, rtol: float, atol: float) -> bool: if isinstance(other, GivensAnsatzOp): if self.norb != other.norb: return False if self.interaction_pairs != other.interaction_pairs: return False if (self.phis is None) != (other.phis is None): return False if (self.phase_angles is None) != (other.phase_angles is None): return False if not np.allclose(self.thetas, other.thetas, rtol=rtol, atol=atol): return False if self.phis is not None and not np.allclose( cast(np.ndarray, self.phis), cast(np.ndarray, other.phis), rtol=rtol, atol=atol, ): return False if self.phase_angles is not None and not np.allclose( cast(np.ndarray, self.phase_angles), cast(np.ndarray, other.phase_angles), rtol=rtol, atol=atol, ): return False return True return NotImplemented
def _brickwork_givens_rotations( interaction_pairs: list[tuple[int, int]], thetas: list[float], phis: list[float], norb: int, ) -> tuple[list[tuple[int, int]], list[float], list[float]]: """Expand a sparse Givens rotation decomposition to a full brickwork pattern.""" # Construct a brickwork pattern of Givens rotations with angles set to zero q, r = divmod(norb, 2) even_layers = [ [((i, i + 1), 0.0, 0.0) for i in range(0, norb - 1, 2)] for _ in range(q + r) ] odd_layers = [ [((i, i + 1), 0.0, 0.0) for i in range(1, norb - 1, 2)] for _ in range(q) ] # even_layer_index[i] is the index of the last even layer acting on orbital i even_layer_index = [-1] * norb # odd_layer_index[i] is the index of the last odd layer acting on orbital i odd_layer_index = [-1] * norb for (i, j), theta, phi in zip(interaction_pairs, thetas, phis): if i > j: # Enforce i < j i, j = j, i theta = -theta phi = -phi if i % 2 == 0: # Even layer # Get the index of the even layer this Givens rotation should go in index = max(odd_layer_index[i], odd_layer_index[j]) + 1 # Add the Givens rotation in the appropriate place even_layers[index][i // 2] = ((i, j), theta, phi) # Update the even layer index even_layer_index[i] = index even_layer_index[j] = index else: # Odd layer # Get the index of the odd layer this Givens rotation should go in index = max(even_layer_index[i], even_layer_index[j]) # Add the Givens rotation in the appropriate place odd_layers[index][i // 2] = ((i, j), theta, phi) # Update the odd layer index odd_layer_index[i] = index odd_layer_index[j] = index # Construct the new Givens rotation decomposition and return new_interaction_pairs = [] new_thetas = [] new_phis = [] for even_layer, odd_layer in itertools.zip_longest( even_layers, odd_layers, fillvalue=() ): for layer in [even_layer, odd_layer]: for pair, theta, phi in layer: new_interaction_pairs.append(pair) new_thetas.append(theta) new_phis.append(phi) return new_interaction_pairs, new_thetas, new_phis