blackboard.py

import numpy as np

# Lattice vectors
ea = np.array([[0, 0], [0, 1], [1, 0], [0, -1], [-1, 0], [1, 1], [1, -1], [-1, -1], [-1, 1]])
# Equilibrium distribution function weights
wa = np.array([4/9, 1/9, 1/9, 1/9, 1/9, 1/36, 1/36, 1/36, 1/36])

def calc_equilibrium_distribution(rho, u):
    ux, uy = u
    feq = np.empty(9, dtype=np.float64)

    u2 = ux*ux + uy*uy

    for i in range(9):
        ex = ea[i, 0]
        ey = ea[i, 1]

        eu = ex*ux + ey*uy
        feq[i] = wa[i] * rho * (1.0 + 3.0*eu + 4.5*eu*eu - 1.5*u2)

    return feq


def calc_equilibrium_distribution2(rho, u):
    f_eq = wa * rho * (1 + 3 * np.dot(ea, u) + 4.5 * (np.dot(ea, u) ** 2) - 1.5 * np.dot(u, u))
    return f_eq

print(calc_equilibrium_distribution(1.0, (0.1, 0.05)))
print(calc_equilibrium_distribution2(1.0, (0.1, 0.05)))