From c450c083f6050934d050611a22f88d5d4f63bf6b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Zolt=C3=A1n=20Kolarovszki?= Date: Thu, 8 Feb 2024 16:49:47 +0100 Subject: [PATCH] feat(fock/pure): `get_tensor_representation` `PureFockSimulator` calculates with a global cutoff (i.e., cutoff in the total number of particles), therefore it does not represent the state vector as a tensor. `get_tensor_representation` is created to reshape the current state vector representation into a tensor. --- piquasso/_backends/calculator.py | 7 ++--- piquasso/_backends/fock/pure/state.py | 18 +++++++++++++ piquasso/_backends/tensorflow/calculator.py | 6 ++--- piquasso/_math/fock.py | 2 +- piquasso/api/calculator.py | 2 +- tests/api/test_calculator.py | 2 +- tests/backends/fock/pure/test_state.py | 29 +++++++++++++++++++++ 7 files changed, 57 insertions(+), 9 deletions(-) diff --git a/piquasso/_backends/calculator.py b/piquasso/_backends/calculator.py index 70d309f1..4989acef 100644 --- a/piquasso/_backends/calculator.py +++ b/piquasso/_backends/calculator.py @@ -50,9 +50,10 @@ def assign(self, array, index, value): return array - def scatter(self, indices, updates, dim): - embedded_matrix = np.zeros((dim,) * 2, dtype=complex) - composite_index = np.array(indices)[:, 0], np.array(indices)[:, 1] + def scatter(self, indices, updates, shape): + embedded_matrix = np.zeros(shape, dtype=complex) + indices_array = np.array(indices) + composite_index = tuple([indices_array[:, i] for i in range(len(shape))]) embedded_matrix[composite_index] = np.array(updates) diff --git a/piquasso/_backends/fock/pure/state.py b/piquasso/_backends/fock/pure/state.py index 41a8b659..62f6befa 100644 --- a/piquasso/_backends/fock/pure/state.py +++ b/piquasso/_backends/fock/pure/state.py @@ -275,3 +275,21 @@ def mean_photon_number(self): accumulator += number * self._np.abs(self._state_vector[index]) ** 2 return accumulator + + def get_tensor_representation(self): + calculator = self._calculator + cutoff = self._config.cutoff + d = self.d + + indices = [] + updates = [] + + for index, number in enumerate(self._space): + indices.append(list(number)) + updates.append(self._state_vector[index]) + + return calculator.scatter( + indices, + updates, + [cutoff] * d, + ) diff --git a/piquasso/_backends/tensorflow/calculator.py b/piquasso/_backends/tensorflow/calculator.py index a554c01e..a9419fca 100644 --- a/piquasso/_backends/tensorflow/calculator.py +++ b/piquasso/_backends/tensorflow/calculator.py @@ -78,8 +78,8 @@ def block(self, arrays): return self.np.stack(output) - def scatter(self, indices, updates, dim): - return self._tf.scatter_nd(indices, updates, (dim, dim)) + def scatter(self, indices, updates, shape): + return self._tf.scatter_nd(indices, updates, shape) def embed_in_identity(self, matrix, indices, dim): tf_indices = [] @@ -100,7 +100,7 @@ def embed_in_identity(self, matrix, indices, dim): tf_indices.append(diagonal_index) updates.append(1.0) - return self.scatter(tf_indices, updates, dim) + return self.scatter(tf_indices, updates, (dim, dim)) def _funm(self, matrix, func): eigenvalues, U = self._tf.linalg.eig(matrix) diff --git a/piquasso/_math/fock.py b/piquasso/_math/fock.py index f01f0b9c..fea80ddf 100644 --- a/piquasso/_math/fock.py +++ b/piquasso/_math/fock.py @@ -261,7 +261,7 @@ def embed_matrix( indices.append(embedded_index) updates.append(matrix[index][0]) - return self._calculator.scatter(indices, updates, dim=self.cardinality) + return self._calculator.scatter(indices, updates, shape=(self.cardinality,) * 2) def get_linear_fock_operator( self, diff --git a/piquasso/api/calculator.py b/piquasso/api/calculator.py index 24f61f0d..9ba1e319 100644 --- a/piquasso/api/calculator.py +++ b/piquasso/api/calculator.py @@ -62,7 +62,7 @@ def loop_hafnian( def assign(self, array, index, value): raise NotImplementedCalculation() - def scatter(self, indices, updates, dim): + def scatter(self, indices, updates, shape): raise NotImplementedCalculation() def embed_in_identity(self, matrix, indices, dim): diff --git a/tests/api/test_calculator.py b/tests/api/test_calculator.py index e4977347..66bcd63c 100644 --- a/tests/api/test_calculator.py +++ b/tests/api/test_calculator.py @@ -84,7 +84,7 @@ def test_BaseCalculator_raises_NotImplementedCalculation_for_scatter( empty_calculator, ): with pytest.raises(NotImplementedCalculation): - empty_calculator.scatter(indices=[], updates=[], dim=2) + empty_calculator.scatter(indices=[], updates=[], shape=(3, 3)) def test_BaseCalculator_raises_NotImplementedCalculation_for_embed_in_identity( diff --git a/tests/backends/fock/pure/test_state.py b/tests/backends/fock/pure/test_state.py index 17420593..dd36611c 100644 --- a/tests/backends/fock/pure/test_state.py +++ b/tests/backends/fock/pure/test_state.py @@ -168,3 +168,32 @@ def test_normalize_if_disabled_in_Config(): norm = state.norm assert not np.isclose(norm, 1.0) + + +def test_PureFockState_get_tensor_representation(): + d = 2 + cutoff = 3 + + with pq.Program() as program: + pq.Q() | pq.StateVector([0, 1]) / 2 + + pq.Q() | pq.StateVector([0, 2]) / 2 + pq.Q() | pq.StateVector([2, 0]) / np.sqrt(2) + + simulator = pq.PureFockSimulator(d=d, config=pq.Config(cutoff=cutoff)) + state = simulator.execute(program).state + + state_tensor = state.get_tensor_representation() + + assert state_tensor.shape == (3,) * 2 + + assert np.allclose( + state_tensor, + np.array( + [ + [0.0 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j], + [0.0 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j], + [0.70710678 + 0.0j, 0.0 + 0.0j, 0.0 + 0.0j], + ] + ), + )