From 2eab952b60f553beb2522b25655f71a88f6a6248 Mon Sep 17 00:00:00 2001
From: Joseph Capriotti <josephrcapriotti@gmail.com>
Date: Thu, 26 Oct 2023 08:11:36 -0600
Subject: [PATCH 01/10] attempt at cython prism ufuncs

---
 .../_extensions/potential_field_prism.c       | 579 ------------------
 .../_extensions/potential_field_prism.pyx     | 314 ++++++++++
 geoana/kernels/_extensions/setup.py           |   6 +-
 3 files changed, 319 insertions(+), 580 deletions(-)
 delete mode 100644 geoana/kernels/_extensions/potential_field_prism.c
 create mode 100644 geoana/kernels/_extensions/potential_field_prism.pyx

diff --git a/geoana/kernels/_extensions/potential_field_prism.c b/geoana/kernels/_extensions/potential_field_prism.c
deleted file mode 100644
index 081d423c..00000000
--- a/geoana/kernels/_extensions/potential_field_prism.c
+++ /dev/null
@@ -1,579 +0,0 @@
-#define PY_SSIZE_T_CLEAN
-#include <Python.h>
-#include "numpy/ndarraytypes.h"
-#include "numpy/ufuncobject.h"
-#include "numpy/npy_3kcompat.h"
-#include <math.h>
-
-static PyMethodDef PrismMethods[] = {
-    {NULL, NULL, 0, NULL}
-};
-
-/* The loop definition must precede the PyMODINIT_FUNC. */
-
-static void double_prism_f(char **args, const npy_intp *dimensions,
-                         const npy_intp *steps, void *data)
-{
-    npy_intp i;
-    npy_intp n = dimensions[0];
-    char *inx = args[0], *iny = args[1], *inz = args[2];
-    char *out = args[3];
-    npy_intp inx_step = steps[0], iny_step = steps[1], inz_step = steps[2];
-    npy_intp out_step = steps[3];
-
-    double x, y, z, v, r, temp;
-
-    for (i = 0; i < n; i++) {
-        /* BEGIN main ufunc computation */
-        x = *(double *)inx;
-        y = *(double *)iny;
-        z = *(double *)inz;
-        v = 0.0;
-        r = sqrt(x * x + y * y + z * z);
-        if (x != 0.0){
-            if (y != 0.0){
-                temp = z + r;
-                // check if x & y were small relative to -z
-                if (temp > 0){
-                    v -= x * y * log(temp);
-                }
-            }
-            v += 0.5 * x * x * atan( y * z / (x * r));
-        }
-        if (y != 0.0){
-            if (z != 0.0){
-                temp = x + r;
-                // check if y & z were small relative to -x
-                if (temp > 0){
-                    v -= y *z * log(temp);
-                }
-            }
-            v += 0.5 * y * y * atan(z * x / (y * r));
-        }
-        if (z != 0.0){
-            if (x != 0.0){
-                temp = y + r;
-                // check if x & z were small relative to -y
-                if (temp > 0){
-                    v -= z * x * log(temp);
-                }
-            }
-            v += 0.5 * z * z * atan(x * y / (z * r));
-        }
-        *((double *)out) = v;
-        /* END main ufunc computation */
-
-        inx += inx_step;
-        iny += iny_step;
-        inz += inz_step;
-        out += out_step;
-    }
-}
-
-/* This a pointer to the above function */
-PyUFuncGenericFunction funcs_prism_f[1] = {&double_prism_f};
-
-static void double_prism_fz(char **args, const npy_intp *dimensions,
-                         const npy_intp *steps, void *data)
-{
-    npy_intp i;
-    npy_intp n = dimensions[0];
-    char *inx = args[0], *iny = args[1], *inz = args[2];
-    char *out = args[3];
-    npy_intp inx_step = steps[0], iny_step = steps[1], inz_step = steps[2];
-    npy_intp out_step = steps[3];
-
-    double x, y, z, v, r, temp;
-
-    for (i = 0; i < n; i++) {
-        /* BEGIN main ufunc computation */
-        x = *(double *)inx;
-        y = *(double *)iny;
-        z = *(double *)inz;
-        v = 0.0;
-        r = sqrt(x * x + y * y + z * z);
-        if (x != 0.0){
-            temp = y + r;
-            // check if x & z were small relative to -y
-            if (temp > 0.0){
-                v += x * log(temp);
-            }
-        }
-        if (y != 0.0){
-            temp = x + r;
-            // check if y & z were small relative to -x
-            if (temp > 0.0){
-                v += y * log(temp);
-            }
-        }
-        if (z != 0.0){
-            v -= z * atan(x * y / (z * r));
-        }
-        *((double *)out) = v;
-        /* END main ufunc computation */
-
-        inx += inx_step;
-        iny += iny_step;
-        inz += inz_step;
-        out += out_step;
-    }
-}
-
-/* This a pointer to the above function */
-PyUFuncGenericFunction funcs_prism_fz[1] = {&double_prism_fz};
-
-static void double_prism_fzz(char **args, const npy_intp *dimensions,
-                         const npy_intp *steps, void *data)
-{
-    npy_intp i;
-    npy_intp n = dimensions[0];
-    char *inx = args[0], *iny = args[1], *inz = args[2];
-    char *out = args[3];
-    npy_intp inx_step = steps[0], iny_step = steps[1], inz_step = steps[2];
-    npy_intp out_step = steps[3];
-
-    double x, y, z, v, r;
-
-    for (i = 0; i < n; i++) {
-        /* BEGIN main ufunc computation */
-        x = *(double *)inx;
-        y = *(double *)iny;
-        z = *(double *)inz;
-        if (z != 0.0){
-            r = sqrt(x * x + y * y + z * z);
-            v = atan(x * y / (z * r));
-        }else{
-            v = 0.0;
-        }
-        *((double *)out) = v;
-        /* END main ufunc computation */
-
-        inx += inx_step;
-        iny += iny_step;
-        inz += inz_step;
-        out += out_step;
-    }
-}
-
-/* This a pointer to the above function */
-PyUFuncGenericFunction funcs_prism_fzz[1] = {&double_prism_fzz};
-
-static void double_prism_fzx(char **args, const npy_intp *dimensions,
-                         const npy_intp *steps, void *data)
-{
-    npy_intp i;
-    npy_intp n = dimensions[0];
-    char *inx = args[0], *iny = args[1], *inz = args[2];
-    char *out = args[3];
-    npy_intp inx_step = steps[0], iny_step = steps[1], inz_step = steps[2];
-    npy_intp out_step = steps[3];
-
-    double x, y, z, v, r;
-    for (i = 0; i < n; i++) {
-        /* BEGIN main ufunc computation */
-        x = *(double *)inx;
-        y = *(double *)iny;
-        z = *(double *)inz;
-        r = sqrt(x * x + y * y + z * z);
-        v = y + r;
-        if (v == 0.0){
-            if (y < 0){
-                v = log(-2 * y);
-            }
-        }else{
-            v = -log(v);
-        }
-        *((double *)out) = v;
-        /* END main ufunc computation */
-
-        inx += inx_step;
-        iny += iny_step;
-        inz += inz_step;
-        out += out_step;
-    }
-}
-
-/* This a pointer to the above function */
-PyUFuncGenericFunction funcs_prism_fzx[1] = {&double_prism_fzx};
-
-static void double_prism_fzy(char **args, const npy_intp *dimensions,
-                         const npy_intp *steps, void *data)
-{
-    npy_intp i;
-    npy_intp n = dimensions[0];
-    char *inx = args[0], *iny = args[1], *inz = args[2];
-    char *out = args[3];
-    npy_intp inx_step = steps[0], iny_step = steps[1], inz_step = steps[2];
-    npy_intp out_step = steps[3];
-
-    double x, y, z, v, r;
-    for (i = 0; i < n; i++) {
-        /* BEGIN main ufunc computation */
-        x = *(double *)inx;
-        y = *(double *)iny;
-        z = *(double *)inz;
-        r = sqrt(x * x + y * y + z * z);
-        v = x + r;
-        if (v == 0.0){
-            if (x < 0){
-                v = log(-2 * x);
-            }
-        }else{
-            v = -log(v);
-        }
-        *((double *)out) = v;
-        /* END main ufunc computation */
-
-        inx += inx_step;
-        iny += iny_step;
-        inz += inz_step;
-        out += out_step;
-    }
-}
-
-/* This a pointer to the above function */
-PyUFuncGenericFunction funcs_prism_fzy[1] = {&double_prism_fzy};
-
-static void double_prism_fzzz(char **args, const npy_intp *dimensions,
-                         const npy_intp *steps, void *data)
-{
-    npy_intp i;
-    npy_intp n = dimensions[0];
-    char *inx = args[0], *iny = args[1], *inz = args[2];
-    char *out = args[3];
-    npy_intp inx_step = steps[0], iny_step = steps[1], inz_step = steps[2];
-    npy_intp out_step = steps[3];
-
-    double x, y, z, v, r, v1, v2;
-    for (i = 0; i < n; i++) {
-        /* BEGIN main ufunc computation */
-        x = *(double *)inx;
-        y = *(double *)iny;
-        z = *(double *)inz;
-        r = sqrt(x * x + y * y + z * z);
-        v1 = x * x + z * z;
-        v2 = y * y + z * z;
-        v = 0.0;
-        if (v1 != 0.0){
-            v += 1.0/v1;
-        }
-        if (v2 != 0.0){
-            v += 1.0/v2;
-        }
-        if (r != 0.0){
-            v *= x * y / r;
-        }
-
-        *((double *)out) = v;
-        /* END main ufunc computation */
-
-        inx += inx_step;
-        iny += iny_step;
-        inz += inz_step;
-        out += out_step;
-    }
-}
-
-/* This a pointer to the above function */
-PyUFuncGenericFunction funcs_prism_fzzz[1] = {&double_prism_fzzz};
-
-static void double_prism_fxxy(char **args, const npy_intp *dimensions,
-                         const npy_intp *steps, void *data)
-{
-    npy_intp i;
-    npy_intp n = dimensions[0];
-    char *inx = args[0], *iny = args[1], *inz = args[2];
-    char *out = args[3];
-    npy_intp inx_step = steps[0], iny_step = steps[1], inz_step = steps[2];
-    npy_intp out_step = steps[3];
-
-    double x, y, z, v, r;
-    for (i = 0; i < n; i++) {
-        /* BEGIN main ufunc computation */
-        x = *(double *)inx;
-        y = *(double *)iny;
-        z = *(double *)inz;
-        if (x != 0.0){
-            v = x * x + y * y;
-            r = sqrt(x * x + y * y + z * z);
-            v = - x * z / (v * r);
-        }else{
-            v = 0.0;
-        }
-        *((double *)out) = v;
-        /* END main ufunc computation */
-
-        inx += inx_step;
-        iny += iny_step;
-        inz += inz_step;
-        out += out_step;
-    }
-}
-
-/* This a pointer to the above function */
-PyUFuncGenericFunction funcs_prism_fxxy[1] = {&double_prism_fxxy};
-
-static void double_prism_fxxz(char **args, const npy_intp *dimensions,
-                         const npy_intp *steps, void *data)
-{
-    npy_intp i;
-    npy_intp n = dimensions[0];
-    char *inx = args[0], *iny = args[1], *inz = args[2];
-    char *out = args[3];
-    npy_intp inx_step = steps[0], iny_step = steps[1], inz_step = steps[2];
-    npy_intp out_step = steps[3];
-
-    double x, y, z, v, r;
-    for (i = 0; i < n; i++) {
-        /* BEGIN main ufunc computation */
-        x = *(double *)inx;
-        y = *(double *)iny;
-        z = *(double *)inz;
-        if (x != 0.0){
-            v = x * x + z * z;
-            r = sqrt(x * x + y * y + z * z);
-            v = - x * y / (v * r);
-        }else{
-            v = 0.0;
-        }
-        *((double *)out) = v;
-        /* END main ufunc computation */
-
-        inx += inx_step;
-        iny += iny_step;
-        inz += inz_step;
-        out += out_step;
-    }
-}
-
-/* This a pointer to the above function */
-PyUFuncGenericFunction funcs_prism_fxxz[1] = {&double_prism_fxxz};
-
-static void double_prism_fxyz(char **args, const npy_intp *dimensions,
-                         const npy_intp *steps, void *data)
-{
-    npy_intp i;
-    npy_intp n = dimensions[0];
-    char *inx = args[0], *iny = args[1], *inz = args[2];
-    char *out = args[3];
-    npy_intp inx_step = steps[0], iny_step = steps[1], inz_step = steps[2];
-    npy_intp out_step = steps[3];
-
-    double x, y, z, v, r;
-    for (i = 0; i < n; i++) {
-        /* BEGIN main ufunc computation */
-        x = *(double *)inx;
-        y = *(double *)iny;
-        z = *(double *)inz;
-        r = sqrt(x * x + y * y + z * z);
-        if (r != 0.0){
-            v = 1.0/r;
-        }else{
-            v = 0.0;
-        }
-        *((double *)out) = v;
-        /* END main ufunc computation */
-
-        inx += inx_step;
-        iny += iny_step;
-        inz += inz_step;
-        out += out_step;
-    }
-}
-
-/* This a pointer to the above function */
-PyUFuncGenericFunction funcs_prism_fxyz[1] = {&double_prism_fxyz};
-
-/* These are the input and return dtypes of the prism functions.*/
-static char types[4] = {NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE, NPY_DOUBLE};
-
-static struct PyModuleDef moduledef = {
-    PyModuleDef_HEAD_INIT,
-    "potential_field_prism",
-    NULL,
-    -1,
-    PrismMethods,
-    NULL,
-    NULL,
-    NULL,
-    NULL
-};
-
-static void *data[1] = {NULL,};
-
-PyMODINIT_FUNC PyInit_potential_field_prism(void)
-{
-    PyObject *m, *d;
-    PyObject *prism_f, *prism_fz, *prism_fzz, *prism_fzx, *prism_fzy;
-    PyObject *prism_fzzz, *prism_fxxy, *prism_fxxz, *prism_fxyz;
-    m = PyModule_Create(&moduledef);
-    if (!m) {
-        return NULL;
-    }
-
-    import_array();
-    import_umath();
-
-    prism_f = PyUFunc_FromFuncAndData(funcs_prism_f, data, types, 1, 3, 1,
-                                    PyUFunc_None, "prism_f",
-"Evaluates the indefinite volume integral for the 1/r kernel.\n\n"
-"This is used to evaluate the gravitational potential of dense prisms.\n\n"
-"Parameters\n"
-"----------\n"
-"x, y, z : (...) numpy.ndarray\n"
-"    The nodal locations to evaluate the function at\n\n"
-"Returns\n"
-"-------\n"
-"(...) numpy.ndarray", 0);
-
-    prism_fz = PyUFunc_FromFuncAndData(funcs_prism_fz, data, types, 1, 3, 1,
-                                    PyUFunc_None, "prism_fz",
-"Evaluates the indefinite volume integral for the d/dz * 1/r kernel.\n\n"
-"This is used to evaluate the gravitational field of dense prisms.\n\n"
-"Parameters\n"
-"----------\n"
-"x, y, z : (...) numpy.ndarray\n"
-"    The nodal locations to evaluate the function at\n\n"
-"Returns\n"
-"-------\n"
-"(...) numpy.ndarray\n\n"
-"Notes\n"
-"-----\n"
-"Can be used to compute other components by cycling the inputs", 0);
-
-    prism_fzz = PyUFunc_FromFuncAndData(funcs_prism_fzz, data, types, 1, 3, 1,
-                                    PyUFunc_None, "prism_fzz",
-"Evaluates the indefinite volume integral for the d**2/dz**2 * 1/r kernel.\n\n"
-"This is used to evaluate the gravitational potential of dense prisms.\n\n"
-"Parameters\n"
-"----------\n"
-"x, y, z : (...) numpy.ndarray\n"
-"    The nodal locations to evaluate the function at\n\n"
-"Returns\n"
-"-------\n"
-"  (...) numpy.ndarray\n\n"
-"Notes\n"
-"-----\n"
-"Can be used to compute other components by cycling the inputs", 0);
-
-    prism_fzx = PyUFunc_FromFuncAndData(funcs_prism_fzx, data, types, 1, 3, 1,
-                                    PyUFunc_None, "prism_fzx",
-"Evaluates the indefinite volume integral for the d**2/(dz*dx) * 1/r kernel.\n\n"
-"This is used to evaluate the gravitational potential of dense prisms.\n\n"
-"Parameters\n"
-"----------\n"
-"x, y, z : (...) numpy.ndarray\n"
-"    The nodal locations to evaluate the function at\n\n"
-"Returns\n"
-"-------\n"
-"(...) numpy.ndarray\n\n"
-"Notes\n"
-"-----\n"
-"Can be used to compute other components by cycling the inputs", 0);
-
-    prism_fzy = PyUFunc_FromFuncAndData(funcs_prism_fzy, data, types, 1, 3, 1,
-                                    PyUFunc_None, "prism_fzy",
-"Evaluates the indefinite volume integral for the d**2/(dz*dx) * 1/r kernel.\n\n"
-"This is used to evaluate the gravitational potential of dense prisms.\n\n"
-"Parameters\n"
-"----------\n"
-"x, y, z : (...) numpy.ndarray\n"
-"    The nodal locations to evaluate the function at\n\n"
-"Returns\n"
-"-------\n"
-"(...) numpy.ndarray\n\n"
-"Notes\n"
-"-----\n"
-"Can be used to compute other components by cycling the inputs", 0);
-
-    prism_fzzz = PyUFunc_FromFuncAndData(funcs_prism_fzzz, data, types, 1, 3, 1,
-                                    PyUFunc_None, "prism_fzzz",
-"Evaluates the indefinite volume integral for the d**3/(dz**3) * 1/r kernel.\n\n"
-"This is used to evaluate the magnetic gradient of susceptible prisms.\n\n"
-"Parameters\n"
-"----------\n"
-"x, y, z : (...) numpy.ndarray\n"
-"    The nodal locations to evaluate the function at\n\n"
-"Returns\n"
-"-------\n"
-"(...) numpy.ndarray\n\n"
-"Notes\n"
-"-----\n"
-"Can be used to compute other components by cycling the inputs", 0);
-
-    prism_fxxy = PyUFunc_FromFuncAndData(funcs_prism_fxxy, data, types, 1, 3, 1,
-                                    PyUFunc_None, "prism_fxxy",
-"Evaluates the indefinite volume integral for the d**3/(dx**2 * dy) * 1/r kernel.\n\n"
-"This is used to evaluate the magnetic gradient of susceptible prisms.\n\n"
-"Parameters\n"
-"----------\n"
-"x, y, z : (...) numpy.ndarray\n"
-"    The nodal locations to evaluate the function at\n\n"
-"Returns\n"
-"-------\n"
-"(...) numpy.ndarray\n\n"
-"Notes\n"
-"-----\n"
-"Can be used to compute other components by cycling the inputs", 0);
-
-    prism_fxxz = PyUFunc_FromFuncAndData(funcs_prism_fxxz, data, types, 1, 3, 1,
-                                    PyUFunc_None, "prism_fxxz",
-"Evaluates the indefinite volume integral for the d**3/(dx**2 * dz) * 1/r kernel.\n\n"
-"This is used to evaluate the magnetic gradient of susceptible prisms.\n\n"
-"Parameters\n"
-"----------\n"
-"x, y, z : (...) numpy.ndarray\n"
-"    The nodal locations to evaluate the function at\n\n"
-"Returns\n"
-"-------\n"
-"(...) numpy.ndarray\n\n"
-"Notes\n"
-"-----\n"
-"Can be used to compute other components by cycling the inputs", 0);
-
-    prism_fxyz = PyUFunc_FromFuncAndData(funcs_prism_fxyz, data, types, 1, 3, 1,
-                                    PyUFunc_None, "prism_fxyz",
-"Evaluates the indefinite volume integral for the d**3/(dx * dy * dz) * 1/r kernel.\n\n"
-"This is used to evaluate the magnetic gradient of susceptible prisms.\n\n"
-"Parameters\n"
-"----------\n"
-"x, y, z : (...) numpy.ndarray\n"
-"    The nodal locations to evaluate the function at\n\n"
-"Returns\n"
-"-------\n"
-"(...) numpy.ndarray\n\n"
-"Notes\n"
-"-----\n"
-"Can be used to compute other components by cycling the inputs", 0);
-
-    d = PyModule_GetDict(m);
-
-    PyDict_SetItemString(d, "prism_f", prism_f);
-    Py_DECREF(prism_f);
-
-    PyDict_SetItemString(d, "prism_fz", prism_fz);
-    Py_DECREF(prism_fz);
-
-    PyDict_SetItemString(d, "prism_fzz", prism_fzz);
-    Py_DECREF(prism_fzz);
-
-    PyDict_SetItemString(d, "prism_fzx", prism_fzx);
-    Py_DECREF(prism_fzx);
-
-    PyDict_SetItemString(d, "prism_fzy", prism_fzy);
-    Py_DECREF(prism_fzy);
-
-    PyDict_SetItemString(d, "prism_fzzz", prism_fzzz);
-    Py_DECREF(prism_fzzz);
-
-    PyDict_SetItemString(d, "prism_fxxy", prism_fxxy);
-    Py_DECREF(prism_fxxy);
-
-    PyDict_SetItemString(d, "prism_fxxz", prism_fxxz);
-    Py_DECREF(prism_fxxz);
-
-    PyDict_SetItemString(d, "prism_fxyz", prism_fxyz);
-    Py_DECREF(prism_fxyz);
-
-    return m;
-}
diff --git a/geoana/kernels/_extensions/potential_field_prism.pyx b/geoana/kernels/_extensions/potential_field_prism.pyx
new file mode 100644
index 00000000..89149046
--- /dev/null
+++ b/geoana/kernels/_extensions/potential_field_prism.pyx
@@ -0,0 +1,314 @@
+cimport cython
+cimport numpy as np
+import numpy as np
+
+from libc.math cimport sqrt, log, atan
+
+# @cython.ufunc
+# @cython.cdivision
+# cdef double prism_f(double x, double y, double z) nogil:
+#     """Evaluates the indefinite volume integral for the 1/r kernel.
+#
+#     This is used to evaluate the gravitational potential of dense prisms.
+#
+#     Parameters
+#     ----------
+#     x, y, z : (...) numpy.ndarray
+#         The nodal locations to evaluate the function at
+#
+#     Returns
+#     -------
+#     (...) numpy.ndarray
+#     """
+#     cdef:
+#         double v = 0.0
+#         double r, temp
+#
+#     r = sqrt(x * x + y * y + z * z)
+#     if x != 0.0 and y != 0.0:
+#         temp = z + r
+#         if temp > 0.0:
+#             v -= x * y * log(temp)
+#         v += 0.5 * x * x * atan( y * z / (x * r))
+#     if y != 0.0 and z != 0.0:
+#         temp = x + r
+#         if temp > 0.0:
+#             v -= y * z * log(temp)
+#         v += 0.5 * y * y * atan(z * x / (y * r))
+#     if z != 0.0 and x != 0.0:
+#         temp = y + r
+#         if temp > 0.0:
+#             v -= z * x * log(temp)
+#         v += 0.5 * z * z * atan(x * y / (z * r))
+#     return v
+
+
+@cython.cdivision
+cdef api double c_prism_fz(double x, double y, double z) nogil:
+    """Evaluates the indefinite volume integral for the d/dz * 1/r kernel.
+
+    This is used to evaluate the gravitational field of dense prisms.
+    
+    Parameters
+    ----------
+    x, y, z : (...) numpy.ndarray
+        The nodal locations to evaluate the function at
+    
+    Returns
+    -------
+    (...) numpy.ndarray
+
+    Notes
+    -----
+    Can be used to compute other components by cycling the inputs
+    """
+    cdef:
+        double v = 0.0
+        double r, temp
+
+    r = sqrt(x * x + y * y + z * z)
+    if x != 0.0:
+        temp = y + r
+        if temp > 0.0:
+            v += x * log(temp)
+    if y != 0.0:
+        temp = x + r
+        if temp > 0.0:
+            v += y * log(temp)
+    if z != 0.0:
+        v -= z * atan(x * y / (z * r))
+    return v
+def prism_fz(x, y, z):
+    cdef double[:] x_, y_, z_, out_
+    out = np.empty_like(x)
+
+    x_ = np.asarray(x).ravel()
+    y_ = np.asarray(y).ravel()
+    z_ = np.asarray(z).ravel()
+    out_ = np.asarray(out).ravel()
+
+    cdef int i
+    for i in range(x_.shape[0]):
+        out_[i] = c_prism_fz(x_[i], y_[i], z_[i])
+    return out
+
+
+
+# @cython.ufunc
+# @cython.cdivision
+# cdef double prism_fzz(double x, double y, double z) nogil:
+#     """Evaluates the indefinite volume integral for the d**2/dz**2 * 1/r kernel.
+#
+#     This is used to evaluate the gravitational gradient of dense prisms.
+#
+#     Parameters
+#     ----------
+#     x, y, z : (...) numpy.ndarray
+#         The nodal locations to evaluate the function at
+#
+#     Returns
+#     -------
+#     (...) numpy.ndarray
+#
+#     Notes
+#     -----
+#     Can be used to compute other components by cycling the inputs
+#     """
+#     cdef:
+#         double v = 0.0
+#         double r
+#
+#     if z != 0.0:
+#         r = sqrt(x * x + y * y + z * z)
+#         v = atan(x * y / (z * r))
+#     return v
+#
+#
+# @cython.ufunc
+# @cython.cdivision
+# cdef double prism_fzx(double x, double y, double z) nogil:
+#     """Evaluates the indefinite volume integral for the d**2/(dz*dx) * 1/r kernel.
+#
+#     This is used to evaluate the gravitational gradient of dense prisms.
+#
+#     Parameters
+#     ----------
+#     x, y, z : (...) numpy.ndarray
+#         The nodal locations to evaluate the function at
+#
+#     Returns
+#     -------
+#     (...) numpy.ndarray
+#
+#     Notes
+#     -----
+#     Can be used to compute other components by cycling the inputs
+#     """
+#     cdef:
+#         double v = 0.0
+#         double r
+#     r = sqrt(x * x + y * y + z * z)
+#     v = y + r
+#     if v == 0.0:
+#         if y < 0:
+#             v = log(-2 * y)
+#     else:
+#         v = -log(v)
+#     return v
+#
+#
+# @cython.ufunc
+# @cython.cdivision
+# cdef double prism_fzy(double x, double y, double z) nogil:
+#     """Evaluates the indefinite volume integral for the d**2/(dz*dy) * 1/r kernel.
+#
+#     This is used to evaluate the gravitational gradient of dense prisms.
+#
+#     Parameters
+#     ----------
+#     x, y, z : (...) numpy.ndarray
+#         The nodal locations to evaluate the function at
+#
+#     Returns
+#     -------
+#     (...) numpy.ndarray
+#
+#     Notes
+#     -----
+#     Can be used to compute other components by cycling the inputs
+#     """
+#     cdef:
+#         double v = 0.0
+#         double r, temp
+#     r = sqrt(x * x + y * y + z * z)
+#     v = y + r
+#     if v == 0.0:
+#         if x < 0:
+#             v = log(-2 * x)
+#     else:
+#         v = -log(v)
+#     return v
+#
+#
+# @cython.ufunc
+# @cython.cdivision
+# cdef double prism_fzzz(double x, double y, double z) nogil:
+#     """Evaluates the indefinite volume integral for the d**3/(dz**3) * 1/r kernel.
+#
+#     This is used to evaluate the magnetic gradient of susceptible prisms.
+#
+#     Parameters
+#     ----------
+#     x, y, z : (...) numpy.ndarray
+#         The nodal locations to evaluate the function at
+#
+#     Returns
+#     -------
+#     (...) numpy.ndarray
+#
+#     Notes
+#     -----
+#     Can be used to compute other components by cycling the inputs
+#     """
+#     cdef:
+#         double v = 0.0
+#         double r, v1, v2
+#     r = sqrt(x * x + y * y + z * z)
+#     v1 = x * x + z * z
+#     v2 = y * y + z * z
+#     if v1 != 0.0:
+#         v += 1.0/v1
+#     if v2 != 0.0:
+#         v += 1.0/v2
+#     if r != 0.0:
+#         v *= x * y / r
+#     return v
+#
+#
+# @cython.ufunc
+# @cython.cdivision
+# cdef double prism_fxxy(double x, double y, double z) nogil:
+#     """Evaluates the indefinite volume integral for the d**3/(dx**2 * dy) * 1/r kernel.
+#
+#     This is used to evaluate the magnetic gradient of susceptible prisms.
+#
+#     Parameters
+#     ----------
+#     x, y, z : (...) numpy.ndarray
+#         The nodal locations to evaluate the function at
+#
+#     Returns
+#     -------
+#     (...) numpy.ndarray
+#
+#     Notes
+#     -----
+#     Can be used to compute other components by cycling the inputs
+#     """
+#     cdef:
+#         double v = 0.0
+#         double r
+#     if x != 0.0:
+#         v = x * x + y * y
+#         r = sqrt(x * x + y * y + z * z)
+#         v = - x * z / (v * r)
+#     return v
+#
+#
+# @cython.ufunc
+# @cython.cdivision
+# cdef double prism_fxxz(double x, double y, double z) nogil:
+#     """Evaluates the indefinite volume integral for the d**3/(dx**2 * dz) * 1/r kernel.
+#
+#     This is used to evaluate the magnetic gradient of susceptible prisms.
+#
+#     Parameters
+#     ----------
+#     x, y, z : (...) numpy.ndarray
+#         The nodal locations to evaluate the function at
+#
+#     Returns
+#     -------
+#     (...) numpy.ndarray
+#
+#     Notes
+#     -----
+#     Can be used to compute other components by cycling the inputs
+#     """
+#     cdef:
+#         double v = 0.0
+#         double r
+#     if x != 0.0:
+#         v = x * x + z * z
+#         r = sqrt(x * x + y * y + z * z)
+#         v = - x * y / (v * r)
+#     return v
+#
+#
+# @cython.ufunc
+# @cython.cdivision
+# cdef double prism_fxyz(double x, double y, double z) nogil:
+#     """Evaluates the indefinite volume integral for the d**3/(dx * dy * dz) * 1/r kernel.
+#
+#     This is used to evaluate the magnetic gradient of susceptible prisms.
+#
+#     Parameters
+#     ----------
+#     x, y, z : (...) numpy.ndarray
+#         The nodal locations to evaluate the function at
+#
+#     Returns
+#     -------
+#     (...) numpy.ndarray
+#
+#     Notes
+#     -----
+#     Can be used to compute other components by cycling the inputs
+#     """
+#     cdef:
+#         double v = 0.0
+#         double r
+#     r = sqrt(x * x + y * y + z * z)
+#     if r != 0.0:
+#         v = 1.0/r
+#     return v
\ No newline at end of file
diff --git a/geoana/kernels/_extensions/setup.py b/geoana/kernels/_extensions/setup.py
index 674461f0..9209270f 100644
--- a/geoana/kernels/_extensions/setup.py
+++ b/geoana/kernels/_extensions/setup.py
@@ -13,6 +13,8 @@ def configuration(parent_package="", top_path=None):
         from Cython.Build import cythonize
 
         cythonize(os.path.join(base_path, "rTE.pyx"))
+
+        cythonize(os.path.join(base_path, "potential_field_prism.pyx"))
     except ImportError:
         pass
 
@@ -22,6 +24,8 @@ def configuration(parent_package="", top_path=None):
     )
 
     config.add_extension(
-        'potential_field_prism', sources=['potential_field_prism.c']
+        'potential_field_prism',
+        sources=['potential_field_prism.c'],
+        include_dirs=[get_numpy_include_dirs()]
     )
     return config

From 0ff961265fdd610af0a480f4cba086cafc922b16 Mon Sep 17 00:00:00 2001
From: Joseph Capriotti <josephrcapriotti@gmail.com>
Date: Thu, 26 Oct 2023 16:58:43 -0600
Subject: [PATCH 02/10] re-enable all the cython ufuncs

---
 .../_extensions/potential_field_prism.pyx     | 536 +++++++++---------
 1 file changed, 260 insertions(+), 276 deletions(-)

diff --git a/geoana/kernels/_extensions/potential_field_prism.pyx b/geoana/kernels/_extensions/potential_field_prism.pyx
index 89149046..b123c493 100644
--- a/geoana/kernels/_extensions/potential_field_prism.pyx
+++ b/geoana/kernels/_extensions/potential_field_prism.pyx
@@ -1,59 +1,57 @@
 cimport cython
-cimport numpy as np
-import numpy as np
 
 from libc.math cimport sqrt, log, atan
 
-# @cython.ufunc
-# @cython.cdivision
-# cdef double prism_f(double x, double y, double z) nogil:
-#     """Evaluates the indefinite volume integral for the 1/r kernel.
-#
-#     This is used to evaluate the gravitational potential of dense prisms.
-#
-#     Parameters
-#     ----------
-#     x, y, z : (...) numpy.ndarray
-#         The nodal locations to evaluate the function at
-#
-#     Returns
-#     -------
-#     (...) numpy.ndarray
-#     """
-#     cdef:
-#         double v = 0.0
-#         double r, temp
-#
-#     r = sqrt(x * x + y * y + z * z)
-#     if x != 0.0 and y != 0.0:
-#         temp = z + r
-#         if temp > 0.0:
-#             v -= x * y * log(temp)
-#         v += 0.5 * x * x * atan( y * z / (x * r))
-#     if y != 0.0 and z != 0.0:
-#         temp = x + r
-#         if temp > 0.0:
-#             v -= y * z * log(temp)
-#         v += 0.5 * y * y * atan(z * x / (y * r))
-#     if z != 0.0 and x != 0.0:
-#         temp = y + r
-#         if temp > 0.0:
-#             v -= z * x * log(temp)
-#         v += 0.5 * z * z * atan(x * y / (z * r))
-#     return v
+@cython.cdivision
+@cython.ufunc
+cdef api double prism_f(double x, double y, double z) nogil:
+    """Evaluates the indefinite volume integral for the 1/r kernel.
+
+    This is used to evaluate the gravitational potential of dense prisms.
 
+    Parameters
+    ----------
+    x, y, z : (...) numpy.ndarray
+        The nodal locations to evaluate the function at
+
+    Returns
+    -------
+    (...) numpy.ndarray
+    """
+    cdef:
+        double v = 0.0
+        double r, temp
+
+    r = sqrt(x * x + y * y + z * z)
+    if x != 0.0 and y != 0.0:
+        temp = z + r
+        if temp > 0.0:
+            v -= x * y * log(temp)
+        v += 0.5 * x * x * atan( y * z / (x * r))
+    if y != 0.0 and z != 0.0:
+        temp = x + r
+        if temp > 0.0:
+            v -= y * z * log(temp)
+        v += 0.5 * y * y * atan(z * x / (y * r))
+    if z != 0.0 and x != 0.0:
+        temp = y + r
+        if temp > 0.0:
+            v -= z * x * log(temp)
+        v += 0.5 * z * z * atan(x * y / (z * r))
+    return v
 
 @cython.cdivision
-cdef api double c_prism_fz(double x, double y, double z) nogil:
+@cython.ufunc
+cdef api double prism_fz(double x, double y, double z) nogil:
     """Evaluates the indefinite volume integral for the d/dz * 1/r kernel.
 
     This is used to evaluate the gravitational field of dense prisms.
-    
+
     Parameters
     ----------
     x, y, z : (...) numpy.ndarray
         The nodal locations to evaluate the function at
-    
+
     Returns
     -------
     (...) numpy.ndarray
@@ -78,237 +76,223 @@ cdef api double c_prism_fz(double x, double y, double z) nogil:
     if z != 0.0:
         v -= z * atan(x * y / (z * r))
     return v
-def prism_fz(x, y, z):
-    cdef double[:] x_, y_, z_, out_
-    out = np.empty_like(x)
-
-    x_ = np.asarray(x).ravel()
-    y_ = np.asarray(y).ravel()
-    z_ = np.asarray(z).ravel()
-    out_ = np.asarray(out).ravel()
-
-    cdef int i
-    for i in range(x_.shape[0]):
-        out_[i] = c_prism_fz(x_[i], y_[i], z_[i])
-    return out
-
-
-
-# @cython.ufunc
-# @cython.cdivision
-# cdef double prism_fzz(double x, double y, double z) nogil:
-#     """Evaluates the indefinite volume integral for the d**2/dz**2 * 1/r kernel.
-#
-#     This is used to evaluate the gravitational gradient of dense prisms.
-#
-#     Parameters
-#     ----------
-#     x, y, z : (...) numpy.ndarray
-#         The nodal locations to evaluate the function at
-#
-#     Returns
-#     -------
-#     (...) numpy.ndarray
-#
-#     Notes
-#     -----
-#     Can be used to compute other components by cycling the inputs
-#     """
-#     cdef:
-#         double v = 0.0
-#         double r
-#
-#     if z != 0.0:
-#         r = sqrt(x * x + y * y + z * z)
-#         v = atan(x * y / (z * r))
-#     return v
-#
-#
-# @cython.ufunc
-# @cython.cdivision
-# cdef double prism_fzx(double x, double y, double z) nogil:
-#     """Evaluates the indefinite volume integral for the d**2/(dz*dx) * 1/r kernel.
-#
-#     This is used to evaluate the gravitational gradient of dense prisms.
-#
-#     Parameters
-#     ----------
-#     x, y, z : (...) numpy.ndarray
-#         The nodal locations to evaluate the function at
-#
-#     Returns
-#     -------
-#     (...) numpy.ndarray
-#
-#     Notes
-#     -----
-#     Can be used to compute other components by cycling the inputs
-#     """
-#     cdef:
-#         double v = 0.0
-#         double r
-#     r = sqrt(x * x + y * y + z * z)
-#     v = y + r
-#     if v == 0.0:
-#         if y < 0:
-#             v = log(-2 * y)
-#     else:
-#         v = -log(v)
-#     return v
-#
-#
-# @cython.ufunc
-# @cython.cdivision
-# cdef double prism_fzy(double x, double y, double z) nogil:
-#     """Evaluates the indefinite volume integral for the d**2/(dz*dy) * 1/r kernel.
-#
-#     This is used to evaluate the gravitational gradient of dense prisms.
-#
-#     Parameters
-#     ----------
-#     x, y, z : (...) numpy.ndarray
-#         The nodal locations to evaluate the function at
-#
-#     Returns
-#     -------
-#     (...) numpy.ndarray
-#
-#     Notes
-#     -----
-#     Can be used to compute other components by cycling the inputs
-#     """
-#     cdef:
-#         double v = 0.0
-#         double r, temp
-#     r = sqrt(x * x + y * y + z * z)
-#     v = y + r
-#     if v == 0.0:
-#         if x < 0:
-#             v = log(-2 * x)
-#     else:
-#         v = -log(v)
-#     return v
-#
-#
-# @cython.ufunc
-# @cython.cdivision
-# cdef double prism_fzzz(double x, double y, double z) nogil:
-#     """Evaluates the indefinite volume integral for the d**3/(dz**3) * 1/r kernel.
-#
-#     This is used to evaluate the magnetic gradient of susceptible prisms.
-#
-#     Parameters
-#     ----------
-#     x, y, z : (...) numpy.ndarray
-#         The nodal locations to evaluate the function at
-#
-#     Returns
-#     -------
-#     (...) numpy.ndarray
-#
-#     Notes
-#     -----
-#     Can be used to compute other components by cycling the inputs
-#     """
-#     cdef:
-#         double v = 0.0
-#         double r, v1, v2
-#     r = sqrt(x * x + y * y + z * z)
-#     v1 = x * x + z * z
-#     v2 = y * y + z * z
-#     if v1 != 0.0:
-#         v += 1.0/v1
-#     if v2 != 0.0:
-#         v += 1.0/v2
-#     if r != 0.0:
-#         v *= x * y / r
-#     return v
-#
-#
-# @cython.ufunc
-# @cython.cdivision
-# cdef double prism_fxxy(double x, double y, double z) nogil:
-#     """Evaluates the indefinite volume integral for the d**3/(dx**2 * dy) * 1/r kernel.
-#
-#     This is used to evaluate the magnetic gradient of susceptible prisms.
-#
-#     Parameters
-#     ----------
-#     x, y, z : (...) numpy.ndarray
-#         The nodal locations to evaluate the function at
-#
-#     Returns
-#     -------
-#     (...) numpy.ndarray
-#
-#     Notes
-#     -----
-#     Can be used to compute other components by cycling the inputs
-#     """
-#     cdef:
-#         double v = 0.0
-#         double r
-#     if x != 0.0:
-#         v = x * x + y * y
-#         r = sqrt(x * x + y * y + z * z)
-#         v = - x * z / (v * r)
-#     return v
-#
-#
-# @cython.ufunc
-# @cython.cdivision
-# cdef double prism_fxxz(double x, double y, double z) nogil:
-#     """Evaluates the indefinite volume integral for the d**3/(dx**2 * dz) * 1/r kernel.
-#
-#     This is used to evaluate the magnetic gradient of susceptible prisms.
-#
-#     Parameters
-#     ----------
-#     x, y, z : (...) numpy.ndarray
-#         The nodal locations to evaluate the function at
-#
-#     Returns
-#     -------
-#     (...) numpy.ndarray
-#
-#     Notes
-#     -----
-#     Can be used to compute other components by cycling the inputs
-#     """
-#     cdef:
-#         double v = 0.0
-#         double r
-#     if x != 0.0:
-#         v = x * x + z * z
-#         r = sqrt(x * x + y * y + z * z)
-#         v = - x * y / (v * r)
-#     return v
-#
-#
-# @cython.ufunc
-# @cython.cdivision
-# cdef double prism_fxyz(double x, double y, double z) nogil:
-#     """Evaluates the indefinite volume integral for the d**3/(dx * dy * dz) * 1/r kernel.
-#
-#     This is used to evaluate the magnetic gradient of susceptible prisms.
-#
-#     Parameters
-#     ----------
-#     x, y, z : (...) numpy.ndarray
-#         The nodal locations to evaluate the function at
-#
-#     Returns
-#     -------
-#     (...) numpy.ndarray
-#
-#     Notes
-#     -----
-#     Can be used to compute other components by cycling the inputs
-#     """
-#     cdef:
-#         double v = 0.0
-#         double r
-#     r = sqrt(x * x + y * y + z * z)
-#     if r != 0.0:
-#         v = 1.0/r
-#     return v
\ No newline at end of file
+
+
+@cython.ufunc
+@cython.cdivision
+cdef api double prism_fzz(double x, double y, double z) nogil:
+    """Evaluates the indefinite volume integral for the d**2/dz**2 * 1/r kernel.
+
+    This is used to evaluate the gravitational gradient of dense prisms.
+
+    Parameters
+    ----------
+    x, y, z : (...) numpy.ndarray
+        The nodal locations to evaluate the function at
+
+    Returns
+    -------
+    (...) numpy.ndarray
+
+    Notes
+    -----
+    Can be used to compute other components by cycling the inputs
+    """
+    cdef:
+        double v = 0.0
+        double r
+
+    if z != 0.0:
+        r = sqrt(x * x + y * y + z * z)
+        v = atan(x * y / (z * r))
+    return v
+
+
+@cython.ufunc
+@cython.cdivision
+cdef api double prism_fzx(double x, double y, double z) nogil:
+    """Evaluates the indefinite volume integral for the d**2/(dz*dx) * 1/r kernel.
+
+    This is used to evaluate the gravitational gradient of dense prisms.
+
+    Parameters
+    ----------
+    x, y, z : (...) numpy.ndarray
+        The nodal locations to evaluate the function at
+
+    Returns
+    -------
+    (...) numpy.ndarray
+
+    Notes
+    -----
+    Can be used to compute other components by cycling the inputs
+    """
+    cdef:
+        double v = 0.0
+        double r
+    r = sqrt(x * x + y * y + z * z)
+    v = y + r
+    if v == 0.0:
+        if y < 0:
+            v = log(-2 * y)
+    else:
+        v = -log(v)
+    return v
+
+
+@cython.ufunc
+@cython.cdivision
+cdef api double prism_fzy(double x, double y, double z) nogil:
+    """Evaluates the indefinite volume integral for the d**2/(dz*dy) * 1/r kernel.
+
+    This is used to evaluate the gravitational gradient of dense prisms.
+
+    Parameters
+    ----------
+    x, y, z : (...) numpy.ndarray
+        The nodal locations to evaluate the function at
+
+    Returns
+    -------
+    (...) numpy.ndarray
+
+    Notes
+    -----
+    Can be used to compute other components by cycling the inputs
+    """
+    cdef:
+        double v = 0.0
+        double r, temp
+    r = sqrt(x * x + y * y + z * z)
+    v = y + r
+    if v == 0.0:
+        if x < 0:
+            v = log(-2 * x)
+    else:
+        v = -log(v)
+    return v
+
+
+@cython.ufunc
+@cython.cdivision
+cdef api double prism_fzzz(double x, double y, double z) nogil:
+    """Evaluates the indefinite volume integral for the d**3/(dz**3) * 1/r kernel.
+
+    This is used to evaluate the magnetic gradient of susceptible prisms.
+
+    Parameters
+    ----------
+    x, y, z : (...) numpy.ndarray
+        The nodal locations to evaluate the function at
+
+    Returns
+    -------
+    (...) numpy.ndarray
+
+    Notes
+    -----
+    Can be used to compute other components by cycling the inputs
+    """
+    cdef:
+        double v = 0.0
+        double r, v1, v2
+    r = sqrt(x * x + y * y + z * z)
+    v1 = x * x + z * z
+    v2 = y * y + z * z
+    if v1 != 0.0:
+        v += 1.0/v1
+    if v2 != 0.0:
+        v += 1.0/v2
+    if r != 0.0:
+        v *= x * y / r
+    return v
+
+
+@cython.ufunc
+@cython.cdivision
+cdef api double prism_fxxy(double x, double y, double z) nogil:
+    """Evaluates the indefinite volume integral for the d**3/(dx**2 * dy) * 1/r kernel.
+
+    This is used to evaluate the magnetic gradient of susceptible prisms.
+
+    Parameters
+    ----------
+    x, y, z : (...) numpy.ndarray
+        The nodal locations to evaluate the function at
+
+    Returns
+    -------
+    (...) numpy.ndarray
+
+    Notes
+    -----
+    Can be used to compute other components by cycling the inputs
+    """
+    cdef:
+        double v = 0.0
+        double r
+    if x != 0.0:
+        v = x * x + y * y
+        r = sqrt(x * x + y * y + z * z)
+        v = - x * z / (v * r)
+    return v
+
+
+@cython.ufunc
+@cython.cdivision
+cdef api double prism_fxxz(double x, double y, double z) nogil:
+    """Evaluates the indefinite volume integral for the d**3/(dx**2 * dz) * 1/r kernel.
+
+    This is used to evaluate the magnetic gradient of susceptible prisms.
+
+    Parameters
+    ----------
+    x, y, z : (...) numpy.ndarray
+        The nodal locations to evaluate the function at
+
+    Returns
+    -------
+    (...) numpy.ndarray
+
+    Notes
+    -----
+    Can be used to compute other components by cycling the inputs
+    """
+    cdef:
+        double v = 0.0
+        double r
+    if x != 0.0:
+        v = x * x + z * z
+        r = sqrt(x * x + y * y + z * z)
+        v = - x * y / (v * r)
+    return v
+
+
+@cython.ufunc
+@cython.cdivision
+cdef api double prism_fxyz(double x, double y, double z) nogil:
+    """Evaluates the indefinite volume integral for the d**3/(dx * dy * dz) * 1/r kernel.
+
+    This is used to evaluate the magnetic gradient of susceptible prisms.
+
+    Parameters
+    ----------
+    x, y, z : (...) numpy.ndarray
+        The nodal locations to evaluate the function at
+
+    Returns
+    -------
+    (...) numpy.ndarray
+
+    Notes
+    -----
+    Can be used to compute other components by cycling the inputs
+    """
+    cdef:
+        double v = 0.0
+        double r
+    r = sqrt(x * x + y * y + z * z)
+    if r != 0.0:
+        v = 1.0/r
+    return v
\ No newline at end of file

From 925a3a4ddeede65ce8cd5e9e8f505a9b048c4c97 Mon Sep 17 00:00:00 2001
From: Joseph Capriotti <josephrcapriotti@gmail.com>
Date: Thu, 26 Oct 2023 17:14:15 -0600
Subject: [PATCH 03/10] add optionally handles to the underlying c functions
 for prism kernels

---
 geoana/kernels/__init__.py | 58 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 58 insertions(+)

diff --git a/geoana/kernels/__init__.py b/geoana/kernels/__init__.py
index 55cc5e2a..6208470a 100644
--- a/geoana/kernels/__init__.py
+++ b/geoana/kernels/__init__.py
@@ -12,3 +12,61 @@
     prism_fxxz,
     prism_fxyz,
 )
+
+try:
+    from numba.extending import get_cython_function_address
+    import ctypes
+
+    def __as_ctypes_func(module, function, argument_types):
+        func_address = get_cython_function_address(module, function)
+        func_type = ctypes.CFUNCTYPE(*argument_types)
+        return func_type(func_address)
+
+    c_prism = __as_ctypes_func(
+        'geoana.kernels._extensions.potential_field_prism',
+        'prism_f',
+        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
+    )
+    c_prism_fz = __as_ctypes_func(
+        'geoana.kernels._extensions.potential_field_prism',
+        'prism_fz',
+        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
+    )
+    c_prism_fzz = __as_ctypes_func(
+        'geoana.kernels._extensions.potential_field_prism',
+        'prism_fzz',
+        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
+    )
+    c_prism_fzx = __as_ctypes_func(
+        'geoana.kernels._extensions.potential_field_prism',
+        'prism_fzx',
+        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
+    )
+    c_prism_fzy = __as_ctypes_func(
+        'geoana.kernels._extensions.potential_field_prism',
+        'prism_fzy',
+        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
+    )
+    c_prism_fzzz = __as_ctypes_func(
+        'geoana.kernels._extensions.potential_field_prism',
+        'prism_fzzz',
+        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
+    )
+    c_prism_fxxy = __as_ctypes_func(
+        'geoana.kernels._extensions.potential_field_prism',
+        'prism_fxxy',
+        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
+    )
+    c_prism_fxxz = __as_ctypes_func(
+        'geoana.kernels._extensions.potential_field_prism',
+        'prism_fxxz',
+        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
+    )
+    c_prism_fxyz = __as_ctypes_func(
+        'geoana.kernels._extensions.potential_field_prism',
+        'prism_fxyz',
+        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
+    )
+
+except ImportError:
+    pass

From d2163706c293f513c5d77bfd4f5f9d6c6a444a2d Mon Sep 17 00:00:00 2001
From: Joseph Capriotti <josephrcapriotti@gmail.com>
Date: Fri, 27 Oct 2023 15:25:39 -0600
Subject: [PATCH 04/10] register jitable prism functions for numba

---
 geoana/kernels/__init__.py             | 58 --------------------------
 geoana/kernels/_extensions/__init__.py | 54 ++++++++++++++++++++++++
 2 files changed, 54 insertions(+), 58 deletions(-)

diff --git a/geoana/kernels/__init__.py b/geoana/kernels/__init__.py
index 6208470a..55cc5e2a 100644
--- a/geoana/kernels/__init__.py
+++ b/geoana/kernels/__init__.py
@@ -12,61 +12,3 @@
     prism_fxxz,
     prism_fxyz,
 )
-
-try:
-    from numba.extending import get_cython_function_address
-    import ctypes
-
-    def __as_ctypes_func(module, function, argument_types):
-        func_address = get_cython_function_address(module, function)
-        func_type = ctypes.CFUNCTYPE(*argument_types)
-        return func_type(func_address)
-
-    c_prism = __as_ctypes_func(
-        'geoana.kernels._extensions.potential_field_prism',
-        'prism_f',
-        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
-    )
-    c_prism_fz = __as_ctypes_func(
-        'geoana.kernels._extensions.potential_field_prism',
-        'prism_fz',
-        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
-    )
-    c_prism_fzz = __as_ctypes_func(
-        'geoana.kernels._extensions.potential_field_prism',
-        'prism_fzz',
-        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
-    )
-    c_prism_fzx = __as_ctypes_func(
-        'geoana.kernels._extensions.potential_field_prism',
-        'prism_fzx',
-        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
-    )
-    c_prism_fzy = __as_ctypes_func(
-        'geoana.kernels._extensions.potential_field_prism',
-        'prism_fzy',
-        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
-    )
-    c_prism_fzzz = __as_ctypes_func(
-        'geoana.kernels._extensions.potential_field_prism',
-        'prism_fzzz',
-        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
-    )
-    c_prism_fxxy = __as_ctypes_func(
-        'geoana.kernels._extensions.potential_field_prism',
-        'prism_fxxy',
-        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
-    )
-    c_prism_fxxz = __as_ctypes_func(
-        'geoana.kernels._extensions.potential_field_prism',
-        'prism_fxxz',
-        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
-    )
-    c_prism_fxyz = __as_ctypes_func(
-        'geoana.kernels._extensions.potential_field_prism',
-        'prism_fxyz',
-        (ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
-    )
-
-except ImportError:
-    pass
diff --git a/geoana/kernels/_extensions/__init__.py b/geoana/kernels/_extensions/__init__.py
index e69de29b..0b1da38d 100644
--- a/geoana/kernels/_extensions/__init__.py
+++ b/geoana/kernels/_extensions/__init__.py
@@ -0,0 +1,54 @@
+try:
+    # register numba jitable versions of the prism functions
+    # if numba is available (and this module is installed).
+    from numba.extending import (
+        overload,
+        get_cython_function_address
+    )
+    from numba import types
+    import ctypes
+
+    from .potential_field_prism import (
+        prism_f,
+        prism_fz,
+        prism_fzz,
+        prism_fzx,
+        prism_fzy,
+        prism_fzzz,
+        prism_fxxy,
+        prism_fxxz,
+        prism_fxyz,
+    )
+    funcs = [
+        prism_f,
+        prism_fz,
+        prism_fzz,
+        prism_fzx,
+        prism_fzy,
+        prism_fzzz,
+        prism_fxxy,
+        prism_fxxz,
+        prism_fxyz,
+    ]
+
+    def _numba_register_prism_func(prism_func):
+        module = 'geoana.kernels._extensions.potential_field_prism'
+        name = prism_func.__name__
+
+        func_address = get_cython_function_address(module, name)
+        func_type = ctypes.CFUNCTYPE(ctypes.c_double, ctypes.c_double, ctypes.c_double, ctypes.c_double)
+        c_func = func_type(func_address)
+
+        @overload(prism_func)
+        def numba_func(x, y, z):
+            if isinstance(x, types.Float):
+                if isinstance(y, types.Float):
+                    if isinstance(z, types.Float):
+                        def f(x, y, z):
+                            return c_func(x, y, z)
+                        return f
+    for func in funcs:
+        _numba_register_prism_func(func)
+
+except ImportError as err:
+    pass
\ No newline at end of file

From 686ca8523ae5ecc74124a8205a90389a57719543 Mon Sep 17 00:00:00 2001
From: Joseph Capriotti <josephrcapriotti@gmail.com>
Date: Mon, 30 Oct 2023 09:57:17 -0600
Subject: [PATCH 05/10] fix fyz func

---
 geoana/kernels/_extensions/potential_field_prism.pyx | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/geoana/kernels/_extensions/potential_field_prism.pyx b/geoana/kernels/_extensions/potential_field_prism.pyx
index b123c493..07134a39 100644
--- a/geoana/kernels/_extensions/potential_field_prism.pyx
+++ b/geoana/kernels/_extensions/potential_field_prism.pyx
@@ -165,7 +165,7 @@ cdef api double prism_fzy(double x, double y, double z) nogil:
         double v = 0.0
         double r, temp
     r = sqrt(x * x + y * y + z * z)
-    v = y + r
+    v = x + r
     if v == 0.0:
         if x < 0:
             v = log(-2 * x)

From cc60e98f23e7eab7e2f7c42fd55b3214fc041efa Mon Sep 17 00:00:00 2001
From: Joseph Capriotti <josephrcapriotti@gmail.com>
Date: Mon, 30 Oct 2023 10:49:08 -0600
Subject: [PATCH 06/10] install all requirements with conda

---
 .github/workflows/test_with_conda.yml | 3 ++-
 requirements_dev.txt                  | 1 -
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/test_with_conda.yml b/.github/workflows/test_with_conda.yml
index e5b5f2bb..79439f1d 100644
--- a/.github/workflows/test_with_conda.yml
+++ b/.github/workflows/test_with_conda.yml
@@ -38,7 +38,8 @@ jobs:
         conda config --show
         conda install --quiet --yes pip numpy scipy "matplotlib>=3.5.0" ipython cython;
         conda install --quiet --yes discretize libdlf setuptools_scm utm pytest pytest-cov graphviz
-        pip install -r requirements_dev.txt
+        conda install --quiet --yes sphinx sphinx-gallery pydata-sphinx-theme numpydoc nbsphinx
+        conda install --quiet --yes pytest pytest-cov
     - name: Install Our Package
       run: |
         export BUILD_GEOANA_EXT=1
diff --git a/requirements_dev.txt b/requirements_dev.txt
index 56c838dc..501ab7fb 100644
--- a/requirements_dev.txt
+++ b/requirements_dev.txt
@@ -1,5 +1,4 @@
 sphinx
-sphinx_rtd_theme
 sphinx-gallery
 sphinx-toolbox
 sphinxcontrib-apidoc

From 6817243423da1b63fa996c8b60ccb4efea3cd622 Mon Sep 17 00:00:00 2001
From: Joseph Capriotti <josephrcapriotti@gmail.com>
Date: Mon, 30 Oct 2023 10:55:23 -0600
Subject: [PATCH 07/10] add sphinx toolbox

---
 .github/workflows/test_with_conda.yml | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/.github/workflows/test_with_conda.yml b/.github/workflows/test_with_conda.yml
index 79439f1d..8229ee18 100644
--- a/.github/workflows/test_with_conda.yml
+++ b/.github/workflows/test_with_conda.yml
@@ -38,8 +38,7 @@ jobs:
         conda config --show
         conda install --quiet --yes pip numpy scipy "matplotlib>=3.5.0" ipython cython;
         conda install --quiet --yes discretize libdlf setuptools_scm utm pytest pytest-cov graphviz
-        conda install --quiet --yes sphinx sphinx-gallery pydata-sphinx-theme numpydoc nbsphinx
-        conda install --quiet --yes pytest pytest-cov
+        conda install --quiet --yes sphinx sphinx-gallery sphinx-toolbox pydata-sphinx-theme numpydoc nbsphinx
     - name: Install Our Package
       run: |
         export BUILD_GEOANA_EXT=1

From 4e646bd73d40aeb89c3ee0ab649181dd431457fe Mon Sep 17 00:00:00 2001
From: Joseph Capriotti <josephrcapriotti@gmail.com>
Date: Mon, 30 Oct 2023 11:00:10 -0600
Subject: [PATCH 08/10] update ignore

---
 .gitignore | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/.gitignore b/.gitignore
index e54e277a..0ed7dd09 100644
--- a/.gitignore
+++ b/.gitignore
@@ -98,7 +98,11 @@ docs/api/generated/*
 
 # Jupyter
 *.ipynb
+
+#Cython generated files
 geoana/kernels/_extensions/rTE.cpp
+geoana/kernels/_extensions/potential_field_prism.c
+geoana/kernels/_extensions/potential_field_prism_api.h
 
 # setuptools_scm
 geoana/version.py

From 772b10a0f22be380a19597f76219e16d7ae958ae Mon Sep 17 00:00:00 2001
From: Joseph Capriotti <josephrcapriotti@gmail.com>
Date: Mon, 30 Oct 2023 13:04:14 -0600
Subject: [PATCH 09/10] add test for jitting

---
 tests/test_potential_prism.py | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/tests/test_potential_prism.py b/tests/test_potential_prism.py
index 949bfe7c..ff602bd3 100644
--- a/tests/test_potential_prism.py
+++ b/tests/test_potential_prism.py
@@ -6,6 +6,7 @@
 import geoana.kernels.potential_field_prism as pf
 import geoana.gravity as grav
 from geoana.em.static import MagneticPrism, MagneticDipoleWholeSpace
+from numba import njit
 
 
 class TestCompiledVsNumpy():
@@ -282,3 +283,35 @@ def test_grav_init_and_errors():
 
     with pytest.raises(TypeError):
         prism.rho = 'abc'
+
+
+@pytest.mark.parametrize('function', [
+    pf.prism_f,
+    pf.prism_fz,
+    pf.prism_fzx,
+    pf.prism_fzy,
+    pf.prism_fzz,
+    pf.prism_fzzz,
+    pf.prism_fxxy,
+    pf.prism_fxxz,
+    pf.prism_fxyz
+])
+def test_numba_jitting_nopython(function):
+
+    # create a vectorized jit function of it:
+
+    x = np.random.rand(10)
+    y = np.random.rand(10)
+    z = np.random.rand(10)
+    @njit
+    def jitted_func(x, y, z):
+        n = len(x)
+        out = np.empty_like(x)
+        for i in range(n):
+            out[i] = function(x[i], y[i], z[i])
+        return out
+
+    v1 = jitted_func(x, y, z)
+    v2 = function(x, y, z)
+
+    np.allclose(v1, v2)

From 57299d8021de7d34620941c2056499b65149abb7 Mon Sep 17 00:00:00 2001
From: Joseph Capriotti <josephrcapriotti@gmail.com>
Date: Mon, 30 Oct 2023 13:04:46 -0600
Subject: [PATCH 10/10] add numba to testing environment

---
 .github/workflows/test_with_conda.yml | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/.github/workflows/test_with_conda.yml b/.github/workflows/test_with_conda.yml
index 19c0600d..30a20df0 100644
--- a/.github/workflows/test_with_conda.yml
+++ b/.github/workflows/test_with_conda.yml
@@ -39,6 +39,7 @@ jobs:
           matplotlib
           jupyter
           utm
+          numba
           pytest
           pytest-cov
           sphinx
@@ -87,6 +88,7 @@ jobs:
           matplotlib
           jupyter
           utm
+          numba
           pytest
           pytest-cov
           sphinx