perf(fock): Caching subspaces

Calculating the subspaces (the indices, essentially) is a costly calculation,
and it is beneficial to have them stored in a cache for later purposes.

piquasso/_backends/fock/pure/calculations.py

@@ -336,12 +336,7 @@ def passive_linear(
         state._calculator, state._config
     ).astype(np.complex128)
 
-    subspace = FockSpace(
-        d=len(interferometer),
-        cutoff=state._space.cutoff,
-        calculator=calculator,
-        config=state._config,
-    )
+    subspace = state._get_subspace(dim=len(interferometer))
 
     subspace_transformations = _get_interferometer_on_fock_space(
         interferometer, subspace, calculator
@@ -553,7 +548,7 @@ def _apply_active_gate_matrix_to_state(
     calculator = state._calculator
     state_vector = state._state_vector
     space = state._space
-    auxiliary_subspace = state._auxiliary_subspace
+    auxiliary_subspace = state._get_subspace(state.d - 1)
 
     @calculator.custom_gradient
     def _apply_active_gate_matrix(state_vector, matrix):

piquasso/_backends/fock/state.py

@@ -31,21 +31,40 @@ def __init__(
     ) -> None:
         super().__init__(calculator=calculator, config=config)
 
-        self._space = fock.FockSpace(
-            d=d,
-            cutoff=self._config.cutoff,
-            calculator=calculator,
-            config=self._config,
-        )
-        # NOTE: This is instantiated here, since it is costly to do so, and is needed
-        # for several, repeating calculations.
-        self._auxiliary_subspace = fock.FockSpace(
-            d=d - 1,
+        self._initialize_subspaces(d)
+
+    def _initialize_subspaces(self, d):
+        """
+        NOTE: Calculating the subspaces (the indices, essentially) is a costly
+        calculation, and it is beneficial to have them stored for later purposes
+        """
+        self._space = self._init_subspace(d)
+
+        self._subspace_cache = {
+            1: self._init_subspace(1),
+            2: self._init_subspace(2),
+            d - 1: self._init_subspace(d - 1),
+            d: self._space,
+        }
+
+    def _init_subspace(self, dim):
+        return fock.FockSpace(
+            d=dim,
             cutoff=self._config.cutoff,
-            calculator=calculator,
+            calculator=self._calculator,
             config=self._config,
         )
 
+    def _get_subspace(self, dim):
+        if dim in self._subspace_cache.keys():
+            return self._subspace_cache[dim]
+
+        subspace = self._init_subspace(dim)
+
+        self._subspace_cache[dim] = subspace
+
+        return subspace
+
     @property
     def d(self) -> int:
         return self._space.d