Merge pull request #29 from InvestmentSystems/25/array-deepcopy

array_deepcopy

Author: flexatone

arraykit.c

@@ -34,10 +34,10 @@
         }\
     } while (0)
 
-// Placeholder of not implemented functions.
-# define AK_NOT_IMPLEMENTED\
+// Placeholder of not implemented pathways / debugging.
+# define AK_NOT_IMPLEMENTED(msg)\
     do {\
-        PyErr_SetNone(PyExc_NotImplementedError);\
+        PyErr_SetString(PyExc_NotImplementedError, msg);\
         return NULL;\
     } while (0)
 
@@ -148,6 +148,65 @@ AK_ResolveDTypeIter(PyObject *dtypes)
     return resolved;
 }
 
+// Numpy implementation: https://github.com/numpy/numpy/blob/a14c41264855e44ebd6187d7541b5b8d59bb32cb/numpy/core/src/multiarray/methods.c#L1557
+PyObject*
+AK_ArrayDeepCopy(PyArrayObject *array, PyObject *memo)
+{
+    PyObject *id = PyLong_FromVoidPtr((PyObject*)array);
+    if (!id) {
+        return NULL;
+    }
+    PyObject *found = PyDict_GetItemWithError(memo, id);
+    if (found) { // found will be NULL if not in dict
+        Py_INCREF(found); // got a borrowed ref, increment first
+        Py_DECREF(id);
+        return found;
+    }
+    else if (PyErr_Occurred()) {
+        goto error;
+    }
+
+    // if dtype is object, call deepcopy with memo
+    PyObject *array_new;
+    PyArray_Descr *dtype = PyArray_DESCR(array); // borrowed ref
+
+    if (PyDataType_ISOBJECT(dtype)) {
+        PyObject *copy = PyImport_ImportModule("copy");
+        if (!copy) {
+            goto error;
+        }
+        PyObject *deepcopy = PyObject_GetAttrString(copy, "deepcopy");
+        Py_DECREF(copy);
+        if (!deepcopy) {
+            goto error;
+        }
+        array_new = PyObject_CallFunctionObjArgs(deepcopy, array, memo, NULL);
+        Py_DECREF(deepcopy);
+        if (!array_new) {
+            goto error;
+        }
+    }
+    else {
+        Py_INCREF(dtype); // PyArray_FromArray steals a reference
+        array_new = PyArray_FromArray(
+                array,
+                dtype,
+                NPY_ARRAY_ENSURECOPY);
+        if (!array_new || PyDict_SetItem(memo, id, array_new)) {
+            Py_XDECREF(array_new);
+            goto error;
+        }
+    }
+    // set immutable
+    PyArray_CLEARFLAGS((PyArrayObject *)array_new, NPY_ARRAY_WRITEABLE);
+    Py_DECREF(id);
+    return array_new;
+error:
+    Py_DECREF(id);
+    return NULL;
+}
+
+
 //------------------------------------------------------------------------------
 // AK module public methods
 //------------------------------------------------------------------------------
@@ -249,6 +308,26 @@ row_1d_filter(PyObject *Py_UNUSED(m), PyObject *a)
     return a;
 }
 
+//------------------------------------------------------------------------------
+// array utility
+
+// Specialized array deepcopy that stores immutable arrays in memo dict.
+static PyObject *
+array_deepcopy(PyObject *Py_UNUSED(m), PyObject *args)
+{
+    PyObject *array, *memo;
+    if (!PyArg_UnpackTuple(args, "array_deepcopy", 2, 2, &array, &memo)) {
+        return NULL;
+    }
+    AK_CHECK_NUMPY_ARRAY(array);
+    if (!PyDict_CheckExact(memo)) {
+        PyErr_Format(PyExc_TypeError, "expected a dict (got %s)",
+                Py_TYPE(memo)->tp_name);
+        return NULL;
+    }
+    return AK_ArrayDeepCopy((PyArrayObject*)array, memo);
+}
+
 //------------------------------------------------------------------------------
 // type resolution
 
@@ -544,6 +623,7 @@ static PyMethodDef arraykit_methods[] =  {
     {"column_2d_filter", column_2d_filter, METH_O, NULL},
     {"column_1d_filter", column_1d_filter, METH_O, NULL},
     {"row_1d_filter", row_1d_filter, METH_O, NULL},
+    {"array_deepcopy", array_deepcopy, METH_VARARGS, NULL},
     {"resolve_dtype", resolve_dtype, METH_VARARGS, NULL},
     {"resolve_dtype_iter", resolve_dtype_iter, METH_O, NULL},
     {NULL},

arraykit.pyi

@@ -25,5 +25,6 @@ def shape_filter(__array: np.array) -> np.ndarray: ...
 def column_2d_filter(__array: np.array) -> np.ndarray: ...
 def column_1d_filter(__array: np.array) -> np.ndarray: ...
 def row_1d_filter(__array: np.array) -> np.ndarray: ...
+def array_deepcopy(__array: np.array, memo: tp.Dict[int, tp.Any]) -> np.ndarray: ...
 def resolve_dtype(__d1: np.dtype, __d2: np.dtype) -> np.dtype: ...
 def resolve_dtype_iter(__dtypes: tp.Iterable[np.dtype]) -> np.dtype: ...

performance/main.py

@@ -15,6 +15,7 @@
 from performance.reference.util import row_1d_filter as row_1d_filter_ref
 from performance.reference.util import resolve_dtype as resolve_dtype_ref
 from performance.reference.util import resolve_dtype_iter as resolve_dtype_iter_ref
+from performance.reference.util import array_deepcopy as array_deepcopy_ref
 
 from performance.reference.array_go import ArrayGO as ArrayGOREF
 
@@ -27,6 +28,7 @@
 from arraykit import row_1d_filter as row_1d_filter_ak
 from arraykit import resolve_dtype as resolve_dtype_ak
 from arraykit import resolve_dtype_iter as resolve_dtype_iter_ak
+from arraykit import array_deepcopy as array_deepcopy_ak
 
 from arraykit import ArrayGO as ArrayGOAK
 
@@ -200,6 +202,33 @@ class ResolveDTypeIterREF(ResolveDTypeIter):
     entry = staticmethod(resolve_dtype_iter_ref)
 
 
+#-------------------------------------------------------------------------------
+class ArrayDeepcopy(Perf):
+    FUNCTIONS = ('memo_new', 'memo_shared')
+    NUMBER = 500
+
+    def pre(self):
+        self.array1 = np.arange(100_000)
+        self.array2 = np.full(100_000, None)
+        self.array2[0] = [np.nan] # add a mutable
+        self.memo = {}
+
+    def memo_new(self):
+        memo = {}
+        self.entry(self.array1, memo)
+        self.entry(self.array2, memo)
+
+    def memo_shared(self):
+        self.entry(self.array1, self.memo)
+        self.entry(self.array2, self.memo)
+
+class ArrayDeepcopyAK(ArrayDeepcopy):
+    entry = staticmethod(array_deepcopy_ak)
+
+class ArrayDeepcopyREF(ArrayDeepcopy):
+    entry = staticmethod(array_deepcopy_ref)
+
+
 #-------------------------------------------------------------------------------
 class ArrayGOPerf(Perf):
     NUMBER = 1000
@@ -240,6 +269,7 @@ def main():
 
     records = [('cls', 'func', 'ak', 'ref', 'ref/ak')]
     for cls_perf in Perf.__subclasses__(): # only get one level
+        print(cls_perf)
         cls_map = {}
         if match and cls_perf.__name__ not in match:
             continue

test/test_util.py

@@ -10,6 +10,7 @@
 from arraykit import row_1d_filter
 from arraykit import mloc
 from arraykit import immutable_filter
+from arraykit import array_deepcopy
 
 from performance.reference.util import mloc as mloc_ref
 
@@ -167,6 +168,61 @@ def test_row_1d_filter_a(self) -> None:
         with self.assertRaises(NotImplementedError):
             row_1d_filter(a1.reshape(1,2,5))
 
+    #---------------------------------------------------------------------------
+
+    def test_array_deepcopy_a1(self) -> None:
+        a1 = np.arange(10)
+        memo = {}
+        a2 = array_deepcopy(a1, memo)
+
+        self.assertNotEqual(id(a1), id(a2))
+        self.assertNotEqual(mloc(a1), mloc(a2))
+        self.assertFalse(a2.flags.writeable)
+        self.assertEqual(a1.dtype, a2.dtype)
+
+    def test_array_deepcopy_a2(self) -> None:
+        a1 = np.arange(10)
+        memo = {}
+        a2 = array_deepcopy(a1, memo)
+
+        self.assertNotEqual(id(a1), id(a2))
+        self.assertNotEqual(mloc(a1), mloc(a2))
+        self.assertIn(id(a1), memo)
+        self.assertEqual(memo[id(a1)].tolist(), a2.tolist())
+        self.assertFalse(a2.flags.writeable)
+
+
+    def test_array_deepcopy_b(self) -> None:
+        a1 = np.arange(10)
+        memo = {id(a1): a1}
+        a2 = array_deepcopy(a1, memo)
+
+        self.assertEqual(mloc(a1), mloc(a2))
+
+
+    def test_array_deepcopy_c1(self) -> None:
+        mutable = [np.nan]
+        memo = {}
+        a1 = np.array((None, 'foo', True, mutable))
+        a2 = array_deepcopy(a1, memo)
+
+        self.assertNotEqual(id(a1), id(a2))
+        self.assertNotEqual(mloc(a1), mloc(a2))
+        self.assertNotEqual(id(a1[3]), id(a2[3]))
+        self.assertFalse(a2.flags.writeable)
+
+    def test_array_deepcopy_c2(self) -> None:
+        memo = {}
+        mutable = [np.nan]
+        a1 = np.array((None, 'foo', True, mutable))
+        a2 = array_deepcopy(a1, memo)
+        self.assertNotEqual(id(a1), id(a2))
+        self.assertNotEqual(mloc(a1), mloc(a2))
+        self.assertNotEqual(id(a1[3]), id(a2[3]))
+        self.assertFalse(a2.flags.writeable)
+        self.assertIn(id(a1), memo)
+
+
 if __name__ == '__main__':
     unittest.main()