DeepXDE and RANS

[gsof1]

Is there a way to introduce RANS with DeepXDE in order to solve simple turbulent flows? An example would be a 2D turbulent channel flow.

[LordPontiag]

Is there a way to introduce RANS with DeepXDE in order to solve simple turbulent flows? An example would be a 2D turbulent channel flow.

Yes it is. Simply write your NS equations beside your turbulence modeling equations. Then for Continuity, and other transport equations create their loss functions inside the PDE defenition and you're done!

[gsof1]

Yes it is. Simply write your NS equations beside your turbulence modeling equations. Then for Continuity, and other transport equations create their loss functions inside the PDE defenition and you're done!

Ok thank you very much for your reply! I tried to do something close to this, but the custom loss function does not seem to work. When i deactivate the custom loss, the code runs but produces incorrect results. Any suggestion would be appreciated! The code is the following:

import deepxde as dde
import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
import h5py

# Load DNS data
data = h5py.File('data_domain.h5')
data_domain = np.array(data['data_domain'])

X_train = data_domain[:, 0:2]  # x, y coordinates
y_train = data_domain[:, 2:4]  # u, v velocities

u_dns = y_train[:, 0]  # u component
v_dns = y_train[:, 1]  # v component

# Constants for the k-omega model
alpha = 5 / 9
beta = 3 / 40
beta_star = 9 / 100
sigma_k = 0.5
sigma_omega = 0.5
nu = 1.0 / 5600  # Non-dimensional viscosity

# Define the PDE system for u, v, p, k, omega
def pde(xy, y):

    # Ensure inputs are TensorFlow tensors
    xy = tf.convert_to_tensor(xy, dtype=tf.float32)
    y = tf.convert_to_tensor(y, dtype=tf.float32)

    if y.shape[1] != 5:
        raise ValueError(f"Expected y to have shape (None, 5), but got {y.shape}")

    
    x, y_pos = xy[:, 0:1], xy[:, 1:2]
    u, v, p, k, omega = y[:, 0:1], y[:, 1:2], y[:, 2:3], y[:, 3:4], y[:, 4:5]

    # Compute turbulent viscosity (non-dimensional)
    nu_t = k / (omega + 1e-8)  # Avoid division by zero

    # Gradients for momentum equations
    du_dx = dde.grad.jacobian(y, xy, i=0, j=0)
    du_dy = dde.grad.jacobian(y, xy, i=0, j=1)
    dv_dx = dde.grad.jacobian(y, xy, i=1, j=0)
    dv_dy = dde.grad.jacobian(y, xy, i=1, j=1)
    dp_dx = dde.grad.jacobian(y, xy, i=2, j=0)
    dp_dy = dde.grad.jacobian(y, xy, i=2, j=1)

    d2u_dx2 = dde.grad.hessian(y, xy, component=0, i=0, j=0)
    d2u_dy2 = dde.grad.hessian(y, xy, component=0, i=1, j=1)
    d2v_dx2 = dde.grad.hessian(y, xy, component=1, i=0, j=0)
    d2v_dy2 = dde.grad.hessian(y, xy, component=1, i=1, j=1)

    # Production term for k (non-dimensional)
    P_k = nu_t * ((du_dx**2) + 2 * (du_dy * dv_dx) + (dv_dy**2))

    # Gradients for k and omega
    dk_dx = dde.grad.jacobian(y, xy, i=3, j=0)
    dk_dy = dde.grad.jacobian(y, xy, i=3, j=1)
    d2k_dx2 = dde.grad.hessian(y, xy, component=3, i=0, j=0)
    d2k_dy2 = dde.grad.hessian(y, xy, component=3, i=1, j=1)

    domega_dx = dde.grad.jacobian(y, xy, i=4, j=0)
    domega_dy = dde.grad.jacobian(y, xy, i=4, j=1)
    d2omega_dx2 = dde.grad.hessian(y, xy, component=4, i=0, j=0)
    d2omega_dy2 = dde.grad.hessian(y, xy, component=4, i=1, j=1)

    # Non-dimensional PDEs
    eq1 = du_dx + dv_dy  # Continuity equation
    eq2 = u * du_dx + v * du_dy + dp_dx - (nu + nu_t) * (d2u_dx2 + d2u_dy2)  # x-momentum
    eq3 = u * dv_dx + v * dv_dy + dp_dy - (nu + nu_t) * (d2v_dx2 + d2v_dy2)  # y-momentum
    eq4 = -dk_dx - dk_dy + (d2k_dx2 + d2k_dy2) * (nu + nu_t / sigma_k) + P_k - beta_star * k * omega  # k-equation
    eq5 = -domega_dx - domega_dy + (d2omega_dx2 + d2omega_dy2) * (nu + nu_t / sigma_omega) + alpha * (omega / k) * P_k - beta * omega**2  # omega-equation

    return [eq1, eq2, eq3, eq4, eq5]

# Define the 2D geometry
x_min, x_max = 0, 4.0 * np.pi  # Streamwise length
y_min, y_max = -1.0, 1.0  # Channel height
geom = dde.geometry.Rectangle([x_min, y_min], [x_max, y_max])

# No-slip walls (u = v = 0 at y = ±1)
def boundary_wall(x, on_boundary):
    return on_boundary and (np.isclose(x[1], y_min) or np.isclose(x[1], y_max))

# Apply no-slip BCs to u and v
bc_u = dde.DirichletBC(geom, lambda x: 0, boundary_wall, component=0)  # u = 0 at y = ±1
bc_v = dde.DirichletBC(geom, lambda x: 0, boundary_wall, component=1)  # v = 0 at y = ±1

# Periodic BCs in x-direction
def boundary_periodic(x, on_boundary):
    return on_boundary and (np.isclose(x[0], x_min) or np.isclose(x[0], x_max))

bc_periodic_u = dde.PeriodicBC(geom, 0, boundary_periodic, component=0)  # Enforce periodicity on u
bc_periodic_v = dde.PeriodicBC(geom, 0, boundary_periodic, component=1)  # Enforce periodicity on v
bc_periodic_p = dde.PeriodicBC(geom, 0, boundary_periodic, component=2)  # Enforce periodicity on p

# PINN data setup
data = dde.data.PDE(
    geom,
    pde,
    [bc_u, bc_v, bc_periodic_u, bc_periodic_v, bc_periodic_p],
    num_domain=5000,  # Number of collocation points in the domain
    num_boundary=500,  # Number of boundary points
)

# Define the neural network
layer_size = [2] + [80] * 5 + [5]  # 2D input, 5 outputs (u, v, p, k, omega)
activation = "tanh"
initializer = "Glorot uniform"
net = dde.nn.FNN(layer_size, activation, initializer)

# Model setup
model = dde.Model(data, net)

X_train_tf = tf.convert_to_tensor(X_train, dtype=tf.float32)


# Define PDE residual loss functions
def loss_pde(y_true, y_pred, component):
    """Compute PDE residuals separately for each component"""
    pde_losses = pde(X_train_tf, y_pred)  # Compute PDE residuals
    return tf.reduce_mean(tf.square(pde_losses[component]))  # Squared residual loss

# Define supervised losses for DNS data (u, v)
def loss_dns_u(y_true, y_pred):
    return tf.reduce_mean(tf.square(u_dns - y_pred[:, 0]))

def loss_dns_v(y_true, y_pred):
    return tf.reduce_mean(tf.square(v_dns - y_pred[:, 1]))

# Boundary condition losses (zero-enforcement for walls and periodicity)
def loss_bc(y_true, y_pred):
    return tf.reduce_mean(tf.square(y_pred - y_true))  # Generic BC loss

# PDE residuals (5 losses)
pde_losses = [
    lambda y_true, y_pred: loss_pde(y_true, y_pred, 0),  # Continuity
    lambda y_true, y_pred: loss_pde(y_true, y_pred, 1),  # x-momentum
    lambda y_true, y_pred: loss_pde(y_true, y_pred, 2),  # y-momentum
    lambda y_true, y_pred: loss_pde(y_true, y_pred, 3),  # k-equation
    lambda y_true, y_pred: loss_pde(y_true, y_pred, 4),  # omega-equation
]

# BC losses (5 losses corresponding to `bc_u`, `bc_v`, `bc_periodic_u`, `bc_periodic_v`, `bc_periodic_p`)
bc_losses = [
    loss_bc,  # Enforce no-slip u=0
    loss_bc,  # Enforce no-slip v=0
    loss_bc,  # Periodic u
    loss_bc,  # Periodic v
    loss_bc,  # Periodic p
]

# Combine losses (Total: 10)
total_losses = pde_losses + bc_losses


# Combine DNS data for u and v
#y_train_dns = np.hstack([u_dns.reshape(-1, 1), v_dns.reshape(-1, 1)])

# Compile the model with custom loss
model.compile(
    "adam",
    lr=0.001,
    loss=total_losses,  # Use the custom loss function
)