Optimize Face Edge Coordinate Construction

[philipc2]

Below are timings taken across multiple grid resolutions on a single NCAR Derecho CPU node.

Resolution	Nodes	Faces	Edges	Grid Load Time (s)	face_edges_xyz (s)
30km	1,310,720	655,362	1,966,080	1.921	0.498
15km	5,242,880	2,621,442	7,864,320	7.458	2.213
7.5km	20,971,520	10,485,762	31,457,280	27.46	6.764
3.75km	83,886,080	41,943,042	125,829,120	113.993	28.683

While the performance here is not bad, it is noticeably slow for the higher resolutions.

@hongyuchen1030 pointed out that we could consider caching this variable for re-use, especially since multiple calls to zonal average could take advantage of this. There is an informative discussion about this in #1180

The current implementation only relies on the face_node_connectivity and uses it to derive the edge information. If the edge_node_connectivity and face_edge_connectivity are available, it may be more efficient to use these to index our coordinates.

uxarray/uxarray/grid/utils.py

Lines 141 to 259 in 3b7e8c4

	def _get_cartesian_face_edge_nodes(
	face_node_conn, n_face, n_max_face_edges, node_x, node_y, node_z
	):
	"""Construct an array to hold the edge Cartesian coordinates connectivity
	for multiple faces in a grid.
	Parameters
	----------
	face_node_conn : np.ndarray
	An array of shape (n_face, n_max_face_edges) containing the node indices for each face. Accessed through `grid.face_node_connectivity.value`.
	n_face : int
	The number of faces in the grid. Accessed through `grid.n_face`.
	n_max_face_edges : int
	The maximum number of edges for any face in the grid. Accessed through `grid.n_max_face_edges`.
	node_x : np.ndarray
	An array of shape (n_nodes,) containing the x-coordinate values of the nodes. Accessed through `grid.node_x`.
	node_y : np.ndarray
	An array of shape (n_nodes,) containing the y-coordinate values of the nodes. Accessed through `grid.node_y`.
	node_z : np.ndarray
	An array of shape (n_nodes,) containing the z-coordinate values of the nodes. Accessed through `grid.node_z`.
	Returns
	-------
	face_edges_cartesian : np.ndarray
	An array of shape (n_face, n_max_face_edges, 2, 3) containing the Cartesian coordinates of the edges
	for each face. It might contain dummy values if the grid has holes.
	Examples
	--------
	>>> face_node_conn = np.array(
	... [
	... [0, 1, 2, 3, 4],
	... [0, 1, 3, 4, INT_FILL_VALUE],
	... [0, 1, 3, INT_FILL_VALUE, INT_FILL_VALUE],
	... ]
	... )
	>>> n_face = 3
	>>> n_max_face_edges = 5
	>>> node_x = np.array([0, 1, 1, 0, 1, 0])
	>>> node_y = np.array([0, 0, 1, 1, 2, 2])
	>>> node_z = np.array([0, 0, 0, 0, 1, 1])
	>>> _get_cartesian_face_edge_nodes(
	... face_node_conn, n_face, n_max_face_edges, node_x, node_y, node_z
	... )
	array([[[[ 0, 0, 0],
	[ 1, 0, 0]],
	[[ 1, 0, 0],
	[ 1, 1, 0]],
	[[ 1, 1, 0],
	[ 0, 1, 0]],
	[[ 0, 1, 0],
	[ 1, 2, 1]],
	[[ 1, 2, 1],
	[ 0, 0, 0]]],
	[[[ 0, 0, 0],
	[ 1, 0, 0]],
	[[ 1, 0, 0],
	[ 0, 1, 0]],
	[[ 0, 1, 0],
	[ 1, 2, 1]],
	[[ 1, 2, 1],
	[ 0, 0, 0]],
	[[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE],
	[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]]],
	[[[ 0, 0, 0],
	[ 1, 0, 0]],
	[[ 1, 0, 0],
	[ 0, 1, 0]],
	[[ 0, 1, 0],
	[ 0, 0, 0]],
	[[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE],
	[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]],
	[[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE],
	[INT_FILL_VALUE, INT_FILL_VALUE, INT_FILL_VALUE]]]])
	"""
	# Shift node connections to create edge connections
	face_node_conn_shift = np.roll(face_node_conn, -1, axis=1)
	# Construct edge connections by combining original and shifted node connections
	face_edge_conn = np.array([face_node_conn, face_node_conn_shift]).T.swapaxes(0, 1)
	# swap the first occurrence of INT_FILL_VALUE with the last value in each sub-array
	face_edge_conn = _swap_first_fill_value_with_last(face_edge_conn)
	# Get the indices of the nodes from face_edge_conn
	face_edge_conn_flat = face_edge_conn.reshape(-1)
	valid_mask = face_edge_conn_flat != INT_FILL_VALUE
	# Get the valid node indices
	valid_edges = face_edge_conn_flat[valid_mask]
	# Create an array to hold the Cartesian coordinates of the edges
	face_edges_cartesian = np.full(
	(len(face_edge_conn_flat), 3), INT_FILL_VALUE, dtype=float
	)
	# Fill the array with the Cartesian coordinates of the edges
	face_edges_cartesian[valid_mask, 0] = node_x[valid_edges]
	face_edges_cartesian[valid_mask, 1] = node_y[valid_edges]
	face_edges_cartesian[valid_mask, 2] = node_z[valid_edges]
	return face_edges_cartesian.reshape(n_face, n_max_face_edges, 2, 3)

[philipc2]

For the MPAS grids above, here is the rough size of each array (with n_max_face_nodes=10)

Cartesian

Resolution	FP64 Values	Size (GB)
30km	39,321,720	0.3146 GB
15km	157,286,520	1.2583 GB
7.5km	629,145,720	5.0332 GB
3.75km	2,516,582,520	20.1327 GB

Spherical

Resolution	FP64 Values	Size (GB)
30km	26,214,480	0.2097 GB
15km	104,857,680	0.8389 GB
7.5km	419,430,480	3.3554 GB
3.75km	1,677,721,680	13.4218 GB

@erogluorhan

These arrays would still fit comfortably in memory if we are working on a single HPC Node (256GB on a Derecho CPU Node).

However, even for the 7.5km resolution, the size of this array could cause issues if we are trying to work on a consumer machine. I wonder if there would be a way to modify our existing Bounds and Zonal Average implementation to not require these arrays to be built in their entirety.

In the zonal average call, instead of constructing the entire array and then indexing, we could instead construct the face_edges for the candidate faces at each step of the loop. This could significantly reduce the maximum memory needed to run the zonal averaging, but may slightly increase the execution time.

uxarray/uxarray/core/zonal.py

Lines 21 to 33 in 3b7e8c4

	faces_edge_nodes_xyz = _get_cartesian_face_edge_nodes(
	uxgrid.face_node_connectivity.values,
	uxgrid.n_face,
	uxgrid.n_max_face_nodes,
	uxgrid.node_x.values,
	uxgrid.node_y.values,
	uxgrid.node_z.values,
	)
	bounds = uxgrid.bounds.values
	for i, lat in enumerate(latitudes):
	face_indices = uxda.uxgrid.get_faces_at_constant_latitude(lat)

[erogluorhan]

Hey @philipc2 , thanks for putting this together, it's very helpful. Please find a few thoughts and questions I have below:

• @hongyuchen1030 pointed out that we could consider caching this variable for re-use, especially since multiple calls to zonal average could take advantage of this. There is an informative discussion about this in DRAFT: Optimized Non-Conservative Zonal Average #1180

  The current implementation only relies on the face_node_connectivity and uses it to derive the edge information. If the edge_node_connectivity and face_edge_connectivity are available, it may be more efficient to use these to index our coordinates.

  When a new data variable is derivable from other attributes of a class (connectivities in this case) and especially when the new ones would introduce comparable or much more memory usage than those attributes (e.g. ~ 4x or 6x more than face_edge_connectivity for instance, though it is not always there in the Grid but can still give us an idea about how much more memory than face_node_connecivity, which would be ~ 2.5x or 4x, respectively, I believe for this MPAS case), it'd be a better choice to avoid added memory consumption (i.e. by not caching / making those vars the new attributes of the class) and focus on making the calculation of those vars as efficient as possible.

  I get that seeing execution times of ~28 seconds for the 3.75 km grid for repeated calculations of such variables would be pretty frustrating though. However, I don't believe caching would be feasible due to the above reasons, and I think we should definitely look into optimizing the performance.

• These arrays would still fit comfortably in memory if we are working on a single HPC Node (256GB on a Derecho CPU Node).

  However, even for the 7.5km resolution, the size of this array could cause issues if we are trying to work on a consumer machine. I wonder if there would be a way to modify our existing Bounds and Zonal Average implementation to not require these arrays to be built in their entirety.

  This is a huge memory consumption. Even regarding the Derecho case, we have to be considerate of an end-to-end workflow, where there will already be memory consumption due to the main tasks such as loading in the grid and data, visualizations, additional analysis variables, etc. Adding these cached variables' memory consumption while they could be derived from the existing Grid attributes does not seem like a feasible way forward.

• In the zonal average call, instead of constructing the entire array and then indexing, we could instead construct the face_edges for the candidate faces at each step of the loop. This could significantly reduce the maximum memory needed to run the zonal averaging, but may slightly increase the execution time.

  This sounds like an interesting path to explore; very cool idea! @hongyuchen1030 may probably have much better thoughts than me about this, but I am happy to further look into the internal code and share my thoughts if needed.

Also some questions:

• Below are timings taken across multiple grid resolutions on a single NCAR Derecho CPU node.

  Sorry for repeatedly asking this type of questions, but I thought grid loading was made much better than this recently (regarding ~114 seconds for 3.75 km grid for instance).

• For the MPAS grids above, here is the rough size of each array (with n_max_face_nodes=10)

  These are only for the variable we are generating (faces_edges_cartesian and faces_edges_spherical), right? If so, could you remind the memory usage for the whole grid and face_node_connectivity again, for 3.75km grid for instance (I believe the latter should be around 5GB)?

[philipc2]

Sorry for repeatedly asking this type of questions, but I thought grid loading was made much better than this recently (regarding ~114 seconds for 3.75 km grid for instance).

No worries! This was without using Dask, so the entire Grid data was loaded into memory, which included all the conversions from UGRID to MPAS.

[philipc2]

@erogluorhan

This is a better representation of the Grid load times.

Resolution	Nodes	Faces	Edges	Grid Load Time (Eager)	Grid Load Time (Dask)	Compute Minimal Variables (node_lat/lon and face_node_conn) (s)
30km	1,310,720	655,362	1,966,080	1.921	0.554	0.401
15km	5,242,880	2,621,442	7,864,320	7.458	0.764	1.613
7.5km	20,971,520	10,485,762	31,457,280	27.46	4.437	5.798
3.75km	83,886,080	41,943,042	125,829,120	113.993	13.772	22.982

[philipc2]

I shared an initial Parallel Numba implementation (which I still need to fully revise to match the current implementation as @hongyuchen1030 pointed out) in #1180.

These are the results for constructing the entire array.

Resolution	Nodes	Faces	Edges	Parallel Numba (s)
30km	1,310,720	655,362	1,966,080	0.05
15km	5,242,880	2,621,442	7,864,320	0.233
7.5km	20,971,520	10,485,762	31,457,280	0.564
3.75km	83,886,080	41,943,042	125,829,120	5.619

About a 5x speedup for the highest resolution, which isn't too shabby.