encoding.py

import numpy as np
import strawberryfields as sf
from thewalrus.symplectic import is_symplectic
import pickle

def load_processed_data(load_path):
    ' input : path to the preprocessed data (a pickle file containing a dictionary with keys "adj" and "targets") '
    ' output : a list of adjacency matrices and a numpy array of targets '
    with open(load_path, 'rb') as f:
        data = pickle.load(f)
    return data['adj'], np.array(data['targets'])

def highest_singular_value(adj):
    ' input : list of adjacency matrices '
    ' output : list of scaled adjacency matrices, each of shape (N, N), where N is the number of nodes of a given graph, and global maximum singular value across all graphs '
    # get the largest singular value across all graphs
    global_s_max = 0 
    for A in adj:
        s = np.linalg.eigvals(A)
        s = np.abs(s)
        s_max_values = np.max(s)
        if s_max_values > global_s_max:
            global_s_max = s_max_values
    
    # scale each matrix
    scaled_adj = []
    for A in adj:
        # we need to scale with (/0.99) here to avoid numerical issues with the svd values of the scaled adjacency matrix of the graph with the global_s_max, which causes problems for the squeezing parameters and the covariance matrix later on (they become non-positive definite)
        A_scaled = A / (global_s_max / 0.99 + 1e-8) 
        scaled_adj.append(A_scaled)

    print(f"Created {len(scaled_adj)} scaled adjacency matrices")
    print(f"Shape of first scaled adjacency: {scaled_adj[0].shape}")
    print(f"Global maximum singular value across all graphs: {global_s_max}")
    
    return scaled_adj, global_s_max

def doubled_adjancencies(scaled_adj):
    ' input : list of scaled adjacency matrices, each of shape (N, N)'
    ' output : list of double encoded adjacency matrices, each of shape (2N, 2N)'
    doubled_adj = []

    for A in scaled_adj:
        zero_block = np.zeros_like(A)
        doubled_adjacency = np.block([[A, zero_block], [zero_block, A]])
        doubled_adj.append(doubled_adjacency)

    print(f"Created {len(doubled_adj)} double encoded adjacency matrices")
    print(f"Shape of first double encoded adjacency: {doubled_adj[0].shape}")

    return doubled_adj

def covariance_from_doubled_adj(doubled_adj):
    ' input : list of double encoded adjacency matrices, each of shape (2N, 2N) '
    ' output : list of covariance matrices, each of shape (2N, 2N) '
    covariances = []

    for i, A in enumerate(doubled_adj):

        N = A.shape[0] // 2

        # X has to be the same size as doubled_adj with blocks of identity matrices and zeros
        X = np.block([[np.zeros((N,N)), np.eye(N)],
                    [np.eye(N), np.zeros((N,N))]])
        
        # check eigenvalues of X @ A
        eigvals = np.linalg.eigvals(X @ A)
        # print(f'Eigenvalues of X @ A for graph {graph_nr}: {eigvals}')
        if np.max(np.abs(eigvals)) >= 1:
            print(f"Warning: |λ_max| = {np.max(np.abs(eigvals))} >= 1")

        # compute Q = (I - X * doubled_adj)^-1
        M = np.eye(A.shape[0]) - X @ A
        Q = np.linalg.inv(M)
        
        # compute covariance matrix: covariance = Q - I/2
        covariance = Q - np.eye(A.shape[0])/2

        # verify that covariance is positive definite
        if np.all(np.linalg.eigvals(covariance) > 0) == False:
            print(f"Covariance matrix for graph {i} is not positive definite. Minimum eigenvalue: {np.min(np.linalg.eigvals(covariance))}")
            break

        covariances.append(covariance)
        
    print(f"Created {len(covariances)} covariance matrices")
    print(f"Shape of first covariance matrix: {covariances[0].shape}")
    print(f"Minimum eigenvalue and maximum eigenvalue of the first covariance matrix: {np.min(np.linalg.eigvals(covariances[0]))}, {np.max(np.linalg.eigvals(covariances[0]))}" )

    return covariances

    
def decomposition(covariances):
    ' input : list of covariance matrices, each of shape (2N, 2N) '
    ' output : lists of the decomposition parameters for each graph: Db (list of diagonal matrices from williamson), S (list of symplectic matrices from williamson),'
    ' S_int (list of symplectic matrices containing only passive transformations from bloch-messiah), squeezing_diagonal (list of diagonal squeezing matrices from bloch-messiah), '
    ' U (list of unitaries extracted from S_int), tList (list of tuples of clement decomposition parameters extracted from U) '
    Db = []
    S = []
    S_int = []
    squeezing_diagonal = []
    U = []
    tList = []

    for i, cov in enumerate(covariances):
        # number of nodes in original graph
        N = cov.shape[0] // 2

        # williamson to decompose covariance matrix by symplectic diagonalization
        Db_k, S_k = sf.decompositions.williamson(cov)
        # check that the symplectic eigenvalues are all 0.5 (using hbar=1) for a pure state
        if not np.allclose(np.diag(Db_k), 0.5, atol=1e-8):
            print(f"Warning: Symplectic eigenvalues for graph {i} are not all 0.5. Minimum symplectic eigenvalue: {np.min(np.diag(Db_k))}, maximum symplectic eigenvalue: {np.max(np.diag(Db_k))}")

        # verify that S is symplectic (should be guaranteed by williamson but there can be numerical errors)
        if not is_symplectic(S_k):
            print(f"S_k is not symplectic for graph {i}. Minimum eigenvalue of covariance matrix: {np.min(np.linalg.eigvals(cov))}")

        # decompose S into symplectic unitaries containing the passive transformations (aka beamsplitters and phase shifters) and a squeezing diagonal matrix using bloch-messiah
        S_int_k, squeezing_diagonal_k, _ = sf.decompositions.bloch_messiah(S_k, tol=1e-8)
        # print(f'S_int_k shape: {S_int_k.shape}')

        # extract unitary from S_int_k
        ReU = S_int_k[:N, :N]
        ImU = S_int_k[:N, N:]
        U_k = ReU - 1j * ImU
        # print(f"U shape: {U.shape}")

        # check unitarity
        deviation = np.max(np.abs(U_k @ U_k.conj().T - np.eye(N)))
        if deviation > 1e-10:
            print(f"Unitarity deviation before correction: {deviation:.2e}")
            
            # QR decomposition
            U_k, _ = np.linalg.qr(U_k)
            
            # verify after correction
            new_deviation = np.max(np.abs(U_k @ U_k.conj().T - np.eye(N)))
            print(f"Unitarity deviation after correction: {new_deviation:.2e}")

        # extract parameters from U using clement decomposition (in StrawberryFields : rectangular_phase_end)
        tList_k = sf.decompositions.rectangular_phase_end(U_k)

        # ensure that the number of beamsplitters is equal to N(N-1)/2 in the covariance matrix, otherwise there is a problem with the decomposition
        if len(tList_k[0]) != N*(N-1)/2:
            print(f"Number of beamsplitters in tList_k for graph {i} is not equal to N(N-1)/2. Number of beamsplitters: {len(tList_k[0])}, N(N-1)/2: {cov.shape[0]*(cov.shape[0]-1)/2}.")
            break

        Db.append(Db_k) # diagonal matrix (how pure is the state, should be identity matrix in our case (pure state))
        S.append(S_k) # symplectic matrix, containg all the transformations
        S_int.append(S_int_k) # symplectic containing only the passive transformations (beamsplitters and phase shifters)
        squeezing_diagonal.append(squeezing_diagonal_k) # diagonal matrix containing the squeezing parameters (how much squeezing is needed for each mode)
        U.append(U_k) # unitary extracted from S_int, containing only the passive transformations
        tList.append(tList_k) # list of tuples of clement decomposition parameters extracted from U, containing the parameters for the beamsplitters and phase shifters in the unitary

    print(f"Decomposition complete. Nr. of graphs decomposed: {len(tList)}")
    print(f"First tList: \n {tList[0][0]} \n (n,m,theta,phi,n_size)")

    return Db, S, S_int, squeezing_diagonal, U, tList


if __name__ == "__main__":
    adj, targets = load_processed_data('data/enzymes.pkl')

    scaled_adj, global_s_max = highest_singular_value(adj)
    doubled_adj = doubled_adjancencies(scaled_adj)
    covariances = covariance_from_doubled_adj(doubled_adj)

    # save the covariance matrices and targets as a pickle file containing a dictionary with keys "covariances" and "targets"
    with open('data/enzymes_covariances.pkl', 'wb') as f:
        pickle.dump({'covariances': covariances, 'targets': targets}, f)

    Db, S, S_int, squeezing_diagonal, U, tList = decomposition(covariances)