perf(fock): Caching subspaces

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

tests/backends/test_backend_equivalence.py

@@ -1439,3 +1439,77 @@ def test_Attenuator_raises_InvalidParam_for_non_zero_mean_thermal_excitation(
         "Non-zero mean thermal excitation is not supported in this backend."
         in error.value.args[0]
     )
+
+
+@pytest.mark.parametrize(
+    "SimulatorClass",
+    (
+        pq.PureFockSimulator,
+        pq.FockSimulator,
+        pq.GaussianSimulator,
+        pq.TensorflowPureFockSimulator,
+    ),
+)
+def test_Interferometer_smaller_than_system_size(SimulatorClass):
+    config = pq.Config(cutoff=3)
+
+    d = 5
+
+    interferometer_matrix = np.array(
+        [
+            [
+                0.67622072 - 0.02632995j,
+                -0.41993552 - 0.43404319j,
+                0.21135489 + 0.36417311j,
+            ],
+            [
+                0.48143126 - 0.18351759j,
+                -0.05256469 + 0.2714796j,
+                -0.7896096 - 0.18600457j,
+            ],
+            [
+                0.47764293 - 0.22007895j,
+                0.74644107 + 0.0402763j,
+                0.35346034 - 0.19922808j,
+            ],
+        ]
+    )
+
+    with pq.Program() as program:
+        pq.Q() | pq.Vacuum()
+
+        for i in range(d):
+            pq.Q(i) | pq.Squeezing(r=0.01 * i)
+
+        pq.Q(0, 4, 2) | pq.Interferometer(interferometer_matrix)
+
+    simulator = SimulatorClass(d=5, config=config)
+
+    state = simulator.execute(program).state
+
+    assert is_proportional(
+        state.fock_probabilities,
+        [
+            0.99850342,
+            0.0,
+            0.0,
+            0.0,
+            0.0,
+            0.0,
+            0.00015807,
+            0.0,
+            0.0002202,
+            0.0,
+            0.00011654,
+            0.00004992,
+            0.0,
+            0.0,
+            0.0,
+            0.00028562,
+            0.0,
+            0.00016609,
+            0.00044906,
+            0.0,
+            0.00005109,
+        ],
+    )