Three-dimensional unsteady Navier-Stokes Equations

Three-dimensional unsteady Navier-Stokes Equations#

Problem Statement#

1. Momentum Equations#

The Navier-Stokes equations describe the conservation of momentum in fluid dynamics. For an incompressible fluid, the equation is:

\[ \frac{\partial \mathbf{u}}{\partial t} + \mathbf{u} \cdot \nabla \mathbf{u} = - \nabla p + \frac{1}{Re} \nabla^2 \mathbf{u} \]

Where:

\(\mathbf{u} = (u, v, w)\) is the velocity field,
\(p\) is the pressure field,
\(\nabla\) is the gradient operator,
\(\nabla^2\) is the Laplacian operator,
\(\mu\) is the dynamic viscosity.

The momentum equations in the code are written for each of the three spatial components (x, y, z). Specifically, the equations correspond to the following:

\(x\)-direction momentum equation (momentum_x in the code):

\[ \rho\left[\frac{\partial u}{\partial t}+\frac{\partial u}{\partial x} u+\frac{\partial u}{\partial y} v+\frac{\partial u}{\partial z} w\right]=-\frac{\partial p}{\partial x}+\mu\left(\frac{\partial^2 u}{\partial x^2}+\frac{\partial^2 u}{\partial y^2}+\frac{\partial^2 u}{\partial z^2}\right)+\rho g_x \]

\(y\)-direction momentum equation (momentum_y in the code):

\[ \rho\left[\frac{\partial v}{\partial t}+\frac{\partial v}{\partial x} u+\frac{\partial v}{\partial y} v+\frac{\partial v}{\partial z} w\right]=-\frac{\partial p}{\partial y}+\mu\left(\frac{\partial^2 v}{\partial x^2}+\frac{\partial^2 v}{\partial y^2}+\frac{\partial^2 v}{\partial z^2}\right)+\rho g_y \]

\(z\)-direction momentum equation (momentum_z in the code):

\[ \rho\left[\frac{\partial w}{\partial t}+\frac{\partial w}{\partial x} u+\frac{\partial w}{\partial y} v+\frac{\partial w}{\partial z} w\right]=-\frac{\partial p}{\partial z}+\mu\left(\frac{\partial^2 w}{\partial x^2}+\frac{\partial^2 w}{\partial y^2}+\frac{\partial^2 w}{\partial z^2}\right)+\rho g_z \]

2. Continuity Equation#

The continuity equation represents the conservation of mass, ensuring that the flow is incompressible (i.e., the divergence of the velocity field is zero). The equation is:

\[ \nabla \cdot \mathbf{u} = 0 \]

In the code, the continuity equation corresponds to:

\[ \frac{\partial u}{\partial x} + \frac{\partial v}{\partial y} + \frac{\partial w}{\partial z} = 0 \]

This guarantees that the volume of the fluid remains constant and the flow is incompressible.

3. Initial and Boundary Conditions (IC and BC)#

The function icbc_cond_func defines the initial conditions (IC) and boundary conditions (BC) for the velocity and pressure fields.

Initial velocity fields: The velocity components are given as functions of spatial variables \(x\), \(y\), and \(z\), as well as time \(t\). Specifically, the velocity components \(u\), \(v\), and \(w\) are defined as:

\[ u = -a \left( e^{a x} \sin(a y + d z) + e^{a z} \cos(a x + d y) \right) e^{-d^2 t} \]

\[ v = -a \left( e^{a y} \sin(a z + d x) + e^{a x} \cos(a y + d z) \right) e^{-d^2 t} \]

\[ w = -a \left( e^{a z} \sin(a x + d y) + e^{a y} \cos(a z + d x) \right) e^{-d^2 t} \]

Initial pressure field: The pressure field \(p\) is given by a more complex expression, involving exponentials and trigonometric functions of the spatial variables \(x\), \(y\), and \(z\):

\[\begin{split} \begin{aligned} p(x, y, z, t)= & -\frac{1}{2} a^2\left[e^{2 a x}+e^{2 a y}+e^{2 a z}\right. \\ & +2 \sin (a x+d y) \cos (a z+d x) e^{a(y+z)} \\ & +2 \sin (a y+d z) \cos (a x+d y) e^{a(z+x)} \\ & \left.+2 \sin (a z+d x) \cos (a y+d z) e^{a(x+y)}\right] e^{-2 d^2 t} \end{aligned} \end{split}\]

4. Final Formulation: 3D Navier-Stokes Equations#

The Navier-Stokes equations, based on the code, are as follows:

Momentum Equation (in three dimensions):#

\[ \frac{\partial \mathbf{u}}{\partial t} + \mathbf{u} \cdot \nabla \mathbf{u} = - \nabla p + \frac{1}{Re} \nabla^2 \mathbf{u} \]

Where:

\(\mathbf{u} = (u, v, w)\) is the velocity field,
\(p\) is the pressure field,
\(Re\) is the Reynolds number,
\(\nabla^2\) is the Laplacian operator.

Continuity Equation (for incompressibility):#

\[ \nabla \cdot \mathbf{u} = 0 \]

This ensures that the flow remains incompressible.

Conclusion#

The Python code essentially implements the 3D Navier-Stokes equations for an incompressible fluid, where the momentum equations are resolved in each spatial direction, and the continuity equation ensures mass conservation. The initial conditions for velocity and pressure are specified, and boundary conditions are likely handled by the methods in icbc_cond_func.

Dimensional Analysis#

Summary of Physical Units

Velocity (\(u, v, w\)): \([u] = [v] = [w] = \text{m/s}\)
Time (\(t\)): \([t] = \text{s}\)
Pressure (\(p\)): \([p] = \text{kg/m} \cdot \text{s}^2\)
Reynolds number (\(Re\)): Dimensionless
Laplacian of velocity (\(\nabla^2 \mathbf{u}\)): \(\text{s}^{-2}\)
Density (\(\rho\)): \(\text{kg/m}^3\)
Dynamic viscosity (\(\mu\)): \(\text{kg/m} \cdot \text{s}\)

The analysis confirms that the Navier-Stokes equations, including the momentum and continuity equations, are dimensionally consistent and properly describe the physical quantities involved in fluid flow.

Code Implementation#

First, we import the necessary libraries for the implementation:

import braintools
import brainstate
import brainunit as u
import jax.tree
import numpy as np

import pinnx

Define the physical units for the problem:

unit_of_space = u.meter
unit_of_speed = u.meter / u.second
unit_of_t = u.second
unit_of_pressure = u.pascal

Define the spatial and temporal domains for the problem:

spatial_domain = pinnx.geometry.Cuboid(xmin=[-1, -1, -1], xmax=[1, 1, 1])
temporal_domain = pinnx.geometry.TimeDomain(0, 1)
spatio_temporal_domain = pinnx.geometry.GeometryXTime(spatial_domain, temporal_domain)
spatio_temporal_domain = spatio_temporal_domain.to_dict_point(
    x=unit_of_space,
    y=unit_of_space,
    z=unit_of_space,
    t=unit_of_t,
)

Define the neural network model for the problem:

net = pinnx.nn.Model(
    pinnx.nn.DictToArray(x=unit_of_space,
                         y=unit_of_space,
                         z=unit_of_space,
                         t=unit_of_t),
    pinnx.nn.FNN([4] + 4 * [50] + [4], "tanh", braintools.init.KaimingUniform()),
    pinnx.nn.ArrayToDict(u_vel=unit_of_speed,
                         v_vel=unit_of_speed,
                         w_vel=unit_of_speed,
                         p=unit_of_pressure),
)

Define the PDE residual function for the Navier-Stokes equations:

a = 1
d = 1
Re = 1
rho = 1 * u.kilogram / u.meter ** 3
mu = 1 * u.pascal * u.second


@brainstate.transform.jit
def pde(x, u):
    jacobian = net.jacobian(x)
    x_hessian = net.hessian(x, y=['u_vel', 'v_vel', 'w_vel'], xi=['x'], xj=['x'])
    y_hessian = net.hessian(x, y=['u_vel', 'v_vel', 'w_vel'], xi=['y'], xj=['y'])
    z_hessian = net.hessian(x, y=['u_vel', 'v_vel', 'w_vel'], xi=['z'], xj=['z'])

    u_vel, v_vel, w_vel, p = u['u_vel'], u['v_vel'], u['w_vel'], u['p']

    du_vel_dx = jacobian['u_vel']['x']
    du_vel_dy = jacobian['u_vel']['y']
    du_vel_dz = jacobian['u_vel']['z']
    du_vel_dt = jacobian['u_vel']['t']
    du_vel_dx_dx = x_hessian['u_vel']['x']['x']
    du_vel_dy_dy = y_hessian['u_vel']['y']['y']
    du_vel_dz_dz = z_hessian['u_vel']['z']['z']

    dv_vel_dx = jacobian['v_vel']['x']
    dv_vel_dy = jacobian['v_vel']['y']
    dv_vel_dz = jacobian['v_vel']['z']
    dv_vel_dt = jacobian['v_vel']['t']
    dv_vel_dx_dx = x_hessian['v_vel']['x']['x']
    dv_vel_dy_dy = y_hessian['v_vel']['y']['y']
    dv_vel_dz_dz = z_hessian['v_vel']['z']['z']

    dw_vel_dx = jacobian['w_vel']['x']
    dw_vel_dy = jacobian['w_vel']['y']
    dw_vel_dz = jacobian['w_vel']['z']
    dw_vel_dt = jacobian['w_vel']['t']
    dw_vel_dx_dx = x_hessian['w_vel']['x']['x']
    dw_vel_dy_dy = y_hessian['w_vel']['y']['y']
    dw_vel_dz_dz = z_hessian['w_vel']['z']['z']

    dp_dx = jacobian['p']['x']
    dp_dy = jacobian['p']['y']
    dp_dz = jacobian['p']['z']

    momentum_x = (
        rho * (du_vel_dt + (u_vel * du_vel_dx + v_vel * du_vel_dy + w_vel * du_vel_dz))
        + dp_dx - mu * (du_vel_dx_dx + du_vel_dy_dy + du_vel_dz_dz)
    )
    momentum_y = (
        rho * (dv_vel_dt + (u_vel * dv_vel_dx + v_vel * dv_vel_dy + w_vel * dv_vel_dz))
        + dp_dy - mu * (dv_vel_dx_dx + dv_vel_dy_dy + dv_vel_dz_dz)
    )
    momentum_z = (
        rho * (dw_vel_dt + (u_vel * dw_vel_dx + v_vel * dw_vel_dy + w_vel * dw_vel_dz))
        + dp_dz - mu * (dw_vel_dx_dx + dw_vel_dy_dy + dw_vel_dz_dz)
    )
    continuity = du_vel_dx + dv_vel_dy + dw_vel_dz

    return [momentum_x, momentum_y, momentum_z, continuity]

Define the initial and boundary conditions for the problem:

@brainstate.transform.jit(static_argnums=1)
def icbc_cond_func(x, include_p: bool = False):
    x = {k: v.mantissa for k, v in x.items()}

    u_ = (
        -a
        * (u.math.exp(a * x['x']) * u.math.sin(a * x['y'] + d * x['z'])
           + u.math.exp(a * x['z']) * u.math.cos(a * x['x'] + d * x['y']))
        * u.math.exp(-(d ** 2) * x['t'])
    )
    v = (
        -a
        * (u.math.exp(a * x['y']) * u.math.sin(a * x['z'] + d * x['x'])
           + u.math.exp(a * x['x']) * u.math.cos(a * x['y'] + d * x['z']))
        * u.math.exp(-(d ** 2) * x['t'])
    )
    w = (
        -a
        * (u.math.exp(a * x['z']) * u.math.sin(a * x['x'] + d * x['y'])
           + u.math.exp(a * x['y']) * u.math.cos(a * x['z'] + d * x['x']))
        * u.math.exp(-(d ** 2) * x['t'])
    )
    p = (
        -0.5
        * a ** 2
        * (
            u.math.exp(2 * a * x['x'])
            + u.math.exp(2 * a * x['y'])
            + u.math.exp(2 * a * x['z'])
            + 2
            * u.math.sin(a * x['x'] + d * x['y'])
            * u.math.cos(a * x['z'] + d * x['x'])
            * u.math.exp(a * (x['y'] + x['z']))
            + 2
            * u.math.sin(a * x['y'] + d * x['z'])
            * u.math.cos(a * x['x'] + d * x['y'])
            * u.math.exp(a * (x['z'] + x['x']))
            + 2
            * u.math.sin(a * x['z'] + d * x['x'])
            * u.math.cos(a * x['y'] + d * x['z'])
            * u.math.exp(a * (x['x'] + x['y']))
        )
        * u.math.exp(-2 * d ** 2 * x['t'])
    )

    r = {'u_vel': u_ * unit_of_speed,
         'v_vel': v * unit_of_speed,
         'w_vel': w * unit_of_speed}
    if include_p:
        r['p'] = p * unit_of_pressure
    return r


bc = pinnx.icbc.DirichletBC(icbc_cond_func)
ic = pinnx.icbc.IC(icbc_cond_func)

Define the problem as a TimePDE object:

problem = pinnx.problem.TimePDE(
    spatio_temporal_domain,
    pde,
    [bc, ic],
    net,
    num_domain=50000,
    num_boundary=5000,
    num_initial=5000,
    num_test=10000,
)

Warning: 283 points required, but 343 points sampled.
Warning: 10000 points required, but 12348 points sampled.

Train the model using the problem data:

model = pinnx.Trainer(problem)

model.compile(braintools.optim.Adam(1e-3)).train(iterations=30000)
model.compile(braintools.optim.LBFGS(1e-3)).train(5000, display_every=200)

Compiling trainer...
'compile' took 0.051972 s

Training trainer...

Step      Train loss                                                                  Test loss                                                                   Test metric 
0         [12.974215 * (kilogram / klitre * (meter / second) / second) ** 2,          [14.700201 * (kilogram / klitre * (meter / second) / second) ** 2,          []          
           24.321922 * (kilogram / klitre * (meter / second) / second) ** 2,           29.931065 * (kilogram / klitre * (meter / second) / second) ** 2,                      
           13.350433 * (kilogram / klitre * (meter / second) / second) ** 2,           17.096483 * (kilogram / klitre * (meter / second) / second) ** 2,                      
           1.2527013 * becquerel2,                                                     1.3801537 * becquerel2,                                                                
           {'ibc0': {'u_vel': 2.5884202 * meter / second,                              {'ibc0': {'u_vel': 2.5884202 * meter / second,                                         
                     'v_vel': 1.5904388 * meter / second,                                        'v_vel': 1.5904388 * meter / second,                                         
                     'w_vel': 1.7298671 * meter / second}},                                      'w_vel': 1.7298671 * meter / second}},                                       
           {'ibc1': {'u_vel': 4.1043954 * meter / second,                              {'ibc1': {'u_vel': 4.1043954 * meter / second,                                         
                     'v_vel': 2.08325 * meter / second,                                          'v_vel': 2.08325 * meter / second,                                           
                     'w_vel': 2.6199307 * meter / second}}]                                      'w_vel': 2.6199307 * meter / second}}]                                       

Verify the results by plotting the loss history and the predicted solution:

x, y, z = np.meshgrid(np.linspace(-1, 1, 10), np.linspace(-1, 1, 10), np.linspace(-1, 1, 10))
t_0 = np.zeros(1000)
t_1 = np.ones(1000)
X_0 = dict(
    x=np.ravel(x) * unit_of_space,
    y=np.ravel(y) * unit_of_space,
    z=np.ravel(z) * unit_of_space,
    t=t_0 * unit_of_t
)
X_1 = dict(
    x=np.ravel(x) * unit_of_space,
    y=np.ravel(y) * unit_of_space,
    z=np.ravel(z) * unit_of_space,
    t=t_1 * unit_of_t
)
output_0 = model.predict(X_0)
output_1 = model.predict(X_1)

out_exact_0 = icbc_cond_func(X_0, True)
out_exact_1 = icbc_cond_func(X_1, True)

f_0 = pde(X_0, output_0)
f_1 = pde(X_1, output_1)
residual_0 = jax.tree.map(lambda x: np.mean(np.absolute(x)), f_0)
residual_1 = jax.tree.map(lambda x: np.mean(np.absolute(x)), f_1)

print("Accuracy at t = 0:")
print("Mean residual:", residual_0)
print("L2 relative error:", pinnx.metrics.l2_relative_error(output_0, out_exact_0))
print("\n")
print("Accuracy at t = 1:")
print("Mean residual:", residual_1)
print("L2 relative error:", pinnx.metrics.l2_relative_error(output_1, out_exact_1))