Merge pull request #55 from mehdiataei/major-refactoring

Added ruff
Autodesk · Aug 2, 2024 · e07976e · e07976e
2 parents 0da352b + aeb1770
commit e07976e
Show file tree

Hide file tree

Showing 78 changed files with 415 additions and 823 deletions.
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
@@ -0,0 +1,27 @@
+name: Lint
+
+on:
+  pull_request:
+    branches:
+      - major-refactoring  # Remember to add main branch later
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Check out code
+      uses: actions/checkout@v2
+
+    - name: Set up Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: '3.11'
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install ruff
+
+    - name: Run Ruff
+      run: ruff check .
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -0,0 +1,6 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.5.6
+    hooks:
+      - id: ruff
+        args: [--fix]
diff --git a/examples/cfd/flow_past_sphere_3d.py b/examples/cfd/flow_past_sphere_3d.py
@@ -29,9 +29,7 @@ def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy):
         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.grid, self.f_0, self.f_1, self.missing_mask, self.boundary_mask = create_nse_fields(grid_shape)
         self.stepper = None
         self.boundary_conditions = []
 
@@ -61,10 +59,7 @@ def define_boundary_indices(self):
         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
+            (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)]
 
@@ -84,23 +79,17 @@ def setup_boundary_masks(self):
             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)
-        )
+        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
-        )
+        self.stepper = IncompressibleNavierStokesStepper(omega, boundary_conditions=self.boundary_conditions)
 
     def run(self, num_steps, post_process_interval=100):
         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_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
 
             if i % post_process_interval == 0 or i == num_steps - 1:
@@ -134,7 +123,5 @@ def post_process(self, i):
     precision_policy = PrecisionPolicy.FP32FP32
     omega = 1.6
 
-    simulation = FlowOverSphere(
-        omega, grid_shape, velocity_set, backend, precision_policy
-    )
+    simulation = FlowOverSphere(omega, grid_shape, velocity_set, backend, precision_policy)
     simulation.run(num_steps=10000, post_process_interval=1000)
diff --git a/examples/cfd/lid_driven_cavity_2d.py b/examples/cfd/lid_driven_cavity_2d.py
@@ -24,9 +24,7 @@ def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy):
         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.grid, self.f_0, self.f_1, self.missing_mask, self.boundary_mask = create_nse_fields(grid_shape)
         self.stepper = None
         self.boundary_conditions = []
 
@@ -42,9 +40,7 @@ def _setup(self, 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]
+            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
@@ -61,23 +57,17 @@ def setup_boundary_masks(self):
             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
-        )
+        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
-        )
+        self.stepper = IncompressibleNavierStokesStepper(omega, boundary_conditions=self.boundary_conditions)
 
     def run(self, num_steps, post_process_interval=100):
         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_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
 
             if i % post_process_interval == 0 or i == num_steps - 1:
@@ -114,7 +104,5 @@ def post_process(self, i):
     precision_policy = PrecisionPolicy.FP32FP32
     omega = 1.6
 
-    simulation = LidDrivenCavity2D(
-        omega, grid_shape, velocity_set, backend, precision_policy
-    )
+    simulation = LidDrivenCavity2D(omega, grid_shape, velocity_set, backend, precision_policy)
     simulation.run(num_steps=5000, post_process_interval=1000)
diff --git a/examples/cfd/lid_driven_cavity_2d_distributed.py b/examples/cfd/lid_driven_cavity_2d_distributed.py
@@ -11,24 +11,24 @@ def __init__(self, omega, grid_shape, velocity_set, backend, precision_policy):
         super().__init__(omega, grid_shape, velocity_set, backend, precision_policy)
 
     def setup_stepper(self, omega):
-        stepper = IncompressibleNavierStokesStepper(
-            omega, boundary_conditions=self.boundary_conditions
-            )
+        stepper = IncompressibleNavierStokesStepper(omega, boundary_conditions=self.boundary_conditions)
         distributed_stepper = distribute(
-            stepper, self.grid, self.velocity_set,
-            )
+            stepper,
+            self.grid,
+            self.velocity_set,
+        )
         self.stepper = distributed_stepper
         return
-    
+
 
 if __name__ == "__main__":
     # Running the simulation
     grid_size = 512
     grid_shape = (grid_size, grid_size)
-    backend = ComputeBackend.JAX            # Must be JAX for distributed multi-GPU computations. Distributed computations on WARP are not supported yet!
+    backend = ComputeBackend.JAX  # Must be JAX for distributed multi-GPU computations. Distributed computations on WARP are not supported yet!
     velocity_set = xlb.velocity_set.D2Q9()
     precision_policy = PrecisionPolicy.FP32FP32
-    omega=1.6
+    omega = 1.6
 
     simulation = LidDrivenCavity2D_distributed(omega, grid_shape, velocity_set, backend, precision_policy)
     simulation.run(num_steps=5000, post_process_interval=1000)
diff --git a/examples/cfd/windtunnel_3d.py b/examples/cfd/windtunnel_3d.py
@@ -19,9 +19,7 @@
 
 
 class WindTunnel3D:
-    def __init__(
-        self, omega, wind_speed, grid_shape, velocity_set, backend, precision_policy
-    ):
+    def __init__(self, omega, wind_speed, grid_shape, velocity_set, backend, precision_policy):
         # initialize backend
         xlb.init(
             velocity_set=velocity_set,
@@ -33,9 +31,7 @@ def __init__(
         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.grid, self.f_0, self.f_1, self.missing_mask, self.boundary_mask = create_nse_fields(grid_shape)
         self.stepper = None
         self.boundary_conditions = []
 
@@ -72,10 +68,8 @@ def define_boundary_indices(self):
         car_length_lbm_unit = grid_size_x / 4
         car_voxelized, pitch = self.voxelize_stl(stl_filename, car_length_lbm_unit)
 
-        car_area = np.prod(car_voxelized.shape[1:])
-        tx, ty, tz = (
-            np.array([grid_size_x, grid_size_y, grid_size_z]) - car_voxelized.shape
-        )
+        # car_area = np.prod(car_voxelized.shape[1:])
+        tx, ty, _ = np.array([grid_size_x, grid_size_y, grid_size_z]) - car_voxelized.shape
         shift = [tx // 4, ty // 2, 0]
         car = np.argwhere(car_voxelized) + shift
         car = np.array(car).T
@@ -97,31 +91,23 @@ def setup_boundary_masks(self):
             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)
-        )
+        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, collision_type="KBC"
-        )
+        self.stepper = IncompressibleNavierStokesStepper(omega, boundary_conditions=self.boundary_conditions, collision_type="KBC")
 
     def run(self, num_steps, print_interval, post_process_interval=100):
         start_time = time.time()
         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_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
 
             if (i + 1) % print_interval == 0:
                 elapsed_time = time.time() - start_time
-                print(
-                    f"Iteration: {i+1}/{num_steps} | Time elapsed: {elapsed_time:.2f}s"
-                )
+                print(f"Iteration: {i + 1}/{num_steps} | Time elapsed: {elapsed_time:.2f}s")
 
                 if i % post_process_interval == 0 or i == num_steps - 1:
                     self.post_process(i)
@@ -178,7 +164,5 @@ def post_process(self, i):
     print(f"Max iterations: {num_steps}")
     print("\n" + "=" * 50 + "\n")
 
-    simulation = WindTunnel3D(
-        omega, wind_speed, grid_shape, velocity_set, backend, precision_policy
-    )
+    simulation = WindTunnel3D(omega, wind_speed, grid_shape, velocity_set, backend, precision_policy)
     simulation.run(num_steps, print_interval, post_process_interval=1000)
diff --git a/examples/cfd_old_to_be_migrated/flow_past_sphere.py b/examples/cfd_old_to_be_migrated/flow_past_sphere.py
@@ -22,8 +22,8 @@
 
 from xlb.operator import Operator
 
-class UniformInitializer(Operator):
 
+class UniformInitializer(Operator):
     def _construct_warp(self):
         # Construct the warp kernel
         @wp.kernel
@@ -149,48 +149,27 @@ def warp_implementation(self, rho, u, vel):
     y = np.arange(nr)
     z = np.arange(nr)
     X, Y, Z = np.meshgrid(x, y, z)
-    indices = np.where(
-        (X - nr // 2) ** 2 + (Y - nr // 2) ** 2 + (Z - nr // 2) ** 2
-        < sphere_radius**2
-    )
+    indices = np.where((X - nr // 2) ** 2 + (Y - nr // 2) ** 2 + (Z - nr // 2) ** 2 < sphere_radius**2)
     indices = np.array(indices).T
     indices = wp.from_numpy(indices, dtype=wp.int32)
 
     # Set boundary conditions on the indices
-    boundary_mask, missing_mask = indices_boundary_masker(
-        indices,
-        half_way_bc.id,
-        boundary_mask,
-        missing_mask,
-        (0, 0, 0)
-    )
+    boundary_mask, missing_mask = indices_boundary_masker(indices, half_way_bc.id, boundary_mask, missing_mask, (0, 0, 0))
 
     # Set inlet bc
     lower_bound = (0, 0, 0)
     upper_bound = (0, nr, nr)
     direction = (1, 0, 0)
     boundary_mask, missing_mask = planar_boundary_masker(
-        lower_bound,
-        upper_bound,
-        direction,
-        equilibrium_bc.id,
-        boundary_mask,
-        missing_mask,
-        (0, 0, 0)
+        lower_bound, upper_bound, direction, equilibrium_bc.id, boundary_mask, missing_mask, (0, 0, 0)
     )
 
     # Set outlet bc
-    lower_bound = (nr-1, 0, 0)
-    upper_bound = (nr-1, nr, nr)
+    lower_bound = (nr - 1, 0, 0)
+    upper_bound = (nr - 1, nr, nr)
     direction = (-1, 0, 0)
     boundary_mask, missing_mask = planar_boundary_masker(
-        lower_bound,
-        upper_bound,
-        direction,
-        do_nothing_bc.id,
-        boundary_mask,
-        missing_mask,
-        (0, 0, 0)
+        lower_bound, upper_bound, direction, do_nothing_bc.id, boundary_mask, missing_mask, (0, 0, 0)
     )
 
     # Set initial conditions
@@ -201,7 +180,7 @@ def warp_implementation(self, rho, u, vel):
     plot_freq = 512
     save_dir = "flow_past_sphere"
     os.makedirs(save_dir, exist_ok=True)
-    #compute_mlup = False # Plotting results
+    # compute_mlup = False # Plotting results
     compute_mlup = True
     num_steps = 1024 * 8
     start = time.time()
@@ -225,4 +204,4 @@ def warp_implementation(self, rho, u, vel):
     end = time.time()
 
     # Print MLUPS
-    print(f"MLUPS: {num_steps*nr**3/(end-start)/1e6}")
+    print(f"MLUPS: {num_steps * nr**3 / (end - start) / 1e6}")