From 77da3295a502c59b326a806d59b9f10d77c2ee59 Mon Sep 17 00:00:00 2001
From: cheginit <cheginit@gmail.com>
Date: Thu, 9 Dec 2021 21:05:11 -0600
Subject: [PATCH] ENH: Add type signatures to the jenksy function to speed it
 up. Also, add a warning if numba is not installed. (#118)

---
 mapclassify/classifiers.py | 26 ++++++++++++++------------
 1 file changed, 14 insertions(+), 12 deletions(-)

diff --git a/mapclassify/classifiers.py b/mapclassify/classifiers.py
index e34530af..60fbbd4e 100644
--- a/mapclassify/classifiers.py
+++ b/mapclassify/classifiers.py
@@ -57,10 +57,11 @@
 FMT = "{:.2f}"
 
 try:
-    from numba import jit
+    from numba import njit
+    HAS_NUMBA = True
 except ImportError:
-
-    def jit(func):
+    HAS_NUMBA = False
+    def njit(func):
         return func
 
 
@@ -507,8 +508,8 @@ def natural_breaks(values, k=5, init=10):
     return (sids, class_ids, fit, cuts)
 
 
-@jit
-def _fisher_jenks_means(values, classes=5, sort=True):
+@njit("f8[:](f8[:], u2)", cache=True)
+def _fisher_jenks_means(values, classes=5):
     """
     Jenks Optimal (Natural Breaks) algorithm implemented in Python.
 
@@ -523,8 +524,6 @@ def _fisher_jenks_means(values, classes=5, sort=True):
     assuring heterogeneity among classes.
 
     """
-    if sort:
-        values.sort()
     n_data = len(values)
     mat1 = np.zeros((n_data + 1, classes + 1), dtype=np.int32)
     mat2 = np.zeros((n_data + 1, classes + 1), dtype=np.float32)
@@ -562,7 +561,7 @@ def _fisher_jenks_means(values, classes=5, sort=True):
         id = int(pivot - 2)
         kclass[countNum - 1] = values[id]
         k = int(pivot - 1)
-    return kclass
+    return np.delete(kclass, 0)
 
 
 class MapClassifier(object):
@@ -1761,8 +1760,8 @@ class FisherJenks(MapClassifier):
     ----------
     y : array
         (n,1), values to classify
-    k : int
-        number of classes required
+    k : int, optional
+        number of classes, defatuls to 5
 
     Attributes
     ----------
@@ -1790,6 +1789,9 @@ class FisherJenks(MapClassifier):
     """
 
     def __init__(self, y, k=K):
+        if not HAS_NUMBA:
+            Warn("Numba not installed. Using slow pure python version.",
+                 UserWarning)
 
         nu = len(np.unique(y))
         if nu < k:
@@ -1799,8 +1801,8 @@ def __init__(self, y, k=K):
         self.name = "FisherJenks"
 
     def _set_bins(self):
-        x = self.y.copy()
-        self.bins = np.array(_fisher_jenks_means(x, classes=self.k)[1:])
+        x = np.sort(self.y).astype("f8")
+        self.bins = _fisher_jenks_means(x, classes=self.k)
 
 
 class FisherJenksSampled(MapClassifier):