Эффективная сортировка и фильтрация JaggedArray по другому - PullRequest
Новая
       19

Эффективная сортировка и фильтрация JaggedArray по другому

1 голос
/ 09 октября 2019

У меня есть JaggedArray (awkward.array.jagged.JaggedArray), который содержит индексы, указывающие на позиции в другом JaggedArray. Оба массива имеют одинаковую длину, но каждый из numpy.ndarrays, который содержит JaggedArrays, может иметь разную длину. Я хотел бы отсортировать второй массив, используя индексы первого массива, в то же время удаляя элементы из второго массива, которые не проиндексированы из первого массива. Первый массив может дополнительно содержать значения -1 (также может быть заменен на None при необходимости, но в настоящее время это не тот случай), что означает, что во втором массиве нет совпадений. В таком случае соответствующая позиция в первом массиве должна быть установлена ​​в значение по умолчанию (например, 0).

Вот практический пример и как я решаю это в данный момент: 

import uproot
import numpy as np
import awkward

def good_index(my_indices, my_values):
    my_list = []
    for index in my_indices:
        if index > -1:
            my_list.append(my_values[index])
        else:
            my_list.append(0)
    return my_list

indices = awkward.fromiter([[0, -1], [3,1,-1], [-1,0,-1]])
values = awkward.fromiter([[1.1, 1.2, 1.3], [2.1,2.2,2.3,2.4], [3.1]])

new_map = awkward.fromiter(map(good_index, indices, values))

Получается new_map: [[1.1 0.0] [2.4 2.2 0.0] [0.0 3.1 0.0]].

Есть ли более эффективный / быстрый способ достижения этого? Я думал, что можно использовать numpy функциональность, такую ​​как numpy.where, но из-за разной длины ndarrays это терпит неудачу по крайней мере для способов, которые я пробовал.

1 Ответ

0 голосов
/ 11 октября 2019

Если все подмассивы в values гарантированно не являются пустыми (так что индексирование с помощью -1 возвращает последний подэлемент, а не ошибку), тогда вы можете сделать это: 

>>> almost = values[indices]       # almost what you want; uses -1 as a real index
>>> almost.content = awkward.MaskedArray(indices.content < 0, almost.content)
>>> almost.fillna(0.0)
<JaggedArray [[1.1 0.0] [2.4 2.2 0.0] [0.0 3.1 0.0]] at 0x7fe54c713c88>

Последний шаг является необязательным, поскольку без него отсутствующими элементами являются None, а не 0.0.

Если некоторые из подмассивов в values пусты, вы можете pad ихубедитесь, что у них есть хотя бы один подэлемент. Все исходные подэлементы индексируются так же, как и раньше, поскольку pad увеличивает длину только, если это необходимо. 

>>> values = awkward.fromiter([[1.1, 1.2, 1.3], [], [2.1, 2.2, 2.3, 2.4], [], [3.1]])
>>> values.pad(1)
<JaggedArray [[1.1 1.2 1.3] [None] [2.1 2.2 2.3 2.4] [None] [3.1]] at 0x7fe54c713978>
...