PML-388: fixed random phase initialization in entangling layer if trainable is False

merlin/builder/circuit_builder.py

@@ -575,6 +575,7 @@ def add_entangling_layer(
         name: str | None = None,
         trainable_inner: bool | None = None,
         trainable_outer: bool | None = None,
+        seed: int | None = None,
     ) -> "CircuitBuilder":
         """Add an entangling layer spanning a range of modes.
 
@@ -595,6 +596,9 @@ def add_entangling_layer(
             Override for the internal phase shifters.
         trainable_outer : bool | None
             Override for the output phase shifters.
+        seed : int | None
+            Optional seed for the random phases assigned to non-trainable
+            phase shifters. ``None`` draws fresh entropy each call.
 
         Returns
         -------
@@ -648,6 +652,7 @@ def add_entangling_layer(
             model=normalized_model,
             trainable_inner=trainable_inner,
             trainable_outer=trainable_outer,
+            seed=seed,
         )
 
         self.circuit.add(component)
@@ -939,16 +944,19 @@ def _mzi_factory(
                             component, "trainable_outer", component.trainable
                         ),
                         base: str = prefix,
+                        block: GenericInterferometer = component,
                     ):
                         """Build a Mach-Zehnder interferometer optionally parameterised per index."""
                         if inner_trainable:
                             phi_inner = pcvl_module.P(f"{base}_li{i}")
                         else:
-                            phi_inner = 0.0
+                            fixed = block.fixed_inner_values
+                            phi_inner = fixed[i] if i < len(fixed) else 0.0
                         if outer_trainable:
                             phi_outer = pcvl_module.P(f"{base}_lo{i}")
                         else:
-                            phi_outer = 0.0
+                            fixed = block.fixed_outer_values
+                            phi_outer = fixed[i] if i < len(fixed) else 0.0
                         return (
                             pcvl_module.BS()
                             // pcvl_module.PS(phi_inner)
@@ -975,16 +983,19 @@ def _bell_factory(
                             component, "trainable_outer", component.trainable
                         ),
                         base: str = prefix,
+                        block: GenericInterferometer = component,
                     ):
                         """Build a Mach-Zehnder interferometer optionally parameterised per index."""
                         if inner_trainable:
                             phi_inner = pcvl_module.P(f"{base}_li{i}")
                         else:
-                            phi_inner = 0.0
+                            fixed = block.fixed_inner_values
+                            phi_inner = fixed[i] if i < len(fixed) else 0.0
                         if outer_trainable:
                             phi_outer = pcvl_module.P(f"{base}_lo{i}")
                         else:
-                            phi_outer = 0.0
+                            fixed = block.fixed_outer_values
+                            phi_outer = fixed[i] if i < len(fixed) else 0.0
 
                         circuit = pcvl_module.Circuit(2)
                         circuit.add(0, pcvl_module.BS())

merlin/core/components.py

@@ -25,7 +25,9 @@
 Components are platform-agnostic and focus on intent rather than implementation.
 """
 
-from dataclasses import dataclass
+import math
+import random
+from dataclasses import dataclass, field
 from enum import Enum
 from typing import Any
 
@@ -189,6 +191,27 @@ class GenericInterferometer:
         Whether inner phase shifters are trainable.
     trainable_outer : bool | None
         Whether outer phase shifters are trainable.
+    seed : int | None
+        Optional seed controlling the random phases drawn for non-trainable
+        inner/outer phase shifters. ``None`` draws fresh entropy without
+        touching global random state.
+    fixed_inner_values : list[float]
+        Fixed (non-trainable) values for the inner phase shifters, drawn
+        randomly unless explicitly provided. When provided, must contain
+        exactly ``span * (span - 1) // 2`` entries.
+    fixed_outer_values : list[float]
+        Fixed (non-trainable) values for the outer phase shifters, drawn
+        randomly unless explicitly provided. When provided, must contain
+        exactly ``span * (span - 1) // 2`` entries.
+
+    Raises
+    ------
+    TypeError
+        If ``model`` is not a string.
+    ValueError
+        If ``model`` is not ``"mzi"`` or ``"bell"``, or if a provided
+        ``fixed_inner_values``/``fixed_outer_values`` list does not contain
+        exactly one entry per phase shifter.
     """
 
     start_mode: int
@@ -198,6 +221,9 @@ class GenericInterferometer:
     model: str = "mzi"
     trainable_inner: bool | None = None
     trainable_outer: bool | None = None
+    seed: int | None = None
+    fixed_inner_values: list[float] = field(default_factory=list)
+    fixed_outer_values: list[float] = field(default_factory=list)
 
     def __post_init__(self):
         """Validate and normalize the interferometer configuration."""
@@ -219,6 +245,35 @@ def __post_init__(self):
         # Normalise the aggregate flag so downstream logic can rely on it
         self.trainable = bool(self.trainable_inner or self.trainable_outer)
 
+        # Non-trainable phase shifters must still be initialised to random
+        # phases (not a constant 0.0), otherwise a "fixed" interferometer
+        # degenerates into a deterministic swap/identity instead of the
+        # random fixed unitary that untrained entangling layers are meant to
+        # provide (e.g. for reservoir-computing style architectures).
+        count = self.span * (self.span - 1) // 2 if self.span > 1 else 0
+        # A partially filled list would silently fall back to 0.0 phases for
+        # the missing indices downstream, so explicit values must cover every
+        # phase shifter.
+        for label, values in (
+            ("fixed_inner_values", self.fixed_inner_values),
+            ("fixed_outer_values", self.fixed_outer_values),
+        ):
+            if values and len(values) != count:
+                raise ValueError(
+                    f"GenericInterferometer {label} must contain exactly "
+                    f"{count} entries for span {self.span}, got {len(values)}"
+                )
+        if count > 0:
+            rng = random.Random(self.seed)
+            if not self.trainable_inner and not self.fixed_inner_values:
+                self.fixed_inner_values = [
+                    rng.uniform(0, 2 * math.pi) for _ in range(count)
+                ]
+            if not self.trainable_outer and not self.fixed_outer_values:
+                self.fixed_outer_values = [
+                    rng.uniform(0, 2 * math.pi) for _ in range(count)
+                ]
+
     def get_params(self) -> dict[str, Any]:
         """Return placeholder names for every internal interferometer parameter.
 

merlin/pcvl_pytorch/noisy_slos.py

@@ -109,6 +109,11 @@ def __init__(
                     for n_i in range(1, (2 * self.n_photons) + 1)
                 ],
             )
+            # Kept for the extra-photon sectors: photon-number-indexed SLOS
+            # graphs spanning up to 2*n_photons, reused by
+            # _augmented_obb_probs to simulate both the (grown) coherent cell
+            # and the one-hot distinguishable cells of each OBB partition.
+            self._regular_slos_graphs = regular_slos_graphs
             self._slos_graphs = [
                 NoisySLOSComputeGraph(
                     noise_groups=noise_groups,
@@ -193,6 +198,98 @@ def _get_extra_photon_combinations(
 
         return output
 
+    def _augmented_obb_probs(
+        self,
+        unitary: torch.Tensor,
+        obb: _InputStateNoisySLOSComputeGraph,
+        extra_vec: torch.Tensor,
+        n_augmented: int,
+    ) -> torch.Tensor:
+        """Compute one g2 extra-photon sector for a fixed emission combination.
+
+        Perceval's source model draws one independent coherent-vs-distinguishable
+        outcome per *original* input photon only -- the same Orthogonal Bad Bits
+        partitions ``obb`` already built for the noiseless ``n_photons`` state --
+        never an extra independent draw per g2-emitted sibling. Whatever the fate
+        of the original photons in a given partition cell, every g2 sibling
+        always joins that cell's coherent (non-excluded) group. This reuses
+        ``obb``'s cached partitions/weights, growing only the coherent cell's
+        Fock state by ``extra_vec`` before convolving with the (unchanged)
+        one-hot distinguishable cells.
+
+        Parameters
+        ----------
+        unitary : torch.Tensor
+            Circuit unitary, batched ``[batch, m, m]``.
+        obb : _InputStateNoisySLOSComputeGraph
+            Cached OBB graph for the base (non-augmented) input state.
+        extra_vec : torch.Tensor
+            Per-mode count of g2-emitted photons for this combination.
+        n_augmented : int
+            Total photon number of this sector (``n_photons`` + emitted).
+
+        Returns
+        -------
+        torch.Tensor
+            Probabilities over the ``n_augmented``-photon Fock basis, shape
+            ``[batch, n_states]``.
+        """
+        batch_size = unitary.size(0)
+        fock_keys = [
+            tuple(row) for row in self._fock_states_per_n[n_augmented].tolist()
+        ]
+        key_to_idx = {key: idx for idx, key in enumerate(fock_keys)}
+        output_probs = torch.zeros(
+            batch_size, len(fock_keys), device=self.device, dtype=self.dtype
+        )
+
+        for order, (cells, counts) in enumerate(obb._partitions):
+            bit_weight = obb._weights[order]
+            for cell, count in zip(cells, counts, strict=True):
+                if order == obb.n_photons:
+                    base_state = extra_vec
+                    bad_states = cell
+                else:
+                    base_state = cell[0] + extra_vec
+                    bad_states = cell[1:]
+
+                base_n = int(base_state.sum().item())
+                _, base_probs = self._regular_slos_graphs[base_n - 1].compute_probs(
+                    unitary, base_state
+                )
+                if base_probs.ndim == 1:
+                    base_probs = base_probs.unsqueeze(0)
+                cell_keys = [self._fock_states_per_n[base_n]]
+                cell_probs = [base_probs]
+                for bad_state in bad_states:
+                    _, bad_probs = self._regular_slos_graphs[0].compute_probs(
+                        unitary, bad_state
+                    )
+                    if bad_probs.ndim == 1:
+                        bad_probs = bad_probs.unsqueeze(0)
+                    cell_keys.append(self._fock_states_per_n[1])
+                    cell_probs.append(bad_probs)
+
+                conv_keys, conv_probs = convolve_distributions(cell_keys, *cell_probs)
+                if conv_probs.ndim == 1:
+                    conv_probs = conv_probs.unsqueeze(0)
+                conv_keys_list = (
+                    [tuple(k.tolist()) for k in conv_keys]
+                    if isinstance(conv_keys, torch.Tensor)
+                    else [tuple(k) for k in conv_keys]
+                )
+
+                for local_idx, key in enumerate(conv_keys_list):
+                    idx = key_to_idx.get(key)
+                    if idx is not None:
+                        output_probs[:, idx] = (
+                            output_probs[:, idx]
+                            + bit_weight * count.item() * conv_probs[:, local_idx]
+                        )
+
+        output_probs = output_probs / output_probs.sum(dim=1, keepdim=True)
+        return output_probs
+
     def compute_probs(
         self,
         unitary: torch.Tensor,
@@ -314,14 +411,19 @@ def compute_probs(
                         probs = reordered_probs
 
                     else:
-                        input_state_to_run = list(input_state)
+                        extra_vec = torch.zeros(
+                            self.m, dtype=torch.int32, device=self.device
+                        )
                         for photon in combination:
-                            input_state_to_run[photon] += 1
-                        probs = cast(
-                            torch.Tensor,
-                            slos_graphs_list[num_photons_added].compute_probs(
-                                unitary, input_state_to_run
-                            ),
+                            extra_vec[photon] += 1
+                        obb = slos_graphs_list[0]._slos_graph_per_input[
+                            tuple(input_state)
+                        ]
+                        probs = self._augmented_obb_probs(
+                            unitary,
+                            obb,
+                            extra_vec,
+                            self.n_photons + num_photons_added,
                         )
 
                     sector.tensor = sector.tensor + weight_k * probs
@@ -367,6 +469,9 @@ def to(self, device: str | torch.device) -> NoisyG2SLOSComputeGraph:
         else:
             for graph in self._slos_graphs:
                 graph.to(self.device)
+            if hasattr(self, "_regular_slos_graphs"):
+                for regular_graph in self._regular_slos_graphs:
+                    regular_graph.to(self.device)
 
         return self