From 0dabaf517dbf378e27cd8ac19cb52998b2c4c01a Mon Sep 17 00:00:00 2001
From: clearya1 <65236277+clearya1@users.noreply.github.com>
Date: Wed, 3 Apr 2024 12:36:24 +0100
Subject: [PATCH] allowing for 2/3 rule smoothing for non square grids

---
 jax_cfd/spectral/utils.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/jax_cfd/spectral/utils.py b/jax_cfd/spectral/utils.py
index b40ee15..97a6b49 100644
--- a/jax_cfd/spectral/utils.py
+++ b/jax_cfd/spectral/utils.py
@@ -127,10 +127,10 @@ def circular_filter_2d(grid: grids.Grid) -> spectral_types.Array:
 
 def brick_wall_filter_2d(grid: grids.Grid):
   """Implements the 2/3 rule."""
-  n, _ = grid.shape
-  filter_ = jnp.zeros((n, n // 2 + 1))
-  filter_ = filter_.at[:int(2 / 3 * n) // 2, :int(2 / 3 * (n // 2 + 1))].set(1)
-  filter_ = filter_.at[-int(2 / 3 * n) // 2:, :int(2 / 3 * (n // 2 + 1))].set(1)
+  n, m = grid.shape
+  filter_ = jnp.zeros((n, m // 2 + 1))
+  filter_ = filter_.at[:int(2 / 3 * n) // 2, :int(2 / 3 * (m // 2 + 1))].set(1)
+  filter_ = filter_.at[-int(2 / 3 * n) // 2:, :int(2 / 3 * (m // 2 + 1))].set(1)
   return filter_