diff options
-rw-r--r-- | numpy/core/arrayprint.py | 7 | ||||
-rw-r--r-- | numpy/core/setup.py | 2 | ||||
-rw-r--r-- | numpy/core/src/common/umathmodule.h | 3 | ||||
-rw-r--r-- | numpy/core/src/multiarray/multiarraymodule.c | 2 | ||||
-rw-r--r-- | numpy/core/src/umath/_scaled_float_dtype.c | 496 | ||||
-rw-r--r-- | numpy/core/tests/test_custom_dtypes.py | 60 |
6 files changed, 569 insertions, 1 deletions
diff --git a/numpy/core/arrayprint.py b/numpy/core/arrayprint.py index f16bcfd39..2a4bef669 100644 --- a/numpy/core/arrayprint.py +++ b/numpy/core/arrayprint.py @@ -420,7 +420,9 @@ def _get_format_function(data, **options): dtype_ = data.dtype dtypeobj = dtype_.type formatdict = _get_formatdict(data, **options) - if issubclass(dtypeobj, _nt.bool_): + if dtypeobj is None: + return formatdict["numpystr"]() + elif issubclass(dtypeobj, _nt.bool_): return formatdict['bool']() elif issubclass(dtypeobj, _nt.integer): if issubclass(dtypeobj, _nt.timedelta64): @@ -1408,6 +1410,9 @@ def dtype_short_repr(dtype): >>> dt = np.int64([1, 2]).dtype >>> assert eval(dtype_short_repr(dt)) == dt """ + if type(dtype).__repr__ != np.dtype.__repr__: + # TODO: Custom repr for user DTypes, logic should likely move. + return repr(dtype) if dtype.names is not None: # structured dtypes give a list or tuple repr return str(dtype) diff --git a/numpy/core/setup.py b/numpy/core/setup.py index 2061eb510..29d309f74 100644 --- a/numpy/core/setup.py +++ b/numpy/core/setup.py @@ -933,6 +933,8 @@ def configuration(parent_package='',top_path=None): join('src', 'umath', 'scalarmath.c.src'), join('src', 'umath', 'ufunc_type_resolution.c'), join('src', 'umath', 'override.c'), + # For testing. Eventually, should use public API and be separate: + join('src', 'umath', '_scaled_float_dtype.c'), ] umath_deps = [ diff --git a/numpy/core/src/common/umathmodule.h b/numpy/core/src/common/umathmodule.h index 6998596ee..5c718a841 100644 --- a/numpy/core/src/common/umathmodule.h +++ b/numpy/core/src/common/umathmodule.h @@ -1,6 +1,9 @@ #include "__umath_generated.c" #include "__ufunc_api.c" +NPY_NO_EXPORT PyObject * +get_sfloat_dtype(PyObject *NPY_UNUSED(mod), PyObject *NPY_UNUSED(args)); + PyObject * add_newdoc_ufunc(PyObject *NPY_UNUSED(dummy), PyObject *args); PyObject * ufunc_frompyfunc(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *NPY_UNUSED(kwds)); int initumath(PyObject *m); diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index 5fde890de..ea9c10543 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -4430,6 +4430,8 @@ static struct PyMethodDef array_module_methods[] = { METH_VARARGS, NULL}, {"_add_newdoc_ufunc", (PyCFunction)add_newdoc_ufunc, METH_VARARGS, NULL}, + {"_get_sfloat_dtype", + get_sfloat_dtype, METH_NOARGS, NULL}, {"_set_madvise_hugepage", (PyCFunction)_set_madvise_hugepage, METH_O, NULL}, {"_reload_guard", (PyCFunction)_reload_guard, diff --git a/numpy/core/src/umath/_scaled_float_dtype.c b/numpy/core/src/umath/_scaled_float_dtype.c new file mode 100644 index 000000000..6ef7ea88f --- /dev/null +++ b/numpy/core/src/umath/_scaled_float_dtype.c @@ -0,0 +1,496 @@ +/* + * This file implements a basic scaled float64 DType. The reason is to have + * a simple parametric DType for testing. It is not meant to be a useful + * DType by itself, but due to the scaling factor has similar properties as + * a Unit DType. + * + * The code here should be seen as a work in progress. Some choices are made + * to test certain code paths, but that does not mean that they must not + * be modified. + * + * NOTE: The tests were initially written using private API and ABI, ideally + * they should be replaced/modified with versions using public API. + */ + +#define _UMATHMODULE +#define _MULTIARRAYMODULE +#define NPY_NO_DEPRECATED_API NPY_API_VERSION +#include "numpy/ndarrayobject.h" +#include "numpy/ufuncobject.h" + +#include "array_method.h" +#include "common.h" +#include "numpy/npy_math.h" +#include "convert_datatype.h" +#include "dtypemeta.h" + + +typedef struct { + PyArray_Descr base; + double scaling; +} PyArray_SFloatDescr; + +static PyArray_DTypeMeta PyArray_SFloatDType; +static PyArray_SFloatDescr SFloatSingleton; + + +static int +sfloat_is_known_scalar_type(PyArray_DTypeMeta *NPY_UNUSED(cls), PyTypeObject *type) +{ + /* Accept only floats (some others may work due to normal casting) */ + if (type == &PyFloat_Type) { + return 1; + } + return 0; +} + + +static PyArray_Descr * +sfloat_default_descr(PyArray_DTypeMeta *NPY_UNUSED(cls)) +{ + Py_INCREF(&SFloatSingleton); + return (PyArray_Descr *)&SFloatSingleton; +} + + +static PyArray_Descr * +sfloat_discover_from_pyobject(PyArray_DTypeMeta *cls, PyObject *NPY_UNUSED(obj)) +{ + return cls->default_descr(cls); +} + + +static PyArray_DTypeMeta * +sfloat_common_dtype(PyArray_DTypeMeta *cls, PyArray_DTypeMeta *other) +{ + if (other->legacy && other->type_num == NPY_DOUBLE) { + Py_INCREF(cls); + return cls; + } + Py_INCREF(Py_NotImplemented); + return (PyArray_DTypeMeta *)Py_NotImplemented; +} + + +static PyArray_Descr * +sfloat_common_instance(PyArray_Descr *descr1, PyArray_Descr *descr2) +{ + PyArray_SFloatDescr *sf1 = (PyArray_SFloatDescr *)descr1; + PyArray_SFloatDescr *sf2 = (PyArray_SFloatDescr *)descr2; + /* We make the choice of using the larger scaling */ + if (sf1->scaling >= sf2->scaling) { + Py_INCREF(descr1); + return descr1; + } + Py_INCREF(descr2); + return descr2; +} + + +/* + * Implement minimal getitem and setitem to make this DType mostly(?) safe to + * expose in Python. + * TODO: This should not use the old-style API, but the new-style is missing! +*/ + +static PyObject * +sfloat_getitem(char *data, PyArrayObject *arr) +{ + PyArray_SFloatDescr *descr = (PyArray_SFloatDescr *)PyArray_DESCR(arr); + double value; + + memcpy(&value, data, sizeof(double)); + return PyFloat_FromDouble(value * descr->scaling); +} + + +static int +sfloat_setitem(PyObject *obj, char *data, PyArrayObject *arr) +{ + if (!PyFloat_CheckExact(obj)) { + PyErr_SetString(PyExc_NotImplementedError, + "Currently only accepts floats"); + return -1; + } + + PyArray_SFloatDescr *descr = (PyArray_SFloatDescr *)PyArray_DESCR(arr); + double value = PyFloat_AsDouble(obj); + value /= descr->scaling; + + memcpy(data, &value, sizeof(double)); + return 0; +} + + +static PyArray_ArrFuncs arrfuncs = { + .getitem = (PyArray_GetItemFunc *)&sfloat_getitem, + .setitem = (PyArray_SetItemFunc *)&sfloat_setitem, +}; + + +static PyArray_SFloatDescr SFloatSingleton = {{ + .elsize = sizeof(double), + .alignment = _ALIGN(double), + .flags = NPY_USE_GETITEM|NPY_USE_SETITEM, + .type_num = -1, + .f = &arrfuncs, + .byteorder = '|', /* do not bother with byte-swapping... */ + }, + .scaling = 1, +}; + + +static PyArray_Descr * +sfloat_scaled_copy(PyArray_SFloatDescr *self, double factor) { + PyArray_SFloatDescr *new = PyObject_New( + PyArray_SFloatDescr, (PyTypeObject *)&PyArray_SFloatDType); + if (new == NULL) { + return NULL; + } + /* Don't copy PyObject_HEAD part */ + memcpy((char *)new + sizeof(PyObject), + (char *)self + sizeof(PyObject), + sizeof(PyArray_SFloatDescr) - sizeof(PyObject)); + + new->scaling = new->scaling * factor; + return (PyArray_Descr *)new; +} + + +PyObject * +python_sfloat_scaled_copy(PyArray_SFloatDescr *self, PyObject *arg) +{ + if (!PyFloat_Check(arg)) { + PyErr_SetString(PyExc_TypeError, + "Scaling factor must be a python float."); + return NULL; + } + double factor = PyFloat_AsDouble(arg); + + return (PyObject *)sfloat_scaled_copy(self, factor); +} + + +static PyObject * +sfloat_get_scaling(PyArray_SFloatDescr *self, PyObject *NPY_UNUSED(args)) +{ + return PyFloat_FromDouble(self->scaling); +} + + +PyMethodDef sfloat_methods[] = { + {"scaled_by", + (PyCFunction)python_sfloat_scaled_copy, METH_O, + "Method to get a dtype copy with different scaling, mainly to " + "avoid having to implement many ways to create new instances."}, + {"get_scaling", + (PyCFunction)sfloat_get_scaling, METH_NOARGS, NULL}, + {NULL, NULL, 0, NULL} +}; + + +static PyObject * +sfloat_new(PyTypeObject *NPY_UNUSED(cls), PyObject *args, PyObject *kwds) +{ + double scaling = 1.; + static char *kwargs_strs[] = {"scaling", NULL}; + + if (!PyArg_ParseTupleAndKeywords( + args, kwds, "|d:_ScaledFloatTestDType", kwargs_strs, &scaling)) { + return NULL; + } + if (scaling == 1.) { + Py_INCREF(&SFloatSingleton); + return (PyObject *)&SFloatSingleton; + } + return (PyObject *)sfloat_scaled_copy(&SFloatSingleton, scaling); +} + + +static PyObject * +sfloat_repr(PyArray_SFloatDescr *self) +{ + PyObject *scaling = PyFloat_FromDouble(self->scaling); + if (scaling == NULL) { + return NULL; + } + PyObject *res = PyUnicode_FromFormat( + "_ScaledFloatTestDType(scaling=%R)", scaling); + Py_DECREF(scaling); + return res; +} + + +static PyArray_DTypeMeta PyArray_SFloatDType = {{{ + PyVarObject_HEAD_INIT(NULL, 0) + .tp_name = "numpy._ScaledFloatTestDType", + .tp_methods = sfloat_methods, + .tp_new = sfloat_new, + .tp_repr = (reprfunc)sfloat_repr, + .tp_str = (reprfunc)sfloat_repr, + .tp_basicsize = sizeof(PyArray_SFloatDescr), + }}, + .type_num = -1, + .abstract = 0, + .legacy = 0, + .parametric = 1, + .f = &arrfuncs, + .scalar_type = NULL, + /* Special methods: */ + .default_descr = &sfloat_default_descr, + .discover_descr_from_pyobject = &sfloat_discover_from_pyobject, + .is_known_scalar_type = &sfloat_is_known_scalar_type, + .common_dtype = &sfloat_common_dtype, + .common_instance = &sfloat_common_instance, +}; + + +/* + * Implement some casts. + */ + +/* + * It would make more sense to test this early on, but this allows testing + * error returns. + */ +static int +check_factor(double factor) { + if (npy_isfinite(factor) && factor != 0.) { + return 0; + } + NPY_ALLOW_C_API_DEF; + NPY_ALLOW_C_API; + PyErr_SetString(PyExc_TypeError, + "error raised inside the core-loop: non-finite factor!"); + NPY_DISABLE_C_API; + return -1; +} + + +static int +cast_sfloat_to_sfloat_unaligned(PyArrayMethod_Context *context, + char *const data[], npy_intp const dimensions[], + npy_intp const strides[], NpyAuxData *NPY_UNUSED(auxdata)) +{ + /* could also be moved into auxdata: */ + double factor = ((PyArray_SFloatDescr *)context->descriptors[0])->scaling; + factor /= ((PyArray_SFloatDescr *)context->descriptors[1])->scaling; + if (check_factor(factor) < 0) { + return -1; + } + + npy_intp N = dimensions[0]; + char *in = data[0]; + char *out = data[1]; + for (npy_intp i = 0; i < N; i++) { + double tmp; + memcpy(&tmp, in, sizeof(double)); + tmp *= factor; + memcpy(out, &tmp, sizeof(double)); + + in += strides[0]; + out += strides[1]; + } + return 0; +} + + +static int +cast_sfloat_to_sfloat_aligned(PyArrayMethod_Context *context, + char *const data[], npy_intp const dimensions[], + npy_intp const strides[], NpyAuxData *NPY_UNUSED(auxdata)) +{ + /* could also be moved into auxdata: */ + double factor = ((PyArray_SFloatDescr *)context->descriptors[0])->scaling; + factor /= ((PyArray_SFloatDescr *)context->descriptors[1])->scaling; + if (check_factor(factor) < 0) { + return -1; + } + + npy_intp N = dimensions[0]; + char *in = data[0]; + char *out = data[1]; + for (npy_intp i = 0; i < N; i++) { + *(double *)out = *(double *)in * factor; + in += strides[0]; + out += strides[1]; + } + return 0; +} + + +static NPY_CASTING +sfloat_to_sfloat_resolve_descriptors( + PyArrayMethodObject *NPY_UNUSED(self), + PyArray_DTypeMeta *NPY_UNUSED(dtypes[2]), + PyArray_Descr *given_descrs[2], + PyArray_Descr *loop_descrs[2]) +{ + loop_descrs[0] = given_descrs[0]; + Py_INCREF(loop_descrs[0]); + + if (given_descrs[1] == NULL) { + loop_descrs[1] = given_descrs[0]; + } + else { + loop_descrs[1] = given_descrs[1]; + } + Py_INCREF(loop_descrs[1]); + + if (((PyArray_SFloatDescr *)loop_descrs[0])->scaling + == ((PyArray_SFloatDescr *)loop_descrs[1])->scaling) { + /* same scaling is just a view */ + return NPY_NO_CASTING | _NPY_CAST_IS_VIEW; + } + else if (-((PyArray_SFloatDescr *)loop_descrs[0])->scaling + == ((PyArray_SFloatDescr *)loop_descrs[1])->scaling) { + /* changing the sign does not lose precision */ + return NPY_EQUIV_CASTING; + } + /* Technically, this is not a safe cast, since over/underflows can occur */ + return NPY_SAME_KIND_CASTING; +} + + +/* + * Casting to and from doubles. + * + * To keep things interesting, we ONLY define the trivial cast with a factor + * of 1. All other casts have to be handled by the sfloat to sfloat cast. + * + * The casting machinery should optimize this step away normally, since we + * flag the this is a view. + */ +static int +cast_float_to_from_sfloat(PyArrayMethod_Context *NPY_UNUSED(context), + char *const data[], npy_intp const dimensions[], + npy_intp const strides[], NpyAuxData *NPY_UNUSED(auxdata)) +{ + npy_intp N = dimensions[0]; + char *in = data[0]; + char *out = data[1]; + for (npy_intp i = 0; i < N; i++) { + *(double *)out = *(double *)in; + in += strides[0]; + out += strides[1]; + } + return 0; +} + + +static NPY_CASTING +float_to_from_sfloat_resolve_descriptors( + PyArrayMethodObject *NPY_UNUSED(self), + PyArray_DTypeMeta *dtypes[2], + PyArray_Descr *NPY_UNUSED(given_descrs[2]), + PyArray_Descr *loop_descrs[2]) +{ + loop_descrs[0] = dtypes[0]->default_descr(dtypes[0]); + if (loop_descrs[0] == NULL) { + return -1; + } + loop_descrs[1] = dtypes[1]->default_descr(dtypes[1]); + if (loop_descrs[1] == NULL) { + return -1; + } + return NPY_NO_CASTING | _NPY_CAST_IS_VIEW; +} + + +static int +init_casts(void) +{ + PyArray_DTypeMeta *dtypes[2] = {&PyArray_SFloatDType, &PyArray_SFloatDType}; + PyType_Slot slots[4] = {{0, NULL}}; + PyArrayMethod_Spec spec = { + .name = "sfloat_to_sfloat_cast", + .nin = 1, + .nout = 1, + .flags = NPY_METH_SUPPORTS_UNALIGNED, + .dtypes = dtypes, + .slots = slots, + /* minimal guaranteed casting */ + .casting = NPY_SAME_KIND_CASTING, + }; + + slots[0].slot = NPY_METH_resolve_descriptors; + slots[0].pfunc = &sfloat_to_sfloat_resolve_descriptors; + + slots[1].slot = NPY_METH_strided_loop; + slots[1].pfunc = &cast_sfloat_to_sfloat_aligned; + + slots[2].slot = NPY_METH_unaligned_strided_loop; + slots[2].pfunc = &cast_sfloat_to_sfloat_unaligned; + + if (PyArray_AddCastingImplementation_FromSpec(&spec, 0)) { + return -1; + } + + spec.name = "float_to_sfloat_cast"; + /* Technically, it is just a copy currently so this is fine: */ + spec.flags = NPY_METH_NO_FLOATINGPOINT_ERRORS; + PyArray_DTypeMeta *double_DType = PyArray_DTypeFromTypeNum(NPY_DOUBLE); + Py_DECREF(double_DType); /* immortal anyway */ + dtypes[0] = double_DType; + + slots[0].slot = NPY_METH_resolve_descriptors; + slots[0].pfunc = &float_to_from_sfloat_resolve_descriptors; + slots[1].slot = NPY_METH_strided_loop; + slots[1].pfunc = &cast_float_to_from_sfloat; + slots[2].slot = 0; + slots[2].pfunc = NULL; + + if (PyArray_AddCastingImplementation_FromSpec(&spec, 0)) { + return -1; + } + + spec.name = "sfloat_to_float_cast"; + dtypes[0] = &PyArray_SFloatDType; + dtypes[1] = double_DType; + + if (PyArray_AddCastingImplementation_FromSpec(&spec, 0)) { + return -1; + } + + return 0; +} + + +/* + * Python entry point, exported via `umathmodule.h` and `multiarraymodule.c`. + * TODO: Should be moved when the necessary API is not internal anymore. + */ +NPY_NO_EXPORT PyObject * +get_sfloat_dtype(PyObject *NPY_UNUSED(mod), PyObject *NPY_UNUSED(args)) +{ + /* Allow calling the function multiple times. */ + static npy_bool initalized = NPY_FALSE; + + if (initalized) { + Py_INCREF(&PyArray_SFloatDType); + return (PyObject *)&PyArray_SFloatDType; + } + + PyArray_SFloatDType.super.ht_type.tp_base = &PyArrayDescr_Type; + + if (PyType_Ready((PyTypeObject *)&PyArray_SFloatDType) < 0) { + return NULL; + } + PyArray_SFloatDType.castingimpls = PyDict_New(); + if (PyArray_SFloatDType.castingimpls == NULL) { + return NULL; + } + + PyObject *o = PyObject_Init( + (PyObject *)&SFloatSingleton, (PyTypeObject *)&PyArray_SFloatDType); + if (o == NULL) { + return NULL; + } + + if (init_casts() < 0) { + return NULL; + } + + initalized = NPY_TRUE; + return (PyObject *)&PyArray_SFloatDType; +} diff --git a/numpy/core/tests/test_custom_dtypes.py b/numpy/core/tests/test_custom_dtypes.py new file mode 100644 index 000000000..fa7cbc047 --- /dev/null +++ b/numpy/core/tests/test_custom_dtypes.py @@ -0,0 +1,60 @@ +import pytest + +import numpy as np +from numpy.testing import assert_array_equal + + +SF = np.core._multiarray_umath._get_sfloat_dtype() + + +@pytest.mark.parametrize("scaling", [1., -1., 2.]) +def test_scaled_float_from_floats(scaling): + a = np.array([1., 2., 3.], dtype=SF(scaling)) + + assert a.dtype.get_scaling() == scaling + assert_array_equal(scaling * a.view(np.float64), np.array([1., 2., 3.])) + + +@pytest.mark.parametrize("scaling", [1., -1., 2.]) +def test_sfloat_from_float(scaling): + a = np.array([1., 2., 3.]).astype(dtype=SF(scaling)) + + assert a.dtype.get_scaling() == scaling + assert_array_equal(scaling * a.view(np.float64), np.array([1., 2., 3.])) + + +def _get_array(scaling, aligned=True): + if not aligned: + a = np.empty(3*8 + 1, dtype=np.uint8)[1:] + a = a.view(np.float64) + a[:] = [1., 2., 3.] + else: + a = np.array([1., 2., 3.]) + + a *= 1./scaling # the casting code also uses the reciprocal. + return a.view(SF(scaling)) + + +@pytest.mark.parametrize("aligned", [True, False]) +def test_sfloat_casts(aligned): + a = _get_array(1., aligned) + + assert np.can_cast(a, SF(-1.), casting="equiv") + assert not np.can_cast(a, SF(-1.), casting="no") + na = a.astype(SF(-1.)) + assert_array_equal(-1 * na.view(np.float64), a.view(np.float64)) + + assert np.can_cast(a, SF(2.), casting="same_kind") + assert not np.can_cast(a, SF(2.), casting="safe") + a2 = a.astype(SF(2.)) + assert_array_equal(2 * a2.view(np.float64), a.view(np.float64)) + + +@pytest.mark.parametrize("aligned", [True, False]) +def test_sfloat_cast_internal_errors(aligned): + a = _get_array(2e300, aligned) + + with pytest.raises(TypeError, + match="error raised inside the core-loop: non-finite factor!"): + a.astype(SF(2e-300)) + |