Merge pull request #49 from hsalehipour/test_new_bc_setup

New way of settig up BC

examples/cfd/flow_past_sphere_3d.py

@@ -1,111 +1,125 @@
 import xlb
 from xlb.compute_backend import ComputeBackend
 from xlb.precision_policy import PrecisionPolicy
-from xlb.helper import create_nse_fields, initialize_eq, assign_bc_id_box_faces
+from xlb.helper import create_nse_fields, initialize_eq
 from xlb.operator.stepper import IncompressibleNavierStokesStepper
 from xlb.operator.boundary_condition import (
     FullwayBounceBackBC,
     EquilibriumBC,
     DoNothingBC,
 )
-from xlb.operator.equilibrium import QuadraticEquilibrium
 from xlb.operator.macroscopic import Macroscopic
+from xlb.operator.boundary_masker import IndicesBoundaryMasker
 from xlb.utils import save_fields_vtk, save_image
 import warp as wp
 import numpy as np
 import jax.numpy as jnp
 
-backend = ComputeBackend.WARP
-velocity_set = xlb.velocity_set.D3Q19()
-precision_policy = PrecisionPolicy.FP32FP32
-
-xlb.init(
-    velocity_set=velocity_set,
-    default_backend=backend,
-    default_precision_policy=precision_policy,
-)
-
-grid_size_x, grid_size_y, grid_size_z = 512, 128, 128
-grid_shape = (grid_size_x, grid_size_y, grid_size_z)
-
-grid, f_0, f_1, missing_mask, boundary_mask = create_nse_fields(grid_shape)
-
-# Velocity on left face (3D)
-boundary_mask, missing_mask = assign_bc_id_box_faces(
-    boundary_mask, missing_mask, grid_shape, EquilibriumBC.id, ["left"]
-)
-
-
-# Wall on all other faces (3D) except right
-boundary_mask, missing_mask = assign_bc_id_box_faces(
-    boundary_mask,
-    missing_mask,
-    grid_shape,
-    FullwayBounceBackBC.id,
-    ["bottom", "right", "front", "back"],
-)
-
-# Do nothing on right face
-boundary_mask, missing_mask = assign_bc_id_box_faces(
-    boundary_mask, missing_mask, grid_shape, DoNothingBC.id, ["right"]
-)
-
-bc_eq = QuadraticEquilibrium()
-bc_left = EquilibriumBC(rho=1.0, u=(0.02, 0.0, 0.0), equilibrium_operator=bc_eq)
-bc_walls = FullwayBounceBackBC()
-bc_do_nothing = DoNothingBC()
-
-
-sphere_radius = grid_size_y // 12
-x = np.arange(grid_size_x)
-y = np.arange(grid_size_y)
-z = np.arange(grid_size_z)
-X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
-indices = np.where(
-    (X - grid_size_x // 6) ** 2
-    + (Y - grid_size_y // 2) ** 2
-    + (Z - grid_size_z // 2) ** 2
-    < sphere_radius**2
-)
-indices = np.array(indices)
-
-# Set boundary conditions on the indices
-indices_boundary_masker = xlb.operator.boundary_masker.IndicesBoundaryMasker(
-    velocity_set=velocity_set,
-    precision_policy=precision_policy,
-    compute_backend=backend,
-)
-
-boundary_mask, missing_mask = indices_boundary_masker(
-    indices, FullwayBounceBackBC.id, boundary_mask, missing_mask, (0, 0, 0)
-)
-
-f_0 = initialize_eq(f_0, grid, velocity_set, backend)
-boundary_conditions = [bc_left, bc_walls, bc_do_nothing]
-omega = 1.8
-
-stepper = IncompressibleNavierStokesStepper(
-    omega, boundary_conditions=boundary_conditions
-)
-
-for i in range(10000):
-    f_1 = stepper(f_0, f_1, boundary_mask, missing_mask, i)
-    f_0, f_1 = f_1, f_0
-
-
-# Write the results. We'll use JAX backend for the post-processing
-if not isinstance(f_0, jnp.ndarray):
-    f_0 = wp.to_jax(f_0)
-
-macro = Macroscopic(compute_backend=ComputeBackend.JAX)
-
-rho, u = macro(f_0)
-
-# remove boundary cells
-u = u[:, 1:-1, 1:-1, 1:-1]
-u_magnitude = (u[0] ** 2 + u[1] ** 2 + u[2] ** 2) ** 0.5
-
-fields = {"u_magnitude": u_magnitude}
-
-save_fields_vtk(fields, timestep=i)
-save_image(fields["u_magnitude"][:, grid_size_y // 2, :], timestep=i)
+class FlowOverSphere:
+    def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy):
+
+        # initialize backend
+        xlb.init(
+            velocity_set=velocity_set,
+            default_backend=backend,
+            default_precision_policy=precision_policy,
+        )
+
+        self.grid_shape = grid_shape
+        self.velocity_set = velocity_set
+        self.backend = backend
+        self.precision_policy = precision_policy
+        self.grid, self.f_0, self.f_1, self.missing_mask, self.boundary_mask = create_nse_fields(grid_shape)
+        self.stepper = None
+        self.boundary_conditions = []
+
+        # Setup the simulation BC, its initial conditions, and the stepper
+        self._setup(omega)
+
+    def _setup(self, omega):
+        self.setup_boundary_conditions()
+        self.setup_boundary_masks()
+        self.initialize_fields()
+        self.setup_stepper(omega)
+
+    def define_boundary_indices(self):
+        inlet = self.grid.boundingBoxIndices['left']
+        outlet = self.grid.boundingBoxIndices['right']
+        walls = [self.grid.boundingBoxIndices['bottom'][i] + self.grid.boundingBoxIndices['top'][i] + 
+                 self.grid.boundingBoxIndices['front'][i] + self.grid.boundingBoxIndices['back'][i] for i in range(self.velocity_set.d)]
+
+        sphere_radius = self.grid_shape[1] // 12
+        x = np.arange(self.grid_shape[0])
+        y = np.arange(self.grid_shape[1])
+        z = np.arange(self.grid_shape[2])
+        X, Y, Z = np.meshgrid(x, y, z, indexing='ij')
+        indices = np.where(
+            (X - self.grid_shape[0] // 6) ** 2
+            + (Y - self.grid_shape[1] // 2) ** 2
+            + (Z - self.grid_shape[2] // 2) ** 2
+            < sphere_radius**2
+        )
+        sphere = [tuple(indices[i]) for i in range(self.velocity_set.d)]
+
+        return inlet, outlet, walls, sphere
+
+    def setup_boundary_conditions(self):
+        inlet, outlet, walls, sphere = self.define_boundary_indices()
+        bc_left = EquilibriumBC(rho=1.0, u=(0.02, 0.0, 0.0), indices=inlet)
+        bc_walls = FullwayBounceBackBC(indices=walls)
+        bc_do_nothing = DoNothingBC(indices=outlet)
+        bc_sphere = FullwayBounceBackBC(indices=sphere)
+        self.boundary_conditions = [bc_left, bc_walls, bc_do_nothing, bc_sphere]
+
+    def setup_boundary_masks(self):
+        indices_boundary_masker = IndicesBoundaryMasker(
+            velocity_set=self.velocity_set,
+            precision_policy=self.precision_policy,
+            compute_backend=self.backend,
+        )
+        self.boundary_mask, self.missing_mask = indices_boundary_masker(
+            self.boundary_conditions, self.boundary_mask, self.missing_mask, (0, 0, 0)
+        )
+
+    def initialize_fields(self):
+        self.f_0 = initialize_eq(self.f_0, self.grid, self.velocity_set, self.backend)
+
+    def setup_stepper(self, omega):
+        self.stepper = IncompressibleNavierStokesStepper(
+            omega, boundary_conditions=self.boundary_conditions
+        )
+
+    def run(self, num_steps):
+        for i in range(num_steps):
+            self.f_1 = self.stepper(self.f_0, self.f_1, self.boundary_mask, self.missing_mask, i)
+            self.f_0, self.f_1 = self.f_1, self.f_0
+
+    def post_process(self, i):
+        # Write the results. We'll use JAX backend for the post-processing
+        if not isinstance(self.f_0, jnp.ndarray):
+            self.f_0 = wp.to_jax(self.f_0)
+
+        macro = Macroscopic(compute_backend=ComputeBackend.JAX)
+        rho, u = macro(self.f_0)
+
+        # remove boundary cells
+        u = u[:, 1:-1, 1:-1, 1:-1]
+        u_magnitude = (u[0] ** 2 + u[1] ** 2 + u[2] ** 2) ** 0.5
+
+        fields = {"u_magnitude": u_magnitude}
+
+        save_fields_vtk(fields, timestep=i)
+        save_image(fields["u_magnitude"][:, self.grid_shape[1] // 2, :], timestep=i)
+
+
+if __name__ == "__main__":
+    # Running the simulation
+    grid_shape = (512, 128, 128)
+    velocity_set = xlb.velocity_set.D3Q19()
+    backend = ComputeBackend.WARP
+    precision_policy = PrecisionPolicy.FP32FP32
+    omega = 1.6
+
+    simulation = FlowOverSphere(omega, grid_shape, velocity_set, backend, precision_policy)
+    simulation.run(num_steps=10000)
+    simulation.post_process(i=10000)

examples/cfd/lid_driven_cavity_2d.py

@@ -1,79 +1,107 @@
 import xlb
 from xlb.compute_backend import ComputeBackend
 from xlb.precision_policy import PrecisionPolicy
-from xlb.helper import create_nse_fields, initialize_eq, assign_bc_id_box_faces
+from xlb.helper import create_nse_fields, initialize_eq
+from xlb.operator.boundary_masker import IndicesBoundaryMasker
 from xlb.operator.stepper import IncompressibleNavierStokesStepper
 from xlb.operator.boundary_condition import FullwayBounceBackBC, EquilibriumBC
-from xlb.operator.equilibrium import QuadraticEquilibrium
 from xlb.operator.macroscopic import Macroscopic
 from xlb.utils import save_fields_vtk, save_image
 import warp as wp
 import jax.numpy as jnp
 
-backend = ComputeBackend.JAX
-velocity_set = xlb.velocity_set.D2Q9()
-precision_policy = PrecisionPolicy.FP32FP32
 
-xlb.init(
-    velocity_set=velocity_set,
-    default_backend=backend,
-    default_precision_policy=precision_policy,
-)
-
-grid_size = 128
-grid_shape = (grid_size, grid_size)
-
-grid, f_0, f_1, missing_mask, boundary_mask = create_nse_fields(grid_shape)
-
-# Velocity on top face (2D)
-boundary_mask, missing_mask = assign_bc_id_box_faces(
-    boundary_mask, missing_mask, grid_shape, EquilibriumBC.id, ["top"]
-)
-
-bc_eq = QuadraticEquilibrium()
-
-bc_top = EquilibriumBC(rho=1.0, u=(0.02, 0.0), equilibrium_operator=bc_eq)
-
-
-# Wall on all other faces (2D)
-boundary_mask, missing_mask = assign_bc_id_box_faces(
-    boundary_mask,
-    missing_mask,
-    grid_shape,
-    FullwayBounceBackBC.id,
-    ["bottom", "left", "right"],
-)
-
-bc_walls = FullwayBounceBackBC()
-
-
-f_0 = initialize_eq(f_0, grid, velocity_set, backend)
-boundary_conditions = [bc_top, bc_walls]
-omega = 1.6
-
-stepper = IncompressibleNavierStokesStepper(
-    omega, boundary_conditions=boundary_conditions
-)
-
-for i in range(500):
-    f_1 = stepper(f_0, f_1, boundary_mask, missing_mask, i)
-    f_0, f_1 = f_1, f_0
-
-
-# Write the results. We'll use JAX backend for the post-processing
-if not isinstance(f_0, jnp.ndarray):
-    f_0 = wp.to_jax(f_0)
-
-macro = Macroscopic(compute_backend=ComputeBackend.JAX)
-
-rho, u = macro(f_0)
-
-# remove boundary cells
-rho = rho[:, 1:-1, 1:-1]
-u = u[:, 1:-1, 1:-1]
-u_magnitude = (u[0] ** 2 + u[1] ** 2) ** 0.5
-
-fields = {"rho": rho[0], "u_x": u[0], "u_y": u[1], "u_magnitude": u_magnitude}
-
-save_fields_vtk(fields, timestep=i, prefix="lid_driven_cavity")
-save_image(fields["u_magnitude"], timestep=i, prefix="lid_driven_cavity")
+class LidDrivenCavity2D:
+    def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy):
+
+        # initialize backend
+        xlb.init(
+            velocity_set=velocity_set,
+            default_backend=backend,
+            default_precision_policy=precision_policy,
+        )
+
+        self.grid_shape = grid_shape
+        self.velocity_set = velocity_set
+        self.backend = backend
+        self.precision_policy = precision_policy
+        self.grid, self.f_0, self.f_1, self.missing_mask, self.boundary_mask = create_nse_fields(grid_shape)
+        self.stepper = None
+        self.boundary_conditions = []
+
+        # Setup the simulation BC, its initial conditions, and the stepper
+        self._setup(omega)
+
+    def _setup(self, omega):
+        self.setup_boundary_conditions()
+        self.setup_boundary_masks()
+        self.initialize_fields()
+        self.setup_stepper(omega)
+
+    def define_boundary_indices(self):
+        lid = self.grid.boundingBoxIndices['top']
+        walls = [self.grid.boundingBoxIndices['bottom'][i] + self.grid.boundingBoxIndices['left'][i] + 
+                 self.grid.boundingBoxIndices['right'][i] for i in range(self.velocity_set.d)]
+        return lid, walls
+
+    def setup_boundary_conditions(self):
+        lid, walls = self.define_boundary_indices()
+        bc_top = EquilibriumBC(rho=1.0, u=(0.02, 0.0), indices=lid)
+        bc_walls = FullwayBounceBackBC(indices=walls)
+        self.boundary_conditions = [bc_top, bc_walls]
+
+    def setup_boundary_masks(self):
+        indices_boundary_masker = IndicesBoundaryMasker(
+            velocity_set=self.velocity_set,
+            precision_policy=self.precision_policy,
+            compute_backend=self.backend,
+        )
+        self.boundary_mask, self.missing_mask = indices_boundary_masker(
+            self.boundary_conditions, self.boundary_mask, self.missing_mask
+        )
+
+    def initialize_fields(self):
+        self.f_0 = initialize_eq(self.f_0, self.grid, self.velocity_set, self.backend)
+
+    def setup_stepper(self, omega):
+        self.stepper = IncompressibleNavierStokesStepper(
+            omega, boundary_conditions=self.boundary_conditions
+        )
+
+    def run(self, num_steps):
+        for i in range(num_steps):
+            self.f_1 = self.stepper(self.f_0, self.f_1, self.boundary_mask, self.missing_mask, i)
+            self.f_0, self.f_1 = self.f_1, self.f_0
+
+    def post_process(self, i):
+        # Write the results. We'll use JAX backend for the post-processing
+        if not isinstance(self.f_0, jnp.ndarray):
+            self.f_0 = wp.to_jax(self.f_0)
+
+        macro = Macroscopic(compute_backend=ComputeBackend.JAX)
+
+        rho, u = macro(self.f_0)
+
+        # remove boundary cells
+        rho = rho[:, 1:-1, 1:-1]
+        u = u[:, 1:-1, 1:-1]
+        u_magnitude = (u[0] ** 2 + u[1] ** 2) ** 0.5
+
+        fields = {"rho": rho[0], "u_x": u[0], "u_y": u[1], "u_magnitude": u_magnitude}
+
+        save_fields_vtk(fields, timestep=i, prefix="lid_driven_cavity")
+        save_image(fields["u_magnitude"], timestep=i, prefix="lid_driven_cavity")
+
+
+if __name__ == "__main__":
+    # Running the simulation
+    grid_size = 128
+    grid_shape = (grid_size, grid_size)
+    backend = ComputeBackend.JAX
+    velocity_set = xlb.velocity_set.D2Q9()
+    precision_policy = PrecisionPolicy.FP32FP32
+    omega = 1.6
+
+    simulation = LidDrivenCavity2D(omega, grid_shape, velocity_set, backend, precision_policy)
+    simulation.run(num_steps=500)
+    simulation.post_process(i=500)