aboutsummaryrefslogtreecommitdiff
path: root/utility
diff options
context:
space:
mode:
Diffstat (limited to 'utility')
-rw-r--r--utility/ndindex.py14
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)