navier stokes with python/numpy

import numpy as np
import matplotlib.pyplot as plt

# Constants
nx = 41  # Number of grid points in x-direction
ny = 41  # Number of grid points in y-direction
nt = 500  # Number of time steps
dt = 0.001  # Time step size
nu = 0.1  # Viscosity
rho = 1.0  # Density

dx = 2 / (nx - 1)
dy = 2 / (ny - 1)

# Initialization
u = np.zeros((ny, nx))
v = np.zeros((ny, nx))
p = np.zeros((ny, nx))
b = np.zeros((ny, nx))

# Function to solve the pressure Poisson equation
def pressure_poisson(p, dx, dy, b):
    for q in range(nt):
        pn = p.copy()
        p[1:-1, 1:-1] = (((pn[1:-1, 2:] + pn[1:-1, :-2]) * dy**2 +
                          (pn[2:, 1:-1] + pn[:-2, 1:-1]) * dx**2) /
                          (2 * (dx**2 + dy**2)) -
                         dx**2 * dy**2 / (2 * (dx**2 + dy**2)) * b[1:-1, 1:-1])

        # Boundary conditions
        p[0, :] = p[1, :]  # p = 0 @ y = 0
        p[ny-1, :] = p[ny-2, :]  # p = 0 @ y = 2
        p[:, 0] = p[:, 1]  # p = 0 @ x = 0
        p[:, nx-1] = 0  # p = 0 @ x = 2

# Main simulation loop
for n in range(nt):
    un = u.copy()
    vn = v.copy()

    b[1:-1, 1:-1] = (rho * (1 / dt *
                           ((u[1:-1, 2:] - u[1:-1, 0:-2]) / (2 * dx) +
                            (v[2:, 1:-1] - v[0:-2, 1:-1]) / (2 * dy)) -
                           ((u[1:-1, 2:] - u[1:-1, 0:-2]) / (2 * dx))**2 -
                           2 * ((u[2:, 1:-1] - u[0:-2, 1:-1]) / (2 * dy) *
                                (v[1:-1, 2:] - v[1:-1, 0:-2]) / (2 * dx)) -
                           ((v[2:, 1:-1] - v[0:-2, 1:-1]) / (2 * dy))**2))

    pressure_poisson(p, dx, dy, b)

    u[1:-1, 1:-1] = (un[1:-1, 1:-1] -
                     un[1:-1, 1:-1] * dt / dx *
                     (un[1:-1, 1:-1] - un[1:-1, 0:-2]) -
                     vn[1:-1, 1:-1] * dt / dy *
                     (un[1:-1, 1:-1] - un[0:-2, 1:-1]) -
                     dt / (2 * rho * dx) * (p[1:-1, 2:] - p[1:-1, 0:-2]) +
                     nu * (dt / dx**2 *
                           (un[1:-1, 2:] - 2 * un[1:-1, 1:-1] + un[1:-1, 0:-2]) +
                           dt / dy**2 *
                           (un[2:, 1:-1] - 2 * un[1:-1, 1:-1] + un[0:-2, 1:-1])))

    v[1:-1, 1:-1] = (vn[1:-1, 1:-1] -
                     un[1:-1, 1:-1] * dt / dx *
                     (vn[1:-1, 1:-1] - vn[1:-1, 0:-2]) -
                     vn[1:-1, 1:-1] * dt / dy *
                     (vn[1:-1, 1:-1] - vn[0:-2, 1:-1]) -
                     dt / (2 * rho * dy) * (p[2:, 1:-1] - p[0:-2, 1:-1]) +
                     nu * (dt / dx**2 *
                           (vn[1:-1, 2:] - 2 * vn[1:-1, 1:-1] + vn[1:-1, 0:-2]) +
                           dt / dy**2 *
                           (vn[2:, 1:-1] - 2 * vn[1:-1, 1:-1] + vn[0:-2, 1:-1])))

    # Boundary conditions
    u[0, :] = 0
    u[ny-1, :] = 0
    u[:, 0] = 0
    u[:, nx-1] = 1  # u = 1 @ x = 2

    v[0, :] = 0
    v[ny-1, :] = 0
    v[:, 0] = 0
    v[:, nx-1] = 0

# Plotting the results
x = np.linspace(0, 2, nx)
y = np.linspace(0, 2, ny)
X, Y = np.meshgrid(x, y)

fig = plt.figure(figsize=(11, 7), dpi=100)
plt.contourf(X, Y, p, alpha=0.5)
plt.colorbar()
plt.contour(X, Y, p)
plt.quiver(X[::2, ::2], Y[::2, ::2], u[::2, ::2], v[::2, ::2])
plt.xlabel('X')
plt.ylabel('Y')
plt.show()
preveen-stack / nodejs

navier stokes with python/numpy #17
