Enable caching of compiled guvectorize functions

Author: tomwhite

sgkit/distance/metrics.py

@@ -13,6 +13,7 @@
         "void(int8[:], int8[:], float64[:])",
     ],
     "(n),(n)->()",
+    cache=True,
 )
 def correlation(x: ArrayLike, y: ArrayLike, out: ArrayLike) -> None:
     """Calculates the correlation between two vectors.
@@ -80,6 +81,7 @@ def correlation(x: ArrayLike, y: ArrayLike, out: ArrayLike) -> None:
         "void(int8[:], int8[:], float64[:])",
     ],
     "(n),(n)->()",
+    cache=True,
 )
 def euclidean(x: ArrayLike, y: ArrayLike, out: ArrayLike) -> None:
     """Calculates the euclidean distance between two vectors.

sgkit/stats/aggregation.py

@@ -24,6 +24,7 @@
     ],
     "(k),(n)->(n)",
     nopython=True,
+    cache=True,
 )
 def count_alleles(g: ArrayLike, _: ArrayLike, out: ArrayLike) -> None:
     """Generalized U-function for computing per sample allele counts.
@@ -61,6 +62,7 @@ def count_alleles(g: ArrayLike, _: ArrayLike, out: ArrayLike) -> None:
     ],
     "(n, k),(n),(c)->(c,k)",
     nopython=True,
+    cache=True,
 )
 def _count_cohort_alleles(
     ac: ArrayLike, cohorts: ArrayLike, _: ArrayLike, out: ArrayLike

sgkit/stats/popgen.py

@@ -121,9 +121,7 @@ def diversity(
 
 # c = cohorts, k = alleles
 @guvectorize(  # type: ignore
-    ["void(int64[:, :], float64[:,:])"],
-    "(c, k)->(c,c)",
-    nopython=True,
+    ["void(int64[:, :], float64[:,:])"], "(c, k)->(c,c)", nopython=True, cache=True
 )
 def _divergence(ac: ArrayLike, out: ArrayLike) -> None:
     """Generalized U-function for computing divergence.
@@ -295,6 +293,7 @@ def divergence(
     ],
     "(c,c)->(c,c)",
     nopython=True,
+    cache=True,
 )
 def _Fst_Hudson(d: ArrayLike, out: ArrayLike) -> None:
     """Generalized U-function for computing Fst using Hudson's estimator.
@@ -326,6 +325,7 @@ def _Fst_Hudson(d: ArrayLike, out: ArrayLike) -> None:
     ],
     "(c,c)->(c,c)",
     nopython=True,
+    cache=True,
 )
 def _Fst_Nei(d: ArrayLike, out: ArrayLike) -> None:
     """Generalized U-function for computing Fst using Nei's estimator.