Network innerconnect

[Asachoo]

This PR aims to improve the performance of innerconnect_s method during Network connection. And solved the issue that the circuit calculation slows down when using auto_reduce=True for small-scale circuits in #1069. Related to #574 #1069

This PR is to achieve acceleration by avoiding repeated matrix calculations and caching intermediate calculation results. In my benchmark, the innerconnect_s method has a 5~6 times performance improvement in MKL/OpenBLAS backend.

Your feedback and suggestions on this optimization are welcome!

Benchmark Code

import numpy as np
from time import time, sleep
from matplotlib.pyplot import plot, show, legend, scatter

def innerconnect_s(A: np.ndarray, k: int, l: int) -> np.ndarray:
    """
    Performance optimized innerconnect_s function.
    """

    if k > A.shape[-1] - 1 or l > A.shape[-1] - 1:
        raise (ValueError("port indices are out of range"))

    nA = A.shape[1]  # num of ports on input s-matrix

    # external ports index
    ext_i = [i for i in range(nA) if i not in (k, l)]

    # Indexing sub-matrices of internal ports (only k, l)
    Akl = 1.0 - A[:, k, l]
    Alk = 1.0 - A[:, l, k]
    Akk = A[:, k, k]
    All = A[:, l, l]

    # Indexing sub-matrices of other external ports
    Ake = A[:, k, ext_i].T
    Ale = A[:, l, ext_i].T
    Aek = A[:, ext_i, k].T
    Ael = A[:, ext_i, l].T

    # Create an suit-sized s-matrix, to store the result
    x, y = np.meshgrid(ext_i, ext_i)
    C = A[:, y, x]

    # create temporary matrices for calculation
    det = (Akl * Alk - Akk * All)
    tmp_a = (Ael * Alk + Aek * All) / det
    tmp_b = (Ael * Akk + Aek * Akl) / det

    # loop through ports and calculates resultant s-parameters
    for i in range(nA - 2):
        for j in range(nA - 2):
            C[:, i, j] += Ake[j] * tmp_a[i] + Ale[j] * tmp_b[i]

    return C

def stand_rst(A: np.ndarray, k: int, l: int) -> np.ndarray:
    """
    Innerconnect_s function in master.
    """
    if k > A.shape[-1] - 1 or l > A.shape[-1] - 1:
        raise (ValueError("port indices are out of range"))

    nA = A.shape[1]  # num of ports on input s-matrix
    # create an empty s-matrix, to store the result
    C = np.zeros(shape=A.shape, dtype="complex")

    # loop through ports and calculates resultant s-parameters
    for i in range(nA):
        for j in range(nA):
            C[:, i, j] = A[:, i, j] + (
                A[:, k, j] * A[:, i, l] * (1 - A[:, l, k])
                + A[:, l, j] * A[:, i, k] * (1 - A[:, k, l])
                + A[:, k, j] * A[:, l, l] * A[:, i, k]
                + A[:, l, j] * A[:, k, k] * A[:, i, l]
            ) / ((1 - A[:, k, l]) * (1 - A[:, l, k]) - A[:, k, k] * A[:, l, l])

    # remove ports that were `connected`
    C = np.delete(C, (k, l), 1)
    C = np.delete(C, (k, l), 2)

    return C


if __name__ == "__main__":
    nfreq = 1001
    fcalls = 200
    a_line = []
    b_line = []

    nports = range(3, 30)
    for nport in nports:
        np.random.seed(123)
        A = np.random.rand(nfreq, nport, nport) + 1j * np.random.rand(
            nfreq, nport, nport
        )
        a, b, r = [], [], []

        for i in range(fcalls):
            k, l = np.random.choice(nport, 2, replace=False)
            a_s = time()
            a_rst = innerconnect_s(A, k, l)
            a_e = time()

            b_s = time()
            b_rst = stand_rst(A, k, l)
            b_e = time()

            a.append((a_e - a_s) * 1000)
            b.append((b_e - b_s) * 1000)
            r.append(np.allclose(a_rst, b_rst))

        print(
            f"innerconnect_s  [nports: {nport}] {np.mean(a):.4f} ms, min: {np.min(a):.4f}, std: {np.std(a):.4f} ms"
        )
        print(
            f"standard method [nports: {nport}] {np.mean(b):.4f} ms, min: {np.min(b):.4f}, std: {np.std(b):.4f} ms"
        )
        print(all(r))
        print()

        a_line.append(np.min(a))
        b_line.append(np.min(b))

        sleep(1)

    scatter(a_line, b_line, label="relationship")
    p = np.polyfit(a_line, b_line, 1)
    f = np.poly1d(p)
    plot(a_line, f(a_line), label=f"linear fit slope: {p[0]:.4f}")
    legend()
    show()

    print(f"linear fit slope: {p[0]:.4f}")

Benchmark in MKL Backend:

Benchmark in OpenBLAS Backend:

[jhillairet]

LGTM, thank you!

I would keep the previous code in comment for pedagogy, as it's generally simpler to understand how a for-loop works than directly its vectorized version.

[FranzForstmayr]

LGTM

[jhillairet]

ok to be merged?

[Asachoo]

Yes, I consider it finished.