Skip to content

Commit

Permalink
feat(fock): Avoid matrix embedding
Browse files Browse the repository at this point in the history 
In `FockSimulator`, the unitary matrices acting on the density operator were
embedded in the Fock space, before directly multiplying with the density
matrix. This can be done faster, by avoiding embedding and modifying only the
relevant parts of the density matrix, as done in `PureFockSimulator` with the
state vector instead of the density matrix. This approach has been implemented
in this patch.
  • Loading branch information
@Kolarovszki
Kolarovszki committed Feb 13, 2024
1 parent 3f77b3e commit 9e98868
Show file tree
Hide file tree
Showing 3 changed files with 121 additions and 77 deletions.
41 changes: 41 additions & 0 deletions piquasso/_backends/fock/calculations.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,20 @@

from scipy.special import comb

from functools import lru_cache

from .state import BaseFockState

from piquasso.api.instruction import Instruction
from piquasso.api.result import Result
from piquasso.api.exceptions import InvalidInstruction, InvalidParameter

from piquasso._math.fock import cutoff_cardinality
from piquasso._math.indices import (
get_index_in_fock_space,
get_auxiliary_modes,
)


def attenuator(state: BaseFockState, instruction: Instruction, shots: int) -> Result:
r"""
Expand Down Expand Up @@ -89,3 +97,36 @@ def attenuator(state: BaseFockState, instruction: Instruction, shots: int) -> Re
new_state._density_matrix = new_density_matrix

return Result(state=new_state)


@lru_cache(maxsize=None)
def calculate_state_index_matrix_list(space, auxiliary_subspace, mode):
d = space.d
cutoff = space.cutoff

state_index_matrix_list = []
auxiliary_modes = get_auxiliary_modes(d, (mode,))
all_occupation_numbers = np.zeros(d, dtype=int)

indices = [cutoff_cardinality(cutoff=n - 1, d=d - 1) for n in range(1, cutoff + 2)]

for n in range(cutoff):
limit = space.cutoff - n
subspace_size = indices[n + 1] - indices[n]

state_index_matrix = np.empty(shape=(limit, subspace_size), dtype=int)

for i, auxiliary_occupation_numbers in enumerate(
auxiliary_subspace[indices[n] : indices[n + 1]]
):
all_occupation_numbers[(auxiliary_modes,)] = auxiliary_occupation_numbers

for j in range(limit):
all_occupation_numbers[mode] = j
state_index_matrix[j, i] = get_index_in_fock_space(
tuple(all_occupation_numbers)
)

state_index_matrix_list.append(state_index_matrix)

return state_index_matrix_list
113 changes: 78 additions & 35 deletions piquasso/_backends/fock/general/calculations.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@
from piquasso.api.instruction import Instruction
from piquasso.api.result import Result

from piquasso._backends.fock.calculations import calculate_state_index_matrix_list


def vacuum(state: FockState, instruction: Instruction, shots: int) -> Result:
state.reset()
Expand Down Expand Up @@ -157,25 +159,84 @@ def cubic_phase(state: FockState, instruction: Instruction, shots: int) -> Resul
gamma = instruction._all_params["gamma"]
hbar = state._config.hbar

operator = state._space.get_single_mode_cubic_phase_operator(
matrix = state._space.get_single_mode_cubic_phase_operator(
gamma=gamma, hbar=hbar, calculator=state._calculator
)
embedded_operator = state._space.embed_matrix(
operator,
modes=instruction.modes,
auxiliary_modes=state._get_auxiliary_modes(modes=instruction.modes),
)
state._density_matrix = (
embedded_operator
@ state._density_matrix
@ embedded_operator.conjugate().transpose()
)
_apply_active_gate_matrix_to_state(state, matrix, instruction.modes[0])

state.normalize()

return Result(state=state)


def _apply_active_gate_matrix_to_state(
state: FockState,
matrix: np.ndarray,
mode: int,
) -> None:
density_matrix = state._density_matrix
space = state._space
auxiliary_subspace = state._get_subspace(state.d - 1)

state_index_matrix_list = calculate_state_index_matrix_list(
space, auxiliary_subspace, mode
)
density_matrix = _calculate_density_matrix_after_apply_active_gate(
density_matrix, matrix, state_index_matrix_list
)

state._density_matrix = density_matrix


def _calculate_density_matrix_after_apply_active_gate(
density_matrix, matrix, state_index_matrix_list
):
"""
Applies the active gate matrix to the density matrix specified by
`state_index_matrix_list`.
"""

new_density_matrix = np.empty_like(density_matrix, dtype=density_matrix.dtype)

cutoff = len(state_index_matrix_list)

for ket in range(cutoff):
state_index_matrix_ket = state_index_matrix_list[ket]
limit_ket = state_index_matrix_ket.shape[0]
sliced_matrix_ket = matrix[:limit_ket, :limit_ket]

for bra in range(max(ket + 1, cutoff)):
state_index_matrix_bra = state_index_matrix_list[bra]
limit_bra = state_index_matrix_bra.shape[0]
sliced_matrix_bra = matrix[:limit_bra, :limit_bra]

for col_ket in range(state_index_matrix_ket.shape[1]):
for col_bra in range(state_index_matrix_bra.shape[1]):
index = np.ix_(
state_index_matrix_ket[:, col_ket],
state_index_matrix_bra[:, col_bra].T,
)
partial_result = np.einsum(
"ij,jk,kl->il",
sliced_matrix_ket,
density_matrix[index],
sliced_matrix_bra.T.conj(),
)

new_density_matrix[index] = partial_result

if bra < ket:
# NOTE: We use the fact that the density matrix is self-adjoint.
index = np.ix_(
state_index_matrix_bra[:, col_bra],
state_index_matrix_ket[:, col_ket].T,
)

new_density_matrix[index] = partial_result.conj().T

return new_density_matrix


def cross_kerr(state: FockState, instruction: Instruction, shots: int) -> Result:
modes = instruction.modes
xi = instruction._all_params["xi"]
Expand All @@ -198,20 +259,11 @@ def cross_kerr(state: FockState, instruction: Instruction, shots: int) -> Result
def displacement(state: FockState, instruction: Instruction, shots: int) -> Result:
r = instruction._all_params["r"]
phi = instruction._all_params["phi"]
mode = instruction.modes[0]

operator = state._space.get_single_mode_displacement_operator(r=r, phi=phi)

embedded_operator = state._space.embed_matrix(
operator,
modes=instruction.modes,
auxiliary_modes=state._get_auxiliary_modes(modes=instruction.modes),
)
matrix = state._space.get_single_mode_displacement_operator(r=r, phi=phi)

state._density_matrix = (
embedded_operator
@ state._density_matrix
@ embedded_operator.conjugate().transpose()
)
_apply_active_gate_matrix_to_state(state, matrix, mode=mode)

state.normalize()

Expand All @@ -221,23 +273,14 @@ def displacement(state: FockState, instruction: Instruction, shots: int) -> Resu
def squeezing(state: FockState, instruction: Instruction, shots: int) -> Result:
r = instruction._all_params["r"]
phi = instruction._all_params["phi"]
mode = instruction.modes[0]

operator = state._space.get_single_mode_squeezing_operator(
matrix = state._space.get_single_mode_squeezing_operator(
r=r,
phi=phi,
)

embedded_operator = state._space.embed_matrix(
operator,
modes=instruction.modes,
auxiliary_modes=state._get_auxiliary_modes(modes=instruction.modes),
)

state._density_matrix = (
embedded_operator
@ state._density_matrix
@ embedded_operator.conjugate().transpose()
)
_apply_active_gate_matrix_to_state(state, matrix, mode=mode)

state.normalize()

Expand Down
44 changes: 2 additions & 42 deletions piquasso/_backends/fock/pure/calculations/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,17 +17,10 @@

from typing import Optional, Tuple, Mapping

from functools import lru_cache

import random
import numpy as np

from piquasso._math.fock import cutoff_cardinality

from piquasso._math.indices import (
get_index_in_fock_space,
get_auxiliary_modes,
)
from ...calculations import calculate_state_index_matrix_list

from ..state import PureFockState
from ..batch_state import BatchPureFockState
Expand Down Expand Up @@ -124,7 +117,7 @@ def _apply_active_gate_matrix(state_vector, matrix):
state_vector = calculator.maybe_convert_to_numpy(state_vector)
matrix = calculator.maybe_convert_to_numpy(matrix)

state_index_matrix_list = _calculate_state_index_matrix_list(
state_index_matrix_list = calculate_state_index_matrix_list(
space, auxiliary_subspace, mode
)
new_state_vector = _calculate_state_vector_after_apply_active_gate(
Expand All @@ -138,39 +131,6 @@ def _apply_active_gate_matrix(state_vector, matrix):
state._state_vector = _apply_active_gate_matrix(state_vector, matrix)


@lru_cache(maxsize=None)
def _calculate_state_index_matrix_list(space, auxiliary_subspace, mode):
d = space.d
cutoff = space.cutoff

state_index_matrix_list = []
auxiliary_modes = get_auxiliary_modes(d, (mode,))
all_occupation_numbers = np.zeros(d, dtype=int)

indices = [cutoff_cardinality(cutoff=n - 1, d=d - 1) for n in range(1, cutoff + 2)]

for n in range(cutoff):
limit = space.cutoff - n
subspace_size = indices[n + 1] - indices[n]

state_index_matrix = np.empty(shape=(limit, subspace_size), dtype=int)

for i, auxiliary_occupation_numbers in enumerate(
auxiliary_subspace[indices[n] : indices[n + 1]]
):
all_occupation_numbers[(auxiliary_modes,)] = auxiliary_occupation_numbers

for j in range(limit):
all_occupation_numbers[mode] = j
state_index_matrix[j, i] = get_index_in_fock_space(
tuple(all_occupation_numbers)
)

state_index_matrix_list.append(state_index_matrix)

return state_index_matrix_list


def _calculate_state_vector_after_apply_active_gate(
state_vector, matrix, state_index_matrix_list
):
Expand Down

0 comments on commit 9e98868

Please sign in to comment.