Merge pull request #13 from FESOM/latlon_support

Add support for lat-lon grids. Fix installation bugs

Author: kacpnowak

examples/lat_long_grid_example.ipynb

@@ -41,8 +41,4 @@ platforms = ["any"]
 zip-safe  = false
 include-package-data = true
 
-[tool.setuptools.packages.find]
-where = ["implicit_filter"]
-include = ["implicit_filter*"]
-
 

pyproject.toml

@@ -7,7 +7,10 @@
 [![arXiv](https://img.shields.io/badge/arXiv-2404.07398-b31b1b.svg)](https://arxiv.org/abs/2404.07398)
 
 The Implicit Filter Python Package provides a collection of classes for filtering data using implicit filtering techniques.
-Currently FESOM and NEMO meshes are supported.
+Currently FESOM, NEMO and longitude-latitude meshes are supported.
+
+For ICON please use this [fork](https://github.com/wienkers/implicit_filter_ICON) made by Aaron Wienkers.
+In the future it will be integrated into the main repository.
 
 For optimal performance usage of Nvidia GPU is highly recommended.
 

readme.md

@@ -1,11 +1,13 @@
 from .jax_filter import JaxFilter
 from .numpy_filter import NumpyFilter
 from .nemo_filter import NemoNumpyFilter
+from .latlon_filter import LatLonNumpyFilter
 from ._auxiliary import make_tri, convert_to_wavenumbers
 
-# If CuPy is not installed this class won't be imported
+# If CuPy isn't installed, this classes won't be imported
 try:
     from .cupy_filter import CuPyFilter
     from .nemo_cupy_filter import NemoCupyFilter
+    from .latlon_cupy_filter import LatLonCupyFilter
 except ModuleNotFoundError:
     pass

implicit_filter/__init__.py src/implicit_filter/__init__.pyimplicit_filter/__init__.py renamed to src/implicit_filter/__init__.py

@@ -154,12 +154,12 @@ def calculate_global_nemo_neighbourhood(e2d: int, nx: int, ny: int, north_adj: p
     xc = nx  # Number of cells in x axis
 
     # Fill ee_pos, arrangement is W;N;E;S
-    ee_pos[:, 0] = np.array([xc * (yc - 1), 1, yc, 0])  # Corner
+    ee_pos[:, 0] = np.array([yc * (xc - 1), 1, yc, 0])  # Corner
     # print(f"x: {0} y: {0} ni: {0}")
     # ids.add(0)
     nza += 3
     for m in range(1, yc - 1):
-        ee_pos[:, m] = [xc * (yc - 1) + m, m + 1, m + yc, m - 1]  # Left border
+        ee_pos[:, m] = [yc * (xc - 1) + m, m + 1, m + yc, m - 1]  # Left border
         # ids.add(m)
         # print(f"x: {0} y: {m} ni: {m}")
         nza += 4
@@ -188,7 +188,7 @@ def calculate_global_nemo_neighbourhood(e2d: int, nx: int, ny: int, north_adj: p
             nza += 4
 
         # print("Top border")
-        # ni = no + (yc - 1)
+        ni = no + (yc - 1)
         # print(f"Northern neighbour of {n + 1} is {north_adj[n + 1]}")
         ee_pos[:, ni] = [ni - yc, north_adj[n + 1] * yc - 1, ni + yc, ni - 1]
 
@@ -221,3 +221,209 @@ def calculate_global_nemo_neighbourhood(e2d: int, nx: int, ny: int, north_adj: p
     nza += e2d
 
     return ee_pos, nza
+
+
+def calculate_global_regular_neighbourhood(e2d: int, nx: int, ny: int) -> Tuple[np.ndarray, int]:
+    """
+    Calculate the neighbourhood of each cell in regular global mesh.
+    The Northern border is ignored
+
+    Parameters:
+    ----------
+    e2d : int
+        The total number of cells (elements) in the mesh.
+
+    nx : int
+        Number of cells in X direction.
+
+    ny : int
+        Number of cells in Y direction.
+
+    Returns:
+    -------
+    ee_pos : np.ndarray
+        2D NumPy array of shape (4, e2d) containing indexes of neighboring cells.
+
+    nza : int
+        The total number of unique elements in ee_pos.
+    """
+
+    ee_pos = np.zeros((4, e2d), dtype=np.int32)
+    # ids = set()
+    # Initialize nza
+    nza = 0
+
+    yc = ny  # Numer of cells in y axis
+    xc = nx  # Number of cells in x axis
+
+    # Fill ee_pos, arrangement is W;N;E;S
+    ee_pos[:, 0] = np.array([yc * (xc - 1), 1, yc, 0])  # Corner
+    # print(f"x: {0} y: {0} ni: {0}")
+    # ids.add(0)
+    nza += 3
+    for m in range(1, yc - 1):
+        ee_pos[:, m] = [yc * (xc - 1) + m, m + 1, m + yc, m - 1]  # Left border
+        # ids.add(m)
+        # print(f"x: {0} y: {m} ni: {m}")
+        nza += 4
+    # print("Left")
+    m = yc - 1
+    # print(f"Northern neighbour of {1} is {north_adj[1]}")
+    ee_pos[:, m] = [xc * yc - 1, m, m + yc, m - 1]  # Second corner
+    # ids.add(m)
+    # print(f"x: {0} y: {m} ni: {ni}")
+    nza += 3
+
+    for n in range(1, xc - 1):  # Center
+        # print("Bottom border")
+        no = yc * n
+        ni = no
+        ee_pos[:, ni] = [ni - yc, ni + 1, ni + yc, ni]
+        # if not ni in ids: ids.add(ni)
+        # print(f"x: {n} y: {0} ni: {ni}")
+        nza += 3
+        # print("Inner center")
+        for m in range(1, yc - 1):
+            ni = no + m
+            ee_pos[:, ni] = [ni - yc, ni + 1, ni + yc, ni - 1]
+            # if not ni in ids: ids.add(ni)
+            # print(f"x: {n} y: {m} ni: {ni}")
+            nza += 4
+
+        # print("Top border")
+        ni = no + (yc - 1)
+        # print(f"Northern neighbour of {n + 1} is {north_adj[n + 1]}")
+        ee_pos[:, ni] = [ni - yc, ni, ni + yc, ni - 1]
+
+        # if not ni in ids: ids.add(ni)
+        # print(f"x: {n} y: {yc-1} ni: {ni}")
+        nza += 3
+
+    no = yc * (xc - 1)
+    ni = no
+    ee_pos[:, ni] = [ni - yc, ni + 1, 0, ni]  # Third corner
+    # if not ni in ids: ids.add(ni)
+    # print(f"x: {xc-1} y: {0} ni: {ni}")
+    nza += 3
+    # print("Right")
+    for m in range(1, yc - 1):  # Right border
+        ni = no + m
+        ee_pos[:, ni] = [ni - yc, ni + 1, m, ni - 1]
+        # if not ni in ids: ids.add(ni)
+        # print(f"x: {xc-1} y: {m} ni: {ni}")
+        nza += 4
+
+    ni = no + (yc - 1)
+    # print(f"Northern neighbour of {xc} is {north_adj[xc]}")
+    ee_pos[:, ni] = [ni - yc, ni, yc - 1, ni - 1]  # Fourth border
+    # if not ni in ids: ids.add(ni)
+    # print(f"x: {xc-1} y: {yc-1} ni: {ni}")
+    nza += 3
+
+    # Final adjustment for nza
+    nza += e2d
+
+    return ee_pos, nza
+
+
+def calculate_local_regular_neighbourhood(e2d: int, nx: int, ny: int) -> Tuple[np.ndarray, int]:
+    """
+    Calculate the neighbourhood of each cell in regular local mesh.
+    All periodic neighbours are ignored
+
+    Parameters:
+    ----------
+    e2d : int
+        The total number of cells (elements) in the mesh.
+
+    nx : int
+        Number of cells in X direction.
+
+    ny : int
+        Number of cells in Y direction.
+
+    Returns:
+    -------
+    ee_pos : np.ndarray
+        2D NumPy array of shape (4, e2d) containing indexes of neighboring cells.
+
+    nza : int
+        The total number of unique elements in ee_pos.
+    """
+
+    ee_pos = np.zeros((4, e2d), dtype=np.int32)
+    # ids = set()
+    # Initialize nza
+    nza = 0
+
+    yc = ny  # Numer of cells in y axis
+    xc = nx  # Number of cells in x axis
+
+    # Fill ee_pos, arrangement is W;N;E;S
+    ee_pos[:, 0] = np.array([0, 1, yc, 0])  # Corner
+    # print(f"x: {0} y: {0} ni: {0}")
+    # ids.add(0)
+    nza += 2
+    for m in range(1, yc - 1):
+        ee_pos[:, m] = [m, m + 1, m + yc, m - 1]  # Left border
+        # ids.add(m)
+        # print(f"x: {0} y: {m} ni: {m}")
+        nza += 3
+    # print("Left")
+    m = yc - 1
+    # print(f"Northern neighbour of {1} is {north_adj[1]}")
+    ee_pos[:, m] = [m, m, m + yc, m - 1]  # Second corner
+    # ids.add(m)
+    # print(f"x: {0} y: {m} ni: {ni}")
+    nza += 2
+
+    for n in range(1, xc - 1):  # Center
+        # print("Bottom border")
+        no = yc * n
+        ni = no
+        ee_pos[:, ni] = [ni - yc, ni + 1, ni + yc, ni]
+        # if not ni in ids: ids.add(ni)
+        # print(f"x: {n} y: {0} ni: {ni}")
+        nza += 3
+        # print("Inner center")
+        for m in range(1, yc - 1):
+            ni = no + m
+            ee_pos[:, ni] = [ni - yc, ni + 1, ni + yc, ni - 1]
+            # if not ni in ids: ids.add(ni)
+            # print(f"x: {n} y: {m} ni: {ni}")
+            nza += 4
+
+        # print("Top border")
+        ni = no + (yc - 1)
+        # print(f"Northern neighbour of {n + 1} is {north_adj[n + 1]}")
+        ee_pos[:, ni] = [ni - yc, ni, ni + yc, ni - 1]
+
+        # if not ni in ids: ids.add(ni)
+        # print(f"x: {n} y: {yc-1} ni: {ni}")
+        nza += 3
+
+    no = yc * (xc - 1)
+    ni = no
+    ee_pos[:, ni] = [ni - yc, ni + 1, ni, ni]  # Third corner
+    # if not ni in ids: ids.add(ni)
+    # print(f"x: {xc-1} y: {0} ni: {ni}")
+    nza += 2
+    # print("Right")
+    for m in range(1, yc - 1):  # Right border
+        ni = no + m
+        ee_pos[:, ni] = [ni - yc, ni + 1, ni, ni - 1]
+        # if not ni in ids: ids.add(ni)
+        # print(f"x: {xc-1} y: {m} ni: {ni}")
+        nza += 3
+
+    ni = no + (yc - 1)
+    # print(f"Northern neighbour of {xc} is {north_adj[xc]}")
+    ee_pos[:, ni] = [ni - yc, ni, ni, ni - 1]  # Fourth border
+    # if not ni in ids: ids.add(ni)
+    # print(f"x: {xc-1} y: {yc-1} ni: {ni}")
+    nza += 2
+
+    # Final adjustment for nza
+    nza += e2d
+
+    return ee_pos, nza