Tal/review

docs/intro.rst

@@ -310,10 +310,10 @@ The `measure_POVM()` method will return tuple with the outcome as the first valu
     m = env.measure_POVM(operators, env.polarization)
 
 
-Reproducability
+Reproducibility
 ===============
 
-For Reproducability **Photon Weave** allows the user to set the seed that is used with random processes. The seed is configured though `Config` class. `Config` class is a singleton class and is consulted at every random process.
+For Reproducibility **Photon Weave** allows the user to set the seed that is used with random processes. The seed is configured though `Config` class. `Config` class is a singleton class and is consulted at every random process.
 
 .. code:: python
 

examples/mach_zehnder_interferometer.py

@@ -10,20 +10,33 @@
 from photon_weave.state.envelope import Envelope
 
 
-def mach_zender_single_shot(phase_shift: float):
+# Generate beam splitter operators
+bs1 = Operation(CompositeOperationType.NonPolarizingBeamSplitter, eta=jnp.pi / 4)
+ps = Operation(FockOperationType.PhaseShift, phi=0)
+bs2 = Operation(CompositeOperationType.NonPolarizingBeamSplitter, eta=jnp.pi / 4)
+
+def mach_zender_single_shot(phase_shift: float) -> list[int]:
+    """Return photon count in each port of a Mach-Zehnder Interferometer.
+
+    Parameters
+    ----------
+    phase_shift : float
+        Phase shift between the two arms of the interferometer.
+
+    Returns
+    -------
+    list[int]
+        Photon count in each port of the interferometer.
+    """
     # Create one envelope
     env1 = Envelope()
     # Create one photon
     env1.fock.state = 1
-
     # Other port will consume vacuum
     env2 = Envelope()
-
-    # Generate operators
-    bs1 = Operation(CompositeOperationType.NonPolarizingBeamSplitter, eta=jnp.pi / 4)
-    ps = Operation(FockOperationType.PhaseShift, phi=phase_shift)
-    bs2 = Operation(CompositeOperationType.NonPolarizingBeamSplitter, eta=jnp.pi / 4)
-
+    # create phase shift operation
+    ps.kwargs["phi"] = phase_shift
+    # Apply operations
     ce = CompositeEnvelope(env1, env2)
     ce.apply_operation(bs1, env1.fock, env2.fock)
     env1.fock.apply_operation(ps)
@@ -35,21 +48,21 @@ def mach_zender_single_shot(phase_shift: float):
 
 
 if __name__ == "__main__":
-    num_shots = 1000
-    angles = jnp.linspace(0, 2 * jnp.pi, 25)
-    results = {float(angle): [] for angle in angles}
-    for angle in angles:
-        for _ in range(num_shots):
+    from tqdm import tqdm
+    import numpy as np
+    num_shots = 100
+    angles = jnp.linspace(0, 2 * jnp.pi, 10)
+    # (num_angles, num_shots , num_ports) Pre allocating this yields a x2 speedup for free
+    results = np.zeros((angles.shape[0], num_shots, 2))
+    pbar = tqdm(total=len(angles) * num_shots, desc="Simulating Mach-Zehnder Interferometer")
+    for i, angle in enumerate(angles):
+        for j in range(num_shots):
             shot_result = mach_zender_single_shot(angle)
-            results[float(angle)].append(shot_result)
-
-    measurements_1 = []
-    measurements_2 = []
-    for angle, shots in results.items():
-        counts_1 = sum(shot[0] for shot in shots) / num_shots
-        counts_2 = sum(shot[1] for shot in shots) / num_shots
-        measurements_1.append(counts_1)
-        measurements_2.append(counts_2)
+            results[i, j, :] = shot_result
+            pbar.update(1)
+    pbar.close()
+    measurements_1 = results.mean(axis=1)[:, 0]
+    measurements_2 = results.mean(axis=1)[:, 1]
 
     # Plot results
     plt.figure(figsize=(10, 6))

photon_weave/operation/polarization_operation.py

@@ -26,12 +26,12 @@ class PolarizationOperationType(Enum):
 
     Constructs an operator, which acts on a single Polarization Space
 
-    Identity (I)
+    Identity (Id)
     ------------
-    Constructs Identity (:math:`\hat{I}`) operator
+    Constructs Identity (:math:`\hat{Id}`) operator
     .. math::
 
-        \hat{I} = \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}
+        \hat{Id} = \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}
 
     Pauli X (X)
     -----------
@@ -48,8 +48,8 @@ class PolarizationOperationType(Enum):
     .. math::
 
         \hat{\sigma_y} = \begin{bmatrix}
-            0 & -i \\
-            i & 0
+            0 & -Id \\
+            Id & 0
             \end{bmatrix}
 
     Pauli Z (Z)
@@ -72,20 +72,20 @@ class PolarizationOperationType(Enum):
     Constructs Phase (:math:`\hat S`) operator
     .. math::
 
-        \hat{S} = \begin{bmatrix} 1 & 0 \\ 0 & i \end{bmatrix}
+        \hat{S} = \begin{bmatrix} 1 & 0 \\ 0 & Id \end{bmatrix}
 
     T Opeator (T)
     -------------
     Constructs T (:math:`\hat T`) operator
     .. math::
-        \hat{T} = \begin{bmatrix} 1 & 0 \\ 0 & e^{i \pi/4} \end{bmatrix}
+        \hat{T} = \begin{bmatrix} 1 & 0 \\ 0 & e^{Id \pi/4} \end{bmatrix}
 
     Sqrt(X) Operator (SX)
     ---------------------
     Constructs SX (:math:`\hat{SX}`) operator
     .. math::
 
-        \hat{SX} = \frac{1}{2} \begin{bmatrix} 1+i & 1-i \\ 1-i & 1+i \end{bmatrix}
+        \hat{SX} = \frac{1}{2} \begin{bmatrix} 1+Id & 1-Id \\ 1-Id & 1+Id \end{bmatrix}
 
     RX Operator (RX)
     ----------------
@@ -95,8 +95,8 @@ class PolarizationOperationType(Enum):
     .. math::
 
         \hat{RX} = \begin{bmatrix}
-        \cos\left(\frac{\theta}{2}\right) & -i\sin\left(\frac{\theta}{2}\right) \\
-        -i\sin\left(\frac{\theta}{2}\right) & \cos\left(\frac{\theta}{2}\right)
+        \cos\left(\frac{\theta}{2}\right) & -Id\sin\left(\frac{\theta}{2}\right) \\
+        -Id\sin\left(\frac{\theta}{2}\right) & \cos\left(\frac{\theta}{2}\right)
         \end{bmatrix}
 
     Usage Example
@@ -126,8 +126,8 @@ class PolarizationOperationType(Enum):
     .. math::
 
         \hat{RZ} = \begin{bmatrix}
-        e^{-i\frac{\theta}{2}} & 0 \\
-        0 & e^{i\frac{\theta}{2}}
+        e^{-Id\frac{\theta}{2}} & 0 \\
+        0 & e^{Id\frac{\theta}{2}}
         \end{bmatrix}
 
     Usage Example
@@ -150,7 +150,7 @@ class PolarizationOperationType(Enum):
     >>> op = Operation(PolarizationOperationType.Custom, operator=operator)
     """
 
-    I: Tuple[bool, List[str], ExpansionLevel, int] = (
+    Id: Tuple[bool, List[str], ExpansionLevel, int] = (
         True,
         [],
         ExpansionLevel.Vector,
@@ -263,7 +263,7 @@ def compute_operator(self, dimensions: List[int], **kwargs: Any) -> jnp.ndarray:
             Returns operator matrix
         """
         match self:
-            case PolarizationOperationType.I:
+            case PolarizationOperationType.Id:
                 return identity_operator()
             case PolarizationOperationType.X:
                 return x_operator()

photon_weave/state/base_state.py

@@ -173,6 +173,7 @@ def __repr__(self) -> str:
             return f"{formatted_matrix}"
         return f"{self.uid}"
 
+    # TODO: identity check is unused. If present for compatibility reasons, it should be commented
     def apply_kraus(
         self,
         operators: List[Union[np.ndarray, jnp.ndarray]],

photon_weave/state/composite_envelope.py

@@ -107,6 +107,7 @@ def reorder(self, *ordered_states: "BaseState") -> None:
 
             # Update the new order
             self.state_objs = list(ordered_states)
+
         elif self.expansion_level == ExpansionLevel.Matrix:
             # Get the state and reshape it
             shape = [os.dimensions for os in self.state_objs] * 2
@@ -129,8 +130,8 @@ def reorder(self, *ordered_states: "BaseState") -> None:
     def measure(
         self,
         *states: "BaseState",
-        separate_measurement: bool = False,
         destructive: bool = True,
+        separate_measurement: bool = False,
     ) -> Dict["BaseState", int]:
         """
         Measures this subspace. If the state is measured partially, then the state
@@ -163,13 +164,13 @@ def measure(
         C = Config()
 
         remaining_states = [s for s in self.state_objs]
-
+        # TODO: Replace this with a match statement to make it more readable
         if self.expansion_level == ExpansionLevel.Vector:
             # Get the state and reshape it into tensor
             shape = [so.dimensions for so in self.state_objs]
             shape.append(1)
             ps = self.state.reshape(shape)
-            for idx, state in enumerate(states):
+            for state in states:
                 # Constructing the einsum str
                 einsum = ESC.measure_vector(remaining_states, [state])
 
@@ -213,17 +214,20 @@ def measure(
                     state.index = None
                     state.expansion_level = ExpansionLevel.Label
                 self.state_objs.remove(state)
+
             if len(self.state_objs) > 0:
                 # Handle reshaping and storing the post measurement product state
                 self.state = ps.reshape(-1, 1)
                 self.state /= jnp.linalg.norm(self.state)
+
             else:
                 # Product state will be deleted
                 self.state = jnp.array([[1]])
+
         elif self.expansion_level == ExpansionLevel.Matrix:
             shape = [so.dimensions for so in self.state_objs] * 2
             ps = self.state.reshape(shape)
-            for idx, state in enumerate(states):
+            for state in states:
                 # Generate einsum string
                 einsum = ESC.measure_matrix(remaining_states, [state])
 
@@ -269,7 +273,7 @@ def measure(
                     state.index = None
                     state.expansion_level = ExpansionLevel.Label
 
-                # Remove the mesaured state from the product state
+                # Remove the measured state from the product state
                 self.state_objs.remove(state)
 
             # Reconstruct the post measurement product state
@@ -643,7 +647,8 @@ def apply_operation(self, operation: Operation, *states: "BaseState") -> None:
             for i, s in enumerate(states):
                 op_type = operation._operation_type
                 assert isinstance(
-                    s, op_type.expected_base_state_types[i]  # type: ignore
+                    s,
+                    op_type.expected_base_state_types[i],  # type: ignore
                 )
             operation.compute_dimensions(
                 [s._num_quanta if isinstance(s, Fock) else 0 for s in states],
@@ -755,6 +760,7 @@ def update_all_indices(self) -> None:
                     ][0]
 
 
+# FIXME: Inherit from Envelope to limit code duplication ?
 class CompositeEnvelope:
     """
     Composite Envelope is a pointer to a container, which includes the state

photon_weave/state/custom_state.py

@@ -133,6 +133,7 @@ def contract(
     def _measured(self) -> bool:
         return False
 
+    # FIXME: Unused argument
     @_measured.setter
     def _measured(self, measured: bool) -> None:
         """
@@ -161,6 +162,7 @@ def set_index(self, minor: int = -1, major: int = -1) -> None:
                 "Either set both parameters (minor, major) or none of them"
             )
 
+    # FIXME: Unused argument
     def measure(
         self, separate_measurement: bool = False, destructive: bool = True
     ) -> Dict[BaseState, int]:
@@ -223,6 +225,7 @@ def measure(
             "Something went wrong, this exception should not be raised"
         )  # pragma: no cover
 
+    # FIXME: Unused argument
     def measure_POVM(
         self,
         operators: List[Union[np.ndarray, jnp.ndarray]],
@@ -281,6 +284,7 @@ def measure_POVM(
             self.contract()
         return (outcome, {})
 
+    # FIXME: Unused argument
     def apply_kraus(
         self,
         operators: List[Union[np.ndarray, jnp.ndarray]],