feat(config): Moving reduction to loop hafnian

The reduction step is moved to the `loop_hafnian_function` in `config`.

piquasso/_backends/gaussian/calculations.py

@@ -32,7 +32,6 @@
 from piquasso.api.instruction import Instruction
 from piquasso.api.errors import InvalidInstruction
 
-from piquasso._math.linalg import reduce_
 from piquasso._math.indices import get_operator_index, get_auxiliary_operator_index
 from piquasso._math.decompositions import decompose_to_pure_and_mixed
 
@@ -413,12 +412,9 @@ def _get_particle_number_choice(
     for n in possible_choices:
         occupation_numbers = previous_sample + (n,)
 
-        B_reduced = reduce_(B, reduce_on=occupation_numbers)
-        gamma_reduced = reduce_(gamma, reduce_on=occupation_numbers)
-
-        np.fill_diagonal(B_reduced, gamma_reduced)
-
-        weight = abs(state._config.loop_hafnian_function(B_reduced)) ** 2 / factorial(n)
+        weight = abs(
+            state._config.loop_hafnian_function(B, gamma, occupation_numbers)
+        ) ** 2 / factorial(n)
         weights = np.append(weights, weight)
 
     weights /= np.sum(weights)

piquasso/_backends/gaussian/probabilities.py

@@ -19,7 +19,7 @@
 
 from scipy.special import factorial
 
-from piquasso._math.linalg import block_reduce, reduce_
+from piquasso._math.linalg import block_reduce
 from piquasso._math.torontonian import torontonian
 
 from piquasso.api.typing import HafnianFunction
@@ -56,23 +56,14 @@ def __init__(
 
         self.loop_hafnian_function = loop_hafnian_function
 
-    def _get_A_reduced(self, reduce_on: Tuple[int, ...]) -> np.ndarray:
-        A_reduced = reduce_(self._A, reduce_on=reduce_on)
-
-        np.fill_diagonal(A_reduced, reduce_(self._gamma, reduce_on=reduce_on))
-
-        return A_reduced
-
     def get_density_matrix_element(
         self, bra: Tuple[int, ...], ket: Tuple[int, ...]
     ) -> float:
         reduce_on = ket + bra
 
-        A_reduced = self._get_A_reduced(reduce_on=reduce_on)
-
         return (
             self._normalization
-            * self.loop_hafnian_function(A_reduced)
+            * self.loop_hafnian_function(self._A, self._gamma, reduce_on)
             / np.sqrt(np.prod(factorial(reduce_on)))
         )
 

piquasso/_math/hafnian.py

@@ -14,14 +14,15 @@
 # limitations under the License.
 
 import math
-from functools import lru_cache, partial
+from functools import lru_cache
 from itertools import combinations_with_replacement
 from typing import List, Callable
 
 import numpy as np
 from scipy.linalg import block_diag
 
 from .combinatorics import powerset
+from .linalg import reduce_
 
 
 @lru_cache()
@@ -149,6 +150,24 @@ def _get_loop_polynom_coefficients(
     return ret
 
 
-hafnian = partial(_hafnian, polynom_function=_get_polynom_coefficients)
+def hafnian(matrix):
+    return _hafnian(matrix, polynom_function=_get_polynom_coefficients)
 
-loop_hafnian = partial(_hafnian, polynom_function=_get_loop_polynom_coefficients)
+
+def loop_hafnian(matrix):
+    return _hafnian(matrix, polynom_function=_get_loop_polynom_coefficients)
+
+
+def _reduce_matrix_with_diagonal(matrix, diagonal, reduce_on):
+    reduced_diagonal = reduce_(diagonal, reduce_on=reduce_on)
+    reduced_matrix = reduce_(matrix, reduce_on=reduce_on)
+
+    np.fill_diagonal(reduced_matrix, reduced_diagonal)
+
+    return reduced_matrix
+
+
+def loop_hafnian_with_reduction(matrix, diagonal, reduce_on):
+    reduced_matrix = _reduce_matrix_with_diagonal(matrix, diagonal, reduce_on)
+
+    return loop_hafnian(reduced_matrix)

piquasso/api/config.py

@@ -21,7 +21,7 @@
 import numpy as np
 
 from piquasso._math.permanent import glynn_gray_permanent
-from piquasso._math.hafnian import loop_hafnian
+from piquasso._math.hafnian import loop_hafnian_with_reduction
 
 from piquasso.api.typing import PermanentFunction, HafnianFunction
 
@@ -41,7 +41,7 @@ def __init__(
         cutoff: int = 4,
         measurement_cutoff: int = 5,
         permanent_function: PermanentFunction = glynn_gray_permanent,
-        loop_hafnian_function: HafnianFunction = loop_hafnian,
+        loop_hafnian_function: HafnianFunction = loop_hafnian_with_reduction,
     ):
         self._original_seed_sequence = seed_sequence
         self.seed_sequence = seed_sequence or int.from_bytes(

piquasso/api/typing.py

@@ -19,4 +19,4 @@
 
 
 PermanentFunction = Callable[[np.ndarray, Tuple[int, ...], Tuple[int, ...]], float]
-HafnianFunction = Callable[[np.ndarray], float]
+HafnianFunction = Callable[[np.ndarray, np.ndarray, Tuple[int, ...]], float]