FEAT: Extracting fock state transition probability amplitudes

.pre-commit-config.yaml

@@ -22,7 +22,7 @@ repos:
     rev: ad3ff374e97e29ca87c94b5dc7eccdd29adc6296
     hooks:
       - id: codespell
-        args: ["-L TE,TE/TM,te,ba,FPR,fpr_spacing,ro,nd,donot,schem,Synopsys"]
+        args: ["-L TE,TE/TM,te,ba,FPR,fpr_spacing,ro,nd,donot,schem,Synopsys,ket"]
         additional_dependencies:
           - tomli
 

docs/examples/05_quantum_integration_basics.py

@@ -1,11 +1,14 @@
 # # Quantum Integration Basics
 
+# $$\newcommand{\ket}[1]{\left|{#1}\right\rangle}$$
+# $$\newcommand{\bra}[1]{\left\langle{#1}\right|}$$
+#
 # One interesting thing to explore would be quantum state evolution through a unitary matrix composed from a physical photonic network that we could model. We will explore in this example how to integrate `sax` and `qutip`.
 
-import gdsfactory as gf  # NOQA : F401
-import sax  # NOQA : F401
-import piel  # NOQA : F401
-import qutip  # NOQA : F401
+import gdsfactory as gf
+import sax
+import piel
+import qutip as qp
 
 # ## Photonics Circuit to Unitary
 
@@ -74,22 +77,20 @@
 
 qutip_qobj.eigenstates
 
-# ## Calculating the Permanent
+# ## Sub Circuit Unitary Analysis
 
-# We can also integrated the `piel` toolchain, with another set of packages for quantum photonic system design such as those provided by `XanaduAI`. We will use their `thewalrus` package to calculate the permanent of our matrix.
+# We can also integrated the `piel` toolchain, with another set of packages for quantum photonic system design such as those provided by `XanaduAI`. We will use their `thewalrus` package to calculate the permanent of our matrix. For example, we can do this for our full circuit unitary:
 
 piel.sax_circuit_permanent(default_state_s_parameters)
 
-# The way this works is straightforward:
-
 # We might want to calculate the permanent of subsections of the larger unitary to calculate certain operations probability:
 
 s_parameters_standard_matrix.shape
 
-# For, example, we might want to just calculate it for the first two input modes. This would be indexed when starting from the first row and column as `start_index` = (0,0) and `stop_index` = (`unitary_size`, `unitary_size`). Note that an error will be raised if a non-unitary matrix is inputted. Some examples are:
+# For, example, we need to just calculate it for the first submatrix component, or a particular switch unitary within a larger circuit. This would be indexed when starting from the first row and column as `start_index` = (0,0) and `stop_index` = (`unitary_size`, `unitary_size`). Note that an error will be raised if a non-unitary matrix is inputted. Some examples are:
 
-our_subunitary = piel.subunitary_selection(
-    s_parameters_standard_matrix, start_index=(0, 0), stop_index=(1, 1)
+our_subunitary = piel.subunitary_selection_on_range(
+    s_parameters_standard_matrix, stop_index=(1, 1), start_index=(0, 0)
 )
 our_subunitary
 
@@ -105,3 +106,85 @@
 # ```python
 # ((-0.2296868-0.18254918j), 0.0)
 # ```
+
+# ## Fock State Evolution Probability
+
+# Say, we want to calculate the evolution of a Fock state input through our photonic circuit. The initial Fock state is defined as $\ket{f_1} = \ket{j_1, j_2, ... j_N}$ and transitions to $\ket{f_2} = \ket{j_1^{'}, j_2^{'}, ... j_N^{'}}$. The evolution of this state through our circuit with unitary $U$ is defined by the subunitary $U_{f_1}^{f_2}$.
+
+# Let us define an example four-mode multi-photon Fock state using `qutip` in the state $\ket{f_1} = \ket{1001}$
+
+initial_fock_state = qp.fock(4, 0) + qp.fock(4, 3)
+initial_fock_state
+
+final_fock_state = qp.fock(4, 1) + qp.fock(4, 2)
+final_fock_state
+
+# The subunitary $U_{f_1}^{f_2}$ is composed from the larger unitary by selecting the rows from the output state Fock state occupation of $\ket{f_2}$, and columns from the input $\ket{f_1}$. In our case, we need to select the columns indexes $(0,3)$ and rows indexes $(1,2)$.
+#
+# If we only consider an photon number of 1 in the particular Fock state, then we can describe the transition probability amplitude to be equivalent to the permanent:
+#
+# $$
+# a(\ket{f_1} \to \ket{f_2}) = \text{per}(U_{f_1}^{f_2})
+# $$
+#
+# If we consider a photon number of more than one for the transition Fock states, then the Permanent needs to be normalised. The probability amplitude for the transition is described as:
+# $$
+# a(\ket{f_1} \to \ket{f_2}) = \frac{\text{per}(U_{f_1}^{f_2})}{\sqrt{(j_1! j_2! ... j_N!)(j_1^{'}! j_2^{'}! ... j_N^{'}!)}}
+# $$
+#
+# TODO review Jeremy's thesis citations
+
+
+# However, we want to explore the probability amplitude of a Fock state transition. Given that our Fock state can have any photon number index, we need to select subsections of the unitary matrix that affect the photon path as described in the algorithm above. Let's implement this functionality based on our `qutip`-defined Fock states.
+
+piel.fock_state_to_photon_number_factorial(initial_fock_state)
+
+# We can analyse this for a multi-photon Fock state:
+
+example_multiphoton_fock_state = (
+    qp.fock(4, 1) + qp.fock(4, 2) + qp.fock(4, 2) + qp.fock(4, 2) + qp.fock(4, 2)
+)
+piel.fock_state_to_photon_number_factorial(example_multiphoton_fock_state)
+
+# In order to implement the algorithm above, we need to determine the indexes we need to extract for the particular Fock state that we are implementing too.
+
+initial_fock_state_indices = piel.fock_state_nonzero_indexes(initial_fock_state)
+initial_fock_state_indices
+
+# ```python
+# (0, 3)
+# ```
+
+final_fock_state_indices = piel.fock_state_nonzero_indexes(final_fock_state)
+final_fock_state_indices
+
+# ```python
+# (1, 2)
+# ```
+
+# We can extract the section of the unitary that corresponds to this Fock state transition. Note that based on (TODO cite Jeremy), the initial Fock state corresponds to the columns of the unitary and the final Fock states corresponds to the rows of the unitary.
+
+piel.subunitary_selection_on_index(
+    unitary_matrix=s_parameters_standard_matrix,
+    rows_index=final_fock_state_indices,
+    columns_index=initial_fock_state_indices,
+)
+
+# ```python~
+# Array([[ 0.40105772+0.49846345j,  0.4000432 +0.38792986j],
+#        [ 0.40004322+0.3879298j , -0.5810837 -0.31133226j]],      dtype=complex64)
+# ```
+
+# We can now extract the transition amplitude probability accordingly:
+
+piel.fock_transition_probability_amplitude(
+    initial_fock_state=initial_fock_state,
+    final_fock_state=final_fock_state,
+    unitary_matrix=s_parameters_standard_matrix,
+)
+
+# ```python
+# Array(-0.06831534-0.10413378j, dtype=complex64)
+# ```
+#
+# TODO this is not numerically right but the functions are fine because we need to verify the unitary-ness of the model composed matrices.

piel/integration/__init__.py

@@ -3,3 +3,4 @@
 from .cocotb_sax import *
 from .sax_thewalrus import *
 from .sax_qutip import *
+from .thewalrus_qutip import *

piel/integration/sax_qutip.py

@@ -1,60 +1,14 @@
 import qutip  # NOQA : F401
 import sax
 import numpy as np
-import jax.numpy as jnp
+from ..tools.qutip.unitary import matrix_to_qutip_qobj
 from piel.tools.sax.utils import sax_to_s_parameters_standard_matrix
 
 __all__ = [
     "sax_to_ideal_qutip_unitary",
-    "standard_s_parameters_to_qutip_qobj",
-    "verify_matrix_is_unitary",
 ]
 
 
-def matrix_to_qutip_qobj(
-    s_parameters_standard_matrix: jnp.ndarray,
-):
-    """
-    This function converts the calculated S-parameters into a standard Unitary matrix topology so that the shape and
-    dimensions of the matrix can be observed.
-
-    I think this means we need to transpose the output of the filtered sax SDense matrix to map it to a QuTip matrix.
-    Note that the documentation and formatting of the standard `sax` mapping to a S-parameter standard notation is
-    already in described in piel/piel/sax/utils.py.
-
-    From this stage we can implement a ``QObj`` matrix accordingly and perform simulations accordingly. https://qutip.org/docs/latest/guide/qip/qip-basics.html#unitaries
-
-    For example, a ``qutip`` representation of an s-gate gate would be:
-
-    ..code-block::
-
-        import numpy as np
-        import qutip
-        # S-Gate
-        s_gate_matrix = np.array([[1.,   0], [0., 1.j]])
-        s_gate = qutip.Qobj(mat, dims=[[2], [2]])
-
-    In mathematical notation, this S-gate would be written as:
-
-    ..math::
-
-        S = \\begin{bmatrix}
-            1 & 0 \\\\
-            0 & i \\\\
-        \\end{bmatrix}
-
-    Args:
-        s_parameters_standard_matrix (nso.ndarray): A dictionary of S-parameters in the form of a SDict from `sax`.
-
-    Returns:
-        qobj_unitary (qutip.Qobj): A QuTip QObj representation of the S-parameters in a unitary matrix.
-
-    """
-    s_parameter_standard_matrix_numpy = np.asarray(s_parameters_standard_matrix)
-    qobj_unitary = qutip.Qobj(s_parameter_standard_matrix_numpy)
-    return qobj_unitary
-
-
 def sax_to_ideal_qutip_unitary(sax_input: sax.SType):
     """
     This function converts the calculated S-parameters into a standard Unitary matrix topology so that the shape and
@@ -100,21 +54,3 @@ def sax_to_ideal_qutip_unitary(sax_input: sax.SType):
     s_parameter_standard_matrix_numpy = np.asarray(s_parameters_standard_matrix)
     qobj_unitary = matrix_to_qutip_qobj(s_parameter_standard_matrix_numpy)
     return qobj_unitary
-
-
-def verify_matrix_is_unitary(matrix: jnp.ndarray) -> bool:
-    """
-    Verify that the matrix is unitary.
-
-    Args:
-        matrix (jnp.ndarray): The matrix to verify.
-
-    Returns:
-        bool: True if the matrix is unitary, False otherwise.
-    """
-    matrix_numpy = np.asarray(matrix)
-    qobj = matrix_to_qutip_qobj(matrix_numpy)
-    return qobj.check_isunitary()
-
-
-standard_s_parameters_to_qutip_qobj = matrix_to_qutip_qobj

piel/integration/sax_thewalrus.py

@@ -1,12 +1,9 @@
 import sax
-import time
-import thewalrus
-import jax.numpy as jnp
 from ..tools.sax import sax_to_s_parameters_standard_matrix
-from typing import Optional
-import numpy as np
+from ..tools.thewalrus import unitary_permanent
 
-__all__ = ["sax_circuit_permanent", "subunitary_selection", "unitary_permanent"]
+
+__all__ = ["sax_circuit_permanent", "unitary_permanent"]
 
 
 def sax_circuit_permanent(
@@ -31,56 +28,3 @@ def sax_circuit_permanent(
     )
     circuit_permanent, computed_time = unitary_permanent(s_parameter_standard_matrix)
     return circuit_permanent, computed_time
-
-
-def subunitary_selection(
-    unitary_matrix: jnp.ndarray,
-    stop_index: tuple,
-    start_index: Optional[tuple] = (0, 0),
-):
-    """
-    This function returns a unitary between the indexes selected, and verifies the indexes are valid by checking that
-    the output matrix is also a unitary.
-
-    TODO implement validation of a 2D matrix.
-    """
-    start_row = start_index[0]
-    end_row = stop_index[0]
-    start_column = start_index[1]
-    end_column = stop_index[1]
-    column_range = jnp.arange(start_column, end_column + 1)
-    row_range = jnp.arange(start_row, end_row + 1)
-    unitary_matrix_row_selection = unitary_matrix.at[row_range, :].get()
-    unitary_matrix_row_column_selection = unitary_matrix_row_selection.at[
-        :, column_range
-    ].get()
-    return unitary_matrix_row_column_selection
-
-
-def unitary_permanent(
-    unitary_matrix: jnp.ndarray,
-) -> tuple:
-    """
-    The permanent of a unitary is used to determine the state probability of combinatorial Gaussian boson samping systems.
-
-    ``thewalrus`` Ryser's algorithm permananet implementation is described here: https://the-walrus.readthedocs.io/en/latest/gallery/permanent_tutorial.html
-
-    Note that this function needs to be as optimised as possible, so we need to minimise our computational complexity of our operation.
-
-    # TODO implement validation
-    # TODO maybe implement subroutine if computation is taking forever.
-    # TODO why two outputs? Understand this properly later.
-
-    Args:
-        unitary_permanent (np.ndarray): The unitary matrix.
-
-    Returns:
-        tuple: The circuit permanent and the time it took to compute it.
-
-    """
-    start_time = time.time()
-    unitary_matrix_numpy = np.asarray(unitary_matrix)
-    circuit_permanent = thewalrus.perm(unitary_matrix_numpy)
-    end_time = time.time()
-    computed_time = end_time - start_time
-    return circuit_permanent, computed_time

piel/integration/thewalrus_qutip.py

@@ -0,0 +1,69 @@
+import jax.numpy as jnp
+import qutip
+
+from ..tools.thewalrus import unitary_permanent
+from ..tools.qutip import (
+    fock_state_nonzero_indexes,
+    fock_state_to_photon_number_factorial,
+    subunitary_selection_on_index,
+)
+
+__all__ = ["fock_transition_probability_amplitude"]
+
+
+def fock_transition_probability_amplitude(
+    initial_fock_state: qutip.Qobj,
+    final_fock_state: qutip.Qobj,
+    unitary_matrix: jnp.ndarray,
+):
+    """
+    This function returns the transition probability amplitude between two Fock states when propagating in between
+    the unitary_matrix which represents a quantum state circuit.
+
+    Note that based on (TODO cite Jeremy), the initial Fock state corresponds to the columns of the unitary and the
+    final Fock states corresponds to the rows of the unitary.
+
+    .. math ::
+        \newcommand{\ket}[1]{\left|{#1}\right\rangle}
+
+    The subunitary :math:`U_{f_1}^{f_2}` is composed from the larger unitary by selecting the rows from the output state
+    Fock state occupation of :math:`\ket{f_2}`, and columns from the input :math:`\ket{f_1}`. In our case, we need to select the
+    columns indexes :math:`(0,3)` and rows indexes :math:`(1,2)`.
+
+    If we consider a photon number of more than one for the transition Fock states, then the Permanent needs to be
+    normalised. The probability amplitude for the transition is described as:
+
+    .. math ::
+        a(\ket{f_1} \to \ket{f_2}) = \frac{\text{per}(U_{f_1}^{f_2})}{\sqrt{(j_1! j_2! ... j_N!)(j_1^{'}! j_2^{'}! ... j_N^{'}!)}}
+
+    Args:
+        initial_fock_state (qutip.Qobj): A QuTip QObj representation of the initial Fock state.
+        final_fock_state (qutip.Qobj): A QuTip QObj representation of the final Fock state.
+        unitary_matrix (jnp.ndarray): A JAX NumPy array representation of the unitary matrix.
+
+    Returns:
+        float: The transition probability amplitude between the initial and final Fock states.
+    """
+    columns_indices = fock_state_nonzero_indexes(initial_fock_state)
+    rows_indices = fock_state_nonzero_indexes(final_fock_state)
+
+    initial_fock_state_photon_number_factorial = fock_state_to_photon_number_factorial(
+        initial_fock_state
+    )
+    final_fock_state_photon_number_factorial = fock_state_to_photon_number_factorial(
+        final_fock_state
+    )
+
+    subunitary_selection = subunitary_selection_on_index(
+        unitary_matrix=unitary_matrix,
+        rows_index=rows_indices,
+        columns_index=columns_indices,
+    )
+
+    transition_probability_amplitude = unitary_permanent(subunitary_selection)[0] / (
+        jnp.sqrt(
+            initial_fock_state_photon_number_factorial
+            * final_fock_state_photon_number_factorial
+        )
+    )
+    return transition_probability_amplitude

piel/tools/__init__.py

@@ -3,3 +3,4 @@
 from .openlane import *
 from .hdl21 import *
 from .sax import *
+from .qutip import *

piel/tools/qutip/__init__.py

@@ -0,0 +1,2 @@
+from .fock import *
+from .unitary import *