refactored numerical_flux.recompute_advective_fluxes; added common operators and slices modules

Author: ray-chew

src/pybella/flow_solver/physics/gas_dynamics/numerical_flux.py

@@ -1,60 +1,30 @@
 # -*- coding: utf-8 -*-
 import numpy as np
-import scipy as sp
 
+from ....utils.operators import create_convolution_kernels, apply_convolution_kernel, apply_u_kernel_convolution, apply_v_kernel_convolution_3d
+from ....utils.slices import get_inner_slice
 
-def recompute_advective_fluxes(flux, Sol, *args, **kwargs):
-    """
-    Recompute the advective fluxes at the cell interfaces, i.e. the faces. This function updates the `flux` container in-place.
-
-    Parameters
-    ----------
-    flux : :py:class:`management.variable.States`
-        Data container for the fluxes at the cell interfaces.
-    Sol : :py:class:`management.variable.States`
-        Data container for the Solution.
-
-    Attention
-    ---------
-    This function is a mess and requires cleaning up.
-
-    """
+def recompute_advective_fluxes(flux, Sol, **kwargs):
+    """Recompute the advective fluxes at the cell interfaces."""
     ndim = Sol.rho.ndim
-    inner_idx = tuple([slice(1, -1)] * ndim)
-
-    if ndim == 2:
-        kernel_u = np.array([[0.5, 1.0, 0.5], [0.5, 1.0, 0.5]])
-        kernel_v = kernel_u.T
-    elif ndim == 3:
-        kernel_u = np.array(
-            [[[1, 2, 1], [2, 4, 2], [1, 2, 1]], [[1, 2, 1], [2, 4, 2], [1, 2, 1]]]
-        )
-        kernel_v = np.swapaxes(kernel_u, 1, 0)
-        kernel_w = np.swapaxes(kernel_u, 2, 0)
-
+    inner_idx = get_inner_slice(ndim)
+    kernels = create_convolution_kernels(ndim)
+    
+    # Handle 3D w-component first
+    if ndim == 3:
         rhoYw = Sol.rhoY * Sol.rhow / Sol.rho
-        flux[2].rhoY[inner_idx] = (
-            sp.signal.fftconvolve(rhoYw, kernel_w, mode="valid") / kernel_w.sum()
-        )
-    else:
-        assert 0, "Unsupported dimension in recompute_advective_flux"
-
+        flux[2].rhoY[inner_idx] = apply_convolution_kernel(rhoYw, kernels['w'])
+    
+    # u-component (all dimensions)
     rhoYu = kwargs.get("u", Sol.rhoY * Sol.rhou / Sol.rho)
-
-    flux[0].rhoY[inner_idx] = np.moveaxis(
-        sp.signal.fftconvolve(rhoYu, kernel_u, mode="valid") / kernel_u.sum(), 0, -1
-    )
-
+    flux[0].rhoY[inner_idx] = apply_u_kernel_convolution(rhoYu, kernels['u'])
+    
+    # v-component (all dimensions)
     rhoYv = kwargs.get("v", Sol.rhoY * Sol.rhov / Sol.rho)
     if ndim == 2:
-        flux[1].rhoY[inner_idx] = (
-            sp.signal.fftconvolve(rhoYv, kernel_v, mode="valid") / kernel_v.sum()
-        )
+        flux[1].rhoY[inner_idx] = apply_convolution_kernel(rhoYv, kernels['v'])
     elif ndim == 3:
-        flux[1].rhoY[inner_idx] = np.moveaxis(
-            sp.signal.fftconvolve(rhoYv, kernel_v, mode="valid") / kernel_v.sum(), -1, 0
-        )
-    # flux[1].rhoY[...,-1] = 0.
+        flux[1].rhoY[inner_idx] = apply_v_kernel_convolution_3d(rhoYv, kernels['v'])
 
 
 def hll_solver(flux, Lefts, Rights, Sol, lmbda, ud, th):

src/pybella/utils/operators.py

@@ -0,0 +1,138 @@
+import numpy as np
+import scipy as sp
+from numba import njit
+from functools import lru_cache
+
+
+@lru_cache(maxsize=2)
+def create_convolution_kernels(ndim):
+    """Create convolution kernels for advective flux computation.
+    
+    Parameters
+    ----------
+    ndim : int
+        Number of dimensions (2 or 3)
+        
+    Returns
+    -------
+    dict
+        Dictionary containing kernels for each direction
+        
+    Notes
+    -----
+    Results are cached since kernels don't change during computation.
+    """
+    if ndim == 2:
+        kernel_u = np.array([[0.5, 1.0, 0.5], [0.5, 1.0, 0.5]])
+        return {
+            'u': kernel_u,
+            'v': kernel_u.T
+        }
+    elif ndim == 3:
+        kernel_u = np.array([
+            [[1, 2, 1], [2, 4, 2], [1, 2, 1]], 
+            [[1, 2, 1], [2, 4, 2], [1, 2, 1]]
+        ])
+        return {
+            'u': kernel_u,
+            'v': np.swapaxes(kernel_u, 1, 0),
+            'w': np.swapaxes(kernel_u, 2, 0)
+        }
+    else:
+        raise ValueError(f"Unsupported dimension: {ndim}")
+
+
+@njit
+def _numba_convolve_2d(data, kernel):
+    """Numba-compiled 2D convolution for better performance."""
+    data_h, data_w = data.shape
+    kernel_h, kernel_w = kernel.shape
+    
+    result_h = data_h - kernel_h + 1
+    result_w = data_w - kernel_w + 1
+    result = np.zeros((result_h, result_w))
+    
+    for i in range(result_h):
+        for j in range(result_w):
+            for ki in range(kernel_h):
+                for kj in range(kernel_w):
+                    result[i, j] += data[i + ki, j + kj] * kernel[ki, kj]
+    
+    return result
+
+
+@njit
+def _numba_convolve_3d(data, kernel):
+    """Numba-compiled 3D convolution for better performance."""
+    data_d, data_h, data_w = data.shape
+    kernel_d, kernel_h, kernel_w = kernel.shape
+    
+    result_d = data_d - kernel_d + 1
+    result_h = data_h - kernel_h + 1
+    result_w = data_w - kernel_w + 1
+    result = np.zeros((result_d, result_h, result_w))
+    
+    for i in range(result_d):
+        for j in range(result_h):
+            for k in range(result_w):
+                for ki in range(kernel_d):
+                    for kj in range(kernel_h):
+                        for kk in range(kernel_w):
+                            result[i, j, k] += data[i + ki, j + kj, k + kk] * kernel[ki, kj, kk]
+    
+    return result
+
+
+def apply_convolution_kernel(data, kernel, normalize=True, axis_swap=None, use_numba=True):
+    """Apply convolution kernel with optional normalization and axis swapping.
+    
+    Parameters
+    ----------
+    data : np.ndarray
+        Input data array
+    kernel : np.ndarray
+        Convolution kernel
+    normalize : bool, default=True
+        Whether to normalize by kernel sum
+    axis_swap : tuple or None, default=None
+        Tuple of (from_axis, to_axis) for np.moveaxis
+    use_numba : bool, default=True
+        Whether to use Numba-compiled convolution (faster for repeated calls)
+        
+    Returns
+    -------
+    np.ndarray
+        Convolved result
+        
+    Notes
+    -----
+    For large arrays or single calls, scipy.signal.fftconvolve might be faster.
+    For repeated calls on smaller arrays, Numba convolution is typically faster.
+    """
+    if use_numba and data.ndim in (2, 3):
+        if data.ndim == 2:
+            result = _numba_convolve_2d(data, kernel)
+        else:  # 3D
+            result = _numba_convolve_3d(data, kernel)
+    else:
+        # Fallback to scipy for other dimensions or when requested
+        result = sp.signal.fftconvolve(data, kernel, mode="valid")
+    
+    if normalize:
+        result = result / kernel.sum()
+    
+    if axis_swap is not None:
+        result = np.moveaxis(result, axis_swap[0], axis_swap[1])
+    
+    return result
+
+
+# Convenience functions for specific operations
+def apply_u_kernel_convolution(data, kernel_u, normalize=True):
+    """Apply u-direction kernel with standard axis swap."""
+    return apply_convolution_kernel(data, kernel_u, normalize=normalize, axis_swap=(0, -1))
+
+
+def apply_v_kernel_convolution_3d(data, kernel_v, normalize=True):
+    """Apply v-direction kernel for 3D with standard axis swap."""
+    return apply_convolution_kernel(data, kernel_v, normalize=normalize, axis_swap=(-1, 0))

src/pybella/utils/slices.py

@@ -0,0 +1,5 @@
+
+
+def get_inner_slice(ndim):
+    """Get slice tuple for inner cells (excluding outermost ghost cells)."""
+    return tuple([slice(1, -1)] * ndim)