Как реализовать векторизацию по циклам блочного умножения с помощью тензорного умножения в Numpy?

Question

У меня есть код, который запускает sample_size последовательности умножения матриц, и для каждой последовательности используется seq_length операций сумм умножения матриц. Но недостатком моего кода является то, что как только seq_length становится выше, чем около 300, алгоритм замедляется, и нет необходимости говорить, что с ростом seq_length весь алгоритм становится все медленнее и медленнее. Поэтому мне было интересно, есть ли какая-либо оптимизация / векторизация, которая может быть реализована так, как я написал свой код, или просто на моем коде в целом.

По сути, здесь я просто определяю набор (2x2) матриц со сложными записями, которые будут использоваться в моем алгоритме позже. Функция cliff_operators() выберет случайную матрицу из всех.

import random
import numpy as np
import matplotlib.pyplot as plt
import time
import sys

init_state = np.array([[1, 0], [0, 0]], dtype=complex)
II = np.identity(2, dtype=complex)
X = np.array([[0, 1], [1, 0]], dtype=complex)
Y = np.array([[0, -1j], [1j, 0]], dtype=complex)
Z = np.array([[1, 0], [0, -1]], dtype=complex)
PPP = (-II + 1j*X + 1j*Y + 1j*Z)/2
PPM = (-II + 1j*X + 1j*Y - 1j*Z)/2
PMM = (-II + 1j*X - 1j*Y - 1j*Z)/2
MMM = (-II - 1j*X - 1j*Y - 1j*Z)/2
MMP = (-II - 1j*X - 1j*Y + 1j*Z)/2
MPP = (-II - 1j*X + 1j*Y + 1j*Z)/2
PMP = (-II + 1j*X - 1j*Y + 1j*Z)/2
MPM = (-II - 1j*X + 1j*Y - 1j*Z)/2

def cliff_operators():
    return random.choice([II, X, Y, Z, PPP, PPM, PMM, MMM, MMP, MPP, PMP, MPM])

Здесь функция compute_channel_operation производит поэлементное матричное произведение точек во входной блочной матрице / тензоре, а затем выполняет поэлементное суммирование всех матриц внутри тензора.

def compute_channel_operation(rho, operators):
    return np.sum(operators@rho@operators.transpose(0, 2, 1).conj(), axis=0)

def depolarizing_error(param):
    XYZ = np.sqrt(param/3)*np.array([X, Y, Z])
    return np.array([np.sqrt(1-param)*II, XYZ[0], XYZ[1], XYZ[2]])

def random_angles(sd):
    return np.random.normal(0, sd, 3)

def unitary_error(params):
    e_1 = np.exp(-1j*(params[0]+params[2])/2)*np.cos(params[1]/2)
    e_2 = np.exp(-1j*(params[0]-params[2])/2)*np.sin(params[1]/2)
    return np.array([[[e_1, e_2], [-e_2.conj(), e_1.conj()]]])


def rb(input_state, seq_length, sample_size, noise_mean,
                            noise_sd, noise2_sd):
    fidelity = []
    for i in range(1, sample_size+1, 1):
        rho = input_state
        sequence = []
        for j in range(1, seq_length+1, 1):
            noise = depolarizing_error(np.random.normal(noise_mean, noise_sd))
            noise_2 = unitary_error(random_angles(noise2_sd))
            unitary = cliff_operators()
            sequence.append(unitary)
            i_ideal_operator = compute_channel_operation(rho,
                                                         np.array([unitary]))

            i_noisy_operator = compute_channel_operation(i_ideal_operator,
                                                         noise)
            i_noisy_operator_2 = compute_channel_operation(i_noisy_operator,
                                                           noise_2)
            sys.stdout.write("\r" + "gate applied: " + str(j))
            rho = i_noisy_operator_2

        # Final random noise
        noise = depolarizing_error(np.random.normal(noise_mean, noise_sd))
        noise_2 = unitary_error(random_angles(noise2_sd))

        # Computes the Hermitian of the forward operators sequence
        unitary_plus_1 = np.linalg.multi_dot(sequence[::-1]).conj().T

        # Final ideal&noisy density operator
        f_ideal_operator = compute_channel_operation(rho,
                                                     np.array([unitary_plus_1]))

        f_noisy_operator = compute_channel_operation(f_ideal_operator, noise)
        f_noisy_operator_2 = compute_channel_operation(f_noisy_operator,
                                                       noise_2)
        fidelity.append(np.trace(input_state@f_noisy_operator_2))

    avg_fidelity = (1/sample_size)*np.sum(fidelity)
    return avg_fidelity


def get_data(rho, seq_length, sample_size, noise_mean, noise_sd, noise2_sd):
    length = []
    fidelity_s = []
    for s in range(2, seq_length, 1):
        avg_fidelity = rb(rho, s, sample_size, noise_mean,
                                               noise_sd, noise2_sd)
        length.append(s)
        fidelity_s.append(avg_fidelity)

    plt.plot(length, fidelity_s)
    plt.title("Fidelity vs Clifford length")
    plt.ylim(0.5, 1)
    plt.xlabel("Clifford length")
    plt.ylabel("Fidelity")
    plt.xlim(0, 100)
    plt.show()

starttime = time.time()
get_data(init_state, 402, 1, 0.005, 0.001, 0.01)
timeElapsed = time.time() - starttime
print(timeElapsed)

Так может ли быть потенциально реализована векторизация для удаления цикла i и j и ускорения его выполнения по мере увеличения seq_length? Можно ли векторизовать циклы над sample_size, чтобы одновременно выполнять n последовательностей в одной большой матрице?