Faster way to extract one numpy array using another

user_13469737 0
Viewed 223

I have two large one-dimensional numpy arrays (~1e8+ elements) of the same size. Array 1 (a1) is an integer indicator value, it tells me the meaning of the values in a2. Array 2 (a2) has measurement values. I want to separate the values in a2, based on the values in a1. I want to iterate over a list of possible values in a1 and select corresponding values (by index) from a2. Because of the data size and the number of different times I have to do this, it needs to be as fast as possible.

Here's a simple example to hopefully explain my problem more clearly:

a1 = [0,   2,   1,   1,   0,   2] # actual values range from 0 to 4095
a2 = [0.5, 2.4, 1.0, 1.2, 0.4, 2.6] # dummy values for example

I would like to separate the values in a2 based on the values in a1. I am saving each separated array into an HDF5 dataset. So, in the end, I need an array for each value in a1

out = [0.5,0.4] # a1 == 0
out = [1.0, 1.2] # a1 == 1
out = [2.4, 2.6] # a1 == 2

Currently, I've tried the following:

import numpy as np

size = int(1e8)
a1 = np.random.randint(0, 4096, size=size)
a2 = np.random.rand(size)

for i in range(0, 4096):
     ind = a1 == i
     out = a2[ind]
     # more code here to save out to h5 file

Based on some timeit testing, np.extract works just slightly faster than this boolean masking approach for large arrays:

import numpy as np

size = int(1e8)
a1 = np.random.randint(0, 4096, size=size)
a2 = np.random.rand(size)

for i in range(0,4096):
     ind = a1 == i
     out = np.extract(ind, a2)
     # more code here to save out to h5 file

I also tried putting this into a pandas series with a1 as the index, but this was much slower.

import pandas as pd
import numpy as np

size = int(1e8)
a1 = np.random.randint(0, 4096, size=size)
a2 = np.random.rand(size)
s = pd.Series(a2, index=a1)

for i in range(0,4096):
     out = s.loc[i] #this is WAY slower than numpy
     # more code here to save out to h5 file

I'm wondering if there's a faster way of doing this? Maybe with sorting and searchsorted? Is it possible to generate a "database-like index" for a numpy array?

numpy python
1 Answers

As I see, your both source arrays are plain pythonic lists. So the first step is to convert them to Numpy arrays:

arr1 = np.array(a1)
arr2 = np.array(a2)

Then, to get your expected result as a list of Numpy arrays, run:

rng = arr1.max() + 1
result = [ arr2[np.nonzero(arr1 == i)] for i in range(rng) ]

Details:

  • np.nonzero(arr1 == i) - generates the list of indices of arr1 where the element contains i,
  • arr2[...] - retrieves elements of arr2 indicated by the above indices,
  • [ ... for i in range(rng) ] - a list comprehension to generate an output array for each value present in arr1.

For your data sample you will get:

[array([0.5, 0.4]), array([1. , 1.2]), array([2.4, 2.6])]

Or if you want your result as a list of plain list, change the above code to:

result = [ arr2[np.nonzero(arr1 == i)].tolist() for i in range(rng) ]

This time you will get:

[[0.5, 0.4], [1.0, 1.2], [2.4, 2.6]]
user_73884770
Related