Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

perf: Rewrite get_fock_space_basis #361

Merged
merged 1 commit into from
Oct 16, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 20 additions & 16 deletions piquasso/_math/combinatorics.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from itertools import chain, combinations
from typing import Tuple, Iterable, Iterator, TypeVar

import numpy as np
import numba as nb


@nb.njit
@nb.njit(cache=True)
def arr_comb(n, k):
n = np.where((n < 0) | (n < k), 0, n)
prod = np.ones(n.shape, dtype=np.int64)
Expand All @@ -34,7 +31,7 @@ def arr_comb(n, k):

@nb.njit(cache=True)
def comb(n, k):
if n < 0 or n < k:
if n < 0 or k < 0 or n < k:
return 0

prod = 1
Expand All @@ -46,14 +43,6 @@ def comb(n, k):
return prod


_T = TypeVar("_T")


def powerset(iterable: Iterable[_T]) -> Iterator[Tuple[_T, ...]]:
s = list(iterable)
return chain.from_iterable(combinations(iterable, r) for r in range(len(s) + 1))


def is_even_permutation(permutation):
permutation = list(permutation)
length = len(permutation)
Expand All @@ -70,15 +59,30 @@ def is_even_permutation(permutation):
return (length - cycles) % 2 == 0


@nb.njit
def partitions(boxes, particles):
@nb.njit(cache=True)
def partitions(boxes, particles, out=None):
r"""
Returns all the possible ways to put a specified number of particles in a specified
number of boxes in anti-lexicographic order.

Args:
boxes: Number of boxes.
particles: Number of particles.
out: Optional output array.
"""

positions = particles + boxes - 1

if positions == -1 or boxes == 0:
return np.empty((1, 0), dtype=np.int32)

size = comb(positions, boxes - 1)
result = np.empty((size, boxes), dtype=np.int32)

if out is None:
result = np.empty((size, boxes), dtype=np.int32)
else:
result = out

separators = np.arange(boxes - 1, dtype=np.int32)
index = size - 1

Expand Down
39 changes: 18 additions & 21 deletions piquasso/_math/fock.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,7 @@
import numpy as np
import numba as nb

from scipy.special import comb
from piquasso._math.combinatorics import comb as nb_comb
from piquasso._math.combinatorics import comb

from piquasso.api.config import Config
from piquasso._math.combinatorics import partitions
Expand All @@ -37,44 +36,42 @@
from piquasso.api.connector import BaseConnector


@functools.lru_cache()
def cutoff_cardinality(*, cutoff: int, d: int) -> int:
@nb.njit(cache=True)
def cutoff_fock_space_dim(cutoff: int, d: int) -> int:
r"""
Calculates the dimension of the cutoff Fock space with the relation

..math::
\sum_{i=0}^{c - 1} {d + i - 1 \choose i} = {d + c - 1 \choose c - 1}.
\sum_{i=0}^{c - 1} {d + i - 1 \choose i} = {d + c - 1 \choose d}.
"""
return comb(d + cutoff - 1, cutoff - 1, exact=True)
return comb(d + cutoff - 1, d)


@nb.njit
def cutoff_cardinality_array(cutoff, d):
@nb.njit(cache=True)
def cutoff_fock_space_dim_array(cutoff, d):
ret = np.empty(cutoff.shape, dtype=np.int32)

for i in range(len(cutoff)):
ret[i] = nb_comb(d + cutoff[i] - 1, d)
ret[i] = comb(d + cutoff[i] - 1, d)

return ret


def symmetric_subspace_cardinality_array(*, d: int, n: int) -> int:
return np.round(comb(d + n - 1, n)).astype(int)
@nb.njit(cache=True)
def symmetric_subspace_cardinality(d: int, n: int) -> int:
return comb(d + n - 1, n)


@nb.njit(cache=True)
def nb_get_fock_space_basis(d: int, cutoff: int) -> np.ndarray:
partitions_list = [partitions(boxes=d, particles=n) for n in range(cutoff)]

total_elements = 0
for part in partitions_list:
total_elements += part.shape[0]
size = cutoff_fock_space_dim(cutoff=cutoff, d=d)

ret = np.empty((total_elements, d), dtype=partitions_list[0].dtype)
ret = np.empty((size, d), dtype=np.int32)
current_row = 0
for part in partitions_list:
num_rows = part.shape[0]
ret[current_row : current_row + num_rows, :] = part
for n in range(cutoff):
num_rows = symmetric_subspace_cardinality(d, n)
out = ret[current_row : current_row + num_rows, :]
_ = partitions(boxes=d, particles=n, out=out)
current_row += num_rows

return ret
Expand Down Expand Up @@ -161,7 +158,7 @@ def get_creation_operator(
modes: Tuple[int, ...], space: np.ndarray, config: Config
) -> np.ndarray:
d = len(space[0])
size = cutoff_cardinality(cutoff=config.cutoff, d=d)
size = cutoff_fock_space_dim(cutoff=config.cutoff, d=d)
operator = np.zeros(shape=(size,) * 2, dtype=config.complex_dtype)

for index, basis in enumerate(space):
Expand Down
14 changes: 7 additions & 7 deletions piquasso/_simulators/fock/calculations.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,8 +29,8 @@
from piquasso.api.exceptions import InvalidInstruction, InvalidParameter

from piquasso._math.fock import (
cutoff_cardinality,
cutoff_cardinality_array,
cutoff_fock_space_dim,
cutoff_fock_space_dim_array,
operator_basis,
get_fock_space_basis,
nb_get_fock_space_basis,
Expand Down Expand Up @@ -116,7 +116,7 @@ def calculate_state_index_matrix_list(d, cutoff, mode):

state_index_matrix_list = []

indices = cutoff_cardinality_array(cutoff=np.arange(cutoff + 1), d=d - 1)
indices = cutoff_fock_space_dim_array(cutoff=np.arange(cutoff + 1), d=d - 1)

for n in range(cutoff):
limit = cutoff - n
Expand Down Expand Up @@ -148,7 +148,7 @@ def calculate_state_index_matrix_list(d, cutoff, mode):
def calculate_interferometer_helper_indices(d, cutoff):
space = nb_get_fock_space_basis(d=d, cutoff=cutoff)

indices = cutoff_cardinality_array(cutoff=np.arange(1, cutoff + 1), d=d)
indices = cutoff_fock_space_dim_array(cutoff=np.arange(1, cutoff + 1), d=d)

subspace_index_tensor = []

Expand Down Expand Up @@ -227,8 +227,8 @@ def calculate_index_list_for_appling_interferometer(
subspace = get_fock_space_basis(d=len(modes), cutoff=cutoff)
auxiliary_subspace = get_fock_space_basis(d=d - len(modes), cutoff=cutoff)

indices = cutoff_cardinality_array(cutoff=np.arange(cutoff + 1), d=len(modes))
auxiliary_indices = cutoff_cardinality_array(
indices = cutoff_fock_space_dim_array(cutoff=np.arange(cutoff + 1), d=len(modes))
auxiliary_indices = cutoff_fock_space_dim_array(
cutoff=np.arange(cutoff + 1), d=d - len(modes)
)
auxiliary_modes = get_auxiliary_modes(d, modes)
Expand Down Expand Up @@ -258,7 +258,7 @@ def get_projection_operator_indices(d, cutoff, modes, basis_vector):

boxes = d - len(modes)

card = cutoff_cardinality(cutoff=new_cutoff, d=boxes)
card = cutoff_fock_space_dim(cutoff=new_cutoff, d=boxes)

basis = np.empty(shape=(card, d), dtype=int)

Expand Down
8 changes: 4 additions & 4 deletions piquasso/_simulators/fock/general/state.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
from piquasso.api.config import Config
from piquasso.api.exceptions import InvalidState, PiquassoException
from piquasso._math.linalg import is_selfadjoint
from piquasso._math.fock import cutoff_cardinality, get_fock_space_basis
from piquasso._math.fock import cutoff_fock_space_dim, get_fock_space_basis
from piquasso._math.indices import (
get_index_in_fock_space,
get_index_in_fock_space_array,
Expand Down Expand Up @@ -53,7 +53,7 @@ def __init__(
self._density_matrix = self._get_empty()

def _get_empty(self) -> np.ndarray:
state_vector_size = cutoff_cardinality(cutoff=self._config.cutoff, d=self.d)
state_vector_size = cutoff_fock_space_dim(cutoff=self._config.cutoff, d=self.d)
return np.zeros(
shape=(state_vector_size,) * 2, dtype=self._config.complex_dtype
)
Expand Down Expand Up @@ -116,7 +116,7 @@ def __eq__(self, other: object) -> bool:

@property
def density_matrix(self) -> np.ndarray:
cardinality = cutoff_cardinality(d=self.d, cutoff=self._config.cutoff)
cardinality = cutoff_fock_space_dim(d=self.d, cutoff=self._config.cutoff)

return self._density_matrix[:cardinality, :cardinality]

Expand Down Expand Up @@ -178,7 +178,7 @@ def reduced(self, modes: Tuple[int, ...]) -> "FockState":
)

for inner_basis in inner_space:
size = cutoff_cardinality(
size = cutoff_fock_space_dim(
cutoff=cutoff - fallback_np.sum(inner_basis), d=outer_size
)

Expand Down
12 changes: 6 additions & 6 deletions piquasso/_simulators/fock/pure/state.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@
from piquasso.api.exceptions import InvalidState, PiquassoException
from piquasso.api.connector import BaseConnector

from piquasso._math.fock import cutoff_cardinality, get_fock_space_basis
from piquasso._math.fock import cutoff_fock_space_dim, get_fock_space_basis
from piquasso._math.linalg import vector_absolute_square
from piquasso._math.indices import (
get_index_in_fock_space,
Expand Down Expand Up @@ -63,11 +63,11 @@ def __init__(
self.state_vector = self._get_empty()

def _get_empty_list(self) -> list:
state_vector_size = cutoff_cardinality(cutoff=self._config.cutoff, d=self.d)
state_vector_size = cutoff_fock_space_dim(cutoff=self._config.cutoff, d=self.d)
return [0.0] * state_vector_size

def _get_empty(self) -> np.ndarray:
state_vector_size = cutoff_cardinality(cutoff=self._config.cutoff, d=self.d)
state_vector_size = cutoff_fock_space_dim(cutoff=self._config.cutoff, d=self.d)

return self._np.zeros(
shape=(state_vector_size,), dtype=self._config.complex_dtype
Expand Down Expand Up @@ -111,7 +111,7 @@ def __eq__(self, other: object) -> bool:

@property
def density_matrix(self) -> np.ndarray:
cardinality = cutoff_cardinality(d=self.d, cutoff=self._config.cutoff)
cardinality = cutoff_fock_space_dim(d=self.d, cutoff=self._config.cutoff)

state_vector = self.state_vector[:cardinality]

Expand Down Expand Up @@ -154,7 +154,7 @@ def get_particle_detection_probability_on_modes(

unordered_modes = fallback_np.concatenate([modes, auxiliary_modes])

card = cutoff_cardinality(d=auxiliary_d, cutoff=auxiliary_cutoff)
card = cutoff_fock_space_dim(d=auxiliary_d, cutoff=auxiliary_cutoff)
auxiliary_basis = get_fock_space_basis(d=auxiliary_d, cutoff=auxiliary_cutoff)

repeated_occupation_numbers = fallback_np.repeat(
Expand Down Expand Up @@ -225,7 +225,7 @@ def _get_mean_position_indices(self, mode):
relevant_column = self._space[:, mode]

nonzero_indices_on_mode = (relevant_column > 0).nonzero()[0]
upper_index = cutoff_cardinality(d=self.d, cutoff=self._config.cutoff - 1)
upper_index = cutoff_fock_space_dim(d=self.d, cutoff=self._config.cutoff - 1)

multipliers = fallback_np.sqrt(
fallback_np.concatenate(
Expand Down
8 changes: 4 additions & 4 deletions piquasso/_simulators/sampling/state.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
from typing import Optional, List
import numpy as np

from piquasso._math.fock import cutoff_cardinality
from piquasso._math.fock import cutoff_fock_space_dim
from piquasso._math.linalg import is_unitary

from piquasso.api.config import Config
Expand Down Expand Up @@ -122,16 +122,16 @@ def state_vector(self):
particle_numbers = fallback_np.sum(self._occupation_numbers, axis=1)

state_vector = np.zeros(
shape=cutoff_cardinality(d=self.d, cutoff=self._config.cutoff),
shape=cutoff_fock_space_dim(d=self.d, cutoff=self._config.cutoff),
dtype=self._config.complex_dtype,
)

for index in range(len(self._occupation_numbers)):
particle_number = particle_numbers[index]

starting_index = cutoff_cardinality(d=self.d, cutoff=particle_number)
starting_index = cutoff_fock_space_dim(d=self.d, cutoff=particle_number)

ending_index = cutoff_cardinality(d=self.d, cutoff=particle_number + 1)
ending_index = cutoff_fock_space_dim(d=self.d, cutoff=particle_number + 1)

coefficient = self._coefficients[index]

Expand Down
4 changes: 2 additions & 2 deletions scripts/profile_cvnn_layer.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
# limitations under the License.

import piquasso as pq
from piquasso._math.fock import cutoff_cardinality
from piquasso._math.fock import cutoff_fock_space_dim
import tensorflow as tf
import time
import numpy as np
Expand Down Expand Up @@ -59,7 +59,7 @@ def create_layer_parameters(d: int):
d = 2
cutoff = 10

target_state_vector = np.zeros(cutoff_cardinality(cutoff=cutoff, d=d), dtype=complex)
target_state_vector = np.zeros(cutoff_fock_space_dim(cutoff=cutoff, d=d), dtype=complex)
target_state_vector[1] = 1.0
target_state = tf.constant(target_state_vector, dtype=tf.complex128)

Expand Down
4 changes: 2 additions & 2 deletions scripts/profile_cvnn_layer_batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
# limitations under the License.

import piquasso as pq
from piquasso._math.fock import cutoff_cardinality
from piquasso._math.fock import cutoff_fock_space_dim
import tensorflow as tf
import time
import numpy as np
Expand Down Expand Up @@ -61,7 +61,7 @@ def create_layer_parameters(d: int):

expected_mean_positions = 0.01 * np.random.rand(number_of_inputs)

target_state_vector = np.zeros(cutoff_cardinality(cutoff=cutoff, d=d), dtype=complex)
target_state_vector = np.zeros(cutoff_fock_space_dim(cutoff=cutoff, d=d), dtype=complex)
target_state_vector[1] = 1.0
target_state = tf.constant(target_state_vector, dtype=tf.complex128)

Expand Down
12 changes: 8 additions & 4 deletions tests/_math/test_torontonian.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,18 @@

import numpy as np

from itertools import chain, combinations


from piquasso._math.torontonian import torontonian, loop_torontonian
from piquasso._math.combinatorics import powerset
from piquasso._math.transformations import from_xxpp_to_xpxp_transformation_matrix


def powerset(iterable):
s = list(iterable)
return chain.from_iterable(combinations(iterable, r) for r in range(len(s) + 1))


def torontonian_naive(A: np.ndarray) -> complex:
"""
Naive torontonian implementation from formulas provided in
Expand Down Expand Up @@ -71,9 +78,6 @@ def loop_torontonian_naive(A, gamma):

d = A.shape[0] // 2

from piquasso._math.transformations import from_xxpp_to_xpxp_transformation_matrix
from piquasso._math.combinatorics import powerset

T = from_xxpp_to_xpxp_transformation_matrix(d)

A = T.T @ A @ T
Expand Down
Loading