DRAFT: Optimized Geometry Functions using Numba

[philipc2]

Closes #XXX

Overview

• Rewrites the following functions to avoid needing to convert between Cartesian and Spherical Coordinates
  • point_within_gca()
  • extreme_gca_latitude()
  • gca_gca_intersection()
• Added compatibility functions with less-strict array and parameter requirements to use for testing.
  • _point_within_gca_cartesian()
  • _extreme_gca_latitude_cartesian()
  • _gca_gca_intersection_cartesian()
• Removed the use of the pyfma package.

[philipc2]

@hongyuchen1030

I still need to finish adjusting all the tests and ensure that everything is passing, but I wanted to quickly updated you on some initial timings.

• 30km (~40,000 mostly hexagonal faces): 4-10 seconds including the initial Numba compilation.

This was without any parallelization. Since the bounding box for each face can be computed interpedently, my next step will be to parallelize the loop that calls _populate_face_latlon_bound

[github-actions]

ASV Benchmarking

Benchmark Comparison Results

Benchmarks that have improved:

Change	Before [b11d011]	After [d928082]	Ratio	Benchmark (Parameter)
-	1.57±0.03s	22.3±0.04ms	0.01	face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
-	222±2ms	7.86±0.06ms	0.04	face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
-	1.98±0.01s	49.6±0.2ms	0.02	face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
-	7.71±0.1ms	3.69±0.03ms	0.48	face_bounds.FaceBounds.time_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
-	466M	372M	0.8	mpas_ocean.Integrate.peakmem_integrate('480km')

Benchmarks that have stayed the same:

Change	Before [b11d011]	After [d928082]	Ratio	Benchmark (Parameter)
	400M	419M	1.05	face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/geoflow-small/grid.nc'))
	445M	412M	0.93	face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/ugrid/quad-hexagon/grid.nc'))
	2.79±0.01s	2.79±0.01s	1.00	import.Imports.timeraw_import_uxarray
	662±20μs	656±7μs	0.99	mpas_ocean.CheckNorm.time_check_norm('120km')
	428±5μs	426±3μs	1.00	mpas_ocean.CheckNorm.time_check_norm('480km')
	645±8ms	631±10ms	0.98	mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('120km')
	41.3±0.2ms	40.6±0.3ms	0.98	mpas_ocean.ConnectivityConstruction.time_face_face_connectivity('480km')
	1.69±0.01ms	1.66±0.05ms	0.98	mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('120km')
	5.73±0.04ms	5.72±0.01ms	1.00	mpas_ocean.ConstructFaceLatLon.time_cartesian_averaging('120km')
	3.37±0.05ms	3.41±0.07ms	1.01	mpas_ocean.ConstructFaceLatLon.time_cartesian_averaging('480km')
	3.54±0.03s	3.54±0s	1.00	mpas_ocean.ConstructFaceLatLon.time_welzl('120km')
	226±1ms	224±2ms	0.99	mpas_ocean.ConstructFaceLatLon.time_welzl('480km')
	1.08±0μs	1.10±0μs	1.02	mpas_ocean.ConstructTreeStructures.time_ball_tree('120km')
	271±5ns	272±7ns	1.00	mpas_ocean.ConstructTreeStructures.time_ball_tree('480km')
	273±8ns	265±3ns	0.97	mpas_ocean.ConstructTreeStructures.time_kd_tree('480km')
	failed	failed	n/a	mpas_ocean.CrossSections.time_constant_lat_fast('120km', 1)
	failed	failed	n/a	mpas_ocean.CrossSections.time_constant_lat_fast('120km', 2)
	failed	failed	n/a	mpas_ocean.CrossSections.time_constant_lat_fast('120km', 4)
	failed	failed	n/a	mpas_ocean.CrossSections.time_constant_lat_fast('120km', 8)
	failed	failed	n/a	mpas_ocean.CrossSections.time_constant_lat_fast('480km', 1)
	failed	failed	n/a	mpas_ocean.CrossSections.time_constant_lat_fast('480km', 2)
	failed	failed	n/a	mpas_ocean.CrossSections.time_constant_lat_fast('480km', 4)
	failed	failed	n/a	mpas_ocean.CrossSections.time_constant_lat_fast('480km', 8)
	122±1ms	123±0.5ms	1.01	mpas_ocean.DualMesh.time_dual_mesh_construction('120km')
	8.49±0.07ms	8.37±0.2ms	0.99	mpas_ocean.DualMesh.time_dual_mesh_construction('480km')
	1.01±0.01s	1.02±0.01s	1.02	mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', False)
	51.2±0.4ms	51.4±0.8ms	1.00	mpas_ocean.GeoDataFrame.time_to_geodataframe('120km', True)
	75.7±0.4ms	76.3±0.3ms	1.01	mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', False)
	5.08±0.1ms	4.87±0.06ms	0.96	mpas_ocean.GeoDataFrame.time_to_geodataframe('480km', True)
	319M	318M	1.00	mpas_ocean.Gradient.peakmem_gradient('120km')
	296M	296M	1.00	mpas_ocean.Gradient.peakmem_gradient('480km')
	2.69±0.01ms	2.68±0.01ms	1.00	mpas_ocean.Gradient.time_gradient('120km')
	305±1μs	307±0.9μs	1.01	mpas_ocean.Gradient.time_gradient('480km')
	failed	failed	n/a	mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('120km')
	failed	failed	n/a	mpas_ocean.HoleEdgeIndices.time_construct_hole_edge_indices('480km')
	389M	389M	1.00	mpas_ocean.Integrate.peakmem_integrate('120km')
	174±0.8ms	174±1ms	1.00	mpas_ocean.Integrate.time_integrate('120km')
	12.1±0.1ms	11.8±0.5ms	0.98	mpas_ocean.Integrate.time_integrate('480km')
	336±3ms	332±2ms	0.99	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'exclude')
	335±3ms	334±4ms	1.00	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'include')
	330±2ms	330±3ms	1.00	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('120km', 'split')
	21.1±0.1ms	21.2±0.1ms	1.00	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'exclude')
	21.6±0.1ms	21.5±0.1ms	0.99	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'include')
	21.4±0.09ms	21.2±0.2ms	0.99	mpas_ocean.MatplotlibConversion.time_dataarray_to_polycollection('480km', 'split')
	55.1±0.2ms	55.3±0.2ms	1.00	mpas_ocean.RemapDownsample.time_inverse_distance_weighted_remapping
	45.0±0.1ms	45.1±0.2ms	1.00	mpas_ocean.RemapDownsample.time_nearest_neighbor_remapping
	361±1ms	360±0.6ms	1.00	mpas_ocean.RemapUpsample.time_inverse_distance_weighted_remapping
	264±0.7ms	265±0.4ms	1.01	mpas_ocean.RemapUpsample.time_nearest_neighbor_remapping
	292M	292M	1.00	quad_hexagon.QuadHexagon.peakmem_open_dataset
	291M	291M	1.00	quad_hexagon.QuadHexagon.peakmem_open_grid
	6.22±0.01ms	6.18±0.06ms	0.99	quad_hexagon.QuadHexagon.time_open_dataset
	5.31±0.01ms	5.29±0.02ms	1.00	quad_hexagon.QuadHexagon.time_open_grid

Benchmarks that have got worse:

Change	Before [b11d011]	After [d928082]	Ratio	Benchmark (Parameter)
+	375M	416M	1.11	face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/mpas/QU/oQU480.231010.nc'))
+	375M	417M	1.11	face_bounds.FaceBounds.peakmem_face_bounds(PosixPath('/home/runner/work/uxarray/uxarray/test/meshfiles/scrip/outCSne8/outCSne8.nc'))
+	489±10μs	550±5μs	1.12	mpas_ocean.ConnectivityConstruction.time_n_nodes_per_face('480km')
+	744±2ns	871±2ns	1.17	mpas_ocean.ConstructTreeStructures.time_kd_tree('120km')