How to convert a list of indices into a cell list (numpy array of lists) in numpy with vectorized implementation?

Question

Cell list is a data structure that maintains lists of data points in an N-D meshgrid. For example, the following list of 2d indices:

ind = [(0, 1), (1, 0), (0, 1), (0, 0), (0, 0), (0, 0), (1, 1)]

is converted to the following 2x2 cell list:

cell = [[[3, 4, 5], [0, 2]],
        [[1, ],     [6, ]]
       ]

using an O(n) algorithm:

# create an empty 2x2 cell list
cell = [[[] for _ in range(2)] for _ in range(2)]
for i in range(len(ind)):
    cell[ind[i][0], ind[i][1]].append(i)

Is there a vectorized way in numpy that can convert the list of indices (ind) into the cell structure described above?

I think you should reconsider storing data like that if you want to use `numpy` — user3483203, May 16 '19 at 01:59
The variable size of the `cell` lists is a good indication that this is not a `numpy` friendly problem. — hpaulj, May 16 '19 at 02:19

score 0 · Answer 1 · answered May 16 '19 at 04:08

I don't think there is a good pure numpy but you can either use pythran or---if you don't want to touch a compiler---scipy.sparse cf. this Q&A which is essentially a 1D version of your problem.

[stb_pthr.py] simplified from Most efficient way to sort an array into bins specified by an index array?

import numpy as np

#pythran export sort_to_bins(int[:], int)

def sort_to_bins(idx, mx=-1):
    if mx==-1:
        mx = idx.max() + 1
    cnts = np.zeros(mx + 1, int)
    for i in range(idx.size):
        cnts[idx[i] + 1] += 1
    for i in range(1, cnts.size):
        cnts[i] += cnts[i-1]
    res = np.empty_like(idx)
    for i in range(idx.size):
        res[cnts[idx[i]]] = i
        cnts[idx[i]] += 1
    return res, cnts[:-1]

Compile: pythran stb_pthr.py

Main script:

import numpy as np
try:
    from stb_pthr import sort_to_bins
    HAVE_PYTHRAN = True
except:
    HAVE_PYTHRAN = False
from scipy import sparse

def fallback(flat, maxind):
    res = sparse.csr_matrix((np.zeros_like(flat),flat,np.arange(len(flat)+1)),
                            (len(flat), maxind)).tocsc()
    return res.indices, res.indptr[1:-1]

if not HAVE_PYTHRAN:
    sort_to_bins = fallback

def to_cell(data, shape=None):
    data = np.asanyarray(data)
    if shape is None:
        *shape, = (1 + c.max() for c in data.T)
    flat = np.ravel_multi_index((*data.T,), shape)
    reord, bnds = sort_to_bins(flat, np.prod(shape))
    return np.frompyfunc(np.split(reord, bnds).__getitem__, 1, 1)(
        np.arange(np.prod(shape)).reshape(shape))

ind = [(0, 1), (1, 0), (0, 1), (0, 0), (0, 0), (0, 0), (1, 1)]

print(to_cell(ind))

from timeit import timeit

ind = np.transpose(np.unravel_index(np.random.randint(0, 100, (1_000_000)), (10, 10)))

if HAVE_PYTHRAN:
    print(timeit(lambda: to_cell(ind), number=10)*100)
    sort_to_bins = fallback # !!! MUST REMOVE THIS LINE AFTER TESTING
print(timeit(lambda: to_cell(ind), number=10)*100)

Sample run, output is answer to OP's toy example and timings (in ms) for the pythran and scipy solutions on a 1,000,000 points example:

[[array([3, 4, 5]) array([0, 2])]
 [array([1]) array([6])]]
11.411489499732852
29.700406698975712

How to convert a list of indices into a cell list (numpy array of lists) in numpy with vectorized implementation?

1 Answers1