Add overloads for float signatures

person142 · Jul 9, 2019 · 0ee76b3 · 0ee76b3
1 parent 1b36e26
commit 0ee76b3
Show file tree

Hide file tree

Showing 7 changed files with 30 additions and 0 deletions.
diff --git a/generate_signatures_json.py b/generate_signatures_json.py
@@ -15,11 +15,13 @@
 
 SPECIAL_DOC_TO_CTYPES = {
     'double': 'c_double',
+    'float': 'c_float',
     'long': 'c_long'
 }
 
 SPECIAL_DOC_TO_CYTHON_SPECIALIZATION = {
     'double': 'double',
+    'float': 'float',
     'long': 'long'
 }
 

diff --git a/numba_special/__init__.py b/numba_special/__init__.py
@@ -284,6 +284,7 @@
 - :data:`scipy.special.expit`::
 
     float64 expit(float64)
+    float32 expit(float32)
 
 - :data:`scipy.special.expm1`::
 
@@ -514,6 +515,7 @@
 - :data:`scipy.special.logit`::
 
     float64 logit(float64)
+    float32 logit(float32)
 
 - :data:`scipy.special.lpmv`::
 

diff --git a/numba_special/function_pointers.pyx b/numba_special/function_pointers.pyx
@@ -86,6 +86,7 @@ functions = {
     'exp2': PyLong_FromVoidPtr(<void *>sc.exp2),
     'expi[double]': PyLong_FromVoidPtr(<void *>sc.expi[double]),
     'expit[double]': PyLong_FromVoidPtr(<void *>sc.expit[double]),
+    'expit[float]': PyLong_FromVoidPtr(<void *>sc.expit[float]),
     'expm1[double]': PyLong_FromVoidPtr(<void *>sc.expm1[double]),
     'expn[double]': PyLong_FromVoidPtr(<void *>sc.expn[double]),
     'expn[long]': PyLong_FromVoidPtr(<void *>sc.expn[long]),
@@ -145,6 +146,7 @@ functions = {
     'log_ndtr[double]': PyLong_FromVoidPtr(<void *>sc.log_ndtr[double]),
     'loggamma[double]': PyLong_FromVoidPtr(<void *>sc.loggamma[double]),
     'logit[double]': PyLong_FromVoidPtr(<void *>sc.logit[double]),
+    'logit[float]': PyLong_FromVoidPtr(<void *>sc.logit[float]),
     'lpmv': PyLong_FromVoidPtr(<void *>sc.lpmv),
     'mathieu_a': PyLong_FromVoidPtr(<void *>sc.mathieu_a),
     'mathieu_b': PyLong_FromVoidPtr(<void *>sc.mathieu_b),

diff --git a/numba_special/generate_overloads.py b/numba_special/generate_overloads.py
@@ -78,11 +78,13 @@ def {FUNCTION}(*args):"""
 
 CTYPES_TO_NUMBA_TYPES = {
     'c_double': 'numba.types.float64',
+    'c_float': 'numba.types.float32',
     'c_long': 'numba_long'
 }
 
 CTYPES_TO_SHORT_NUMBA_TYPES = {
     'c_double': 'float64',
+    'c_float': 'float32',
     'c_long': 'long'
 }
 

diff --git a/numba_special/numba_overloads.py b/numba_special/numba_overloads.py
@@ -95,6 +95,7 @@
     'exp2': ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double),
     'expi[double]': ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double),
     'expit[double]': ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double),
+    'expit[float]': ctypes.CFUNCTYPE(ctypes.c_float, ctypes.c_float),
     'expm1[double]': ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double),
     'expn[double]': ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double, ctypes.c_double),
     'expn[long]': ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_long, ctypes.c_double),
@@ -154,6 +155,7 @@
     'log_ndtr[double]': ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double),
     'loggamma[double]': ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double),
     'logit[double]': ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double),
+    'logit[float]': ctypes.CFUNCTYPE(ctypes.c_float, ctypes.c_float),
     'lpmv': ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double),
     'mathieu_a': ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double, ctypes.c_double),
     'mathieu_b': ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double, ctypes.c_double),
@@ -738,6 +740,10 @@ def expit(*args):
         f = get_scalar_function('expit[double]')
         return lambda *args: f(*args)
 
+    if args == (numba.types.float32,):
+        f = get_scalar_function('expit[float]')
+        return lambda *args: f(*args)
+
 
 @numba.extending.overload(sc.expm1)
 def expm1(*args):
@@ -1145,6 +1151,10 @@ def logit(*args):
         f = get_scalar_function('logit[double]')
         return lambda *args: f(*args)
 
+    if args == (numba.types.float32,):
+        f = get_scalar_function('logit[float]')
+        return lambda *args: f(*args)
+
 
 @numba.extending.overload(sc.lpmv)
 def lpmv(*args):

diff --git a/numba_special/signatures.json b/numba_special/signatures.json
@@ -535,6 +535,10 @@
         "double": [
             "c_double",
             "c_double"
+        ],
+        "float": [
+            "c_float",
+            "c_float"
         ]
     },
     "expm1": {
@@ -929,6 +933,10 @@
         "double": [
             "c_double",
             "c_double"
+        ],
+        "float": [
+            "c_float",
+            "c_float"
         ]
     },
     "lpmv": {

diff --git a/numba_special/tests/test_functions.py b/numba_special/tests/test_functions.py
@@ -16,6 +16,10 @@
 
 CTYPES_TO_TEST_POINTS = {
     'c_double': [-100.0, -10.0, -1.0, -0.1, 0.0, 0.1, 1.0, 10.0, 100.0],
+    'c_float': np.array(
+        [-100.0, -10.0, -1.0, -0.1, 0.0, 0.1, 1.0, 10.0, 100.0],
+        dtype=np.float32
+    ),
     'c_long': [-100, -10, -1, 0, 1, 10, 100],
 }