Calculating partial transposition.

partial_transposition.py

import numpy as np
from math import sqrt

def partial_transposition(density_matrix: np.ndarray, basis_num_1: int=0, basis_num_2: int=0):
    if basis_num_1 == 0 and basis_num_2 == 0:
        basis_num_1 = int(sqrt(density_matrix.shape[0]))
        basis_num_2 = basis_num_1
    elif basis_num_1 == 0:
        basis_num_1 = int((density_matrix.shape[0])/basis_num_2)
    elif basis_num_2 == 0:
        basis_num_2 = int((density_matrix.shape[0])/basis_num_1)
    if basis_num_1 * basis_num_2 != int(density_matrix.shape[0]):
        raise ValueError
    for i in range(basis_num_1):
        for j in range(basis_num_1):
            density_matrix[basis_num_2*i:basis_num_2*(i+1),basis_num_2*j:basis_num_2*(j+1)] = density_matrix[basis_num_2*i:basis_num_2*(i+1),basis_num_2*j:basis_num_2*(j+1)].T
    return density_matrix

if __name__ == "__main__":
    rho = np.array([[10, 11, 12, 13, 14, 15],
                    [16, 17, 18, 19, 20, 21],
                    [22, 23, 24, 25, 26, 27],
                    [28, 29, 30, 31, 32, 33],
                    [34, 35, 36, 37, 38, 39],
                    [40, 41, 42, 43, 44, 45]])
    rho_pt = partial_transposition(rho, 3, 2)
    print(rho_pt)