diff options
Diffstat (limited to 'numpy/core/numeric.py')
| -rw-r--r-- | numpy/core/numeric.py | 28 |
1 files changed, 18 insertions, 10 deletions
diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py index a609ff3f8..6ca47a87b 100644 --- a/numpy/core/numeric.py +++ b/numpy/core/numeric.py @@ -478,26 +478,34 @@ set_string_function(array_repr, 1) little_endian = (sys.byteorder == 'little') + def indices(dimensions, dtype=int): """Returns an array representing a grid of indices with row-only, and column-only variation. """ - tmp = ones(dimensions, dtype) - lst = [] - for i in range(len(dimensions)): - lst.append( add.accumulate(tmp, i, dtype)-1 ) - return array(lst) + dimensions = tuple(dimensions) + N = len(dimensions) + if N == 0: + return array([],dtype=dtype) + res = empty((N,)+dimensions, dtype=dtype) + for i, dim in enumerate(dimensions): + tmp = arange(dim,dtype=dtype) + tmp.shape = (1,)*i + (dim,)+(1,)*(N-i-1) + newdim = dimensions[:i] + (1,)+ dimensions[i+1:] + val = zeros(newdim, dtype) + add(tmp, val, res[i]) + return res def fromfunction(function, shape, **kwargs): """Returns an array constructed by calling a function on a tuple of number grids. - The function should accept as many arguments as the length of shape and work - on array inputs. The shape argument is a sequence of numbers indicating the - length of the desired output for each axis. + The function should accept as many arguments as the length of shape and + work on array inputs. The shape argument is a sequence of numbers + indicating the length of the desired output for each axis. - The function can also accept keyword arguments (except dtype), which will be - passed through fromfunction to the function itself. The dtype argument + The function can also accept keyword arguments (except dtype), which will + be passed through fromfunction to the function itself. The dtype argument (default float) determines the data-type of the index grid passed to the function. """ |