Оптимизированный метод для разделения двумерного массива

Question

Я пытаюсь разделить двумерный массив numpy на 2 отдельных массива numpy, основываясь на содержимом определенного столбца. Это мой код:

import numpy as np
import pandas as pd

@profile
def partition_data(arr,target_colm):
    total_colms = arr.shape[1]
    target_data = arr[:,target_colm]
    type1_data = []
    type2_data = []
    for i in range(arr.shape[0]):
        if target_data[i]==0:  # if value==0, put in another array
            type1_data = np.append(type1_data,arr[i])
        else:
            type2_data = np.append(type2_data,arr[i])
    type1_data = np.array(type1_data).reshape(int(len(type1_data)/total_colms),total_colms)
    type2_data = np.array(type2_data).reshape(int(len(type2_data)/total_colms),total_colms)
    return type1_data, type2_data

d = pd.read_csv('data.csv').values
x,y = partition_data(d,7)  # check values of 7th column

Примечание: для моего эксперимента я использовал массив (14359,42) элементов.

Теперь, когда я профилирую эту функцию с помощью линейного профилировщика kernprof, я получаю следующие результаты.

Wrote profile results to code.py.lprof
Timer unit: 1e-06 s
Total time: 7.3484 s
File: code2.py
Function: part_data at line 8
Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
     8                                           @profile
     9                                           def part_data(arr,target_col):
    10         1          7.0      7.0      0.0      total_colms = arr.shape[1]
    11         1         14.0     14.0      0.0      target_data = arr[:,target_col]
    12         1          2.0      2.0      0.0      type1_data = []
    13         1          1.0      1.0      0.0      type2_data = []
    14      5161      40173.0      7.8      0.5      for i in range(arr.shape[0]):
    15      5160      39225.0      7.6      0.5          if target_data[i]==6:
    16      4882    7231260.0   1481.2     98.4              type1_data = np.append(type1_data,arr[i])
    17                                                   else:
    18       278      33915.0    122.0      0.5              type2_data = np.append(type2_data,arr[i])
    19         1       3610.0   3610.0      0.0      type1_data = np.array(type1_data).reshape(int(len(type1_data)/total_colms),total_colms)
    20         1        187.0    187.0      0.0      type2_data = np.array(type2_data).reshape(int(len(type2_data)/total_colms),total_colms)
    21         1          3.0      3.0      0.0      return type1_data, type2_data

Здесь одна строка-16 занимает значительное время. В будущем реальный размер данных, с которыми я буду работать, будет намного больше.

Кто-нибудь может предложить более быстрый способ разбиения массива NumPy?

Answer 1

Это должно сделать это намного быстрее:

def partition_data_vectorized(arr, target_colm):
    total_colms = arr.shape[1]
    target_data = arr[:,target_colm]
    mask = target_data == 0
    type1_data = arr[mask, :]
    type2_data = arr[~mask, :]
    return (
        type1_data.reshape(int(type1_data.size / total_colms), total_colms), 
        type2_data.reshape(int(type2_data.size / total_colms), total_colms))

Некоторые сроки:

# Generate some sample inputs:
arr = np.random.rand(10000, 42)
arr[:, 7] = np.random.randint(0, 10, 10000)

%timeit c, d = partition_data_vectorized(arr, 7)
# 2.09 ms ± 200 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
%timeit a, b = partition_data(arr, 7)
# 4.07 s ± 102 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Это в 2000 раз быстрее, чем не векторизованный расчет!

Сравнивая результаты:

np.all(b == d)
# Out: True
np.all(a == c)
# Out: True

Таким образом, результаты верны, и это в 2000 раз быстрее, просто заменив цикл for и повторное создание массива на np.append векторизованными операциями.