xnumpy.py 4.68 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59
# -*- coding: utf-8 -*-
# NumPy-related utilities
# Copyright (C) 2018  Nexedi SA and Contributors.
#                     Kirill Smelkov <kirr@nexedi.com>
#
# This program is free software: you can Use, Study, Modify and Redistribute
# it under the terms of the GNU General Public License version 3, or (at your
# option) any later version, as published by the Free Software Foundation.
#
# You can also Link and Combine this program with other software covered by
# the terms of any of the Free Software licenses or any of the Open Source
# Initiative approved licenses and Convey the resulting work. Corresponding
# source of such a combination shall include the source code for all other
# software used.
#
# This program is distributed WITHOUT ANY WARRANTY; without even the implied
# warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#
# See COPYING file for full licensing terms.
# See https://www.nexedi.com/licensing for rationale and options.
import numpy as np
from numpy.lib import stride_tricks as npst

# _as_strided is similar to numpy.lib.stride_tricks.as_strided, but allows to
# set all shape/stridev/dtype in one go.
#
# It must be used with extreme care, because if there is math error in the
# arguments, the resulting array can cover wrong memory. Bugs here thus can
# lead to mysterious crashes.
def _as_strided(arr, shape, stridev, dtype):
    # the code below is very close to
    #
    #   a = stride_tricks.as_strided(arr, shape=shape, strides=stridev)
    #
    # but we don't use as_strided() because we also have to change dtype
    # with shape and strides in one go - else changing dtype after either
    # via a.dtype = ..., or via a.view(dtype=...) can raise errors like
    #
    #   "When changing to a larger dtype, its size must be a
    #    divisor of the total size in bytes of the last axis
    #    of the array."
    aiface = dict(arr.__array_interface__)
    aiface['shape']   = shape
    aiface['strides'] = stridev
    # type: for now we only care that itemsize is the same
    aiface['typestr'] = '|V%d' % dtype.itemsize
    aiface['descr']   = [('', aiface['typestr'])]

    a = np.asarray(npst.DummyArray(aiface, base=arr))

    # restore full dtype - it should not raise here, since itemsize is the same
    a.dtype = dtype

    # restore full array type (mimics subok=True)
    if type(a) is not type(arr):
        a = a.view(type=type(arr))

    # we are done
    return a
60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126


# structured creates view of the array interpreting its minor axis as fully covered by dtype.
#
# The minor axis of the array must be C-contiguous and be of dtype.itemsize in size.
#
# Structured is similar to arr.view(dtype) + corresponding reshape, but does
# not have limitations of ndarray.view(). For example:
#
#   In [1]: a = np.arange(3*3, dtype=np.int32).reshape((3,3))
#
#   In [2]: a
#   Out[2]:
#   array([[0, 1, 2],
#          [3, 4, 5],
#          [6, 7, 8]], dtype=int32)
#
#   In [3]: b = a[:2,:2]
#
#   In [4]: b
#   Out[4]:
#   array([[0, 1],
#          [3, 4]], dtype=int32)
#
#   In [5]: dtxy = np.dtype([('x', np.int32), ('y', np.int32)])
#
#   In [6]: dtxy
#   Out[6]: dtype([('x', '<i4'), ('y', '<i4')])
#
#   In [7]: b.view(dtxy)
#   ---------------------------------------------------------------------------
#   ValueError                                Traceback (most recent call last)
#   <ipython-input-66-af98529aa150> in <module>()
#   ----> 1 b.view(dtxy)
#
#   ValueError: To change to a dtype of a different size, the array must be C-contiguous
#
#   In [8]: structured(b, dtxy)
#   Out[8]: array([(0, 1), (3, 4)], dtype=[('x', '<i4'), ('y', '<i4')])
#
# Structured always creates view and never copies data.
def structured(arr, dtype):
    dtype = np.dtype(dtype) # convenience

    atype = arr.dtype
    # m* denotes minor *
    maxis   = np.argmin(np.abs(arr.strides))
    mstride = arr.strides[maxis]
    if mstride < 0:
        raise ValueError("minor-axis is not C-contiguous: stride (%d) < 0" % mstride)
    if mstride != atype.itemsize:
        raise ValueError("minor-axis is not C-contiguous: stride (%d) != itemsize (%d)" % (mstride, atype.itemsize))

    # verify dtype fully covers whole minor axis
    mnelem = arr.shape[maxis]
    msize  = mnelem * atype.itemsize
    if dtype.itemsize != msize:
        raise ValueError("dtype.itemsize (%d) != sizeof(minor-axis) (%d)" % (dtype.itemsize, msize))

    # ok to go
    shape   = arr.shape[:maxis] + arr.shape[maxis+1:]
    stridev = arr.strides[:maxis] + arr.strides[maxis+1:]

    # NOTE we cannot use just np.ndarray because if arr is a slice it can give:
    #   TypeError: expected a single-segment buffer object
    #return np.ndarray.__new__(type(arr), shape, dtype, buffer(arr), 0, stridev)
    return _as_strided(arr, shape, stridev, dtype)