Exploration of Unstructured Grid Face Node Indexing using Numba

[philipc2]

Exploration of Unstructured Grid Face Node Indexing using Numba

Overview

The connectivity used to represent the face node connectivity of a mixed-topology grid in the UGRID conventions is stored in a dense array, with fill values used to pad faces with less than the maximum number of nodes. This means that applying an operation (or simply accessing the indices of the nodes of each face) requires the manipulation and indexing of the face_node_connectivity variable. Below are various ways of accessing the nodes for each face, with an arbitrary result array used to store the average of index value, meant to simulate some operation.

Datasets

Element / Resolution	30km	15km	7.5km	3.75km
Faces	655,362	2,621,442	10,485,762	41,943,042
Nodes	1,310,720	5,242,880	20,971,520	83,886,080

Each face has a maximum of 6 nodes.

Machine

I ran the following tests on an Macbook Pro with an M1 pro chip and 32gb of memory. The data was able to fit comfortably into memory for each dataset.

Functions

Pure Python Loops

import numpy as np
def _index(face_node_conn, n_nodes_per_face, n_face):
    result = np.empty(n_face, dtype=np.float64)
    for i in range(n_face):
        face_nodes = face_node_conn[i, :n_nodes_per_face[i]]
        result[i] = face_nodes.mean()
    return result

Numba (@njit)

import numpy as np
from numba import njit

@njit
def _numba_index(face_node_conn, n_nodes_per_face, n_face):
    result = np.empty(n_face, dtype=np.float64)
    for i in range(n_face):
        face_nodes = face_node_conn[i, :n_nodes_per_face[i]]
        result[i] = face_nodes.mean()
    return result

Parallel Numba (@njit(parallel =True)

import numpy as np
from numba import njit

@njit(parallel=True)
def _numba_index_par(face_node_conn, n_nodes_per_face, n_face):
    result = np.empty(n_face, dtype=np.float64)
    for i in prange(n_face):
        face_nodes = face_node_conn[i, :n_nodes_per_face[i]]
        result[i] = face_nodes.mean()
    return result

Vectorized Numpy with Partitioned Connectivity

This implementation uses a pre-partitioned face_node_connectivity to avoid needing to use the n_nodes_per_face variable. The implementation is being worked on in #978

import numpy as np 
def _vectorized_index(partitioned_face_node_connectivity, n_face):    
    result = np.empty(n_face, dtype=np.float64) 
    for (
        cur_face_node_partition,
        cur_original_face_indices,
    ) in partitioned_face_node_connectivity:
        
        result[cur_original_face_indices] = cur_face_node_partition.mean(axis=1)
    return result

Results

[philipc2]

@hongyuchen1030

It would seem that if we can entirely write certain spherical geometry helpers using pure Numba, we would be able to use the parallel for loop with prange to massively speed up the computations, since we can apply some operation onto each face interpedently of each other.

The implementation could look something like:

1. Have a general wrapper that can handle accessing grid information (i.e. uxgrid.face_node_connectivity and calling functions
2. Have a parallel for loop using prange that calls purely numba-decorated helper functions (i.e. gca_gca_intersection or others). All the functions called here should be in pure Numba

[aaronzedwick]

Is the vectorized numpy compatible with the njit wrapper?

[philipc2]

.mean(axis=1) is unsupported in numba since it doesn't accept additional parameters.

numba/numba#1269