plot_basis.py

import numpy as np
from matplotlib import pyplot as plt

N = 3

def get_basis(idx):
    x = (idx // (N - 1)) + 1
    y = (idx % (N - 1)) + 1

    sin = np.sin(2 * np.pi * x * np.arange(N**2) / N**2)
    cos = np.cos(2 * np.pi * y * np.arange(N**2) / N**2)

    sin = sin / np.linalg.norm(sin)
    cos = cos / np.linalg.norm(cos)

    sin = np.reshape(sin, (N**2, 1))
    cos = np.reshape(cos, (1, N**2))

    return sin @ cos

plt.subplot(2, 2, 1)
plt.imshow(get_basis(0), cmap="gray", interpolation="nearest")

plt.subplot(2, 2, 2)
plt.imshow(get_basis(1), cmap="gray", interpolation="nearest")

plt.subplot(2, 2, 3)
plt.imshow(get_basis(2), cmap="gray", interpolation="nearest")

plt.subplot(2, 2, 4)
plt.imshow(get_basis(3), cmap="gray", interpolation="nearest")


plt.show()