Merge pull request #19 from hsalehipour/benchmark

added benchmark examples and interpolating Bounceback BC

examples/CFD/cavity3d.py

@@ -4,7 +4,7 @@
 
 In this example you'll be introduced to the following concepts:
 
-1. Lattice: The simulation employs a D2Q9 lattice. It's a 2D lattice model with nine discrete velocity directions, which is typically used for 2D simulations.
+1. Lattice: The simulation employs a D3Q27 lattice. It's a 3D lattice model with 27 discrete velocity directions.
 
 2. Boundary Conditions: The code implements two types of boundary conditions:
 
@@ -14,74 +14,98 @@
 4. Visualization: The simulation outputs data in VTK format for visualization. The data can be visualized using software like Paraview.
 
 """
-
-import os
-
 # Use 8 CPU devices
 # os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
 
 import numpy as np
 from src.utils import *
 from jax.config import config
+import json, codecs
 
 from src.models import BGKSim, KBCSim
 from src.lattice import LatticeD3Q19, LatticeD3Q27
 from src.boundary_conditions import *
 
+
+config.update('jax_enable_x64', True)
+
 class Cavity(KBCSim):
+    # Note: We have used BGK with D3Q19 (or D3Q27) for Re=(1000, 3200) and KBC with D3Q27 for Re=10,000
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
     def set_boundary_conditions(self):
+        # Note:
+        # We have used halfway BB for Re=(1000, 3200) and regularized BC for Re=10,000
+
+        # apply inlet boundary condition to the top wall
+        moving_wall = self.boundingBoxIndices['top']
+        vel_wall = np.zeros(moving_wall.shape, dtype=self.precisionPolicy.compute_dtype)
+        vel_wall[:, 0] = prescribed_vel
+        # self.BCs.append(BounceBackHalfway(tuple(moving_wall.T), self.gridInfo, self.precisionPolicy, vel_wall))
+        self.BCs.append(Regularized(tuple(moving_wall.T), self.gridInfo, self.precisionPolicy, 'velocity', vel_wall))
+
         # concatenate the indices of the left, right, and bottom walls
         walls = np.concatenate(
             (self.boundingBoxIndices['left'], self.boundingBoxIndices['right'],
              self.boundingBoxIndices['front'], self.boundingBoxIndices['back'],
              self.boundingBoxIndices['bottom']))
         # apply bounce back boundary condition to the walls
-        self.BCs.append(BounceBack(tuple(walls.T), self.gridInfo, self.precisionPolicy))
-
-        # apply inlet equilibrium boundary condition to the top wall
-        moving_wall = self.boundingBoxIndices['top']
-
-        rho_wall = np.ones((moving_wall.shape[0], 1), dtype=self.precisionPolicy.compute_dtype)
-        vel_wall = np.zeros(moving_wall.shape, dtype=self.precisionPolicy.compute_dtype)
-        vel_wall[:, 0] = prescribed_vel
-        self.BCs.append(EquilibriumBC(tuple(moving_wall.T), self.gridInfo, self.precisionPolicy, rho_wall, vel_wall))
+        # self.BCs.append(BounceBackHalfway(tuple(walls.T), self.gridInfo, self.precisionPolicy))
+        vel_wall = np.zeros(walls.shape, dtype=self.precisionPolicy.compute_dtype)
+        self.BCs.append(Regularized(tuple(walls.T), self.gridInfo, self.precisionPolicy, 'velocity', vel_wall))
+        return
 
     def output_data(self, **kwargs):
         # 1: -1 to remove boundary voxels (not needed for visualization when using full-way bounce-back)
-        rho = np.array(kwargs['rho'][1:-1, 1:-1, 1:-1])
-        u = np.array(kwargs['u'][1:-1, 1:-1, 1:-1, :])
+        rho = np.array(kwargs['rho'])
+        u = np.array(kwargs['u'])
         timestep = kwargs['timestep']
-        u_prev = kwargs['u_prev'][1:-1, 1:-1, 1:-1, :]
+        u_prev = kwargs['u_prev']
 
         u_old = np.linalg.norm(u_prev, axis=2)
         u_new = np.linalg.norm(u, axis=2)
 
         err = np.sum(np.abs(u_old - u_new))
         print('error= {:07.6f}'.format(err))
         fields = {"rho": rho[..., 0], "u_x": u[..., 0], "u_y": u[..., 1], "u_z": u[..., 2]}
-        save_fields_vtk(timestep, fields)
+        # save_fields_vtk(timestep, fields)
+
+        # output profiles of velocity at mid-plane for benchmarking
+        output_filename = "./profiles_" + f"{timestep:07d}.json"
+        ux_mid = 0.5*(u[nx//2, ny//2, :, 0] + u[nx//2+1, ny//2+1, :, 0])
+        uz_mid = 0.5*(u[:, ny//2, nz//2, 2] + u[:, ny//2+1, nz//2+1, 2])
+        ldc_ref_result = {'ux(x=y=0)': list(ux_mid/prescribed_vel),
+                          'uz(z=y=0)': list(uz_mid/prescribed_vel)}
+        json.dump(ldc_ref_result, codecs.open(output_filename, 'w', encoding='utf-8'),
+                separators=(',', ':'),
+                sort_keys=True,
+                indent=4)
+
         # Calculate the velocity magnitude
-        u_mag = np.linalg.norm(u, axis=2)
+        # u_mag = np.linalg.norm(u, axis=2)
         # live_volume_randering(timestep, u_mag)
 
 if __name__ == '__main__':
-    nx = 101
-    ny = 101
-    nz = 101
+    # Note: 
+    # We have used BGK with D3Q19 (or D3Q27) for Re=(1000, 3200) and KBC with D3Q27 for Re=10,000
+    precision = 'f64/f64'
+    lattice = LatticeD3Q27(precision)
 
-    Re = 50000.0
-    prescribed_vel = 0.1
-    clength = nx - 1
+    nx = 256
+    ny = 256
+    nz = 256
 
-    precision = 'f32/f32'
-    lattice = LatticeD3Q27(precision)
+    Re = 10000.0
+    prescribed_vel = 0.06
+    clength = nx - 2
+
+    # characteristic time
+    tc = prescribed_vel/clength
+    niter_max = int(500//tc)
 
     visc = prescribed_vel * clength / Re
     omega = 1.0 / (3. * visc + 0.5)
-
     os.system("rm -rf ./*.vtk && rm -rf ./*.png")
 
     kwargs = {
@@ -91,9 +115,9 @@ def output_data(self, **kwargs):
         'ny': ny,
         'nz': nz,
         'precision': precision,
-        'io_rate': 100,
-        'print_info_rate': 100,
-        'downsampling_factor': 2
+        'io_rate': int(10//tc),
+        'print_info_rate': int(10//tc),
+        'downsampling_factor': 1
     }
     sim = Cavity(**kwargs)
-    sim.run(2000)
+    sim.run(niter_max)

examples/CFD/cylinder2d.py

@@ -18,6 +18,7 @@
 
 """
 import os
+import json
 import jax
 from time import time
 from jax.config import config
@@ -33,7 +34,7 @@
 # os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
 jax.config.update('jax_enable_x64', True)
 
-class Cylinder(KBCSim):
+class Cylinder(BGKSim):
     def __init__(self, **kwargs):
         super().__init__(**kwargs)
 
@@ -44,83 +45,99 @@ def set_boundary_conditions(self):
         cx, cy = 2.*diam, 2.*diam
         cylinder = (xx - cx)**2 + (yy-cy)**2 <= (diam/2.)**2
         cylinder = coord[cylinder]
-        self.BCs.append(BounceBackHalfway(tuple(cylinder.T), self.gridInfo, self.precisionPolicy))
-        # wall = np.concatenate([cylinder, self.boundingBoxIndices['top'], self.boundingBoxIndices['bottom']])
-        # self.BCs.append(BounceBack(tuple(wall.T), self.gridInfo, self.precisionPolicy))
+        implicit_distance = np.reshape((xx - cx)**2 + (yy-cy)**2 - (diam/2.)**2, (self.nx, self.ny))
+        self.BCs.append(InterpolatedBounceBackBouzidi(tuple(cylinder.T), implicit_distance, self.gridInfo, self.precisionPolicy))
 
+        # Outflow BC
         outlet = self.boundingBoxIndices['right']
         rho_outlet = np.ones(outlet.shape[0], dtype=self.precisionPolicy.compute_dtype)
         self.BCs.append(ExtrapolationOutflow(tuple(outlet.T), self.gridInfo, self.precisionPolicy))
-        # self.BCs.append(Regularized(tuple(outlet.T), self.gridInfo, self.precisionPolicy, 'pressure', rho_outlet))
 
+        # Inlet BC
         inlet = self.boundingBoxIndices['left']
         rho_inlet = np.ones((inlet.shape[0], 1), dtype=self.precisionPolicy.compute_dtype)
         vel_inlet = np.zeros(inlet.shape, dtype=self.precisionPolicy.compute_dtype)
         yy_inlet = yy.reshape(self.nx, self.ny)[tuple(inlet.T)]
         vel_inlet[:, 0] = poiseuille_profile(yy_inlet,
                                              yy_inlet.min(),
                                              yy_inlet.max()-yy_inlet.min(), 3.0 / 2.0 * prescribed_vel)
-        # self.BCs.append(EquilibriumBC(tuple(inlet.T), self.gridInfo, self.precisionPolicy, rho_inlet, vel_inlet))
         self.BCs.append(Regularized(tuple(inlet.T), self.gridInfo, self.precisionPolicy, 'velocity', vel_inlet))
 
+        # No-slip BC for top and bottom
         wall = np.concatenate([self.boundingBoxIndices['top'], self.boundingBoxIndices['bottom']])
-        self.BCs.append(BounceBack(tuple(wall.T), self.gridInfo, self.precisionPolicy))
+        vel_wall = np.zeros(wall.shape, dtype=self.precisionPolicy.compute_dtype)
+        self.BCs.append(Regularized(tuple(wall.T), self.gridInfo, self.precisionPolicy, 'velocity', vel_wall))
 
     def output_data(self, **kwargs):
-        # 1:-1 to remove boundary voxels (not needed for visualization when using full-way bounce-back)
+        # 1:-1 to remove boundary voxels (not needed for visualization when using bounce-back)
         rho = np.array(kwargs["rho"][..., 1:-1, :])
         u = np.array(kwargs["u"][..., 1:-1, :])
         timestep = kwargs["timestep"]
         u_prev = kwargs["u_prev"][..., 1:-1, :]
 
-        # compute lift and drag over the cyliner
-        cylinder = self.BCs[0]
-        boundary_force = cylinder.momentum_exchange_force(kwargs['f_poststreaming'], kwargs['f_postcollision'])
-        boundary_force = np.sum(boundary_force, axis=0)
-        drag = boundary_force[0]
-        lift = boundary_force[1]
-        cd = 2. * drag / (prescribed_vel ** 2 * diam)
-        cl = 2. * lift / (prescribed_vel ** 2 * diam)
-
-        u_old = np.linalg.norm(u_prev, axis=2)
-        u_new = np.linalg.norm(u, axis=2)
-        err = np.sum(np.abs(u_old - u_new))
-        print('error= {:07.6f}, CL = {:07.6f}, CD = {:07.6f}'.format(err, cl, cd))
-        save_image(timestep, u)
+        if timestep == 0:
+            self.CL_max = 0.0
+            self.CD_max = 0.0
+        if timestep > 0.8 * t_max:
+            # compute lift and drag over the cyliner
+            cylinder = self.BCs[0]
+            boundary_force = cylinder.momentum_exchange_force(kwargs['f_poststreaming'], kwargs['f_postcollision'])
+            boundary_force = np.sum(np.array(boundary_force), axis=0)
+            drag = boundary_force[0]
+            lift = boundary_force[1]
+            cd = 2. * drag / (prescribed_vel ** 2 * diam)
+            cl = 2. * lift / (prescribed_vel ** 2 * diam)
+
+            u_old = np.linalg.norm(u_prev, axis=2)
+            u_new = np.linalg.norm(u, axis=2)
+            err = np.sum(np.abs(u_old - u_new))
+            self.CL_max = max(self.CL_max, cl)
+            self.CD_max = max(self.CD_max, cd)
+            print('error= {:07.6f}, CL = {:07.6f}, CD = {:07.6f}'.format(err, cl, cd))
+            save_image(timestep, u)
 
 # Helper function to specify a parabolic poiseuille profile
 poiseuille_profile  = lambda x,x0,d,umax: np.maximum(0.,4.*umax/(d**2)*((x-x0)*d-(x-x0)**2))
 
 if __name__ == '__main__':
     precision = 'f64/f64'
-    lattice = LatticeD2Q9(precision)
-
-    prescribed_vel = 0.005
-    diam = 80
-
-    nx = int(22*diam)
-    ny = int(4.1*diam)
-
-    Re = 100.0
-    visc = prescribed_vel * diam / Re
-    omega = 1.0 / (3. * visc + 0.5)
-
-    print('omega = ', omega)
-    print("Mesh size: ", nx, ny)
-    print("Number of voxels: ", nx * ny)
-
-    os.system('rm -rf ./*.vtk && rm -rf ./*.png')
-
-    kwargs = {
-        'lattice': lattice,
-        'omega': omega,
-        'nx': nx,
-        'ny': ny,
-        'nz': 0,
-        'precision': precision,
-        'io_rate': 500,
-        'print_info_rate': 500,
-        'return_fpost': True    # Need to retain fpost-collision for computation of lift and drag
-    }
-    sim = Cylinder(**kwargs)
-    sim.run(1000000)
+    # diam_list = [10, 20, 30, 40, 60, 80]
+    diam_list = [80]
+    CL_list, CD_list = [], []
+    result_dict = {}
+    result_dict['resolution_list'] = diam_list
+    for diam in diam_list:
+        scale_factor = 80 / diam
+        prescribed_vel = 0.003 * scale_factor
+        lattice = LatticeD2Q9(precision)
+
+        nx = int(22*diam)
+        ny = int(4.1*diam)
+
+        Re = 100.0
+        visc = prescribed_vel * diam / Re
+        omega = 1.0 / (3. * visc + 0.5)
+
+        os.system('rm -rf ./*.vtk && rm -rf ./*.png')
+
+        kwargs = {
+            'lattice': lattice,
+            'omega': omega,
+            'nx': nx,
+            'ny': ny,
+            'nz': 0,
+            'precision': precision,
+            'io_rate': int(500 / scale_factor),
+            'print_info_rate': int(10000 / scale_factor),
+            'return_fpost': True    # Need to retain fpost-collision for computation of lift and drag
+        }
+        sim = Cylinder(**kwargs)
+        t_max = int(1000000 / scale_factor)
+        sim.run(t_max)
+        CL_list.append(sim.CL_max)
+        CD_list.append(sim.CD_max)
+
+    result_dict['CL'] = CL_list
+    result_dict['CD'] = CD_list
+    with open('data.json', 'w') as fp:
+        json.dump(result_dict, fp)

examples/CFD/taylor_green_vortex.py

@@ -50,7 +50,7 @@ def initialize_populations(self, rho, u):
         ADE = AdvectionDiffusionBGK(**kwargs)
         ADE.initialize_macroscopic_fields = self.initialize_macroscopic_fields
         print("Initializing the distribution functions using the specified macroscopic fields....")
-        f = ADE.run(int(20000*32/nx))
+        f = ADE.run(int(20000*nx/32))
         return f
 
     def output_data(self, **kwargs):
@@ -83,7 +83,6 @@ def output_data(self, **kwargs):
         ErrL2ResListRho = []
         result_dict[precision] = dict.fromkeys(['vel_error', 'rho_error'])
         for nx in resList:
-            print("Running at nx = ny = {:07.6f}".format(nx))
             ny = nx
             twopi = 2.0 * np.pi
             coord = np.array([(i, j) for i in range(nx) for j in range(ny)])

src/base.py

@@ -11,17 +11,16 @@
 
 # JAX-related imports
 from jax import jit, lax, vmap
-from jax.config import config
 from jax.experimental import mesh_utils
 from jax.experimental.multihost_utils import process_allgather
 from jax.experimental.shard_map import shard_map
 from jax.sharding import NamedSharding, PartitionSpec, PositionalSharding, Mesh
 import orbax.checkpoint as orb
+
 # functools imports
 from functools import partial
 
 # Local/Custom Libraries
-import src.models
 from src.utils import downsample_field
 
 jax.config.update("jax_spmd_mode", 'allow_all')