pinn_bfs.py

import jax
import jax.numpy as jnp
from jax import grad, vmap, jit
import optax
import numpy as np
import matplotlib.pyplot as plt

plt.ion()  # Enable interactive mode for immediate plot display

# -------------------------------------------------------------
# Geometry & Hyperparameters
# -------------------------------------------------------------
Re = 2000.0  # Turbulent flow regime (Re > 1000 for BFS)
nu = 1.0 / Re
layers = [2, 128, 128, 128, 128, 3]  # Deeper network for turbulent flow

# -------------------------------------------------------------
# MLP Model
# -------------------------------------------------------------
def init_mlp(layers, key):
    params = []
    keys = jax.random.split(key, len(layers))
    for m, n, k in zip(layers[:-1], layers[1:], keys):
        W = jax.random.normal(k, (n, m)) * jnp.sqrt(2.0/m)
        b = jnp.zeros((n,))
        params.append((W, b))
    return params

def mlp(params, x):
    for W, b in params[:-1]:
        x = jnp.tanh(W @ x + b)
    W, b = params[-1]
    return W @ x + b

def model(params, X):
    return vmap(lambda x: mlp(params, x))(X)

# -------------------------------------------------------------
# PDE Residuals
# -------------------------------------------------------------
def pde_residual(params, xy):
    def f(xy):
        u, v, p = mlp(params, xy)

        grads = grad(lambda r: mlp(params, r)[0])(xy)
        u_x, u_y = grads[0], grads[1]

        grads_v = grad(lambda r: mlp(params, r)[1])(xy)
        v_x, v_y = grads_v[0], grads_v[1]

        grads_p = grad(lambda r: mlp(params, r)[2])(xy)
        p_x, p_y = grads_p[0], grads_p[1]

        lap_u = jnp.trace(jax.hessian(lambda r: mlp(params, r)[0])(xy))
        lap_v = jnp.trace(jax.hessian(lambda r: mlp(params, r)[1])(xy))

        r_cont = u_x + v_y
        r_u = u*u_x + v*u_y + p_x - nu*lap_u
        r_v = u*v_x + v*v_y + p_y - nu*lap_v

        return r_cont, r_u, r_v

    return vmap(f)(xy)

# -------------------------------------------------------------
# Boundary Conditions
# -------------------------------------------------------------
def inlet_bc(xy):
    x, y = xy
    return jnp.array([1.0*(y >= 0), 0.0, 0.0])

def loss_fn(params, Xf, Xi, Xw, Xo):
    r_cont, r_u, r_v = pde_residual(params, Xf)
    loss_pde = jnp.mean(r_cont**2) + jnp.mean(r_u**2) + jnp.mean(r_v**2)

    u_in = model(params, Xi)
    loss_in = jnp.mean((u_in - vmap(inlet_bc)(Xi))**2)

    u_w = model(params, Xw)
    loss_wall = jnp.mean((u_w[:, :2])**2)

    p_out = model(params, Xo)[:, 2]
    loss_out = jnp.mean((p_out - 0.0)**2)

    return loss_pde + loss_in + loss_wall + loss_out

# -------------------------------------------------------------
# Training Setup
# -------------------------------------------------------------
key = jax.random.PRNGKey(42)
params = init_mlp(layers, key)

optimizer = optax.adam(5e-4)  # Slightly lower learning rate for turbulent flow
opt_state = optimizer.init(params)

@jit
def train_step(params, opt_state, Xf, Xi, Xw, Xo):
    loss, grads = jax.value_and_grad(loss_fn)(params, Xf, Xi, Xw, Xo)
    updates, opt_state = optimizer.update(grads, opt_state)
    params = optax.apply_updates(params, updates)
    return params, opt_state, loss

# -------------------------------------------------------------
# Generate Collocation & Boundary Points
# -------------------------------------------------------------
Nf = 8000  # More collocation points for turbulent flow
Xf = jax.random.uniform(key, (Nf, 2),
                        minval=jnp.array([0., -1.]),
                        maxval=jnp.array([4., 1.]))

key, *subkeys = jax.random.split(key, 5)
Xi = jnp.stack([jnp.zeros(400),  # More boundary points for turbulent flow
                jax.random.uniform(subkeys[0], (400,), minval=0., maxval=1.)], axis=1)

Xw1 = jnp.stack([jax.random.uniform(subkeys[1], (500,), minval=0., maxval=4.),
                 -jnp.ones(500)], axis=1)
Xw2 = jnp.stack([jax.random.uniform(subkeys[2], (500,), minval=1., maxval=4.),
                 jnp.ones(500)], axis=1)
Xw = jnp.concatenate([Xw1, Xw2], axis=0)

Xo = jnp.stack([jnp.ones(400)*4,
                jax.random.uniform(subkeys[3], (400,), minval=-1., maxval=1.)], axis=1)

# -------------------------------------------------------------
# Train the Model
# -------------------------------------------------------------
for epoch in range(5000):  # More epochs for turbulent flow
    params, opt_state, loss = train_step(params, opt_state, Xf, Xi, Xw, Xo)
    if epoch % 200 == 0:
        print(f"Epoch {epoch}, Loss = {loss:.6f}")

# -------------------------------------------------------------
# Prediction & Visualization
# -------------------------------------------------------------
nx, ny = 200, 100
x_vals = np.linspace(0, 4, nx)
y_vals = np.linspace(-1, 1, ny)
Xg = jnp.array([[x, y] for x in x_vals for y in y_vals])

pred = model(params, Xg)
u_pred = np.array(pred[:,0]).reshape(nx, ny)
v_pred = np.array(pred[:,1]).reshape(nx, ny)
p_pred = np.array(pred[:,2]).reshape(nx, ny)
vel_mag = np.sqrt(u_pred**2 + v_pred**2)  # Velocity magnitude

X, Y = np.meshgrid(x_vals, y_vals, indexing="xy")

# Velocity magnitude plot
plt.figure(figsize=(10,4))
plt.contourf(X, Y, vel_mag.T, 50, cmap='viridis')
plt.colorbar(label="Velocity Magnitude")
plt.title(f"Velocity Magnitude (PINN - BFS, Re={Re:.0f}, Turbulent)")
plt.xlabel("x"); plt.ylabel("y")
plt.show(block=False)

# Velocity streamlines plot
plt.figure(figsize=(10,4))
plt.streamplot(X, Y, u_pred.T, v_pred.T, density=1.3)
plt.title(f"Velocity Streamlines (PINN - BFS, Re={Re:.0f}, Turbulent)")
plt.xlabel("x"); plt.ylabel("y")
plt.show(block=False)

# Pressure contour plot
plt.figure(figsize=(10,4))
plt.contourf(X, Y, p_pred.T, 50, cmap='jet')
plt.colorbar(label="Pressure")
plt.title(f"Pressure Contour (PINN - BFS, Re={Re:.0f}, Turbulent)")
plt.xlabel("x"); plt.ylabel("y")
plt.show(block=False)

input("Press Enter to close all plots...")  # Keep plots open until user presses Enter
plt.close('all')