Trajectree 0.0.0__py3-none-any.whl

trajectree/experimental/sparse.py

@@ -0,0 +1,115 @@
+from ..fock_optics.measurement import *
+from ..fock_optics.outputs import *
+
+import numpy as np
+import scipy.sparse as sp
+
+# Support functions:
+def create_op(left_indices, op, right_indices, N):
+    if left_indices == 0:
+        return sp.kron(op, sp.eye(N**right_indices))
+    elif right_indices == 0:
+        return sp.kron(sp.eye(N**left_indices), op)
+    else:    
+        out_op = sp.kron(sp.eye(N**left_indices), op)
+        return sp.kron(out_op, sp.eye(N**right_indices))
+def _find_mat_exp(mat):
+    ans = sp.eye(mat.shape[0])
+    intermediate = 1
+    for i in range(1, 50+1):
+        intermediate *= mat/i
+        intermediate.eliminate_zeros()
+        ans += intermediate
+    return ans 
+def read_quantum_state_sparse(sparse_state, N):
+    temp_sparse_state = sp.csr_matrix(sparse_state)
+    temp_sparse_state.data = np.round(temp_sparse_state.data, 10)
+    temp_sparse_state.eliminate_zeros()
+    labels = generate_labels(4,N)
+    state = temp_sparse_state.nonzero()[0]
+    print(f"{len(state)} non-zero elements Corresponding Basis terms:")
+    for k in state: print(labels[k],"-",k,"-",temp_sparse_state[k].data)
+
+def extend_state_sparse(state):
+    return sp.kron(state, state)
+# TMSV_state_dense = extend_state_sparse(TMSV_state)
+
+def bell_state_measurement_sparse(TMSV_state_dense, N, efficiency, a_dag, is_dm = False):
+    # BSM BS implementation
+    BSM_H_0_Mode_op = create_op(2, a_dag, 5, N)
+    print(BSM_H_0_Mode_op.shape, len(BSM_H_0_Mode_op.nonzero()[0]), len(BSM_H_0_Mode_op.nonzero()[1]))
+    BSM_V_0_Mode_op = create_op(3, a_dag, 4, N)
+    BSM_H_1_Mode_op = create_op(6, a_dag, 1, N)
+    BSM_V_1_Mode_op = create_op(7, a_dag, 0, N)
+
+    hamiltonian_BS_H = -np.pi/4 * ( BSM_H_0_Mode_op.T@BSM_H_1_Mode_op - BSM_H_0_Mode_op@BSM_H_1_Mode_op.T )
+    unitary_BS_H = _find_mat_exp(hamiltonian_BS_H)
+
+    hamiltonian_BS_V = -np.pi/4 * ( BSM_V_0_Mode_op.T@BSM_V_1_Mode_op - BSM_V_0_Mode_op@BSM_V_1_Mode_op.T )
+    unitary_BS_V = _find_mat_exp(hamiltonian_BS_V)
+
+
+    # BSM povm implementation
+    povm_op_1 = sp.csr_matrix(create_threshold_POVM_OP_Dense(efficiency, 1, N))
+    povm_op_0 = sp.csr_matrix(create_threshold_POVM_OP_Dense(efficiency, 0, N))
+
+    BSM_povm = create_op(2, povm_op_1, 0, N)
+    BSM_povm = create_op(0, sp.kron(BSM_povm, povm_op_0), 2, N)
+    BSM_povm = sp.kron(BSM_povm, sp.kron(povm_op_0, povm_op_1))
+
+    # print(unitary_BS_V.shape, unitary_BS_H.shape, TMSV_state_dense.shape)
+    
+    if is_dm:
+        post_BS_State = unitary_BS_V @ unitary_BS_H @ TMSV_state_dense @ (unitary_BS_V @ unitary_BS_H).conj().T
+        post_BSM_State = BSM_povm @ post_BS_State @ BSM_povm.conj().T
+    else:
+        post_BS_State = unitary_BS_V @ unitary_BS_H @ TMSV_state_dense
+        post_BSM_State = BSM_povm @ post_BS_State
+
+    # post_BSM_State.data = np.round(post_BSM_State.data, 10)
+    # post_BSM_State.eliminate_zeros()
+
+    return post_BSM_State
+# post_BSM_State = bell_state_measurement_sparse(TMSV_state_dense, N, efficiency)
+
+def rotate_and_measure_sparse(post_BSM_State, N, efficiency, a_dag):
+    # Polarization rotators mode operators
+    rotator_H_0_Mode_op = create_op(0, a_dag, 7, N)
+    rotator_V_0_Mode_op = create_op(1, a_dag, 6, N)
+    rotator_H_1_Mode_op = create_op(4, a_dag, 3, N)
+    rotator_V_1_Mode_op = create_op(5, a_dag, 2, N)
+
+    povm_op_1 = sp.csr_matrix(create_threshold_POVM_OP_Dense(efficiency, 1, N))
+
+    # polarization analysis detector POVMs
+    pol_analyzer_povm = create_op(0, povm_op_1, 3, N)
+    pol_analyzer_povm = create_op(0, sp.kron(pol_analyzer_povm, povm_op_1), 3, N)
+
+    # Applying rotations and measuring
+
+    signal_angles = np.linspace(0, np.pi, 10)
+    # idler_angles = np.linspace(0, np.pi, 20)
+    idler_angles = [0]
+    coincidence = []
+
+    for i, idler_angle in enumerate(idler_angles):
+        coincidence_probs = []
+
+        hamiltonian_rotator_1 = -idler_angle * ( rotator_H_1_Mode_op.T@rotator_V_1_Mode_op - rotator_H_1_Mode_op@rotator_V_1_Mode_op.T )
+        unitary_rotator_1 = _find_mat_exp(hamiltonian_rotator_1)
+        post_idler_detection_state = unitary_rotator_1 @ post_BSM_State
+        # post_idler_detection_state = post_BSM_State
+        
+        for j, angle in enumerate(signal_angles):
+            # print("idler:", i, "signal:", j)
+        
+            hamiltonian_rotator_0 = -angle * ( rotator_H_0_Mode_op.T@rotator_V_0_Mode_op - rotator_H_0_Mode_op@rotator_V_0_Mode_op.T )
+            unitary_rotator_0 = _find_mat_exp(hamiltonian_rotator_0)
+            post_rotations_state = unitary_rotator_0 @ post_idler_detection_state
+
+            measured_state = pol_analyzer_povm @ post_rotations_state
+
+            coincidence_probs.append(sp.linalg.norm(measured_state)**2)
+        coincidence.append(coincidence_probs)
+    return coincidence, idler_angles
+# coincidence, idler_angles = rotate_and_measure_sparse(post_BSM_State, N, efficiency)

trajectree/fock_optics/devices.py

@@ -0,0 +1,58 @@
+from scipy.linalg import expm
+
+import numpy as np
+from numpy import kron
+
+from quimb.tensor import MatrixProductOperator as mpo #type: ignore
+
+import qutip as qt
+
+# Beamsplitter transformation
+def create_BS_MPO(site1, site2, theta, total_sites, N, tag = 'BS'): 
+
+    a = qt.destroy(N).full()
+    a_dag = a.T
+    I = np.eye(N)
+    
+    # This corresponds to the BS hamiltonian:
+
+    hamiltonian_BS = -theta * ( kron(I, a_dag)@kron(a, I) - kron(I, a)@kron(a_dag, I) )
+    unitary_BS = expm(hamiltonian_BS)
+
+    # print("unitary_BS", unitary_BS)
+
+    BS_MPO = mpo.from_dense(unitary_BS, dims = N, sites = (site1,site2), L=total_sites, tags=tag)
+    # BS_MPO = BS_MPO.fill_empty_sites(mode = "full")
+    return BS_MPO
+
+
+def generalized_mode_mixer(site1, site2, theta, phi, psi, lamda, total_sites, N, tag = 'MM'): 
+
+    a = qt.destroy(N).full()
+    a_dag = a.T
+    I = np.eye(N)
+    
+    # This corresponds to the BS hamiltonian: This is a different difinition from the one in 
+    # create_BS_MPO. This is because of how the generalized beamsplitter is defined in DOI: 10.1088/0034-4885/66/7/203 . 
+    hamiltonian_BS = theta * (kron(a_dag, I)@kron(I, a) + kron(a, I)@kron(I, a_dag))
+    unitary_BS = expm(-1j * hamiltonian_BS)
+
+    # print("unitary_BS\n", np.round(unitary_BS, 4))
+
+    pre_phase_shifter = np.kron(phase_shifter(N, phi[0]/2), phase_shifter(N, phi[1]/2))
+    post_phase_shifter = np.kron(phase_shifter(N, psi[0]/2), phase_shifter(N, psi[1]/2))
+    global_phase_shifter = np.kron(phase_shifter(N, lamda[0]/2), phase_shifter(N, lamda[1]/2))
+
+    # This construction for the generalized beamsplitter is based on the description in paper DOI: 10.1088/0034-4885/66/7/203
+    generalized_BS = global_phase_shifter @  (pre_phase_shifter @ unitary_BS @ post_phase_shifter)
+
+    # print("generalized_BS\n", np.round(generalized_BS, 4))
+
+    BS_MPO = mpo.from_dense(generalized_BS, dims = N, sites = (site1,site2), L=total_sites, tags=tag)
+    # BS_MPO = BS_MPO.fill_empty_sites(mode = "full")
+    return BS_MPO
+
+
+def phase_shifter(N, theta):
+    diag = [np.exp(1j * theta * i) for i in range(N)]
+    return np.diag(diag, k=0)

trajectree/fock_optics/light_sources.py

@@ -0,0 +1,123 @@
+from .utils import create_MPO
+from .devices import create_BS_MPO
+
+from scipy import sparse as sp
+from scipy.linalg import expm
+
+import numpy as np
+from numpy.linalg import matrix_power
+from numpy import kron, sqrt
+
+from quimb.tensor.tensor_arbgeom import tensor_network_apply_op_vec #type: ignore
+from quimb.tensor.tensor_1d_compress import enforce_1d_like #type: ignore
+
+import qutip as qt
+from math import factorial
+
+
+def create_TMSV_OP_Dense(N, mean_photon_num):
+    a = qt.destroy(N).full()
+    a_dag = a.T
+    truncation = (N-1)   
+
+    op = expm(1j * mean_photon_num * (kron(a_dag, a_dag) + kron(a, a)))
+
+    return op
+
+
+
+########## Light Source ###########
+
+def light_source(vacuum, N, mean_photon_num, num_modes, error_tolerance, TMSV_indices = ((0,2),(5,7)), compress = True, contract = True):
+
+    psi = vacuum.copy()
+    psi.add_tag("L0")
+    site_tags = psi.site_tags
+
+    # Creating TMSV ops:
+    TMSV_op_dense = create_TMSV_OP_Dense(N, mean_photon_num)
+
+    TMSV_MPO_H = create_MPO(site1 = TMSV_indices[0][0], site2 = TMSV_indices[0][1], total_sites = num_modes, op = TMSV_op_dense, N = N, tag = r"$TMSV_H$")
+    # TMSV_MPO_H.draw()
+    # print("sites present in light_source:", TMSV_MPO_H.sites)
+    enforce_1d_like(TMSV_MPO_H, site_tags=site_tags, inplace=True)
+    # print("sites present in light_source:", TMSV_MPO_H.sites)
+    TMSV_MPO_H.add_tag("L1")
+
+    TMSV_MPO_V = create_MPO(site1 = TMSV_indices[1][0], site2 = TMSV_indices[1][1], total_sites = num_modes, op = TMSV_op_dense, N = N, tag = r"$TMSV_V$")
+    enforce_1d_like(TMSV_MPO_V, site_tags=site_tags, inplace=True)
+    TMSV_MPO_V.add_tag("L1")
+
+    # Creating PBS ops:
+    U_PBS_H_Signal = create_BS_MPO(site1 = 2, site2 = 6, theta=np.pi/2, total_sites = num_modes, N = N, tag = r"$PBS_S$")
+    enforce_1d_like(U_PBS_H_Signal, site_tags=site_tags, inplace=True)
+    U_PBS_H_Signal.add_tag("L1")
+
+    U_PBS_H_Idler = create_BS_MPO(site1 = 0, site2 = 4, theta=np.pi/2, total_sites = num_modes, N = N, tag = r"$PBS_I$")
+    enforce_1d_like(U_PBS_H_Idler, site_tags=site_tags, inplace=True)
+    U_PBS_H_Signal.add_tag("L1")
+
+    # Create entangled state:
+    psi = tensor_network_apply_op_vec(TMSV_MPO_H, psi, compress=compress, contract = contract, cutoff = error_tolerance)
+    psi = tensor_network_apply_op_vec(TMSV_MPO_V, psi, compress=compress, contract = contract, cutoff = error_tolerance)
+    psi = tensor_network_apply_op_vec(U_PBS_H_Idler, psi, compress=compress, contract = contract, cutoff = error_tolerance)
+    psi = tensor_network_apply_op_vec(U_PBS_H_Signal, psi, compress=compress, contract = contract, cutoff = error_tolerance)
+
+    psi.normalize()
+
+    # print("trace is:", np.linalg.norm(psi.to_dense()))
+
+    for _ in range(4):
+        psi.measure(0, remove = True, renorm = True, inplace = True)
+
+    # Not used for TN implermentation. Used for validating impelmentation with dense version
+    TMSV_state = psi.to_dense()
+    TMSV_state = np.reshape(TMSV_state.data, (-1, 1), order = 'C')
+    TMSV_state = sp.csr_matrix(TMSV_state)
+    TMSV_state.data = np.round(TMSV_state.data, 10)
+    TMSV_state.eliminate_zeros()
+
+    return psi, TMSV_state
+
+
+# Generate truncation filter MPO 
+# TODO: Make a function to renormalize a quantum state. How: find the projection of the quantum state onto itself and calculate the 
+# probability. Next, take the square root of this number, divide it by the number nodes in the quantum state and multiply it with 
+# all the states in the MPS. For density matrices, simply find the trace directly and do the same thing as the previous example except
+# for not taking the square root.  The truncation filter would not work without the renormalization 
+def create_truncation_filter_Dense(truncation):
+    # This is only the projection operator. The states need to be normalized first. 
+    N = truncation+1
+    vacuum = np.zeros(N**2)
+    vacuum[0] = 1
+
+    a = qt.destroy(N).full()
+    a_dag = a.T
+    I = np.eye(N)
+
+    # # debug
+    # labels = generate_labels(1,N)
+
+    op = 0
+    for trunc in range(truncation, -1, -1):
+        state = kron(matrix_power(a_dag, trunc), I) @ vacuum / sqrt(factorial(trunc) * factorial(0))
+        op+=np.outer(state, state)
+        coeffs = [trunc+1, 0]
+
+        # # Debug
+        # state_inds = state.nonzero()[0]
+        # print("TMSV state:", [labels[i] for i in state_inds], "Val:", state[state_inds[0]])
+        # print("coeffs", coeffs)
+
+        for i in range(trunc):
+            coeffs = [coeffs[0]-1, coeffs[1]+1]
+            state = kron(a, a_dag) @ state / sqrt((coeffs[0]) * (coeffs[1]))
+            op += np.outer(state, state)
+
+
+            # # debug
+            # state_inds = state.nonzero()[0]
+            # print("TMSV state:", [labels[i] for i in state_inds], "Val:", state[state_inds[0]])
+            # print("coeffs", coeffs)
+
+    return op

trajectree/fock_optics/measurement.py

@@ -0,0 +1,236 @@
+from .devices import generalized_mode_mixer, create_BS_MPO
+from ..trajectory import quantum_channel
+from .noise_models import single_mode_bosonic_noise_channels
+
+from scipy.linalg import sqrtm
+from scipy import sparse as sp
+
+import numpy as np
+from numpy.linalg import matrix_power
+from numpy import sqrt
+
+from quimb.tensor import MatrixProductOperator as mpo #type: ignore
+from quimb.tensor.tensor_arbgeom import tensor_network_apply_op_vec #type: ignore
+from quimb.tensor.tensor_1d_compress import enforce_1d_like #type: ignore
+
+import qutip as qt
+from math import factorial
+
+from functools import lru_cache
+
+
+# This is the actual function that generates the POVM operator.
+def create_threshold_POVM_OP_Dense(efficiency, outcome, N):
+    a = qt.destroy(N).full()
+    a_dag = a.T
+    create0 = a_dag * sqrt(efficiency)
+    destroy0 = a * sqrt(efficiency)
+    series_elem_list = [((-1)**i) * matrix_power(create0, (i+1)) @ matrix_power(destroy0, (i+1)) / factorial(i+1) for i in range(N-1)] # (-1)^i * a_dag^(i+1) @ a^(i+1) / (i+1)! = (-1)^(i+2) * a_dag^(i+1) @ a^(i+1) / (i+1)! since goes from 0->n
+    # print(series_elem_list[0])
+    dense_op = sum(series_elem_list)
+
+    if outcome == 0:
+        dense_op = np.eye(dense_op.shape[0]) - dense_op
+    # print(sqrtm(dense_op))
+    return dense_op
+
+@lru_cache(maxsize=20)
+def factorial(x):
+    n = 1
+    for i in range(2, x+1):
+        n *= i
+    return n
+
+@lru_cache(maxsize=20)
+def comb(n, k):
+    return factorial(n) / (factorial(k) * factorial(n - k))
+
+@lru_cache(maxsize=20)
+def projector(n, N):
+    state = np.zeros(N)
+    state[n] = 1
+    return np.outer(state, state)
+
+# Testing stuff out here. 
+def create_PNR_POVM_OP_Dense(eff, outcome, N, debug = False):
+    a_dag = qt.create(N).full()
+    vacuum = np.zeros(N)
+    vacuum[0] = 1
+
+    @lru_cache(maxsize=20)
+    def create_povm_list(eff, N):
+        povms = []
+        # m is the outcome here
+        for m in range(N-1):
+            op = 0
+            for n in range(m, N):
+                op += comb(n,m) * eff**m * (1-eff)**(n-m) * projector(n, N)
+            povms.append(op)
+
+        povms.append(np.eye(N) - sum(povms))
+        return povms
+    
+    povms = create_povm_list(eff, N)
+    if debug:
+        return povms[outcome], povms
+    return povms[outcome]
+
+
+
+def generate_sqrt_POVM_MPO(sites, outcome, total_sites, efficiency, N, pnr = False, tag = "POVM"):
+    if pnr:
+        dense_op = sqrtm(create_PNR_POVM_OP_Dense(efficiency, outcome, N)).astype(np.complex128)
+    else:
+        dense_op = sqrtm(create_threshold_POVM_OP_Dense(efficiency, outcome, N)).astype(np.complex128)
+
+    sqrt_POVM_MPOs = []
+    for i in sites:
+        sqrt_POVM_MPOs.append(mpo.from_dense(dense_op, dims = N, sites = (i,), L=total_sites, tags=tag))
+
+    return sqrt_POVM_MPOs
+
+
+def bell_state_measurement(psi, N, site_tags, num_modes, efficiencies, dark_counts_gain,  error_tolerance, beamsplitters = [[2,6],[3,7]], measurements = {0:(2,7), 1:(3,6)}, pnr = False, det_outcome = 1, use_trajectory = False, return_MPOs = False, compress = True, contract = True):
+
+    """Perform Bell state measrement or return the MPOs used in the measurement.
+    Args:
+        psi (mps): The input state to be measured.
+        N (int): local Hilbert space dimension
+        site_tags (list): The tags for the sites in the MPS.
+        num_modes (int): The number of modes in the MPS.
+        efficiencies list[float]: The efficiencies of the (pairs of) detectors in the BSM.
+        error_tolerance (float): The error tolerance for the tensor network.
+        measurements (dict): The sites for the measurements. Default is {1:(2,7), 0:(3,6)}.
+        pnr (bool): Whether to use photon number resolving measurement. Default is False.
+        pnr_outcome (int): The outcome for the photon number resolving measurement. Default is 1. When not using PNR, this can be anything other than 1 since threshold detectors don't distinguish between photon numbers. 
+        return_MPOs (bool): Whether to return the MPOs used in the measurement. Default is False.
+        compress (bool): Whether to compress the MPS after applying the MPOs. Default is True.
+        contract (bool): Whether to contract the MPS after applying the MPOs. Default is True.
+        
+        Returns:
+            mps: The measured state after the Bell state measurement.
+            
+    """
+
+    U_BS_H = create_BS_MPO(site1 = beamsplitters[0][0], site2 = beamsplitters[0][1], theta=np.pi/4, total_sites = num_modes, N = N, tag = r"$U_{BS_H}$")
+    enforce_1d_like(U_BS_H, site_tags=site_tags, inplace=True)
+    U_BS_H.add_tag("L2")
+
+    U_BS_V = create_BS_MPO(site1 = beamsplitters[1][0], site2 = beamsplitters[1][1], theta=np.pi/4, total_sites = num_modes, N = N, tag = r"$U_{BS_V}$")
+    enforce_1d_like(U_BS_V, site_tags=site_tags, inplace=True)
+    U_BS_V.add_tag("L3")
+
+    # Note that these are not used if using trajectree to implement detector inefficiency. 
+    BSM_POVM_1_OPs = generate_sqrt_POVM_MPO(sites=measurements[1], outcome = det_outcome, total_sites=num_modes, efficiency=efficiencies[0], N=N, pnr = pnr)
+    BSM_POVM_1_OPs.extend(generate_sqrt_POVM_MPO(sites=measurements[0], outcome = 0, total_sites=num_modes, efficiency=efficiencies[1], N=N, pnr = pnr))
+
+    if return_MPOs:
+        returned_MPOs = [U_BS_H, U_BS_V]
+        if use_trajectory:
+            quantum_channel_list = [quantum_channel(N = N, num_modes = num_modes, formalism = "closed", unitary_MPOs = BSM_MPO, name = "BSM") for BSM_MPO in returned_MPOs]
+
+            damping_kraus_ops_0 = single_mode_bosonic_noise_channels(noise_parameter = 1-efficiencies[0], N = N)
+            damping_kraus_ops_1 = single_mode_bosonic_noise_channels(noise_parameter = 1-efficiencies[1], N = N)
+            two_mode_kraus_ops_0 = [sp.kron(op1, op2) for op1 in damping_kraus_ops_0 for op2 in damping_kraus_ops_0]
+            two_mode_kraus_ops_1 = [sp.kron(op1, op2) for op1 in damping_kraus_ops_1 for op2 in damping_kraus_ops_1]
+            quantum_channel_list.append(quantum_channel(N = N, num_modes = num_modes, formalism = "kraus", kraus_ops_tuple = ((2,3), two_mode_kraus_ops_0))) # The tuples in this list are defined as (sites, kraus_ops). The sites are the sites where the Kraus ops are applied.
+            quantum_channel_list.append(quantum_channel(N = N, num_modes = num_modes, formalism = "kraus", kraus_ops_tuple = ((6,7), two_mode_kraus_ops_1))) # The tuples in this list are defined as (sites, kraus_ops). The sites are the sites where the Kraus ops are applied.
+
+            amplification_kraus_ops_0 = single_mode_bosonic_noise_channels(noise_parameter = dark_counts_gain[0], N = N)
+            amplification_kraus_ops_1 = single_mode_bosonic_noise_channels(noise_parameter = dark_counts_gain[1], N = N)
+            two_mode_kraus_ops_0 = [sp.kron(op1, op2) for op1 in amplification_kraus_ops_0 for op2 in amplification_kraus_ops_0]
+            two_mode_kraus_ops_1 = [sp.kron(op1, op2) for op1 in amplification_kraus_ops_1 for op2 in amplification_kraus_ops_1]
+            quantum_channel_list.append(quantum_channel(N = N, num_modes = num_modes, formalism = "kraus", kraus_ops_tuple = ((2,3), two_mode_kraus_ops_0))) # The tuples in this list are defined as (sites, kraus_ops). The sites are the sites where the Kraus ops are applied.
+            quantum_channel_list.append(quantum_channel(N = N, num_modes = num_modes, formalism = "kraus", kraus_ops_tuple = ((6,7), two_mode_kraus_ops_1))) # The tuples in this list are defined as (sites, kraus_ops). The sites are the sites where the Kraus ops are applied.
+
+            BSM_POVM_1_OPs = generate_sqrt_POVM_MPO(sites=measurements[1], outcome = det_outcome, total_sites=num_modes, efficiency=1, N=N, pnr = pnr)
+            BSM_POVM_1_OPs.extend(generate_sqrt_POVM_MPO(sites=measurements[0], outcome = 0, total_sites=num_modes, efficiency=1, N=N, pnr = pnr))
+
+            det_quantum_channels = [quantum_channel(N = N, num_modes = num_modes, formalism = "closed", unitary_MPOs = DET_MPO, name = "DET") for DET_MPO in BSM_POVM_1_OPs]
+            quantum_channel_list.extend(det_quantum_channels)
+    
+            return quantum_channel_list
+
+        returned_MPOs.extend(BSM_POVM_1_OPs) # Collect all the MPOs in a list and return them. The operators are ordered as such: 
+
+        quantum_channel_list = [quantum_channel(N = N, num_modes = num_modes, formalism = "closed", unitary_MPOs = BSM_MPO, name = "BSM") for BSM_MPO in returned_MPOs]
+
+        return quantum_channel_list
+
+    psi = tensor_network_apply_op_vec(U_BS_H, psi, compress=compress, contract = contract, cutoff = error_tolerance)
+    psi = tensor_network_apply_op_vec(U_BS_V, psi, compress=compress, contract = contract, cutoff = error_tolerance)
+
+    for POVM_OP in BSM_POVM_1_OPs:
+        POVM_OP.add_tag("L4")
+        psi = tensor_network_apply_op_vec(POVM_OP, psi, compress=compress, contract = contract, cutoff = error_tolerance)
+
+    return psi
+
+
+
+def rotate_and_measure(psi, N, site_tags, num_modes, efficiency, error_tolerance, idler_angles, signal_angles, rotations = {"signal":(4,5), "idler":(0,1)}, measurements = {1:(0,4), 0:(1,5)}, pnr = False, det_outcome = 1, return_MPOs = False, compress = True, contract = True, draw = False):
+    # idler_angles = [0]
+    # angles = [np.pi/4]
+
+    coincidence = []
+
+    POVM_1_OPs = generate_sqrt_POVM_MPO(sites = measurements[1], outcome = det_outcome, total_sites=num_modes, efficiency=efficiency, N=N, pnr = pnr)
+    POVM_0_OPs = generate_sqrt_POVM_MPO(sites = measurements[0], outcome = 0, total_sites=num_modes, efficiency=efficiency, N=N, pnr = pnr)
+    # POVM_0_OPs = generate_sqrt_POVM_MPO(sites=(0,4), outcome = 0, total_sites=num_modes, efficiency=efficiency, N=N, pnr = pnr)
+    # enforce_1d_like(POVM_OP, site_tags=site_tags, inplace=True)
+
+    meas_ops = POVM_1_OPs
+    meas_ops.extend(POVM_0_OPs)
+
+    for i, idler_angle in enumerate(idler_angles):
+        coincidence_probs = []
+
+        # rotator_node_1 = create_BS_MPO(site1 = rotations["idler"][0], site2 = rotations["idler"][1], theta=idler_angle, total_sites = num_modes, N = N, tag = r"$Rotator_I$")
+        ######################
+        # We make this correction here since the rotator hamiltonian is 1/2(a_v b_h + a_h b_v), which does not show up in the bs unitary, whose function we are reusing to 
+        # rotate the state.
+        rotator_node_1 = generalized_mode_mixer(site1 = rotations["idler"][0], site2 = rotations["idler"][1], theta = -idler_angle/2, phi = [0,0], psi = [0,0], lamda = [0,0], total_sites = num_modes, N = N, tag = 'MM')
+
+        
+        enforce_1d_like(rotator_node_1, site_tags=site_tags, inplace=True)
+        rotator_node_1.add_tag("L5")
+        if not return_MPOs: # If the user wants the MPOs, we don't need to apply the rotator to the state.
+            idler_rotated_psi = tensor_network_apply_op_vec(rotator_node_1, psi, compress=compress, contract = contract, cutoff = error_tolerance)
+
+
+        for j, angle in enumerate(signal_angles):
+            # print("idler:", i, "signal:", j)
+        
+            # rotator_node_2 = create_BS_MPO(site1 = rotations["signal"][0], site2 = rotations["signal"][1], theta=angle, total_sites = num_modes, N = N, tag = r"$Rotator_S$")
+            ##########################
+            # We make this correction here since the rotator hamiltonian is 1/2(a_v b_h + a_h b_v), which does not show up in the bs unitary, whose function we are reusing to 
+            # rotate the state.
+            rotator_node_2 = generalized_mode_mixer(site1 = rotations["signal"][0], site2 = rotations["signal"][1], theta = -angle/2, phi = [0,0], psi = [0,0], lamda = [0,0], total_sites = num_modes, N = N, tag = 'MM') 
+            
+            
+            enforce_1d_like(rotator_node_2, site_tags=site_tags, inplace=True)
+
+            if return_MPOs:
+                meas_ops = [rotator_node_1, rotator_node_2] + meas_ops # Collect all the MPOs in a list and return them
+                return meas_ops
+        
+            # Rotate and measure:
+            rotator_node_2.add_tag("L5")
+            rho_rotated = tensor_network_apply_op_vec(rotator_node_2, idler_rotated_psi, compress=compress, contract = contract, cutoff = error_tolerance)
+
+            # read_quantum_state(psi)
+            # read_quantum_state(rho_rotated)
+
+            for POVM_OP in meas_ops:
+                POVM_OP.add_tag("L6")
+                rho_rotated = tensor_network_apply_op_vec(POVM_OP, rho_rotated, compress=compress, contract = contract, cutoff = error_tolerance)
+        
+            if draw:
+                # only for drawing the TN. Not used otherwise
+                fix = {(f"L{j}",f"I{num_modes - i-1}"):(3*j,i+5) for j in range(10) for i in range(10)}
+                rho_rotated.draw(color = [r'$HH+VV$', r'$U_{BS_H}$', r"$U_{BS_V}$", 'POVM', r'$Rotator_I$', r'$Rotator_S$'], title = "Polarization entanglement swapping MPS", fix = fix, show_inds = True, show_tags = False)
+                # rho_rotated.draw_tn()
+            coincidence_probs.append((rho_rotated.norm())**2)
+        coincidence.append(coincidence_probs)
+    
+    return np.array(coincidence)

trajectree/fock_optics/noise_models.py

@@ -0,0 +1,41 @@
+from scipy import sparse as sp
+from scipy.linalg import expm
+
+import numpy as np
+
+import qutip as qt
+from math import factorial
+
+def single_mode_bosonic_noise_channels(noise_parameter, N):
+    """This function produces the Kraus operatorsd for the single mode bosonic noise channels. This includes pure loss and 
+    pure gain channels. The pure gain channel is simply the transpose of the pure loss channel.
+    
+    Args:
+        noise_parameter (float): The noise parameter, (loss for pure loss and gain for pure gain channels). For the pure loss channel, this 
+                                 parameter is the dimensionless noise term: 1-transmissivity (of beamsplitter in beamsplitter model of attenuation).
+                                 For a fiber, transmissivity = e**(-chi), where chi = l/l_att, where l is the length of the fiber and 
+                                 l_att is the attenuation length. If the noise_parameter is greater than 1, it is assumed to be a gain channel.
+        N (int): local Hilbert space dimension being considered.
+    """
+    a = qt.destroy(N).full()
+    a_dag = qt.create(N).full()
+    n = a_dag @ a
+    
+    # TODO: Theoretically, verify these
+    normalization = 1
+    gain_channel = False
+
+    if noise_parameter > 1: 
+        gain_channel = True
+        normalization = np.sqrt(1/noise_parameter)
+        noise_parameter = (noise_parameter-1)/(noise_parameter) # Convert gain to loss parameter
+
+    kraus_ops = []
+    for l in range(N): # you can lose anywhere from 0 to N-1 (=trunc) photons in the truncated Hilbert space. 
+        kraus_ops.append(sp.csr_array(normalization * np.sqrt(1/factorial(l) * (noise_parameter/(1-noise_parameter))**l) * (np.linalg.matrix_power(a, l) @ expm(n/2 * np.log(1-noise_parameter)))))
+
+    if gain_channel: 
+        for l in range(N):
+            kraus_ops[l] = kraus_ops[l].T.conjugate()
+
+    return kraus_ops

trajectree/fock_optics/outputs.py

@@ -0,0 +1,65 @@
+from scipy import sparse as sp
+
+import numpy as np
+
+from matplotlib import pyplot as plt
+
+# Generating labels for reading state. 
+def generate_labels(num_systems, N):
+    dim = N**2
+    labels = []
+    state_labels = []
+    for i in range(dim):
+        state_labels.append(f"{i//N}H{i%N}V")
+    # print("sates:", self.state_labels)
+    for i in range(dim**num_systems):
+        new_label = ""
+        for j in range(num_systems-1, -1, -1):
+            # print("appending to labels:", f"{self.state_labels[(i//self.dim**j)%self.dim]}_{chr(65+j)} ")
+            new_label += f"{state_labels[(i//dim**j)%dim]}_{chr(65+j)} "
+        labels.append(new_label[:-1])
+    return labels
+
+def read_quantum_state(TN_state, N, num_states = 4, return_dense = False, precision = 10):
+    dense_state = TN_state.to_dense()
+    if return_dense: return dense_state
+    dense_state = np.reshape(dense_state.data, (-1, 1), order = 'C')
+    dense_state = sp.csr_matrix(dense_state)
+    dense_state.data = np.round(dense_state.data, precision)
+    dense_state.eliminate_zeros()
+
+    print_quantum_state(N, dense_state, num_states)
+
+def print_quantum_state(N, dense_state, num_states = 4):
+    labels = generate_labels(num_states,N)
+    state = dense_state.nonzero()[0]
+    print("Corresponding Basis terms:")
+    for k in state: print(labels[k],"-",k,"-",dense_state[k].data)
+
+
+
+def plot_coincidences(coincidence, idler_angles, signal_angles, title = ''):
+    visibilities = []
+    for i in range(len(coincidence)):
+        visibility = (max(coincidence[i]) - min(coincidence[i])) / (max(coincidence[i]) + min(coincidence[i]))
+        visibilities.append(visibility)
+        # print(visibility, coincidence[i])
+
+    idler_angles = np.array(list(map(float, idler_angles)))/np.pi
+
+    plt.figure()
+    plt.grid(True)
+    for i in range(len(idler_angles)):
+        # print(fringe_real[i])
+        plt.plot(signal_angles, coincidence[i], label=r'{:.2f}$\pi$'.format(idler_angles[i]))
+    plt.title(title)
+    plt.ylabel("Coincidence probability")
+    plt.xlabel(r"$\alpha$ (rad)")    
+    plt.legend(title = "$\delta$")
+
+    plt.figure()
+    plt.grid(True)
+    plt.plot(idler_angles*np.pi, visibilities)
+    plt.title("Visiblilities")
+    plt.ylabel("Visibility")
+    plt.xlabel(r"$\delta$")    

trajectree/fock_optics/utils.py

@@ -0,0 +1,180 @@
+import numpy as np
+from numpy import sqrt
+
+from quimb.tensor import MatrixProductState as mps #type: ignore
+from quimb.tensor import MatrixProductOperator as mpo #type: ignore
+from quimb.tensor.tensor_arbgeom import tensor_network_apply_op_vec #type: ignore
+from quimb.tensor.tensor_core import new_bond #type: ignore
+from quimb.tensor.tensor_1d_compress import enforce_1d_like #type: ignore
+from quimb.tensor.tensor_1d import TensorNetwork1DOperator #type: ignore
+
+import qutip as qt
+import re
+
+
+###### SUPPORT FUNCTIONS ######
+
+# Vacuum state creation
+def fill_fn(shape):
+    arr = np.zeros(shape)
+    idx = tuple([0]*(len(shape)))
+    arr[idx] = 1
+    return arr
+def create_vacuum_state(num_modes, N, bond_dim = 2):
+    return mps.from_fill_fn(
+                fill_fn,
+                L=num_modes,
+                bond_dim=bond_dim,
+                phys_dim=N,
+                cyclic=False,
+                tags="In"
+            )
+
+def create_ladder_MPO(site, total_sites, N, tag="$Ladder$"):
+    a = qt.destroy(N).full()
+    a_dag = a.T
+    TMSV_MPO = mpo.from_dense(a_dag, dims = N, sites = (site,), L=total_sites, tags=tag) 
+    # return TMSV_MPO.fill_empty_sites(mode = "minimal")
+    return TMSV_MPO
+
+def create_MPO(site1, site2, total_sites, op, N, tag):
+    MPO = mpo.from_dense(op, dims = N, sites = (site1,site2), L=total_sites, tags=tag)    
+    return MPO
+
+###### POVM OPERATORS #######
+
+
+########## TMSV Operator ############
+
+
+
+
+########## EXTEND MPS ###########
+
+def extend_MPS(psi, psi_second = None):
+    # print("inside extend_MPS")
+    # psi_second.draw()
+    # print(psi_second)
+    
+    psi.permute_arrays('lrp')
+
+    # psi_second.draw()
+    # print(psi_second)
+
+    # This is supposed to be passed as the second MPS to extend the first MPS with. 
+    if psi_second == None:
+        psi_second = psi.copy()
+    else:
+        psi_second.permute_arrays('lrp')
+    
+    psi_num_modes = len(psi.site_tags)
+    psi2_num_modes = len(psi_second.site_tags)
+
+    psi_second.reindex({f"k{i}":f"k{i+psi_num_modes}" for i in range(psi2_num_modes)}, inplace = True)
+    psi_second.retag({f"I{i}":f"I{i+psi_num_modes}" for i in range(psi2_num_modes)}, inplace = True)
+
+    psi = psi.combine(psi_second)
+
+    psi_last_tensor = psi.select_tensors(f"I{psi_num_modes-1}", which='any')[0]
+    psi2_first_tensor = psi.select_tensors(f"I{psi_num_modes}", which='any')[0]
+
+    new_bond(psi2_first_tensor, psi_last_tensor, axis1=0, axis2=1)
+
+    # Simply find the tags for the input modes. 
+    pattern = re.compile(r"I[0-9][0-9]*")
+    tags = []
+    for tag_list in [t.tags for t in psi]:
+        for tag in tag_list:
+            match = re.search(pattern, tag)
+            if match:
+                tags.append(match.string)
+                break
+            
+    sorted_arrays = [array for array, _ in sorted( zip(psi.arrays, tags), key = lambda pair: int(pair[1][1:]) )]
+
+    psi = mps(sorted_arrays)
+    return psi
+
+
+def calc_fidelity_swapping(state, reference_state, N, error_tolerance):
+    reference_mps = create_bimode_bell_state(reference_state, N)
+    projector_mpo = outer_product_mps(reference_mps)
+
+    projector_mpo.reindex({"k0":"k0","k1":"k1","k2":"k4","k3":"k5"}, inplace = True)
+    projector_mpo.reindex({"b0":"b0","b1":"b1","b2":"b4","b3":"b5"}, inplace = True)
+    projector_mpo.retag({"I0":"I0","I1":"I1","I2":"I4","I3":"I5"}, inplace = True)
+
+    # print("sites present in projector_mpo:", projector_mpo.sites)
+    enforce_1d_like(projector_mpo, site_tags=state.site_tags, inplace=True)
+
+    state = tensor_network_apply_op_vec(projector_mpo, state, compress=True, contract = True, cutoff = error_tolerance)
+    # state.draw()
+    return state.norm()**2
+
+    
+    
+    # Calculate and return fidelity of the projected state. 
+
+
+def create_bimode_bell_state(bell_state, N, error_tolerance = 1e-12):
+    I = np.eye(N)
+
+    a_dag = qt.create(N).full()
+    a = qt.destroy(N).full()
+
+    vacuum_state = np.zeros((N,1))
+    vacuum_state[0] = 1
+    vac_projector = np.outer(vacuum_state, vacuum_state)
+
+    one_state = a_dag @ vacuum_state # For now, we're defining the 1 state as having only one photon. This could be changed to have any number of non-zero photons.
+    # print("one_state:", one_state)   # This is because the ideal case is having exactly one photon for the 1 state. 
+    one_projector = np.outer(one_state, one_state)                                 
+
+    NOT_gate = vacuum_state @ one_state.conj().T + one_state @ vacuum_state.conj().T
+    H_gate = (1/sqrt(2)) * ((vacuum_state - one_state) @ one_state.conj().T + (vacuum_state + one_state) @ vacuum_state.conj().T)
+    C_NOT_close = np.kron(vac_projector, I) + np.kron(one_projector, NOT_gate)
+    C_NOT_open = np.kron(one_projector, I) + np.kron(vac_projector, NOT_gate)
+
+    NOT_MPO_0 = mpo.from_dense(NOT_gate, dims = N, sites = (0,), L=4, tags="a_dag")
+    NOT_MPO_1 = mpo.from_dense(NOT_gate, dims = N, sites = (1,), L=4, tags="a_dag")
+    H_MPO = mpo.from_dense(H_gate, dims = N, sites = (0,), L=4, tags="H")
+    C_NOT_close_MPO_1 = mpo.from_dense(C_NOT_close, dims = N, sites = (0,1), L=4, tags="C_NOT_close_1")
+    C_NOT_close_MPO_2 = mpo.from_dense(C_NOT_close, dims = N, sites = (1,2), L=4, tags="C_NOT_close_2")
+    C_NOT_open_MPO = mpo.from_dense(C_NOT_open, dims = N, sites = (2,3), L=4, tags="C_create_open")
+    
+    vacuum = create_vacuum_state(4, N, bond_dim = 2)
+
+    if bell_state == "psi_minus":
+        psi = tensor_network_apply_op_vec(NOT_MPO_0, vacuum, compress=True, contract = True, cutoff = error_tolerance)
+        psi = tensor_network_apply_op_vec(NOT_MPO_1, psi, compress=True, contract = True, cutoff = error_tolerance)
+    elif bell_state == "psi_plus":
+        psi = tensor_network_apply_op_vec(NOT_MPO_1, vacuum, compress=True, contract = True, cutoff = error_tolerance)
+    elif bell_state == "phi_plus":
+        psi = vacuum
+    elif bell_state == "phi_minus":
+        psi = tensor_network_apply_op_vec(NOT_MPO_0, vacuum, compress=True, contract = True, cutoff = error_tolerance)
+
+    
+    psi = tensor_network_apply_op_vec(H_MPO, psi, compress=True, contract = True, cutoff = error_tolerance)
+    # read_quantum_state(psi, N, num_states = 2)
+    psi = tensor_network_apply_op_vec(C_NOT_close_MPO_1, psi, compress=True, contract = True, cutoff = error_tolerance)
+    # read_quantum_state(psi, N, num_states = 2)
+    psi = tensor_network_apply_op_vec(C_NOT_close_MPO_2, psi, compress=True, contract = True, cutoff = error_tolerance)
+    psi = tensor_network_apply_op_vec(C_NOT_open_MPO, psi, compress=True, contract = True, cutoff = error_tolerance)
+    
+    return psi
+
+
+def outer_product_mps(psi):
+    psi_H = psi.H
+    psi_H.retag_({'In': 'Out'})
+    psi_H.site_ind_id = 'b{}'
+    rho = (psi_H | psi)
+    for i in range(rho.L):
+        rho ^= f"I{i}"   
+    rho = TensorNetwork1DOperator(rho)
+    rho._upper_ind_id = psi.site_ind_id
+    rho._lower_ind_id = psi_H.site_ind_id
+    rho = rho.fuse_multibonds()
+    rho_MPO = rho.view_as_(mpo, cyclic = False, L = 8) # L is important. Its hard coded now, but must be configutrable based on the input state. 
+    return rho_MPO

trajectree/optical_quant_info.py

@@ -0,0 +1,137 @@
+from trajectree.fock_optics.measurement import *
+from trajectree.fock_optics.utils import *
+from trajectree.fock_optics.light_sources import *
+from trajectree.fock_optics.devices import *
+from trajectree.trajectory import *
+
+from trajectree.protocols.swap import perform_swapping_simulation
+
+import numpy as np
+
+def quantum_encoder(mean_photon_num, N, psi_control, control_indices, error_tolerance):
+    """This function performs the quantum encoder, basically "copy" one state into two modes. Obviously, you aren't copying shit. You are simply entangling another 
+    state (modes (0,1) or mode b in the paper) with the control state, creating a "copy". The bell state required to make the copy is added before the control MPS 
+    ((0,1):(H,V){b,d} and (2,3):(H,V){a,c}). {a,c} modes are measured out at the end of the function and hence, only 2 modes are pre-added to the retured MPS.   
+    
+    """
+    # Entangled state from EPS
+    vacuum = create_vacuum_state(num_modes=8, N=N)
+    bell_state, _ = light_source(vacuum, N, mean_photon_num, 8, error_tolerance, compress=True, contract=True)
+
+    # psi_control.draw()
+    # print(psi_control)
+
+    psi = extend_MPS(bell_state, psi_control)
+
+    # psi.draw()
+    # print(psi)
+
+    # PBS op: (The V mode is transmitted and not reflected)
+    U_PBS_V = create_BS_MPO(site1 = 3, site2 = bell_state.L+control_indices[1], theta=np.pi/2, total_sites = psi.L, N = N, tag = r"$PBS$")
+    enforce_1d_like(U_PBS_V, site_tags=psi.site_tags, inplace=True)
+    psi = tensor_network_apply_op_vec(U_PBS_V, psi, compress=True, contract = True, cutoff = error_tolerance)
+
+    # Measuring D_d
+    # This is meant to change the basis from HV -> FS: (See https://doi.org/10.1103/PhysRevA.64.062311)
+    U_PBS_FS = create_BS_MPO(site1 = 2, site2 = 3, theta=np.pi/4, total_sites = psi.L, N = N, tag = r"$rotator$")
+    enforce_1d_like(U_PBS_FS, site_tags=psi.site_tags, inplace=True)
+    psi = tensor_network_apply_op_vec(U_PBS_FS, psi, compress=True, contract = True, cutoff = error_tolerance)
+
+    # Performing measurements:
+    BSM_POVM_1_OPs = generate_sqrt_POVM_MPO(sites=[2], outcome = 1, total_sites=psi.L, efficiency=1, N=N, pnr = True)
+    BSM_POVM_1_OPs.extend(generate_sqrt_POVM_MPO(sites=[3], outcome = 0, total_sites=psi.L, efficiency=1, N=N, pnr = True))
+
+    for POVM_OP in BSM_POVM_1_OPs:
+        psi = tensor_network_apply_op_vec(POVM_OP, psi, compress=True, contract = True, cutoff = error_tolerance)
+
+    return psi
+
+
+def destructive_CNOT(control_b_sites, target_sites, psi, N, error_tolerance):
+
+    # print("control_b_sites:", control_b_sites)
+    # print("target_sites", target_sites)
+
+    # Rotaing bases of encoded control's b mode and the target mode:
+    U_rotator_FS = create_BS_MPO(site1 = control_b_sites[0], site2 = control_b_sites[1], theta=np.pi/4, total_sites = psi.L, N = N, tag = r"$rotator$")
+    enforce_1d_like(U_rotator_FS, site_tags=psi.site_tags, inplace=True)
+    psi = tensor_network_apply_op_vec(U_rotator_FS, psi, compress=True, contract = True, cutoff = error_tolerance)
+
+    U_rotator_FS = create_BS_MPO(site1 = target_sites[0], site2 = target_sites[1], theta=np.pi/4, total_sites = psi.L, N = N, tag = r"$rotator$")
+    enforce_1d_like(U_rotator_FS, site_tags=psi.site_tags, inplace=True)
+    psi = tensor_network_apply_op_vec(U_rotator_FS, psi, compress=True, contract = True, cutoff = error_tolerance)
+
+    # Applying PBS in rotated basis (only the V modes are reflected and the H modes are transmitted. Hence, the V modes undergo pi/2 rotation and the H modes undergo no rotation):
+    
+    # Implementation using SWAP operator. This does not generalize to higher truncations. 
+    # SWAP = qt.qip.operations.swap().full()
+    # U_PBS_F = create_MPO(1, 3, psi.L, SWAP, N, r"$PBS$")
+    U_PBS_F = create_BS_MPO(site1 = target_sites[1], site2 = control_b_sites[1], theta=np.pi/2, total_sites = psi.L, N = N, tag = r"$PBS$")
+    enforce_1d_like(U_PBS_F, site_tags=psi.site_tags, inplace=True)
+    psi = tensor_network_apply_op_vec(U_PBS_F, psi, compress=True, contract = True, cutoff = error_tolerance)
+
+    # Undoing rotations:
+    U_inverse_rotator_FS = create_BS_MPO(site1 = target_sites[0], site2 = target_sites[1], theta=-np.pi/4, total_sites = psi.L, N = N, tag = r"$rotator$")
+    enforce_1d_like(U_inverse_rotator_FS, site_tags=psi.site_tags, inplace=True)
+    psi = tensor_network_apply_op_vec(U_inverse_rotator_FS, psi, compress=True, contract = True, cutoff = error_tolerance)
+
+    U_inverse_rotator_FS = create_BS_MPO(site1 = control_b_sites[0], site2 = control_b_sites[1], theta=-np.pi/4, total_sites = psi.L, N = N, tag = r"$rotator$")
+    enforce_1d_like(U_inverse_rotator_FS, site_tags=psi.site_tags, inplace=True)
+    psi = tensor_network_apply_op_vec(U_inverse_rotator_FS, psi, compress=True, contract = True, cutoff = error_tolerance)
+
+    # Measuring the b mode (name d after the PBS)
+    BSM_POVM_1_OPs = generate_sqrt_POVM_MPO(sites=[control_b_sites[0]], outcome = 1, total_sites=psi.L, efficiency=1, N=N, pnr = True)
+    BSM_POVM_1_OPs.extend(generate_sqrt_POVM_MPO(sites=[control_b_sites[1]], outcome = 0, total_sites=psi.L, efficiency=1, N=N, pnr = True))
+
+    for POVM_OP in BSM_POVM_1_OPs:
+        psi = tensor_network_apply_op_vec(POVM_OP, psi, compress=True, contract = True, cutoff = error_tolerance)
+
+    return psi
+
+def CNOT(psi_control_modes, psi_target_modes, psi_control, psi_target, N, mean_photon_num, error_tolerance):
+    """Pass psi_target as None if the same MPS has both the control and target modes.
+    Args:
+        psi_control_modes (list): List of control modes (H,V).
+        psi_target_modes (list): List of target modes (H,V).
+        psi_control (MPS): MPS for the control modes.
+        psi_target (MPS): MPS for the target modes, can be None if the same MPS as target is used.
+        N (int): Number of photons.
+        mean_photon_num (float): Mean photon number for the EPS used in implemeting the CNOT gate.
+        error_tolerance (float): Tolerance for numerical errors in tensor network operations.
+    Returns:
+        MPS: The resulting MPS after applying the CNOT operation.
+    """
+
+    psi_encoded_control = quantum_encoder(mean_photon_num, N, psi_control, psi_control_modes, error_tolerance)
+
+    # read_quantum_state(psi_encoded_control,  N, num_states = 6)
+
+    if not psi_target == None:
+        psi = extend_MPS(psi_target, psi_encoded_control)
+        psi_control_b_modes = [psi_target.L, psi_target.L+1] # [site + psi_target.L for site in psi_control_modes]
+    
+    else: 
+        psi_control_b_modes = [0,1]
+        psi_target_modes = [4+site for site in psi_target_modes] # We add 4 since the additional modes from the EPS are pre-added to the MPS. 
+        psi = psi_encoded_control
+
+    psi = destructive_CNOT(psi_control_b_modes, psi_target_modes, psi, N, error_tolerance)
+
+    # read_quantum_state(psi, N, num_states = 6)
+
+    norm = psi.normalize()
+    for _ in range(4):
+        psi.measure(0, remove = True, renorm = True, inplace = True)
+    psi[-1].modify(data=psi[-1].data * norm**0.5)
+
+    return psi
+
+
+def H(psi, sites, N, error_tolerance):
+    # TODO: This function does not work for N > 2.
+    # This definition is based on the paper: https://arxiv.org/pdf/quant-ph/9706022
+    H = generalized_mode_mixer(sites[0], sites[1], -np.pi/4, [0,-np.pi], [0,-np.pi], [0,0], psi.L, N)
+    # H = generalized_mode_mixer(0, 1, np.pi/4, 0, 0, 0, 2, N)
+    enforce_1d_like(H, site_tags=psi.site_tags, inplace=True)
+    psi = tensor_network_apply_op_vec(H, psi, compress=True, contract = True, cutoff = error_tolerance)
+    return psi

trajectree/sequence/swap.py

@@ -0,0 +1,76 @@
+from ..fock_optics.noise_models import *
+from ..fock_optics.measurement import *
+from ..fock_optics.utils import *
+from ..fock_optics.light_sources import *
+
+from ..trajectory import *
+
+import time 
+import numpy as np
+import copy
+
+def generate_swapping_circuit(N, num_modes, site_tags, bsm_det_effs, bsm_dark_counts_gain, bsm_measurements,  channel_loss, error_tolerance):
+    quantum_channel_list = []
+
+    # Amplitude damping due to fibers
+    damping_kraus_ops = single_mode_bosonic_noise_channels(noise_parameter = channel_loss, N = N)
+    two_mode_kraus_ops = [sp.kron(op, op) for op in damping_kraus_ops]
+    quantum_channel_list.append(quantum_channel(N = N, num_modes = num_modes, formalism = "kraus", kraus_ops_tuple = ((2,3), two_mode_kraus_ops))) # The tuples in this list are defined as (sites, kraus_ops). The sites are the sites where the Kraus ops are applied.
+    quantum_channel_list.append(quantum_channel(N = N, num_modes = num_modes, formalism = "kraus", kraus_ops_tuple = ((6,7), two_mode_kraus_ops))) # The tuples in this list are defined as (sites, kraus_ops). The sites are the sites where the Kraus ops are applied.
+
+    # Quantum channel for the Bell state measurement
+    # BSM_MPOs = bell_state_measurement(None, N, site_tags, num_modes, bsm_det_effs, error_tolerance, measurements = bsm_measurements, pnr = False, use_trajectory = True, return_MPOs = True, compress=True, contract=True)
+    # BSM_quantum_channels = [quantum_channel(N = N, num_modes = num_modes, formalism = "closed", unitary_MPOs = BSM_MPO, name = "BSM") for BSM_MPO in BSM_MPOs]
+    BSM_quantum_channels = bell_state_measurement(None, N, site_tags, num_modes, bsm_det_effs, bsm_dark_counts_gain, error_tolerance, measurements = bsm_measurements, pnr = False, use_trajectory = True, return_MPOs = True, compress=True, contract=True)
+    quantum_channel_list.extend(BSM_quantum_channels)
+
+    return quantum_channel_list
+
+def analyze_entanglement(quantum_channel_list, N, site_tags, num_modes, efficiency, error_tolerance, idler_angles, signal_angles):
+    PA_MPOs = rotate_and_measure(None, N, site_tags, num_modes, efficiency, error_tolerance, idler_angles, signal_angles, return_MPOs = True)
+    PA_quantum_channels = [quantum_channel(N = N, num_modes = num_modes, formalism = "closed", unitary_MPOs = PA_MPO) for PA_MPO in PA_MPOs]
+    print("num pa quantum channels:", len(PA_quantum_channels))
+    quantum_channel_list.extend(PA_quantum_channels)
+
+
+def create_swapping_initial_state(num_modes, N, mean_photon_num, error_tolerance):
+    # Create Vacuum state:
+    vacuum = create_vacuum_state(num_modes=num_modes, N=N)
+
+    # Entangled state from EPS
+    psi, TMSV_state = light_source(vacuum, N, mean_photon_num, num_modes, error_tolerance, compress=True, contract=True)
+
+    psi = extend_MPS(psi)
+    return psi
+
+def perform_swapping_simulation(N, num_modes, num_simulations, params, error_tolerance = 1e-10):
+
+    psi = create_swapping_initial_state(num_modes, N, params["chi"], error_tolerance)
+
+    quantum_channels = generate_swapping_circuit(N, num_modes, psi.site_tags, [params["BSM_det_loss_1"], params["BSM_det_loss_2"]], [params["BSM_dark_counts_1"], params["BSM_dark_counts_2"]], params["BSM_meas"], params["channel_loss"], error_tolerance)
+
+    if params["if_analyze_entanglement"]:
+        analyze_entanglement(quantum_channels, N, psi.site_tags, num_modes, params["PA_det_loss"], error_tolerance, params["alpha_list"], params["delta_list"])
+
+    t_eval = trajectory_evaluator(quantum_channels)
+
+    fidelities = []
+    probabilities = []
+
+    for i in range(num_simulations): 
+        start = time.time()
+        psi_iter = copy.deepcopy(t_eval.perform_simulation(psi, error_tolerance, normalize = False))
+
+        probabilities.append(psi_iter.normalize())
+        
+        if params["calc_fidelity"]:
+            fidelity = np.abs(calc_fidelity_swapping(psi_iter, "psi_minus", N, error_tolerance))
+            fidelities.append(fidelity)
+    
+        time_taken = time.time() - start
+        # print("time taken for simulation", i, ":", time_taken)
+
+    print("completed set", "cache_hits:", t_eval.cache_hit, "cache_partial_hits:", t_eval.cache_partial_hit, "cache_misses:", t_eval.cache_miss,  "time taken:", time_taken)
+
+    return fidelities, probabilities, t_eval
+

trajectree/trajectory.py

@@ -0,0 +1,210 @@
+import numpy as np
+from quimb.tensor import MatrixProductOperator as mpo #type: ignore
+from quimb.tensor.tensor_arbgeom import tensor_network_apply_op_vec #type: ignore
+
+
+
+class quantum_channel:
+    def __init__(self, N, num_modes, formalism, kraus_ops_tuple = None, unitary_MPOs = None, name = "quantum_channel"):
+        self.N = N
+        self.name = name
+        self.num_modes = num_modes
+        self.formalism = formalism
+        if self.formalism == 'kraus':
+            # Calculate the MPOs of the Kraus operators
+            self.kraus_MPOs = quantum_channel.find_quantum_channels_MPOs(kraus_ops_tuple, N, num_modes)
+        elif self.formalism == 'closed':
+            self.unitary_MPOs = unitary_MPOs
+
+    def get_MPOs(self):
+        if self.formalism == 'closed':
+            return self.unitary_MPOs
+        elif self.formalism == 'kraus':
+            return self.kraus_MPOs
+
+    @staticmethod
+    def find_quantum_channels_MPOs(ops_tuple, N, num_modes):
+        (sites, ops) = ops_tuple
+        quantum_channels = quantum_channel.calc_mpos(ops, N, sites, num_modes)
+        return quantum_channels
+
+    # Just a function which calcualte the MPOs of the Kraus ops
+    @staticmethod
+    def calc_mpos(ops, N, sites, num_modes):
+        MPOs = []
+        for op in ops:
+            MPO = mpo.from_dense(op.todense(), dims = N, sites = sites, L=num_modes, tags="op")
+            MPOs.append(MPO)
+        return MPOs
+
+
+class trajectree_node:
+    def __init__(self, probs, trajectories, trajectory_indices):
+        self.probs = probs
+        self.trajectories = trajectories
+        self.trajectory_indices = trajectory_indices
+
+class trajectory_evaluator():
+    def __init__(self, quantum_channels, cache_size = 2):
+        self.quantum_channels = quantum_channels
+        self.kraus_channels = []
+        for quantum_channel in self.quantum_channels:
+            if quantum_channel.formalism == 'kraus':
+                self.kraus_channels.append(quantum_channel)
+
+        self.trajectree = [{} for i in range(len(self.kraus_channels)+1)] # +1 because you also cache the end of the simulation so you prevent doing the final unitary operations multiple times. 
+        self.traversed_nodes = ()
+        self.cache_size = cache_size
+
+        # for debugging only:
+        self.cache_hit = 0
+        self.cache_miss = 0
+        self.cache_partial_hit = 0
+
+
+    def apply_kraus(self, psi, kraus_MPOs, error_tolerance, normalize = True):
+        trajectory_probs = np.array([])
+        trajectories = np.array([])
+        for kraus_MPO in kraus_MPOs:
+
+            trajectory = tensor_network_apply_op_vec(kraus_MPO, psi, compress=True, contract = True, cutoff = error_tolerance)
+            trajectory_prob = np.real(trajectory.H @ trajectory)
+            
+            if trajectory_prob < 1e-25: # Using 1e-25 arbitrarily. Trajectories with probability less than this are pruned.  
+                continue
+
+            if normalize:
+                trajectory.normalize()
+            trajectory_probs = np.append(trajectory_probs, trajectory_prob)
+            trajectories = np.append(trajectories, trajectory)
+
+        return trajectories, trajectory_probs
+
+
+    def cache_trajectree_node(self, trajectory_probs, trajectories):
+        sorted_indices = np.argsort(trajectory_probs)
+
+        print("trajectory_probs", trajectory_probs)
+
+        cached_trajectory_indices = sorted_indices[-self.cache_size:]
+        cached_trajectories = np.array(trajectories)[cached_trajectory_indices]
+
+        new_node = trajectree_node(trajectory_probs, cached_trajectories, cached_trajectory_indices)
+        self.trajectree[len(self.traversed_nodes)][self.traversed_nodes] = new_node
+
+        self.last_cached_node = new_node
+
+        return cached_trajectory_indices
+
+
+    def discover_trajectree_node(self, psi, kraus_MPOs, error_tolerance, normalize = True, selected_trajectory_index = None):
+        
+        trajectories, trajectory_probs = self.apply_kraus(psi, kraus_MPOs, error_tolerance, normalize)
+
+        cached_trajectory_indices = self.cache_trajectree_node(trajectory_probs, trajectories) # cached_trajectory_indices is returned only for debugging. 
+
+        if selected_trajectory_index == None:
+            selected_trajectory_index = np.random.choice(a = len(trajectory_probs), p = trajectory_probs/sum(trajectory_probs))
+        
+        self.traversed_nodes = self.traversed_nodes + (selected_trajectory_index,)
+
+        return trajectories[selected_trajectory_index]
+
+
+    def query_trajectree(self, psi, kraus_MPOs, error_tolerance, cache = True, selected_trajectory_index = None, normalize = True):
+        self.skip_unitary = False
+        self.cache_unitary = False
+
+        if cache == False:
+            psi = tensor_network_apply_op_vec(self.kraus_channels[len(self.traversed_nodes)].get_MPOs()[selected_trajectory_index], psi, compress=True, contract = True, cutoff = error_tolerance)
+            self.traversed_nodes = self.traversed_nodes + (selected_trajectory_index,)
+            return psi
+
+        if self.traversed_nodes in self.trajectree[len(self.traversed_nodes)]: # Check if the dictionary at level where the traversal is now, i.e., len(self.traversed_nodes)
+                                                                               # has the path that the present traversal has taken. 
+            node = self.trajectree[len(self.traversed_nodes)][self.traversed_nodes] # If found, index that node into the node object to call the probabilities and trajectories cached inside it.
+            if selected_trajectory_index == None:
+                selected_trajectory_index = np.random.choice(a = len(node.probs), p = node.probs/sum(node.probs)) # The cached nodes have all the probabilities, but not all the trajectories cache. So, we can select
+                                                                                                  # what trajecory our traversal takes and later see if the actual trajectory has been cached or needs to be retrieved. 
+            self.cache_unitary = False # If the node has been found, we do not cache the unitary. The unitary is either already cached or we don't need to cache it at all.
+
+            if selected_trajectory_index in node.trajectory_indices: # See if the selected trajectory's MPS has been cached or not. 
+                self.skip_unitary = True # If we're skipping the unitary entirely, it just does not matter whether we cache the unitary or not.
+                self.cache_hit += 1
+                psi = node.trajectories[np.where(node.trajectory_indices == selected_trajectory_index)[0][0]]
+            else: 
+                self.skip_unitary = False # If the trajectory has not been cached, we will have to apply the unitary to it.
+                self.cache_partial_hit += 1
+                psi = tensor_network_apply_op_vec(self.kraus_channels[len(self.traversed_nodes)].get_MPOs()[selected_trajectory_index], psi, compress=True, contract = True, cutoff = error_tolerance) # If not, simply calculate that trajectory. 
+                                                                                                                                                                                             # You don't need to cache it since we have already cached what we had to.  
+                if normalize:
+                    psi.normalize()
+            self.traversed_nodes = self.traversed_nodes + (selected_trajectory_index,)
+
+
+        else: # If the node has not been discovered, we'll have to find all probabilities and cache the results. 
+            self.skip_unitary = False
+            self.cache_unitary = True
+            self.cache_miss += 1
+            psi = self.discover_trajectree_node(psi, kraus_MPOs, error_tolerance, normalize, selected_trajectory_index = selected_trajectory_index)
+
+        return psi
+
+    def apply_unitary_MPOs(self, psi, unitary_MPOs, error_tolerance):
+        return tensor_network_apply_op_vec(unitary_MPOs, psi, compress=True, contract = True, cutoff = error_tolerance)
+
+
+    def calculate_density_matrix(self, psi, error_tolerance):
+        dm = 0
+        trajectree_indices_list = [[]]
+        for quantum_channel in self.quantum_channels:
+            if quantum_channel.formalism == 'kraus':
+                trajectree_indices_list = [[*i, j] for i in trajectree_indices_list for j in range(len(quantum_channel.get_MPOs()))]
+        for trajectree_indices in trajectree_indices_list:
+            psi_new_dense = self.perform_simulation(psi, error_tolerance, cache = True, trajectree_indices = trajectree_indices, normalize = False).to_dense()
+            dm += psi_new_dense @ psi_new_dense.conj().T
+        return dm
+
+    def update_cached_node(self, unitary_MPOs, last_cached_node, error_tolerance):
+        for kraus_idx in range(len(last_cached_node.trajectories)):
+            last_cached_node.trajectories[kraus_idx] = self.apply_unitary_MPOs(last_cached_node.trajectories[kraus_idx], unitary_MPOs, error_tolerance)
+
+
+
+    def perform_simulation(self, psi, error_tolerance, cache = True, trajectree_indices = None, normalize = True):
+        self.traversed_nodes = ()
+        self.skip_unitary = False
+        self.cache_unitary = False
+        for quantum_channel in self.quantum_channels:
+            if quantum_channel.formalism == 'kraus':
+                kraus_MPOs = quantum_channel.get_MPOs()
+                if not trajectree_indices == None: # If the list of trajectoery indices is provided, we will use that to traverse the trajectree. The random number generators will not be used.
+                    psi = self.query_trajectree(psi, kraus_MPOs, error_tolerance, cache, trajectree_indices.pop(0), normalize)
+                else: # In this branch, you actually select the trajectory redomly and perform realistic simulations. 
+                    psi = self.query_trajectree(psi, kraus_MPOs, error_tolerance, cache = cache, normalize = normalize)
+
+            elif quantum_channel.formalism == 'closed' and not self.skip_unitary:
+                unitary_MPOs = quantum_channel.get_MPOs()
+                                
+                if not cache: # If we aren't aching the trajectories at all, simply apply the unitary MPOs to the state.
+                    psi = self.apply_unitary_MPOs(psi, unitary_MPOs, error_tolerance)
+                    continue
+
+                last_cached_node = self.trajectree[len(self.traversed_nodes)-1][self.traversed_nodes[:-1]]
+                
+                if self.cache_unitary:
+                    self.update_cached_node(unitary_MPOs, last_cached_node, error_tolerance)
+
+                # This is where we are checking if the psi is cached or not. If it is, simply use the last cached node 
+                # node to update psi. If not, apply the unitary MPOs to psi.
+                traj_idx = np.where(last_cached_node.trajectory_indices == self.traversed_nodes[-1])    
+                if traj_idx[0].size > 0:
+                    psi = last_cached_node.trajectories[traj_idx[0][0]]
+                else:
+                    psi = self.apply_unitary_MPOs(psi, unitary_MPOs, error_tolerance)
+            
+            else:
+                # print("unitary skipped:", self.traversed_nodes)
+                pass
+
+        return psi

trajectree-0.0.0.dist-info/METADATA

@@ -0,0 +1,18 @@
+Metadata-Version: 2.4
+Name: Trajectree
+Version: 0.0.0
+Summary: Trajectree is a quantum trajectory theory and tensor network based quantum optics simulator.
+Author-email: Ansh Singal <asingal@u.northwestern.edu>
+License-Expression: MIT
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: quimb
+Requires-Dist: numpy
+Requires-Dist: scipy
+Requires-Dist: matplotlib
+Dynamic: license-file
+
+Trajectree is a quantum trajectory theory and tensor network based quantum optics simulator. 

trajectree-0.0.0.dist-info/RECORD

@@ -0,0 +1,16 @@
+trajectree/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+trajectree/optical_quant_info.py,sha256=27LWg0jIWT2ZXWVK8HEd-2gqAJl0GzmsyYi7sxyYzvA,7589
+trajectree/trajectory.py,sha256=yqY68605Qxf49LOvK4IRv9w1lvX_5kLfnFHAK0GMUmE,11262
+trajectree/experimental/sparse.py,sha256=05aFsMcBSCUoBop1LJG_lB5kTULRSYhZhGT9RD-UH4Q,4851
+trajectree/fock_optics/devices.py,sha256=1gwl4AIktqu3obXmNhR1EoKR4iWulP4LfWD4xenKfx8,2130
+trajectree/fock_optics/light_sources.py,sha256=qzBevi1Uqcxb6KuYMuVQ9AkJbHCCigBf21N6s01x-V0,4823
+trajectree/fock_optics/measurement.py,sha256=avTx8LD5OZFg1JCFkEJRZmDFAzlaAMN7XPxU0WqVn-s,13099
+trajectree/fock_optics/noise_models.py,sha256=nyU_jNqjJkkcafq7vIXI8IRDiwthiOr4ZQqf55knZMM,1895
+trajectree/fock_optics/outputs.py,sha256=Dg19FvodfKWDTB7B6IZ_EYUtC3vgsjlBugyrF4_tBvc,2270
+trajectree/fock_optics/utils.py,sha256=SqMTDHe7QcGD7PeSb5Pj5MlxlSiXbjgSsGwQUArQfxk,7098
+trajectree/sequence/swap.py,sha256=g7yWJg6Ow9D07GIHOxiINL6ao3jaAniGAhnW1AR_2ps,4288
+trajectree-0.0.0.dist-info/licenses/LICENSE,sha256=7EI8xVBu6h_7_JlVw-yPhhOZlpY9hP8wal7kHtqKT_E,1074
+trajectree-0.0.0.dist-info/METADATA,sha256=NA9LiSQFm_64Ixwemj7k1gtOjt23DivWEed92Rq5vZQ,621
+trajectree-0.0.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+trajectree-0.0.0.dist-info/top_level.txt,sha256=6x9i8aAZcVn5tZ9J-2IVQMTGdn4bw6DiSd8pR3v4VR8,11
+trajectree-0.0.0.dist-info/RECORD,,

trajectree-0.0.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

trajectree-0.0.0.dist-info/licenses/LICENSE

@@ -0,0 +1,19 @@
+Copyright (c) 2018 The Python Packaging Authority
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

trajectree-0.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+trajectree