Merge pull request #8 from FESOM/numpy

Numpy

Author: kacpnowak

docs/implicit_filter.rst

@@ -21,6 +21,14 @@ implicit\_filter.cupy_filter module
    :undoc-members:
    :show-inheritance:
 
+implicit\_filter.numpy_filter module
+------------------------------------
+
+.. automodule:: implicit_filter.numpy_filter
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 implicit\_filter.filter module
 ------------------------------
 

implicit_filter/__init__.py

@@ -1,4 +1,5 @@
 from .jax_filter import JaxFilter
+from .numpy_filter import NumpyFilter
 from ._auxiliary import make_tri, convert_to_wavenumbers
 
 # If CuPy is not installed this class won't be imported

implicit_filter/_numpy_functions.py

@@ -0,0 +1,115 @@
+import numpy as np
+
+
+def make_smooth(elem_area: np.ndarray, dx: np.ndarray, dy: np.ndarray, nn_num: np.ndarray,
+                nn_pos: np.ndarray, tri: np.ndarray, n2d: int, e2d: int):
+    """
+    Calculate the smoothness matrix and metric matrix for a given mesh.
+    It does not support metric terms
+
+    Parameters:
+    ----------
+    elem_area : np.ndarray
+        A 1D NumPy array containing the area of each triangle.
+
+    dx : np.ndarray
+        A 2D NumPy array of shape (e2d, 3) containing x-derivatives of P1 basis functions.
+
+    dy : np.ndarray
+        A 2D NumPy array of shape (e2d, 3) containing y-derivatives of P1 basis functions.
+
+    nn_num : np.ndarray
+        A 1D JAX NumPy array of shape (n2d,) containing the number of neighboring nodes for each node.
+
+    nn_pos : jnp.ndarray
+        A 2D JAX NumPy array of shape (max_neighboring_nodes, n2d) containing positions of neighboring nodes.
+
+    tri : jnp.ndarray
+        A 2D JAX NumPy array representing the connectivity of triangles in the mesh.
+
+    n2d : int
+        The total number of nodes in the mesh.
+
+    e2d : int
+        The total number of triangles (elements) in the mesh.
+
+    full : bool, optional
+        A flag indicating whether to use the 'full' calculation including metric factors (True) or not (False).
+        Default is True.
+
+    Returns:
+    -------
+    smooth_m : np.ndarray
+        A 2D NumPy array of shape (max_neighboring_nodes, n2d) containing the smoothness matrix.
+
+    """
+
+    smooth_m = np.zeros(nn_pos.shape)  # Place for non - zero entries
+    aux = np.zeros([n2d], dtype=int)  # auxiliary array
+    for j in range(e2d):
+        enodes = tri[j, :]  # vertices of triangle
+        hh = np.sqrt(2 * elem_area[j])
+        # This expression corresponds to the specific mesh above obtained
+        # by splitting quads. A slightly different expression is needed
+        # for equilateral triangles.
+
+        for n in range(3):  # Each contributes to rows
+            row = enodes[n]  # and columns
+            cc = nn_num[row]  # num of neighbors
+            tmp = nn_pos[0:cc, row]
+            aux[tmp] = np.arange(0, cc)  # Map their order
+            for m in range(3):
+                col = enodes[m]
+                pos = aux[col]  # Position of column among neighbors
+                tmp_x = dx[j, m] * dx[j, n]
+                tmp_y = dy[j, n] * dy[j, m]
+                smooth_m[pos, row] += (tmp_x + tmp_y) * elem_area[j]
+
+    return smooth_m
+
+
+def make_smat(nn_pos: np.ndarray, nn_num: np.ndarray, smooth_m: np.ndarray, n2d: int, nza: int):
+    """
+    Convert the smoothness matrix into a redundant sparse form (s(k), i(k), j(k)) as required by scipy.
+
+    Parameters:
+    ----------
+    nn_pos : np.ndarray
+        A 2D NumPy array of shape (max_neighboring_nodes, n2d) containing positions of neighboring nodes.
+
+    nn_num : np.ndarray
+        A 1D NumPy array of shape (n2d,) containing the number of neighboring nodes for each node.
+
+    smooth_m : np.ndarray
+        A 2D NumPy array of shape (max_neighboring_nodes, n2d) containing the smoothness matrix.
+
+    n2d : int
+        The total number of nodes in the mesh.
+
+    nza : int
+        The total number of nonzero elements.
+
+    Returns:
+    -------
+    ss : np.ndarray
+        A 1D NumPy array of shape (nza,) containing the nonzero entries of the sparse matrix.
+
+    ii : np.ndarray
+        A 1D NumPy array of shape (nza,) containing the row indices of the nonzero entries.
+
+    jj : np.ndarray
+        A 1D NumPy array of shape (nza,) containing the column indices of the nonzero entries.
+    """
+    nza = np.sum(nn_num)  # The number of nonzero elements:
+    ss = np.zeros([nza])  # Place for nonzero entries
+    ii = np.zeros([nza])  # Place for their rows
+    jj = np.zeros([nza])  # Place for their columns
+    nz = 0
+    for n in range(n2d):
+        for m in range(nn_num[n]):
+            ss[nz] = smooth_m[m, n]
+            ii[nz] = n
+            jj[nz] = nn_pos[m, n]
+            nz += 1
+
+    return ss, ii, jj

implicit_filter/jax_filter.py

@@ -165,7 +165,7 @@ def prepare(self, n2d: int, e2d: int, tri: np.ndarray, xcoord: np.ndarray, ycoor
         self._ne_pos = jnp.array(ne_pos)
         self._area = jnp.array(area)
 
-        smooth, metric = make_smooth(jMt, self._elem_area, self._dx, self._dy, jnn_num, jnn_pos, jtri, n2d, e2d, False)
+        smooth, metric = make_smooth(jMt, self._elem_area, self._dx, self._dy, jnn_num, jnn_pos, jtri, n2d, e2d, full)
 
         smooth = vmap(lambda n: smooth[:, n] / self._area[n])(jnp.arange(0, n2d)).T
         metric = vmap(lambda n: metric[:, n] / self._area[n])(jnp.arange(0, n2d)).T

implicit_filter/numpy_filter.py

@@ -0,0 +1,86 @@
+import math
+from typing import Tuple
+
+import numpy as np
+from scipy.sparse import csc_matrix, identity
+from scipy.sparse.linalg import cg
+
+from implicit_filter._auxiliary import neighbouring_nodes, neighboring_triangles, areas
+from implicit_filter._numpy_functions import make_smooth, make_smat
+from implicit_filter._utils import VeryStupidIdeaError, SolverNotConvergedError
+from implicit_filter.filter import Filter
+
+
+class NumpyFilter(Filter):
+    """
+    A class for filtering data using JAX-based implicit filtering techniques.
+    Extends the base Filter class.
+    """
+
+    def _check_filter_order(self, n: int) -> None:
+        if n < 1:
+            raise ValueError("Filter order must be positive")
+        elif n > 2:
+            raise VeryStupidIdeaError("Filter order too large", ["It really shouldn't be larger than 2"])
+
+    def compute(self, n: int, k: float, data: np.ndarray) -> np.ndarray:
+        self._check_filter_order(n)
+        return self._compute(n, k, data)
+
+    def compute_velocity(self, n: int, k: float, ux: np.ndarray, vy: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        raise NotImplementedError("Filtering non-scalar values are not supported with NumPy Filer")
+
+    def prepare(self, n2d: int, e2d: int, tri: np.ndarray, xcoord: np.ndarray, ycoord: np.ndarray, meshtype: str = 'r',
+                carthesian: bool = False, cyclic_length: float = 360.0 * math.pi / 180.0, full: bool = False):
+        if full:
+            raise NotImplementedError("Computation including metric terms are not supported with NumPy Filer")
+
+        ne_num, ne_pos = neighboring_triangles(n2d, e2d, tri)
+        nn_num, nn_pos = neighbouring_nodes(n2d, tri, ne_num, ne_pos)
+        area, elem_area, dx, dy, Mt = areas(n2d, e2d, tri, xcoord, ycoord, ne_num, ne_pos, meshtype, carthesian,
+                                            cyclic_length)
+
+        self._elem_area = elem_area
+        self._dx = dx
+        self._dy = dy
+        self._ne_num = ne_num
+        self._ne_pos = ne_pos
+        self._area = area
+
+        smooth = make_smooth(self._elem_area, self._dx, self._dy, nn_num, nn_pos, tri, n2d, e2d)
+
+        for i in range(n2d):
+            smooth[:, i] /= self._area[i]
+
+        self._ss, self._ii, self._jj = make_smat(nn_pos, nn_num, smooth, n2d, int(np.sum(nn_num)))
+        self._n2d = n2d
+        self._full = full
+
+    def __transform_atribute(self, atr: str, lmbd, fill=None):
+        """
+        If atribute atr exists then transform it using given Callable lmbd, otherwise it set with fill value
+        """
+        if hasattr(self, atr):
+            setattr(self, atr, lmbd(getattr(self, atr)))
+        else:
+            setattr(self, atr, fill)
+
+    def __init__(self, *initial_data, **kwargs):
+        super().__init__(initial_data, kwargs)
+        # Transform from Numpy array
+        self.__transform_atribute("_n2d", lambda x: int(x), 0)
+        self.__transform_atribute("_full", lambda x: bool(x), False)
+
+    def _compute(self, n, kl, ttu, tol=1e-6, maxiter=150000):
+        Smat1 = csc_matrix((self._ss * (1.0 / np.square(kl)), (self._ii, self._jj)), shape=(self._n2d, self._n2d))
+        Smat = identity(self._n2d) + 0.5 * (Smat1 ** n)
+
+        ttw = ttu - Smat @ ttu  # Work with perturbations
+
+        tts, code = cg(Smat, ttw, tol=tol, maxiter=maxiter)
+        if code != 0:
+            raise SolverNotConvergedError("Solver has not converged without metric terms",
+                                          [f"output code with code: {code}"])
+
+        tts += ttu
+        return np.array(tts)