Skip to content Skip to sidebar Skip to footer

How To Recover Original Indices For A Flattened Numpy Array?

I've got a multidimensional numpy array that I'm trying to stick into a pandas data frame. I'd like to flatten the array, and create a pandas index that reflects the pre-flattened

Solution 1:

You could use pd.MultiIndex.from_product:

import numpy as np
import pandas as pd
import string

def using_multiindex(A, columns):
    shape = A.shape
    index = pd.MultiIndex.from_product([range(s)for s in shape], names=columns)
    df = pd.DataFrame({'A': A.flatten()}, index=index).reset_index()
    return df

A = np.array([[[ 0.43793885,  0.40078139,  0.48078691,  0.05334248],
    [ 0.76331509,  0.82514441,  0.86169078,  0.86496111],
    [ 0.75572665,  0.80860943,  0.79995337,  0.63123724]],

   [[ 0.20648946,  0.57042315,  0.71777265,  0.34155005],
    [ 0.30843717,  0.39381407,  0.12623462,  0.93481552],
    [ 0.3267771 ,  0.64097038,  0.30405215,  0.57726629]]])

df = using_multiindex(A, list('ZYX'))

yields

    Z  Y  X         A
0   0  0  0  0.437939
1   0  0  1  0.400781
2   0  0  2  0.480787
3   0  0  3  0.053342
...
21  1  2  1  0.640970
22  1  2  2  0.304052
23  1  2  3  0.577266

Or if performance is a top priority, consider using senderle's cartesian_product. (See the code, below.)


Here is a benchmark for A with shape (100, 100, 100):

In [321]: %timeit  using_cartesian_product(A, columns)
100 loops, best of 3: 13.8 ms per loop

In [318]: %timeit using_multiindex(A, columns)
10 loops, best of 3: 35.6 ms per loop

In [320]: %timeit indices_merged_arr_generic(A, columns)
10 loops, best of 3: 29.1 ms per loop

In [319]: %timeit using_product(A)
1 loop, best of 3: 461 ms per loop

This is the setup I used for the benchmark:

import numpy as np
import pandas as pd
import functools
import itertools as IT
import string
product = IT.product

defcartesian_product_broadcasted(*arrays):
    """
    http://stackoverflow.com/a/11146645/190597 (senderle)
    """
    broadcastable = np.ix_(*arrays)
    broadcasted = np.broadcast_arrays(*broadcastable)
    dtype = np.result_type(*arrays)
    rows, cols = functools.reduce(np.multiply, broadcasted[0].shape), len(broadcasted)
    out = np.empty(rows * cols, dtype=dtype)
    start, end = 0, rows
    for a in broadcasted:
        out[start:end] = a.reshape(-1)
        start, end = end, end + rows
    return out.reshape(cols, rows).T

defusing_cartesian_product(A, columns):
    shape = A.shape
    coords = cartesian_product_broadcasted(*[np.arange(s, dtype='int') for s in shape])
    df = pd.DataFrame(coords, columns=columns)
    df['A'] = A.flatten()
    return df

defusing_multiindex(A, columns):
    shape = A.shape
    index = pd.MultiIndex.from_product([range(s)for s in shape], names=columns)
    df = pd.DataFrame({'A': A.flatten()}, index=index).reset_index()
    return df

defindices_merged_arr_generic(arr, columns):
    n = arr.ndim
    grid = np.ogrid[tuple(map(slice, arr.shape))]
    out = np.empty(arr.shape + (n+1,), dtype=arr.dtype)
    for i inrange(n):
        out[...,i] = grid[i]
    out[...,-1] = arr
    out.shape = (-1,n+1)
    df = pd.DataFrame(out, columns=['A']+columns)
    return df

defusing_product(A):
    x, y, z = A.shape
    x_, y_, z_ = zip(*product(range(x), range(y), range(z)))
    df = pd.DataFrame(A.flatten()).assign(x=x_, y=y_, z=z_)
    return df

A = np.random.random((100,100,100))
shape = A.shape
columns = list(string.ascii_uppercase[-len(shape):][::-1])

Solution 2:

from itertools import product

np.random.seed(0)
A = np.random.rand(2, 3, 4)
x, y, z = A.shape
x_, y_, z_ = zip(*product(range(x), range(y), range(z)))
df = pd.DataFrame(A.flatten()).assign(x=x_, y=y_, z=z_)
>>> df

           0  x  y  z
00.54881400010.71518900120.60276300230.54488300340.42365501050.64589401160.43758701270.89177301380.96366302090.383442021100.791725022110.528895023120.568045100130.925597101140.071036102150.087129103160.020218110170.832620111180.778157112190.870012113200.978618120210.799159121220.461479122230.780529123

Solution 3:

My solution is based on this this answer by Divakar involving np.ogrid. This function should work for any array of any dimension.

def indices_merged_arr(arr):
    n = arr.ndim
    grid = np.ogrid[tuple(map(slice, arr.shape))]
    out = np.empty(arr.shape + (n+1,), dtype=arr.dtype)
    for i in range(n):
        out[...,i+1] = grid[i]
    out[...,0] = arr
    out.shape = (-1,n+1)
    return out

A = np.array([[[ 0.43793885,  0.40078139,  0.48078691,  0.05334248],
               [ 0.76331509,  0.82514441,  0.86169078,  0.86496111],
               [ 0.75572665,  0.80860943,  0.79995337,  0.63123724]],

              [[ 0.20648946,  0.57042315,  0.71777265,  0.34155005],
               [ 0.30843717,  0.39381407,  0.12623462,  0.93481552],
               [ 0.3267771 ,  0.64097038,  0.30405215,  0.57726629]]])

df = pd.DataFrame(indices_merged_arr(A), columns=list('Axyz'))
df

           A    x    y    z
00.4379390.00.00.010.4007810.00.01.020.4807870.00.02.030.0533420.00.03.040.7633150.01.00.050.8251440.01.01.060.8616910.01.02.070.8649610.01.03.080.7557270.02.00.090.8086090.02.01.0100.7999530.02.02.0110.6312370.02.03.0120.2064891.00.00.0130.5704231.00.01.0140.7177731.00.02.0150.3415501.00.03.0160.3084371.01.00.0170.3938141.01.01.0180.1262351.01.02.0190.9348161.01.03.0200.3267771.02.00.0210.6409701.02.01.0220.3040521.02.02.0230.5772661.02.03.0

Solution 4:

As hpaulj pointed out in a comment, I could add indexing=='ij' to the meshgrid call:

A = np.random.rand(2,3,4)
dimnames = ['z', 'y', 'x']
ranges   = [ np.arange(x) for x in A.shape ]
ix       = [ x.flatten()  for x in np.meshgrid(*ranges, indexing='ij') ]
for name, col inzip(dimnames, ix):
    df[name] = col
df = df.set_index(dimnames).squeeze()

# Compare the resultsfor ix, val in df.iteritems():
    print ix, val == A[ix]
(0, 0, 0) True
(0, 0, 1) True
(0, 0, 2) True
(0, 0, 3) True
(0, 1, 0) True
(0, 1, 1) True
(0, 1, 2) True
(0, 1, 3) True
(0, 2, 0) True
(0, 2, 1) True
(0, 2, 2) True
(0, 2, 3) True
(1, 0, 0) True
(1, 0, 1) True
(1, 0, 2) True
(1, 0, 3) True
(1, 1, 0) True
(1, 1, 1) True
(1, 1, 2) True
(1, 1, 3) True
(1, 2, 0) True
(1, 2, 1) True
(1, 2, 2) True
(1, 2, 3) True

Solution 5:

Another possibility, although others maybe faster...

x,y,z = np.indices(A.shape)

df = pd.DataFrame(np.array([p.flatten() for p in [x,y,z,A]]).T
                  ,columns=['x','y','z',0])

Post a Comment for "How To Recover Original Indices For A Flattened Numpy Array?"