Merge pull request #117 from markflorisson88/_newaxis

Newaxis support for memoryview slices & avoid taking address of memoryviews for utility functions

Cython/Compiler/ExprNodes.py

@@ -2523,26 +2523,32 @@ class IndexNode(ExprNode):
             import MemoryView
 
             skip_child_analysis = True
+            newaxes = [newaxis for newaxis in indices if newaxis.is_none]
             have_slices, indices = MemoryView.unellipsify(indices,
+                                                          newaxes,
                                                           self.base.type.ndim)
-            self.memslice_index = len(indices) == self.base.type.ndim
+
+            self.memslice_index = (not newaxes and
+                                   len(indices) == self.base.type.ndim)
             axes = []
 
             index_type = PyrexTypes.c_py_ssize_t_type
             new_indices = []
 
-            if len(indices) > self.base.type.ndim:
+            if len(indices) - len(newaxes) > self.base.type.ndim:
                 self.type = error_type
                 return error(indices[self.base.type.ndim].pos,
                              "Too many indices specified for type %s" %
                                                         self.base.type)
 
-            suboffsets_dim = -1
+            axis_idx = 0
             for i, index in enumerate(indices[:]):
                 index.analyse_types(env)
-                access, packing = self.base.type.axes[i]
+                if not index.is_none:
+                    access, packing = self.base.type.axes[axis_idx]
+                    axis_idx += 1
+
                 if isinstance(index, SliceNode):
-                    suboffsets_dim = i
                     self.memslice_slice = True
                     if index.step.is_none:
                         axes.append((access, packing))
@@ -2558,6 +2564,11 @@ class IndexNode(ExprNode):
                             setattr(index, attr, value)
                             new_indices.append(value)
 
+                elif index.is_none:
+                    self.memslice_slice = True
+                    new_indices.append(index)
+                    axes.append(('direct', 'strided'))
+
                 elif index.type.is_int or index.type.is_pyobject:
                     if index.type.is_pyobject and not self.warned_untyped_idx:
                         warning(index.pos, "Index should be typed for more "
@@ -2633,6 +2644,11 @@ class IndexNode(ExprNode):
             self.index = None
             self.is_temp = True
             self.use_managed_ref = True
+
+            if not MemoryView.validate_axes(self.pos, axes):
+                self.type = error_type
+                return
+
             self.type = PyrexTypes.MemoryViewSliceType(
                             self.base.type.dtype, axes)
 

Cython/Compiler/MemoryView.py

@@ -277,8 +277,8 @@ class MemoryViewSliceBufferEntry(Buffer.BufferEntry):
         """
         Slice a memoryviewslice.
 
-        indices     - list of index nodes. If not a SliceNode, then it must be
-                      coercible to Py_ssize_t
+        indices     - list of index nodes. If not a SliceNode, or NoneNode,
+                      then it must be coercible to Py_ssize_t
 
         Simply call __pyx_memoryview_slice_memviewslice with the right
         arguments.
@@ -307,28 +307,14 @@ class MemoryViewSliceBufferEntry(Buffer.BufferEntry):
         code.putln("%(dst)s.memview = %(src)s.memview;" % locals())
         code.put_incref_memoryviewslice(dst)
 
-        for dim, index in enumerate(indices):
+        dim = -1
+        for index in indices:
             error_goto = code.error_goto(index.pos)
-
-            if not isinstance(index, ExprNodes.SliceNode):
-                # normal index
-                idx = index.result()
-
+            if not index.is_none:
+                dim += 1
                 access, packing = self.type.axes[dim]
-                if access == 'direct':
-                    indirect = False
-                else:
-                    indirect = True
-                    generic = (access == 'full')
-                    if new_ndim != 0:
-                        return error(index.pos,
-                                     "All preceding dimensions must be "
-                                     "indexed and not sliced")
-
-                d = locals()
-                code.put(load_slice_util("SliceIndex", d))
-            else:
 
+            if isinstance(index, ExprNodes.SliceNode):
                 # slice, unspecified dimension, or part of ellipsis
                 d = locals()
                 for s in "start stop step".split():
@@ -344,7 +330,6 @@ class MemoryViewSliceBufferEntry(Buffer.BufferEntry):
                     not d['have_step']):
                     # full slice (:), simply copy over the extent, stride
                     # and suboffset. Also update suboffset_dim if needed
-                    access, packing = self.type.axes[dim]
                     d['access'] = access
                     code.put(load_slice_util("SimpleSlice", d))
                 else:
@@ -352,6 +337,31 @@ class MemoryViewSliceBufferEntry(Buffer.BufferEntry):
 
                 new_ndim += 1
 
+            elif index.is_none:
+                # newaxis
+                attribs = [('shape', 1), ('strides', 0), ('suboffsets', -1)]
+                for attrib, value in attribs:
+                    code.putln("%s.%s[%d] = %d;" % (dst, attrib, new_ndim, value))
+
+                new_ndim += 1
+
+            else:
+                # normal index
+                idx = index.result()
+
+                if access == 'direct':
+                    indirect = False
+                else:
+                    indirect = True
+                    generic = (access == 'full')
+                    if new_ndim != 0:
+                        return error(index.pos,
+                                     "All preceding dimensions must be "
+                                     "indexed and not sliced")
+
+                d = locals()
+                code.put(load_slice_util("SliceIndex", d))
+
         if not no_suboffset_dim:
             code.funcstate.release_temp(suboffset_dim)
 
@@ -361,11 +371,13 @@ def empty_slice(pos):
     return ExprNodes.SliceNode(pos, start=none,
                                stop=none, step=none)
 
-def unellipsify(indices, ndim):
+def unellipsify(indices, newaxes, ndim):
     result = []
     seen_ellipsis = False
     have_slices = False
 
+    n_indices = len(indices) - len(newaxes)
+
     for index in indices:
         if isinstance(index, ExprNodes.EllipsisNode):
             have_slices = True
@@ -374,16 +386,19 @@ def unellipsify(indices, ndim):
             if seen_ellipsis:
                 result.append(full_slice)
             else:
-                nslices = ndim - len(indices) + 1
+                nslices = ndim - n_indices + 1
                 result.extend([full_slice] * nslices)
                 seen_ellipsis = True
         else:
-            have_slices = have_slices or isinstance(index, ExprNodes.SliceNode)
+            have_slices = (have_slices or
+                           isinstance(index, ExprNodes.SliceNode) or
+                           index.is_none)
             result.append(index)
 
-    if len(result) < ndim:
+    result_length = len(result) - len(newaxes)
+    if result_length < ndim:
         have_slices = True
-        nslices = ndim - len(result)
+        nslices = ndim - result_length
         result.extend([empty_slice(indices[-1].pos)] * nslices)
 
     return have_slices, result
@@ -708,6 +723,14 @@ def get_axes_specs(env, axes):
 
     return axes_specs
 
+def validate_axes(pos, axes):
+    if len(axes) >= Options.buffer_max_dims:
+        error(pos, "More dimensions than the maximum number"
+                   " of buffer dimensions were used.")
+        return False
+
+    return True
+
 def all(it):
     for item in it:
         if not item:

Cython/Compiler/Nodes.py

@@ -883,10 +883,13 @@ class MemoryViewSliceTypeNode(CBaseTypeNode):
             self.type = PyrexTypes.ErrorType()
             return self.type
 
-        MemoryView.validate_memslice_dtype(self.pos, base_type)
-        self.type = PyrexTypes.MemoryViewSliceType(base_type, axes_specs)
+        if not MemoryView.validate_axes(self.pos, axes_specs):
+            self.type = error_type
+        else:
+            MemoryView.validate_memslice_dtype(self.pos, base_type)
+            self.type = PyrexTypes.MemoryViewSliceType(base_type, axes_specs)
+            self.use_memview_utilities(env)
 
-        self.use_memview_utilities(env)
         return self.type
 
     def use_memview_utilities(self, env):

Cython/Compiler/PyrexTypes.py

@@ -645,7 +645,7 @@ class MemoryViewSliceType(PyrexType):
 
         tup = (obj.result(), self.ndim, to_py_func, from_py_func,
                self.dtype.is_pyobject)
-        return "__pyx_memoryview_fromslice(&%s, %s, %s, %s, %d);" % tup
+        return "__pyx_memoryview_fromslice(%s, %s, %s, %s, %d);" % tup
 
     def dtype_object_conversion_funcs(self, env):
         import MemoryView, Code

Cython/Utility/MemoryView.pyx

@@ -707,8 +707,17 @@ cdef memoryview memview_slice(memoryview memview, object indices):
 
     for dim, index in enumerate(indices):
         if PyIndex_Check(index):
-            slice_memviewslice(p_src, p_dst, dim, new_ndim, p_suboffset_dim,
-                               index, 0, 0, 0, 0, 0, False)
+            slice_memviewslice(
+                p_dst, p_src.shape[dim], p_src.strides[dim], p_src.suboffsets[dim],
+                dim, new_ndim, p_suboffset_dim,
+                index, 0, 0, # start, stop, step
+                0, 0, 0, # have_{start,stop,step}
+                False)
+        elif index is None:
+            p_dst.shape[new_ndim] = 1
+            p_dst.strides[new_ndim] = 0
+            p_dst.suboffsets[new_ndim] = -1
+            new_ndim += 1
         else:
             start = index.start or 0
             stop = index.stop or 0
@@ -718,18 +727,21 @@ cdef memoryview memview_slice(memoryview memview, object indices):
             have_stop = index.stop is not None
             have_step = index.step is not None
 
-            slice_memviewslice(p_src, p_dst, dim, new_ndim, p_suboffset_dim,
-                               start, stop, step, have_start, have_stop, have_step,
-                               True)
+            slice_memviewslice(
+                p_dst, p_src.shape[dim], p_src.strides[dim], p_src.suboffsets[dim],
+                dim, new_ndim, p_suboffset_dim,
+                start, stop, step,
+                have_start, have_stop, have_step,
+                True)
             new_ndim += 1
 
     if isinstance(memview, _memoryviewslice):
-        return memoryview_fromslice(&dst, new_ndim,
+        return memoryview_fromslice(dst, new_ndim,
                                     memviewsliceobj.to_object_func,
                                     memviewsliceobj.to_dtype_func,
                                     memview.dtype_is_object)
     else:
-        return memoryview_fromslice(&dst, new_ndim, NULL, NULL,
+        return memoryview_fromslice(dst, new_ndim, NULL, NULL,
                                     memview.dtype_is_object)
 
 
@@ -754,18 +766,13 @@ cdef extern from "pystate.h":
     PyObject *PyErr_Format(PyObject *exc, char *msg, ...) nogil
 
 @cname('__pyx_memoryview_slice_memviewslice')
-cdef int slice_memviewslice({{memviewslice_name}} *src,
-                              {{memviewslice_name}} *dst,
-                              int dim,
-                              int new_ndim,
-                              int *suboffset_dim,
-                              Py_ssize_t start,
-                              Py_ssize_t stop,
-                              Py_ssize_t step,
-                              int have_start,
-                              int have_stop,
-                              int have_step,
-                              bint is_slice) nogil except -1:
+cdef int slice_memviewslice(
+        {{memviewslice_name}} *dst,
+        Py_ssize_t shape, Py_ssize_t stride, Py_ssize_t suboffset,
+        int dim, int new_ndim, int *suboffset_dim,
+        Py_ssize_t start, Py_ssize_t stop, Py_ssize_t step,
+        int have_start, int have_stop, int have_step,
+        bint is_slice) nogil except -1:
     """
     Create a new slice dst given slice src.
 
@@ -776,14 +783,8 @@ cdef int slice_memviewslice({{memviewslice_name}} *src,
                       where slicing offsets should be added
     """
 
-    cdef:
-        Py_ssize_t shape, stride, suboffset
-        Py_ssize_t new_shape
-        bint negative_step
-
-    shape = src.shape[dim]
-    stride = src.strides[dim]
-    suboffset = src.suboffsets[dim]
+    cdef Py_ssize_t new_shape
+    cdef bint negative_step
 
     if not is_slice:
         # index is a normal integer-like index
@@ -958,7 +959,7 @@ cdef class _memoryviewslice(memoryview):
 
 
 @cname('__pyx_memoryview_fromslice')
-cdef memoryview_fromslice({{memviewslice_name}} *memviewslice,
+cdef memoryview_fromslice({{memviewslice_name}} memviewslice,
                           int ndim,
                           object (*to_object_func)(char *),
                           int (*to_dtype_func)(char *, object) except 0,
@@ -975,8 +976,8 @@ cdef memoryview_fromslice({{memviewslice_name}} *memviewslice,
 
     result = _memoryviewslice(None, 0, dtype_is_object)
 
-    result.from_slice = memviewslice[0]
-    __PYX_INC_MEMVIEW(memviewslice, 1)
+    result.from_slice = memviewslice
+    __PYX_INC_MEMVIEW(&memviewslice, 1)
 
     result.from_object = (<memoryview> memviewslice.memview).base
     result.typeinfo = memviewslice.memview.typeinfo
@@ -1055,7 +1056,7 @@ cdef memoryview_copy_from_slice(memoryview memview, {{memviewslice_name}} *memvi
         to_object_func = NULL
         to_dtype_func = NULL
 
-    return memoryview_fromslice(memviewslice, memview.view.ndim,
+    return memoryview_fromslice(memviewslice[0], memview.view.ndim,
                                 to_object_func, to_dtype_func,
                                 memview.dtype_is_object)
 

Cython/Utility/MemoryView_C.c

@@ -655,6 +655,8 @@ int {{set_function}}(const char *itemp, PyObject *obj); /* proto */
 /////////////// MemviewDtypeToObject ///////////////
 {{#__pyx_memview_<dtype_name>_to_object}}
 
+/* Convert a dtype to or from a Python object */
+
 {{if to_py_function}}
 PyObject *{{get_function}}(const char *itemp) {
     return (PyObject *) {{to_py_function}}(*({{dtype}} *) itemp);
@@ -672,6 +674,7 @@ int {{set_function}}(const char *itemp, PyObject *obj) {
 {{endif}}
 
 /////////////// MemviewObjectToObject.proto ///////////////
+/* Function callbacks (for memoryview object) for dtype object */
 PyObject *{{get_function}}(const char *itemp); /* proto */
 int {{set_function}}(const char *itemp, PyObject *obj); /* proto */
 
@@ -693,8 +696,8 @@ int {{set_function}}(const char *itemp, PyObject *obj) {
 /* Dimension is indexed with 'start:stop:step' */
 
 if (unlikely(__pyx_memoryview_slice_memviewslice(
-    &{{src}},
     &{{dst}},
+    {{src}}.shape[{{dim}}], {{src}}.strides[{{dim}}], {{src}}.suboffsets[{{dim}}],
     {{dim}},
     {{new_ndim}},
     &{{suboffset_dim}},
@@ -793,6 +796,8 @@ __pyx_fill_slice_{{dtype_name}}({{type_decl}} *p, Py_ssize_t extent, Py_ssize_t
 
 ////////// FillStrided1DScalar //////////
 /* Fill a slice with a scalar value. The dimension is direct and strided or contiguous */
+/* This can be used as a callback for the memoryview object to efficienty assign a scalar */
+/* Currently unused */
 static void
 __pyx_fill_slice_{{dtype_name}}({{type_decl}} *p, Py_ssize_t extent, Py_ssize_t stride,
                                 size_t itemsize, void *itemp)

docs/src/userguide/memoryviews.rst

@@ -76,6 +76,22 @@ the slice can be transposed in the same way that numpy slices can be transposed:
 
 This gives a new, transposed, view on the data.
 
+Newaxis
+-------
+New axes can be introduced by indexing an array with ``None`` ::
+
+    cdef double[:] myslice = np.linspace(0, 10, num=50)
+
+    # 2D array with shape (1, 50)
+    myslice[None] # or
+    myslice[None, :]
+
+    # 2D array with shape (50, 1)
+    myslice[:, None]
+
+One may mix new axis indexing with all other forms of indexing and slicing.
+See also an example_.
+
 Specifying data layout
 ======================
 
@@ -230,3 +246,4 @@ Unlike object attributes of extension classes, memoryview slices are not initial
 to None.
 
 .. _NumPy: http://docs.scipy.org/doc/numpy/reference/arrays.ndarray.html#memory-layout
+.. _example: http://www.scipy.org/Numpy_Example_List#newaxis

tests/errors/memview_declarations.pyx

@@ -54,8 +54,15 @@ cdef struct Invalid:
 cdef Valid[:] validslice
 cdef Invalid[:] invalidslice
 
+cdef int[:, :, :, :] four_D
+four_D[None, None, None, None]
+four_D[None, None, None, None, None]
+
+cdef int[:, :, :, :, :, :, :, :] eight_D = object()
+
 # These are VALID
 cdef int[::view.indirect_contiguous, ::view.contiguous] a9
+four_D[None, None, None]
 
 _ERRORS = u'''
 11:25: Cannot specify an array that is both C and Fortran contiguous.
@@ -79,4 +86,7 @@ _ERRORS = u'''
 46:35: Can only create cython.array from pointer or array
 47:24: Cannot assign type 'double' to 'Py_ssize_t'
 55:13: Invalid base type for memoryview slice: Invalid
+58:6: More dimensions than the maximum number of buffer dimensions were used.
+59:6: More dimensions than the maximum number of buffer dimensions were used.
+61:9: More dimensions than the maximum number of buffer dimensions were used.
 '''

tests/run/memslice.pyx

@@ -49,6 +49,14 @@ def testcase(func):
 include "mockbuffers.pxi"
 include "cythonarrayutil.pxi"
 
+def _print_attributes(memview):
+    print "shape: " + " ".join(map(str, memview.shape))
+    print "strides: " + " ".join([str(stride // memview.itemsize)
+                                      for stride in memview.strides])
+    print "suboffsets: " + " ".join(
+        [str(suboffset if suboffset < 0 else suboffset // memview.itemsize)
+             for suboffset in memview.suboffsets])
+
 #
 # Buffer acquire and release tests
 #
@@ -2217,3 +2225,63 @@ def test_inplace_assignment():
 
     m[0] = get_int()
     print m[0]
+
+@testcase
+def test_newaxis(int[:] one_D):
+    """
+    >>> A = IntMockBuffer("A", range(6))
+    >>> test_newaxis(A)
+    acquired A
+    3
+    3
+    3
+    3
+    released A
+    """
+    cdef int[:, :] two_D_1 = one_D[None]
+    cdef int[:, :] two_D_2 = one_D[None, :]
+    cdef int[:, :] two_D_3 = one_D[:, None]
+    cdef int[:, :] two_D_4 = one_D[..., None]
+
+    print two_D_1[0, 3]
+    print two_D_2[0, 3]
+    print two_D_3[3, 0]
+    print two_D_4[3, 0]
+
+@testcase
+def test_newaxis2(int[:, :] two_D):
+    """
+    >>> A = IntMockBuffer("A", range(6), shape=(3, 2))
+    >>> test_newaxis2(A)
+    acquired A
+    shape: 3 1 1
+    strides: 2 0 0
+    suboffsets: -1 -1 -1
+    <BLANKLINE>
+    shape: 1 2 1
+    strides: 0 1 0
+    suboffsets: -1 -1 -1
+    <BLANKLINE>
+    shape: 3 1 1 1
+    strides: 2 0 1 0
+    suboffsets: -1 -1 -1 -1
+    <BLANKLINE>
+    shape: 1 2 2 1
+    strides: 0 2 1 0
+    suboffsets: -1 -1 -1 -1
+    released A
+    """
+    cdef int[:, :, :] a = two_D[..., None, 1, None]
+    cdef int[:, :, :] b = two_D[None, 1, ..., None]
+    cdef int[:, :, :, :] c = two_D[..., None, 1:, None]
+    cdef int[:, :, :, :] d = two_D[None, 1:, ..., None]
+
+    _print_attributes(a)
+    print
+    _print_attributes(b)
+    print
+    _print_attributes(c)
+    print
+    _print_attributes(d)
+
+