From 93cf4d9fd6d29fa8ed26d527ebe1724264da86d8 Mon Sep 17 00:00:00 2001
From: Philip Chmielowiec <67855069+philipc2@users.noreply.github.com>
Date: Fri, 17 May 2024 23:20:33 -0500
Subject: [PATCH] Re-introduce Numba to face area computations  (#794)

* re-introduce numba to face area calculation

* update coordinate conversion
---
 uxarray/grid/area.py        | 11 +++++++++++
 uxarray/grid/coordinates.py |  3 +++
 2 files changed, 14 insertions(+)

diff --git a/uxarray/grid/area.py b/uxarray/grid/area.py
index 6e1692feb..1dc6de52e 100644
--- a/uxarray/grid/area.py
+++ b/uxarray/grid/area.py
@@ -2,7 +2,13 @@
 
 from uxarray.grid.coordinates import _lonlat_rad_to_xyz
 
+from numba import njit, config
+from uxarray.constants import ENABLE_JIT_CACHE, ENABLE_JIT
 
+config.DISABLE_JIT = not ENABLE_JIT
+
+
+@njit(cache=ENABLE_JIT_CACHE)
 def calculate_face_area(
     x, y, z, quadrature_rule="gaussian", order=4, coords_type="spherical"
 ):
@@ -92,6 +98,7 @@ def calculate_face_area(
     return area, jacobian
 
 
+@njit(cache=ENABLE_JIT_CACHE)
 def get_all_face_area_from_coords(
     x,
     y,
@@ -166,6 +173,7 @@ def get_all_face_area_from_coords(
     return area, jacobian
 
 
+@njit(cache=ENABLE_JIT_CACHE)
 def calculate_spherical_triangle_jacobian(node1, node2, node3, dA, dB):
     """Calculate Jacobian of a spherical triangle. This is a helper function
     for calculating face area.
@@ -255,6 +263,7 @@ def calculate_spherical_triangle_jacobian(node1, node2, node3, dA, dB):
     return dJacobian
 
 
+@njit(cache=ENABLE_JIT_CACHE)
 def calculate_spherical_triangle_jacobian_barycentric(node1, node2, node3, dA, dB):
     """Calculate Jacobian of a spherical triangle. This is a helper function
     for calculating face area.
@@ -333,6 +342,7 @@ def calculate_spherical_triangle_jacobian_barycentric(node1, node2, node3, dA, d
     return 0.5 * dJacobian
 
 
+@njit(cache=ENABLE_JIT_CACHE)
 def get_gauss_quadratureDG(nCount):
     """Gauss Quadrature Points for integration.
 
@@ -577,6 +587,7 @@ def get_gauss_quadratureDG(nCount):
     return dG, dW
 
 
+@njit(cache=ENABLE_JIT_CACHE)
 def get_tri_quadratureDG(nOrder):
     """Triangular Quadrature Points for integration.
 
diff --git a/uxarray/grid/coordinates.py b/uxarray/grid/coordinates.py
index 7032d8d22..ba0b27e57 100644
--- a/uxarray/grid/coordinates.py
+++ b/uxarray/grid/coordinates.py
@@ -8,7 +8,10 @@
 
 from typing import Union
 
+from numba import njit
 
+
+@njit(cache=True)
 def _lonlat_rad_to_xyz(
     lon: Union[np.ndarray, float],
     lat: Union[np.ndarray, float],