diff options
Diffstat (limited to 'utility')
| -rw-r--r-- | utility/ndindex.py | 14 | 
1 files changed, 14 insertions, 0 deletions
| diff --git a/utility/ndindex.py b/utility/ndindex.py new file mode 100644 index 0000000..0c5ed9b --- /dev/null +++ b/utility/ndindex.py @@ -0,0 +1,14 @@ +import numpy +from numpy.lib.stride_tricks import as_strided +import numpy.core.numeric as _nx + +class ndindex(numpy.ndindex): +    pass + +    def __init__(self, *shape, order): +        if len(shape) == 1 and isinstance(shape[0], tuple): +            shape = shape[0] +        x = as_strided(_nx.zeros(1), shape=shape, +                       strides=_nx.zeros_like(shape)) +        self._it = _nx.nditer(x, flags=['multi_index', 'zerosize_ok'], +                              order=order) | 
