refactor numerical_flux.py

reworked numerical_flux.recompute_advective_fluxes function too; all tests passed.

Author: ray-chew

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

@@ -100,7 +100,7 @@ def _compute_flux_and_recovery(mem, ud, lmbda, split_step, tag=None):
 
     Lefts, Rights = gd_recovery.do(mem, ud, lmbda, split_step, tag)
 
-    flux = gd_flux.hll_solver(flux, Lefts, Rights, mem.sol, lmbda, ud, mem.th)
+    flux = gd_flux.hll_solver(mem, flux, Lefts, Rights)
 
     return flux
 

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

@@ -1,104 +1,85 @@
 # -*- coding: utf-8 -*-
 import numpy as np
 
-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
+from ....utils.operators import create_convolution_kernels, apply_directional_convolution
+from ....utils.slices import get_inner_slice, get_interface_indices, get_last_dim_inner_slice
 
 def recompute_advective_fluxes(flux, Sol, **kwargs):
-    """Recompute the advective fluxes at the cell interfaces."""
+    """Recompute the advective fluxes at the cell interfaces.
+    
+    Parameters
+    ----------
+    flux : list
+        List of flux objects for each direction
+    Sol : object
+        Solution object containing rho, rhoY, rhou, rhov, rhow fields
+    **kwargs
+        Optional pre-computed velocity components ('u', 'v', 'w')
+    """
     ndim = Sol.rho.ndim
     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] = apply_convolution_kernel(rhoYw, kernels['w'])
+    # Define the component order and corresponding flux indices
+    components = ['u', 'v'] if ndim == 2 else ['u', 'v', 'w']
+    rho_components = ['rhou', 'rhov'] if ndim == 2 else ['rhou', 'rhov', 'rhow']
 
-    # u-component (all dimensions)
-    rhoYu = kwargs.get("u", Sol.rhoY * Sol.rhou / Sol.rho)
-    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] = apply_convolution_kernel(rhoYv, kernels['v'])
-    elif ndim == 3:
-        flux[1].rhoY[inner_idx] = apply_v_kernel_convolution_3d(rhoYv, kernels['v'])
-
+    for i, (comp, rho_comp) in enumerate(zip(components, rho_components)):
+        # Use provided velocity or compute from momentum
+        if comp in kwargs:
+            rhoY_vel = kwargs[comp]
+        else:
+            momentum = getattr(Sol, rho_comp)
+            rhoY_vel = Sol.rhoY * momentum / Sol.rho
+        
+        # Apply directional convolution
+        flux[i].rhoY[inner_idx] = apply_directional_convolution(
+            rhoY_vel, kernels[comp], comp, ndim
+        )
 
-def hll_solver(flux, Lefts, Rights, Sol, lmbda, ud, th):
+def hll_solver(mem, flux, Lefts, Rights):
     """
     HLL solver for the Riemann problem. Chooses the advected quantities from `Lefts` or `Rights` based on the direction given by `flux`.
 
-    Parameters
-    ----------
-    flux : :py:class:`management.variable.States`
-        Data container for fluxes.
-    Lefts : :py:class:`management.variable.States`
-        Container for the quantities on the left of the cell interfaces.
-    Rights : :py:class:`management.variable.States`
-        Container for the quantities on the right of the cell interfaces.
-    Sol : :py:class:`management.variable.Vars`
-        Solution data container.
-    lmbda : float
-        :math:`\\frac{dt}{dx}`, where :math:`dx` is the grid-size in the direction of the substep.
-    ud : :py:class:`inputs.user_data.UserDataInit`
-        Class container for the initial condition.
-    th : :py:class:`physics.gas_dynamics.thermodynamic.init`
-        Class container for the thermodynamical constants.
-
     Returns
     -------
     :py:class:`management.variable.States`
         `flux` data container with the solution of the Riemann problem.
+    
     """
-    # flux: index 1 to end = Left[inner_idx]: index 0 to -1 = Right[inner_idx]: index 1 to end
+    def _compute_flux_component(flux_attr, state_attr=None, state_value=1.0):
+        """Helper function to compute a single flux component."""
+        left_weight = upl[left_idx] / Lefts.Y[left_idx]
+        right_weight = upr[right_idx] / Rights.Y[right_idx]
+        
+        if state_attr is not None:
+            left_val = getattr(Lefts, state_attr)[left_idx]
+            right_val = getattr(Rights, state_attr)[right_idx]
+        else:
+            left_val = right_val = state_value
+        
+        getattr(flux, flux_attr)[remove_cols_idx] = flux.rhoY[remove_cols_idx] * (
+            left_weight * left_val + right_weight * right_val
+        )
 
-    ndim = Sol.rho.ndim
-    left_idx, right_idx, remove_cols_idx = (
-        [slice(None)] * ndim,
-        [slice(None)] * ndim,
-        [slice(None)] * ndim,
-    )
+    ndim = mem.sol.rho.ndim
+    left_idx, right_idx, _ = get_interface_indices(ndim)
+    remove_cols_idx = get_last_dim_inner_slice(ndim)
 
-    remove_cols_idx[-1] = slice(1, -1)
-    left_idx[-1] = slice(0, -1)
-    right_idx[-1] = slice(1, None)
-
-    left_idx, right_idx, remove_cols_idx = (
-        tuple(left_idx),
-        tuple(right_idx),
-        tuple(remove_cols_idx),
-    )
-
-    Lefts.primitives(th)
-    Rights.primitives(th)
+    # Compute primitive variables
+    Lefts.primitives(mem.th)
+    Rights.primitives(mem.th)
 
+    # Compute upwind weights
     upwind = 0.5 * (1.0 + np.sign(flux.rhoY))
     upl = upwind[right_idx]
     upr = 1.0 - upwind[left_idx]
 
-    flux.rhou[remove_cols_idx] = flux.rhoY[remove_cols_idx] * (
-        upl[left_idx] / Lefts.Y[left_idx] * Lefts.u[left_idx]
-        + upr[right_idx] / Rights.Y[right_idx] * Rights.u[right_idx]
-    )
-    flux.rho[remove_cols_idx] = flux.rhoY[remove_cols_idx] * (
-        upl[left_idx] / Lefts.Y[left_idx] * 1.0
-        + upr[right_idx] / Rights.Y[right_idx] * 1.0
-    )
-
-    flux.rhov[remove_cols_idx] = flux.rhoY[remove_cols_idx] * (
-        upl[left_idx] / Lefts.Y[left_idx] * Lefts.v[left_idx]
-        + upr[right_idx] / Rights.Y[right_idx] * Rights.v[right_idx]
-    )
-    flux.rhow[remove_cols_idx] = flux.rhoY[remove_cols_idx] * (
-        upl[left_idx] / Lefts.Y[left_idx] * Lefts.w[left_idx]
-        + upr[right_idx] / Rights.Y[right_idx] * Rights.w[right_idx]
-    )
-    flux.rhoX[remove_cols_idx] = flux.rhoY[remove_cols_idx] * (
-        upl[left_idx] / Lefts.Y[left_idx] * Lefts.X[left_idx]
-        + upr[right_idx] / Rights.Y[right_idx] * Rights.X[right_idx]
-    )
+    # Compute all flux components
+    _compute_flux_component('rhou', 'u')
+    _compute_flux_component('rho')  # Uses default state_value=1.0
+    _compute_flux_component('rhov', 'v')
+    _compute_flux_component('rhow', 'w')
+    _compute_flux_component('rhoX', 'X')
 
-    return flux
+    return flux

src/pybella/utils/operators.py

@@ -3,7 +3,6 @@
 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.
@@ -127,12 +126,49 @@ def apply_convolution_kernel(data, kernel, normalize=True, axis_swap=None, use_n
     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))
-
+# Configuration for directional convolutions
+DIRECTION_CONFIG = {
+    2: {
+        'u': {'axis_swap': (0, -1)},
+        'v': {'axis_swap': None}
+    },
+    3: {
+        'u': {'axis_swap': (0, -1)},
+        'v': {'axis_swap': (-1, 0)},
+        'w': {'axis_swap': None}
+    }
+}
 
-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))
+def apply_directional_convolution(data, kernel, direction, ndim, normalize=True, use_numba=True):
+    """Apply convolution kernel for a specific direction with appropriate axis swapping.
+    
+    Parameters
+    ----------
+    data : np.ndarray
+        Input data array
+    kernel : np.ndarray
+        Convolution kernel for the direction
+    direction : str
+        Direction ('u', 'v', or 'w')
+    ndim : int
+        Number of dimensions (2 or 3)
+    normalize : bool, default=True
+        Whether to normalize by kernel sum
+    use_numba : bool, default=True
+        Whether to use Numba-compiled convolution
+        
+    Returns
+    -------
+    np.ndarray
+        Convolved result with appropriate axis swapping
+    """
+    if direction not in DIRECTION_CONFIG[ndim]:
+        raise ValueError(f"Direction '{direction}' not supported for {ndim}D")
+    
+    config = DIRECTION_CONFIG[ndim][direction]
+    return apply_convolution_kernel(
+        data, kernel, 
+        normalize=normalize, 
+        axis_swap=config['axis_swap'], 
+        use_numba=use_numba
+    )

src/pybella/utils/slices.py

@@ -14,6 +14,12 @@ def get_inner_slice(ndim):
     """Get slice tuple for inner cells (excluding outermost ghost cells)."""
     return tuple([slice(1, -1)] * ndim)
 
+def get_last_dim_inner_slice(ndim):
+    """Get slice for inner faces in the last dimension only."""
+    idx = [slice(None)] * ndim
+    idx[-1] = slice(1, -1)
+    return tuple(idx)
+
 
 def get_interface_indices(ndim):
     """