adding ufuncs

numpy · ngoldbaum · Aug 16, 2024 · Jul 29, 2024 · Jul 29, 2024 · Jul 31, 2024
commit 96ee075cdc6058d0363612b678c22433a4544cf7
diff --git a/quaddtype/meson.build b/quaddtype/meson.build
@@ -5,7 +5,7 @@ py = py_mod.find_installation()
 
 c = meson.get_compiler('c')
 
-sleef_dep = c.find_library('sleef')
+sleef_dep = c.find_library('sleef', dirs:['/usr/local/lib'])
 sleefquad_dep = c.find_library('sleefquad')
 
 incdir_numpy = run_command(py,

diff --git a/quaddtype/quaddtype/src/ops.hpp b/quaddtype/quaddtype/src/ops.hpp
@@ -0,0 +1,33 @@
+#include <sleef.h>
+#include <sleefquad.h>
+
+typedef int (*unary_op_def)(Sleef_quad *, Sleef_quad *);
+
+static inline int
+quad_negative(Sleef_quad *op, Sleef_quad *out)
+{
+    *out = Sleef_negq1(*op);
+    return 0;
+}
+
+static inline int
+quad_absolute(Sleef_quad *op, Sleef_quad *out)
+{
+    *out = Sleef_fabsq1(*op);
+    return 0;
+}
+
+// binary ops
+typedef int (*binop_def)(Sleef_quad *, Sleef_quad *, Sleef_quad *);
+
+static inline int quad_add(Sleef_quad *out, Sleef_quad *in1, Sleef_quad *in2)
+{
+    *out = Sleef_addq1_u05(*in1, *in2);
+    return 0;
+}
+
+static inline int quad_sub(Sleef_quad *out, Sleef_quad *in1, Sleef_quad *in2)
+{
+    *out = Sleef_subq1_u05(*in1, *in2);
+    return 0;
+}
diff --git a/quaddtype/quaddtype/src/umath.cpp b/quaddtype/quaddtype/src/umath.cpp
@@ -0,0 +1,224 @@
+#define PY_ARRAY_UNIQUE_SYMBOL QuadPrecType_ARRAY_API
+#define PY_UFUNC_UNIQUE_SYMBOL QuadPrecType_UFUNC_API
+#define NPY_NO_DEPRECATED_API NPY_2_0_API_VERSION
+#define NPY_TARGET_VERSION NPY_2_0_API_VERSION
+#define NO_IMPORT_ARRAY
+#define NO_IMPORT_UFUNC
+
+extern "C" {
+    #include <Python.h>
+
+    #include "numpy/arrayobject.h"
+    #include "numpy/ndarraytypes.h"
+    #include "numpy/ufuncobject.h"
+
+    #include "numpy/dtype_api.h"
+}
+
+#include "scalar.h"
+#include "dtype.h"
+#include "umath.h"
+#include "ops.hpp"
+
+template <unary_op_def unary_op>
+int quad_generic_unary_op_strided_loop(PyArrayMethod_Context *context,
+        char *const data[], npy_intp const dimensions[],
+        npy_intp const strides[], NpyAuxData *auxdata)
+{
+    npy_intp N = dimensions[0];
+    char *in_ptr = data[0];
+    char *out_ptr = data[1];
+    npy_intp in_stride = strides[0];
+    npy_intp out_stride = strides[1];
+
+    while (N--) 
+    {
+        unary_op((Sleef_quad *)in_ptr, (Sleef_quad *)out_ptr);
+        in_ptr += in_stride;
+        out_ptr += out_stride;
+    }
+    return 0;
+}
+
+static NPY_CASTING
+quad_unary_op_resolve_descriptors(PyObject *self,
+        PyArray_DTypeMeta *dtypes[], QuadPrecDTypeObject *given_descrs[],
+        QuadPrecDTypeObject *loop_descrs[], npy_intp *unused)
+{
+    Py_INCREF(given_descrs[0]);
+    loop_descrs[0] = given_descrs[0];
+
+    if (given_descrs[1] == NULL) {
+        Py_INCREF(given_descrs[0]);
+        loop_descrs[1] = given_descrs[0];
+        return NPY_NO_CASTING;
+    }
+    Py_INCREF(given_descrs[1]);
+    loop_descrs[1] = given_descrs[1];
+
+    return NPY_NO_CASTING;  // Quad precision is always the same precision
+}
+
+template <unary_op_def unary_op>
+int create_quad_unary_ufunc(PyObject *numpy, const char *ufunc_name)
+{
+    PyObject *ufunc = PyObject_GetAttrString(numpy, ufunc_name);
+    if (ufunc == NULL) {
+        return -1;
+    }
+
+    PyArray_DTypeMeta *dtypes[2] = {
+       &QuadPrecDType, &QuadPrecDType};
+
+    PyType_Slot slots[] = {
+       {NPY_METH_resolve_descriptors, (void *)&quad_unary_op_resolve_descriptors},
+       {NPY_METH_strided_loop, (void *)&quad_generic_unary_op_strided_loop<unary_op>},
+       {0, NULL}
+    };
+
+    PyArrayMethod_Spec Spec = {
+        .name = "quad_unary_op",
+        .nin = 1,
+        .nout = 1,
+        .casting = NPY_NO_CASTING,
+        .flags = (NPY_ARRAYMETHOD_FLAGS)0,
+        .dtypes = dtypes,
+        .slots = slots,
+    };
+
+    if (PyUFunc_AddLoopFromSpec(ufunc, &Spec) < 0) {
+        return -1;
+    }
+
+    return 0;
+}
+
+int init_quad_unary_ops(PyObject *numpy)
+{
+    if (create_quad_unary_ufunc<quad_negative>(numpy, "negative") < 0) {
+        return -1;
+    }
+    if (create_quad_unary_ufunc<quad_absolute>(numpy, "absolute") < 0) {
+        return -1;
+    }
+    return 0;
+}
+
+// Binary ufuncs
+
+template <binop_def binop>
+int quad_generic_binop_strided_loop(PyArrayMethod_Context *context,
+        char *const data[], npy_intp const dimensions[],
+        npy_intp const strides[], NpyAuxData *auxdata)
+{
+    npy_intp N = dimensions[0];
+    char *in1_ptr = data[0], *in2_ptr = data[1];
+    char *out_ptr = data[2];
+    npy_intp in1_stride = strides[0];
+    npy_intp in2_stride = strides[1];
+    npy_intp out_stride = strides[2];
+
+    while (N--) {
+        binop((Sleef_quad *)out_ptr, (Sleef_quad *)in1_ptr, (Sleef_quad *)in2_ptr);
+
+        in1_ptr += in1_stride;
+        in2_ptr += in2_stride;
+        out_ptr += out_stride;
+    }
+    return 0;
+}
+
+static NPY_CASTING
+quad_binary_op_resolve_descriptors(PyObject *self,
+        PyArray_DTypeMeta *dtypes[], QuadPrecDTypeObject *given_descrs[],
+        QuadPrecDTypeObject *loop_descrs[], npy_intp *unused)
+{
+    Py_INCREF(given_descrs[0]);
+    loop_descrs[0] = given_descrs[0];
+    Py_INCREF(given_descrs[1]);
+    loop_descrs[1] = given_descrs[1];
+
+    if (given_descrs[2] == NULL) {
+        Py_INCREF(given_descrs[0]);
+        loop_descrs[2] = given_descrs[0];
+    }
+    else {
+        Py_INCREF(given_descrs[2]);
+        loop_descrs[2] = given_descrs[2];
+    }
+
+    return NPY_NO_CASTING;  // Quad precision is always the same precision
+}
+
+// todo: skipping the promoter for now, since same type operation will be requried
+
+template <binop_def binop>
+int create_quad_binary_ufunc(PyObject *numpy, const char *ufunc_name)
+{
+    PyObject *ufunc = PyObject_GetAttrString(numpy, ufunc_name);
+    if (ufunc == NULL) {
+        return -1;
+    }
+
+    PyArray_DTypeMeta *dtypes[3] = {
+       &QuadPrecDType, &QuadPrecDType, &QuadPrecDType};
+
+    PyType_Slot slots[] = {
+       {NPY_METH_resolve_descriptors,
+            (void *)&quad_binary_op_resolve_descriptors},
+       {NPY_METH_strided_loop,
+            (void *)&quad_generic_binop_strided_loop<binop>},
+       {0, NULL}
+    };
+
+    PyArrayMethod_Spec Spec = {
+        .name = "quad_binop",
+        .nin = 2,
+        .nout = 1,
+        .casting = NPY_NO_CASTING,
+        .flags = (NPY_ARRAYMETHOD_FLAGS)0,
+        .dtypes = dtypes,
+        .slots = slots,
+    };
+
+    if (PyUFunc_AddLoopFromSpec(ufunc, &Spec) < 0) {
+        return -1;
+    }
+
+    return 0;
+}
+
+int init_quad_binary_ops(PyObject *numpy)
+{
+    if (create_quad_binary_ufunc<quad_add>(numpy, "add") < 0) {
+        return -1;
+    }
+    if (create_quad_binary_ufunc<quad_sub>(numpy, "subtract") < 0) {
+        return -1;
+    }
+
+    return 0;
+}
+
+int init_quad_umath(void)
+{
+    PyObject * numpy = PyImport_ImportModule("numpy");
+    if (!numpy)
+        return -1;
+
+    if (init_quad_unary_ops(numpy) < 0) {
+        goto err;
+    }
+
+    if (init_quad_binary_ops(numpy) < 0) {
+        goto err;
+    }
+
+    Py_DECREF(numpy);
+    return 0;
+
+    err:
+        Py_DECREF(numpy);
+        return -1;
+
+}
diff --git a/quaddtype/quaddtype/src/umath.h b/quaddtype/quaddtype/src/umath.h
@@ -0,0 +1,15 @@
+#ifndef _QUADDTYPE_UMATH_H
+#define _QUADDTYPE_UMATH_H
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+int
+init_quad_umath(void);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif
diff --git a/quaddtype/tests/test_arr.py → quaddtype/tests/test_quaddtype.py b/quaddtype/tests/test_arr.py → quaddtype/tests/test_quaddtype.py
diff --git a/temp.py b/temp.py
@@ -0,0 +1,111 @@
+import numpy as np
+from quaddtype import QuadPrecDType, QuadPrecision
+import matplotlib.pyplot as plt
+
+def get_color(t, interior_t):
+    epsilon = QuadPrecision("1e-10")
+
+    if abs(t - QuadPrecision(1.0)) < epsilon:
+        value = int(255 * float(interior_t))
+        return np.array([value, value, value], dtype=np.uint8)
+
+    t = np.power(t, 0.5)
+    t = np.mod(t * 20, 1.0)
+
+    if t < 0.16:
+        return np.array([0, int(255 * (t / 0.16)), int(128 + 127 * (t / 0.16))], dtype=np.uint8)
+    elif t < 0.33:
+        return np.array([0, 255, int(255 * (1 - (t - 0.16) / 0.17))], dtype=np.uint8)
+    elif t < 0.5:
+        return np.array([int(255 * ((t - 0.33) / 0.17)), 255, 0], dtype=np.uint8)
+    elif t < 0.66:
+        return np.array([255, int(255 * (1 - (t - 0.5) / 0.16)), 0], dtype=np.uint8)
+    elif t < 0.83:
+        return np.array([255, 0, int(255 * ((t - 0.66) / 0.17))], dtype=np.uint8)
+    else:
+        return np.array([int(255 * (1 - (t - 0.83) / 0.17)), 0, int(128 * ((t - 0.83) / 0.17))], dtype=np.uint8)
+
+def iterate_and_compute_derivatives(c, max_iter):
+    z = 0
+    dz = 1
+    dc = 0
+    dzdz = 0
+
+    for _ in range(max_iter):
+        dzdz = 2 * (z * dzdz + dz * dz)
+        dz = 2 * z * dz + dc
+        z = z * z + c
+        dc = 1
+
+    return z, dz, dc, dzdz
+
+def estimate_interior_distance(c, max_iter):
+    z, dz, dc, dzdz = iterate_and_compute_derivatives(c, max_iter)
+
+    dz_abs_sq = np.abs(dz) ** 2
+    numerator = 1 - dz_abs_sq
+
+    denominator = np.abs(dc * dz + dzdz * z * dc)
+
+    return numerator / denominator
+
+def mandelbrot(c, max_iter, radius2):
+    z = 0
+    for i in range(max_iter):
+        z = z * z + c
+        if np.abs(z) ** 2 > radius2:
+            log_zn = np.log(np.abs(z))
+            nu = np.log(log_zn / np.log(2)) / np.log(2)
+            return i + 1 - nu, z
+    return max_iter, z
+
+def mandelbrot_set(width, height, max_iter, center_r, center_i, zoom):
+    radius = 2.0
+    radius2 = radius * radius
+    zoom_q = 1 / zoom
+
+    x = np.linspace(center_r - radius / zoom, center_r + radius / zoom, width)
+    y = np.linspace(center_i - radius / zoom, center_i + radius / zoom, height)
+    c = x[np.newaxis, :] + 1j * y[:, np.newaxis]
+
+    smooth_iter, final_z = np.frompyfunc(lambda c: mandelbrot(c, max_iter, radius2), 1, 2)(c)
+    smooth_iter = smooth_iter.astype(np.float64)
+    final_z = final_z.astype(np.complex128)
+
+    img = np.zeros((height, width, 3), dtype=np.uint8)
+
+    interior_mask = smooth_iter == max_iter
+    interior_c = c[interior_mask]
+    interior_distance = np.frompyfunc(lambda c: estimate_interior_distance(c, max_iter), 1, 1)(interior_c)
+    interior_distance = interior_distance.astype(np.float64)
+    interior_t = interior_distance - np.floor(interior_distance)
+
+    exterior_mask = ~interior_mask
+    t = smooth_iter[exterior_mask] / max_iter
+
+    interior_colors = np.array(list(map(lambda t: get_color(1.0, t), interior_t)))
+    exterior_colors = np.array(list(map(lambda t: get_color(t, 0.0), t)))
+
+    img[interior_mask] = interior_colors
+    img[exterior_mask] = exterior_colors
+
+    return img
+
+def plot_mandelbrot(width, height, max_iter, center_r, center_i, zoom):
+    img_array = mandelbrot_set(width, height, max_iter, center_r, center_i, zoom)
+
+    plt.figure(figsize=(10, 10))
+    plt.imshow(img_array)
+    plt.axis('off')
+    plt.title(f'Mandelbrot Set (zoom: {zoom}, center: {center_r} + {center_i}i, iterations: {max_iter}, dtype: numpy.float64)')
+    plt.show()
+
+if __name__ == "__main__":
+    width = 800
+    height = 800
+    max_iter = 1000
+    center_r = -0.75
+    center_i = 0.0
+    zoom = 1.0
+
+    plot_mandelbrot(width, height, max_iter, center_r, center_i, zoom)