Merge pull request #16 from hsalehipour/main

modified two benchmark examples and fixed a bug in moving bounce-back

.gitignore

@@ -10,6 +10,7 @@
 *zraw
 *mhd
 *png
+*pdf
 # Exclusions
 !jax-lbm.png
 !airfoil.png

examples/CFD/channel3d.py

@@ -62,7 +62,7 @@ def __init__(self, **kwargs):
     def set_boundary_conditions(self):
         # top and bottom sides of the channel are no-slip and the other directions are periodic
         wall = np.concatenate((self.boundingBoxIndices['bottom'], self.boundingBoxIndices['top']))
-        self.BCs.append(BounceBack(tuple(wall.T), self.gridInfo, self.precisionPolicy))
+        self.BCs.append(Regularized(tuple(wall.T), self.gridInfo, self.precisionPolicy, 'velocity', np.zeros((wall.shape[0], 3))))
         return
 
     def initialize_macroscopic_fields(self):
@@ -107,8 +107,11 @@ def output_data(self, **kwargs):
         dns_dic = get_dns_data()
         plt.clf()
         plt.semilogx(yplus, uplus,'r.', yplus, uplus_loglaw, 'k:', dns_dic['y+'], dns_dic['Umean'], 'b-')
-        fname = "uplus_" + str(timestep).zfill(4) + '.png'
-        plt.savefig(fname)
+        ax = plt.gca()
+        ax.set_xlim([0.1, 300])
+        ax.set_ylim([0, 20])
+        fname = "uplus_" + str(timestep//10000).zfill(5) + '.pdf'
+        plt.savefig(fname, format='pdf')
         fields = {"rho": rho[..., 0], "u_x": u[..., 0], "u_y": u[..., 1], "u_z": u[..., 2]}
         save_fields_vtk(timestep, fields)
 
@@ -117,17 +120,19 @@ def output_data(self, **kwargs):
 if __name__ == "__main__":
     precision = "f64/f64"
     lattice = LatticeD3Q27(precision)
+
     # h: channel half-width
-    h = 10
+    h = 50
+
     # Define channel geometry based on h
     nx = 6*h
     ny = 3*h
     nz = 2*h
 
     # Define flow regime
     Re_tau = 180
-    u_tau = 0.004
-    # DeltaPlus = Re_tau/h    # DeltaPlus = u_tau / nu * Delta where u_tau / nu = Re_tau/h
+    u_tau = 0.001
+    DeltaPlus = Re_tau/h    # DeltaPlus = u_tau / nu * Delta where u_tau / nu = Re_tau/h
     visc = u_tau * h / Re_tau
     omega = 1.0 / (3.0 * visc + 0.5)
 
@@ -146,8 +151,8 @@ def output_data(self, **kwargs):
         'ny': ny,
         'nz': nz,
         'precision': precision,
-        'io_rate': 20000,
-        'print_info_rate': 20000
+        'io_rate': 500000,
+        'print_info_rate': 100000
     }
     sim = turbulentChannel(**kwargs)
-    sim.run(4000000)
+    sim.run(10000000)

examples/CFD/taylor_green_vortex.py

@@ -6,16 +6,13 @@
 
 
 from src.boundary_conditions import *
-from jax.config import config
 from src.utils import *
 import numpy as np
 from src.lattice import LatticeD2Q9
 from src.models import BGKSim, KBCSim, AdvectionDiffusionBGK
-import jax.numpy as jnp
-from jax.experimental import mesh_utils
-from jax.sharding import PositionalSharding
 import os
 import matplotlib.pyplot as plt
+import json
 
 # Use 8 CPU devices
 # os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
@@ -25,8 +22,8 @@
 jax.config.update('jax_enable_x64', True)
 
 def taylor_green_initial_fields(xx, yy, u0, rho0, nu, time):
-    ux = -u0 * np.cos(xx) * np.sin(yy) * np.exp(-2 * nu * time)
-    uy = u0 * np.sin(xx) * np.cos(yy) * np.exp(-2 * nu * time)
+    ux = u0 * np.sin(xx) * np.cos(yy) * np.exp(-2 * nu * time)
+    uy = -u0 * np.cos(xx) * np.sin(yy) * np.exp(-2 * nu * time)
     rho = 1.0 - rho0 * u0 ** 2 / 12. * (np.cos(2. * xx) + np.cos(2. * yy)) * np.exp(-4 * nu * time)
     return ux, uy, np.expand_dims(rho, axis=-1)
 
@@ -51,7 +48,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(20000)
+        f = ADE.run(int(20000*32/nx))
         return f
 
     def output_data(self, **kwargs):
@@ -64,49 +61,67 @@ def output_data(self, **kwargs):
         time = timestep * (kx**2 + ky**2)/2.
         ux_th, uy_th, rho_th = taylor_green_initial_fields(xx, yy, vel_ref, 1, visc, time)
         vel_err_L2 = np.sqrt(np.sum((u[..., 0]-ux_th)**2 + (u[..., 1]-uy_th)**2) / np.sum(ux_th**2 + uy_th**2))
-        print("error= {:07.6f}".format(vel_err_L2))
+        rho_err_L2 = np.sqrt(np.sum((rho - rho_th)**2) / np.sum(rho_th**2))
+        print("Vel error= {:07.6f}, Pressure error= {:07.6f}".format(vel_err_L2, rho_err_L2))
         if timestep == endTime:
             ErrL2ResList.append(vel_err_L2)
+            ErrL2ResListRho.append(rho_err_L2)
         # save_image(timestep, u)
 
 
 if __name__ == "__main__":
-    precision = "f64/f64"
-    lattice = LatticeD2Q9(precision)
-
-    resList = [32, 64, 128, 256, 512]
-    ErrL2ResList = []
-
-    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)])
-        xx, yy = coord[:, 0], coord[:, 1]
-        kx, ky = twopi / nx, twopi / ny
-        xx = xx.reshape((nx, ny)) * kx
-        yy = yy.reshape((nx, ny)) * ky
-
-        Re = 1000.0
-        vel_ref = 0.04*32/nx
-
-        visc = vel_ref * nx / Re
-        omega = 1.0 / (3.0 * visc + 0.5)
-        print("omega = ", omega)
-        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
-        }
-        sim = TaylorGreenVortex(**kwargs)
-        endTime = int(20000*nx/32.0)
-        sim.run(endTime)
-    plt.loglog(resList, ErrL2ResList, '-o')
-    plt.loglog(resList, 1e-3*(np.array(resList)/128)**(-2), '--')
-    plt.savefig('Error.png'); plt.savefig('Error.pdf', format='pdf')
+    precision_list = ["f32/f32", "f64/f32", "f64/f64"]
+    resList = [32, 64, 128, 256, 512, 1024]
+    result_dict = dict.fromkeys(precision_list)
+    result_dict['resolution_list'] = resList
+
+    for precision in precision_list:
+        lattice = LatticeD2Q9(precision)
+        ErrL2ResList = []
+        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)])
+            xx, yy = coord[:, 0], coord[:, 1]
+            kx, ky = twopi / nx, twopi / ny
+            xx = xx.reshape((nx, ny)) * kx
+            yy = yy.reshape((nx, ny)) * ky
+
+            Re = 1600.0
+            vel_ref = 0.04*32/nx
+
+            visc = vel_ref * nx / Re
+            omega = 1.0 / (3.0 * visc + 0.5)
+            print("omega = ", omega)
+            os.system("rm -rf ./*.vtk && rm -rf ./*.png")
+            kwargs = {
+                'lattice': lattice,
+                'omega': omega,
+                'nx': nx,
+                'ny': ny,
+                'nz': 0,
+                'precision': precision,
+                'io_rate': 5000,
+                'print_info_rate': 1000
+            }
+            sim = TaylorGreenVortex(**kwargs)
+            tc = 2.0/(2. * visc * (kx**2 + ky**2))
+            endTime = int(0.05*tc)
+            sim.run(endTime)
+        result_dict[precision]['vel_error'] = ErrL2ResList
+        result_dict[precision]['rho_error'] = ErrL2ResListRho
+
+    with open('data.json', 'w') as fp:
+        json.dump(result_dict, fp)
+
+    # plt.loglog(resList, ErrL2ResList, '-o')
+    # plt.loglog(resList, 1e-3*(np.array(resList)/128)**(-2), '--')
+    # plt.savefig('ErrorVel.png'); plt.savefig('ErrorVel.pdf', format='pdf')
+
+    # plt.figure()
+    # plt.loglog(resList, ErrL2ResListRho, '-o')
+    # plt.loglog(resList, 1e-3*(np.array(resList)/128)**(-2), '--')
+    # plt.savefig('ErrorRho.png'); plt.savefig('ErrorRho.pdf', format='pdf')

src/boundary_conditions.py

@@ -445,8 +445,8 @@ def apply(self, fout, fin, time):
         """
         indices, vel = self.update_function(time)
         c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype)
-        cu = 3.0 * jnp.dot(vel, c)
-        return fout.at[indices].set(fin[indices][..., self.lattice.opp_indices] - 6.0 * self.lattice.w * cu)
+        cu = 6.0 * self.lattice.w * jnp.dot(vel, c)
+        return fout.at[indices].set(fin[indices][..., self.lattice.opp_indices] - cu)
 
 
 class BounceBackHalfway(BoundaryCondition):
@@ -467,13 +467,16 @@ class BounceBackHalfway(BoundaryCondition):
         Whether the boundary condition needs extra configuration before it can be applied. For this class, it is True.
     isSolid : bool
         Whether the boundary condition represents a solid boundary. For this class, it is True.
+    vel : array-like
+        The prescribed value of velocity vector for the boundary condition. No-slip BC is assumed if vel=None (default).
     """
-    def __init__(self, indices, gridInfo, precision_policy):
+    def __init__(self, indices, gridInfo, precision_policy, vel=None):
         super().__init__(indices, gridInfo, precision_policy)
         self.name = "BounceBackHalfway"
         self.implementationStep = "PostStreaming"
         self.needsExtraConfiguration = True
         self.isSolid = True
+        self.vel = vel
 
     def configure(self, boundaryBitmask):
         """
@@ -523,7 +526,12 @@ def apply(self, fout, fin):
         nbd = len(self.indices[0])
         bindex = np.arange(nbd)[:, None]
         fbd = fout[self.indices]
-        fbd = fbd.at[bindex, self.imissing].set(fin[self.indices][bindex, self.iknown])
+        if self.vel is not None:
+            c = jnp.array(self.lattice.c, dtype=self.precisionPolicy.compute_dtype)
+            cu = 6.0 * self.lattice.w * jnp.dot(self.vel, c)
+            fbd = fbd.at[bindex, self.imissing].set(fin[self.indices][bindex, self.iknown] - cu[bindex, self.iknown])
+        else:
+            fbd = fbd.at[bindex, self.imissing].set(fin[self.indices][bindex, self.iknown])
 
         return fbd
 
@@ -926,8 +934,8 @@ def configure(self, boundaryBitmask):
 
         Parameters
         ----------
-        connectivity_bitmask : np.ndarray
-            The connectivity bitmask for the lattice.
+        boundaryBitmask : np.ndarray
+            The connectivity bitmask for the boundary voxels.
         """        
         shiftDir = ~boundaryBitmask[:, self.lattice.opp_indices]
         idx = np.array(self.indices).T