Balanced_Truncation.py

import numpy as np
import matplotlib.pyplot as plt
from scipy.linalg import solve_continuous_lyapunov as clyap
from scipy.linalg import svd

# --> Utility function.
vec2array = lambda x, ny, nz : x.reshape(ny, nz)
array2vec = lambda x : x.flatten()

# --> Differential operators.
def laplacian_1D(nx, dx, periodic=False):
    # --> Diagonals of the finite-difference approximation.
    d, u = nx*[-2/dx**2], (nx-1)*[1/dx**2]

    # --> Laplacian matrix.
    L = np.diag(d) + np.diag(u, k=1) + np.diag(u, k=-1)

    # --> Periodic boundary conditions?
    if periodic:
        L[0, -1] = 1/dx**2
        L[-1, 0] = 1/dx**2

    return L

def derivative(nx, dx, periodic=False):
    # --> Diagonals of the finite-difference approximation.
    u = (nx-1)*[1/(2*dx)]
    l = (nx-1)*[-1/(2*dx)]

    # --> Differentiation matrix.
    D = np.diag(u, k=1) + np.diag(l, k=-1)

    # --> Periodic boundary conditions?
    if periodic:
        D[0, -1] = -1/(2*dx)
        D[-1, 0] = 1/(2*dx)
    return D

if __name__ == "__main__":
    # --> Parameters.
    ny, nz = 31, 32
    Ly, Lz = 2.0, 2*np.pi
    dy, dz = Ly/(ny+1), Lz/nz
    Re = 100.0

    # --> Differential operators.
    Dz = derivative(nz, dz, periodic=True)
    D2y, D2z = laplacian_1D(ny, dy), laplacian_1D(nz, dz, periodic=True)
    Iy, Iz = np.eye(ny), np.eye(nz)
    D2 = np.kron(D2y, Iz) + np.kron(Iy, D2z)

    # --> Mesh.
    y = np.linspace(-Ly/2, Ly/2, ny+2)[1:-1]
    z = np.linspace(0, Lz, nz+1)[:-1]

    # --> Baseflow velocity profile.
    u = 1 - y**2
    du = -2*y

    # --> Orr-Sommerfeld Squire operators.
    OS = D2/Re
    S  = D2/Re
    C = -np.kron(np.diag(du), Dz)

    # --> State-space model.
    A = np.block([[OS, np.zeros_like(OS)], [C, S]])
    Bv = np.exp(-y**2/0.05).reshape(-1, 1) @ np.exp(-(z-Lz/2)**2/0.05).reshape(1, -1)
    Beta = np.zeros_like(Bv)
    B = np.r_[Bv.flatten(), Beta.flatten()].reshape(-1, 1)
    C = np.eye(len(A))

    # --> Lyapunov equations.
    P = clyap(A, -B @ B.T) ; Q = clyap(A.T, -C.T @ C)

    # --> SVD of Gramian.
    Up, Sp, _ = np.linalg.svd(P)
    Uq, Sq, _ = np.linalg.svd(Q)

    # --> Matrix square-roots.
    Psqrt = Up @ np.diag(np.sqrt(Sp)) #NOTE : In practice we would have a @ Up.T but it ain't needed for the rest.
    Qsqrt = Uq @ np.diag(np.sqrt(Sq))

    # --> Balancing transformation.
    W, S, V = svd(Psqrt.T @ Qsqrt, full_matrices=False) ; V = V.T
    T = np.diag(np.sqrt(1.0/S)) @ V.T @ Qsqrt.T
    Tinv = Psqrt @ W @ np.diag(np.sqrt(1.0/S))

    # --> Plot figures.
    fig, ax = plt.subplots(1, 1, figsize=(4, 2))
    ax.semilogy(S)
    ax.set(xlim=(0, 20), xlabel="Rank")
    ax.set(ylabel=r"$\sigma_i$")

    for i in range(3):
        fig, ax = plt.subplots(2, 2, figsize=(8, 4))

        #####
        #####     BALANCED MODES / BPOD MODES
        #####

        ax[0, 0].pcolormesh(
            z, y, data := vec2array(Tinv[:ny*nz, i], ny, nz),
            vmin=-abs(data).max(), vmax=abs(data).max(),
            shading="gouraud", cmap="RdBu")

        ax[0, 0].set_aspect("equal")

        ax[1, 0].pcolormesh(
            z, y, data := vec2array(Tinv[ny*nz:, i], ny, nz),
            vmin=-abs(data).max(), vmax=abs(data).max(),
            shading="gouraud", cmap="RdBu")

        ax[1, 0].set_aspect("equal")

        #####
        #####     ADJOINT BALANCED MODES / ADJOINT BPOD MODES
        #####

        ax[0, 1].pcolormesh(
            z, y, data := vec2array(T.T[:ny*nz, i], ny, nz),
            vmin=-abs(data).max(), vmax=abs(data).max(),
            shading="gouraud", cmap="RdBu")

        ax[0, 1].set_aspect("equal")

        ax[1, 1].pcolormesh(
            z, y, data := vec2array(T.T[ny*nz:, i], ny, nz),
            vmin=-abs(data).max(), vmax=abs(data).max(),
            shading="gouraud", cmap="RdBu")

        ax[1, 1].set_aspect("equal")

    plt.show()

Yes, I got mine to work as well. The problem was the "order = 'F'" in my reshape which put the forcing at the wrong place. Thanks for sharing!