fixing a bug in array shape's compatibility in ZouHe BC

Autodesk · Feb 1, 2024 · bc346c8 · bc346c8
1 parent 3a7de9c
commit bc346c8
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/examples/CFD/cylinder2d.py b/examples/CFD/cylinder2d.py
@@ -50,8 +50,9 @@ def set_boundary_conditions(self):
 
         # Outflow BC
         outlet = self.boundingBoxIndices['right']
-        rho_outlet = np.ones(outlet.shape[0], dtype=self.precisionPolicy.compute_dtype)
+        rho_outlet = np.ones((outlet.shape[0], 1), dtype=self.precisionPolicy.compute_dtype)
         self.BCs.append(ExtrapolationOutflow(tuple(outlet.T), self.gridInfo, self.precisionPolicy))
+        # self.BCs.append(ZouHe(tuple(outlet.T), self.gridInfo, self.precisionPolicy, 'pressure', rho_outlet))
 
         # Inlet BC
         inlet = self.boundingBoxIndices['left']

diff --git a/src/boundary_conditions.py b/src/boundary_conditions.py
@@ -702,11 +702,11 @@ def calculate_vel(self, fpop, rho):
         """
         Calculate velocity based on the prescribed pressure/density (Zou/He BC)
         """
-        unormal = -1. + 1. / rho * (jnp.sum(fpop[self.indices] * self.imiddleMask, axis=1) +
-                               2. * jnp.sum(fpop[self.indices] * self.iknownMask, axis=1))
+        unormal = -1. + 1. / rho * (jnp.sum(fpop[self.indices] * self.imiddleMask, axis=1, keepdims=True) +
+                               2. * jnp.sum(fpop[self.indices] * self.iknownMask, axis=1, keepdims=True))
 
         # Return the above unormal as a normal vector which sets the tangential velocities to zero
-        vel = unormal[:, None] * self.normals
+        vel = unormal * self.normals
         return vel
 
     @partial(jit, static_argnums=(0,), inline=True)