Optimize point_within_gca by Eliminating Redundant Lat/Lon Conversions

[amberchen122]

Closes #936

Overview

This PR addresses a performance issue in the point_within_gca function, which unnecessarily re-converted lat/lon coordinates from Cartesian coordinates. Since Uxarray already stores lat/lon coordinates, they can be passed directly to avoid redundant computation. This change reduces time and space complexity while better aligning with the architecture.

Proposed changes:

• Root cause: change point_within_gca to accept lat_lon coordinates
• Upsteam functions that called point_within_gca:

  • gca_gca_intersection

  • gca_constLat_intersection

  • _populate_face_latlon_bound

  • test files

PR Checklist

General

• An issue is linked created and linked
• Add appropriate labels
• Filled out Overview and Expected Usage (if applicable) sections

Testing

• Adequate tests are created if there is new functionality
• Tests cover all possible logical paths in your function
• Tests are not too basic (such as simply calling a function and nothing else)

Documentation

• Docstrings have been added to all new functions
• Docstrings have updated with any function changes
• Internal functions have a preceding underscore (_) and have been added to docs/internal_api/index.rst
• User functions have been added to docs/user_api/index.rst

Examples

• Any new notebook examples added to docs/examples/ folder
• Clear the output of all cells before committing
• New notebook files added to docs/examples.rst toctree
• New notebook files added to new entry in docs/gallery.yml with appropriate thumbnail photo in docs/_static/thumbnails/

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[amberchen122]

@philipc2 As we discussed in today's meeting, could you please suggest the new function API/signature for point_within_gca?

I will then proceed with changing the function bodies :)

[philipc2]

Hi @amberchen122

This is one of the main sections we'd need to update:

uxarray/uxarray/grid/arcs.py

Lines 170 to 177 in 82e46b8

	pt_lonlat = np.array(_xyz_to_lonlat_rad_no_norm(pt[0], pt[1], pt[2]))
	GCRv0_lonlat = np.array(
	_xyz_to_lonlat_rad_no_norm(gca_cart[0][0], gca_cart[0][1], gca_cart[0][2])
	)
	GCRv1_lonlat = np.array(
	_xyz_to_lonlat_rad_no_norm(gca_cart[1][0], gca_cart[1][1], gca_cart[1][2])
	)
	gca_cart = np.asarray(gca_cart)

We should update the helper functions to take in both the Cartesian and LatLon versions of each point/gca. This may look like the following:

point_within_gca(pt_cart, pt_latlon_rad, gca_cart, gca_latlon_rad, is_directed=False)

In addition to this PR, I will look into getting #878 finished this week, which will ensure that the coordinates will always be normalized. I'll request a review from you once this is prepared.

[philipc2]

I'd like to also mention that we should try to avoid as many conversions between np.ndarray and list as possible. This could be something to considering when optimizing.

[hongyuchen1030]

@amberchen122 To answer your question

uxarray/uxarray/grid/intersections.py

Line 13 in 85748f1

def gca_gca_intersection(gca1_cart, gca1_rad, gca2_cart, gca2_rad):

1. When the function is meant to calculate the scalar values, gca1_cart : [1, 3] ， gca2_cart: [1, 3], where gca1 and gca2 mean the two vertices of the GCR respectively. And I was trying this function vectorizable, which means if gca1_cart : [n, 3] ， gca2_cart: [n, 3]then it should be able to calculate n intersection point at the same time . (I forgot if the current implementation met that goal yet)
2. Don't delete the fma flag, it might be very useful in the future

[amberchen122]

@amberchen122 To answer your question

uxarray/uxarray/grid/intersections.py

Line 13 in 85748f1

def gca_gca_intersection(gca1_cart, gca1_rad, gca2_cart, gca2_rad):

1. When the function is meant to calculate the scalar values, gca1_cart : [1, 3] ， gca2_cart: [1, 3], where gca1 and gca2 mean the two vertices of the GCR respectively. And I was trying this function vectorizable, which means if gca1_cart : [n, 3] ， gca2_cart: [n, 3]then it should be able to calculate n intersection point at the same time . (I forgot if the current implementation met that goal yet)
2. Don't delete the fma flag, it might be very useful in the future

Thank you for the clarification!

1. The current gca_gca_intersection implementation is not vectorized yet. Therefore, it currently only deals with gca1_cart : [1, 3] ， gca2_cart: [1, 3] inputs.
2. The current gca_gca_intersection implementation is not utilizing the fma flag.

For this PR, would it be ok for me to update the function docstring to reflect what they are currently implementing? I will put a note in the comment about the modified parts(gca dimension, and fma flag) for future reference.

[hongyuchen1030]

@amberchen122 To answer your question

uxarray/uxarray/grid/intersections.py

Line 13 in 85748f1

def gca_gca_intersection(gca1_cart, gca1_rad, gca2_cart, gca2_rad):

1. When the function is meant to calculate the scalar values, gca1_cart : [1, 3] ， gca2_cart: [1, 3], where gca1 and gca2 mean the two vertices of the GCR respectively. And I was trying this function vectorizable, which means if gca1_cart : [n, 3] ， gca2_cart: [n, 3]then it should be able to calculate n intersection point at the same time . (I forgot if the current implementation met that goal yet)
2. Don't delete the fma flag, it might be very useful in the future

Thank you for the clarification!

1. The current gca_gca_intersection implementation is not vectorized yet. Therefore, it currently only deals with gca1_cart : [1, 3] ， gca2_cart: [1, 3] inputs.
2. The current gca_gca_intersection implementation is not utilizing the fma flag.

For this PR, would it be ok for me to update the function docstring to reflect what they are currently implementing? I will put a note in the comment about the modified parts(gca dimension, and fma flag) for future reference.

sounds good! Yeah always keep the docstring updated with the current implementaton

[github-actions]

ASV Benchmarking

Benchmark Comparison Results

Benchmarks that have improved:

Change	Before [c66286b]	After [38548ec]	Ratio	Benchmark (Parameter)
-	1.95±0.07ms	1.57±0.1ms	0.81	mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('120km')
-	408M	354M	0.87	mpas_ocean.Integrate.peakmem_integrate('480km')

Benchmarks that have stayed the same:

Change	Before [c66286b]	After [38548ec]	Ratio	Benchmark (Parameter)
	380M	379M	1	face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
	381M	380M	1	face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
	386M	383M	0.99	face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
	384M	381M	0.99	face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
	1.38±0s	1.26±0s	0.91	face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
	181±3ms	167±1ms	0.92	face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
	1.64±0.01s	1.54±0s	0.94	face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
	8.07±0.2ms	7.44±0.09ms	0.92	face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
	1.75±0.01s	1.78±0.04s	1.02	import.Imports.timeraw_import_uxarray
	649±3ms	655±10ms	1.01	mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('120km')
	40.8±0.5ms	41.4±0.8ms	1.01	mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('480km')
	506±10μs	510±20μs	1.01	mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('480km')
	1.27±0μs	1.25±0μs	0.98	mpas_ocean.ConstructTreeStructures.time_ball_tree('120km')
	306±8ns	306±5ns	1	mpas_ocean.ConstructTreeStructures.time_ball_tree('480km')
	858±7ns	825±4ns	0.96	mpas_ocean.ConstructTreeStructures.time_kd_tree('120km')
	296±3ns	288±4ns	0.97	mpas_ocean.ConstructTreeStructures.time_kd_tree('480km')
	399M	399M	1	mpas_ocean.GeoDataFrame.peakmem_to_geodataframe('120km', False)
	389M	388M	1	mpas_ocean.GeoDataFrame.peakmem_to_geodataframe('120km', True)
	362M	362M	1	mpas_ocean.GeoDataFrame.peakmem_to_geodataframe('480km', False)
	361M	361M	1	mpas_ocean.GeoDataFrame.peakmem_to_geodataframe('480km', True)
	1.04±0.01s	1.07±0.03s	1.03	mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', False)
	58.2±0.2ms	59.3±3ms	1.02	mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', True)
	78.8±0.3ms	80.5±0.5ms	1.02	mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', False)
	5.61±0.2ms	5.84±0.1ms	1.04	mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', True)
	265M	263M	0.99	mpas_ocean.Gradient.peakmem_gradient('120km')
	241M	241M	1	mpas_ocean.Gradient.peakmem_gradient('480km')
	2.70±0.01ms	2.72±0.02ms	1.01	mpas_ocean.Gradient.time_gradient('120km')
	287±2μs	288±2μs	1.01	mpas_ocean.Gradient.time_gradient('480km')
	231±20μs	248±5μs	1.07	mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('120km')
	121±2μs	120±0.9μs	1	mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('480km')
	370M	370M	1	mpas_ocean.Integrate.peakmem_integrate('120km')
	177±1ms	174±2ms	0.98	mpas_ocean.Integrate.time_integrate('120km')
	11.9±0.02ms	11.9±0.02ms	1	mpas_ocean.Integrate.time_integrate('480km')
	350±4ms	356±2ms	1.02	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'exclude')
	351±5ms	353±2ms	1.01	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'include')
	355±5ms	360±10ms	1.01	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'split')
	23.2±0.3ms	23.6±0.6ms	1.02	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'exclude')
	22.9±0.7ms	23.0±0.3ms	1	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'include')
	23.2±0.3ms	23.4±0.2ms	1.01	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'split')
	54.7±0.08ms	54.8±0.4ms	1	mpas_ocean.RemapDownsample.time_inverse_distance_weighted_remapping
	44.4±0.08ms	44.6±0.3ms	1	mpas_ocean.RemapDownsample.time_nearest_neighbor_remapping
	360±0.8ms	359±0.8ms	1	mpas_ocean.RemapUpsample.time_inverse_distance_weighted_remapping
	264±1ms	263±0.6ms	1	mpas_ocean.RemapUpsample.time_nearest_neighbor_remapping
	238M	241M	1.01	quad_hexagon.QuadHexagon.peakmem_open_dataset
	236M	236M	1	quad_hexagon.QuadHexagon.peakmem_open_grid
	6.57±0.04ms	6.66±0.2ms	1.01	quad_hexagon.QuadHexagon.time_open_dataset
	5.59±0.07ms	5.76±0.2ms	1.03	quad_hexagon.QuadHexagon.time_open_grid