Доступ к элементам неудобного массива, которые не являются переданным индексом

jet_lepton_indices = [[0, 2], [1], [2, 3]] 
lepton_flags = [[0, 10, 20, 30], [0, 10, 20, 30], [0, 10, 20, 30, 40]]

flags_in_tot  = []
flags_out_tot = []

for index in range(len(jet_lepton_indices)):
    indices = jet_lepton_indices[index]
    flags = lepton_flags[index]
    flags_in = []
    flags_out = []
    for f in range(len(flags)):
        if f in indices:
            flags_in.append(flags[f])
        else:
            flags_out.append(flags[f])
    flags_in_tot.append(flags_in)
    flags_out_tot.append(flags_out)

flags_in_tot
[[0, 20], [10], [20, 30]]

flags_out_tot
[[10, 30], [0, 20, 30], [0, 10, 40]]

(«TL;DR» находится внизу, под горизонтальной линией.)

~ (побитовое нет) не работало с вашим массивом целых чисел, потому что оно просто инвертировало биты в целых числах:

>>> jet_lepton_indices
<Array [[0, 2], [1], [2, 3]] type='3 * var * int64'>
>>> ~jet_lepton_indices
<Array [[-1, -3], [-2], [-3, -4]] type='3 * var * int64'>

В конечном итоге вам нужно преобразовать срез целочисленного массива в срез логического массива. Как срезы, целочисленные массивы содержат строго больше информации, чем логические массивы: они могут дублировать и изменять порядок элементов срезанного массива в дополнение к простому удалению элементов. Таким образом, срезы целочисленного массива всегда можно преобразовать в срезы логического массива, но не наоборот. Фактически, был запрос на такую ​​функцию, #497, и в этой проблеме описано несколько способов ее получения, каждый из которых отличается от того, который я разработал ниже. (Я все еще собираюсь показать пример, который я только что разработал, потому что он проще и демонстрирует общую схему: cartesian чтобы увеличить измерения, сделайте что-то в новом измерении, затем агрегируйте его, чтобы вернуться к старому количеству измерений. )

Еще один факт о срезах логических массивов заключается в том, что они должны согласовываться с длиной списка массива, который они срезают. (Сноска: чтобы инвертировать выборку, вам нужно знать универсальный набор, поэтому инвертировать можно только срез логического массива.) Следовательно, чтобы преобразовать срез целочисленного массива в срез логического массива, нам нужно использовать массив нарезать, lepton_flags. Мы можем использовать ak.local_index, чтобы создать целочисленный массив целочисленных индексов для всех элементов, существующих в lepton_flags:

>>> all_indices = ak.local_index(lepton_flags)
>>> all_indices
<Array [[0, 1, 2, 3], [0, ..., 3], [0, 1, 2, 3, 4]] type='3 * var * int64'>

Теперь цель будет состоять в том, чтобы найти логические значения для каждого из этих индексов, которые говорят, находится ли индекс в jet_lepton_indices или нет. Такой вопрос имеет форму: «Для каждого элемента в X (локальный индекс) существует ли какой-либо элемент в Y (jet_lepton_indices), для которого Z (они равны)?» «Для каждого» одного массива с другим массивом обрабатывается ak.cartesian , и поскольку мы хотим агрегировать все, что связано с одним элементом X («является ли ak.any item равным? ") нам понадобится nested=True, чтобы создать новое измерение, чтобы позже агрегировать его.

>>> pairs = ak.cartesian([all_indices, jet_lepton_indices], nested=True)
>>> pairs.show(type=True)
type: 3 * var * var * (
    int64,
    int64
)
[[[(0, 0), (0, 2)], [(1, 0), (1, 2)], [(2, ...), ...], [(3, 0), (3, 2)]],
 [[(0, 1)], [(1, 1)], [(2, 1)], [(3, 1)]],
 [[(0, 2), (0, 3)], [(1, 2), (1, 3)], ..., [(3, ...), ...], [(4, 2), (4, 3)]]]

Пары более глубоко вложены (var * var *), чем all_indices и jet_lepton_indices (var *), потому что мы попросили сгруппировать результаты по одинаковому первому индексу (nested=True).

Левый элемент в каждой из этих пар — от all_indices, правый — от jet_lepton_indices для всех комбинаций. Чтобы разделить их, используйте ak.unzip:

>>> whole_set, in_set = ak.unzip(pairs)
>>> whole_set
<Array [[[0, 0], [1, 1], [...], [3, 3]], ...] type='3 * var * var * int64'>
>>> in_set
<Array [[[0, 2], [0, 2], [...], [0, 2]], ...] type='3 * var * var * int64'>

whole_set и in_set совпадают, потому что они происходят от одного и того же pairs. Поскольку они выстраиваются в линию, мы можем использовать к ним ==, чтобы получить логическое значение True тогда и только тогда, когда член whole_set находится в in_set.

>>> whole_set == in_set
<Array [[[True, False], ..., [False, ...]], ...] type='3 * var * var * bool'>

Если какой-либо (ak.any) из этих самых внутренних списков (axis=-1) равен True, то мы хотим сказать, что вся группа, представляющая элемент из all_indices, находится в jet_lepton_indices.

>>> jet_lepton_boolean = ak.any(whole_set == in_set, axis=-1)
>>> jet_lepton_boolean
<Array [[True, False, True, False], ..., [False, ...]] type='3 * var * bool'>

Это jet_lepton_boolean — логический массив, который можно использовать как срез для создания тех же элементов, что и jet_lepton_indices. Как логическое значение, его можно отрицать с помощью ~.

>>> lepton_flags[~jet_lepton_boolean]
<Array [[10, 30], [0, 20, 30], [0, 10, 40]] type='3 * var * int64'>

Это выбор lepton_flags, который вам нужен: это все, кроме того, что было в

>>> lepton_flags[jet_lepton_indices]
<Array [[0, 20], [10], [20, 30]] type='3 * var * int64'>

В качестве альтернативы вы могли бы создать отрицательные логические значения напрямую, используя != вместо ==.

Вот краткое описание этого метода как функции:

def indices_to_booleans(indices, array_to_slice):
    whole_set, in_set = ak.unzip(ak.cartesian([
        ak.local_index(array_to_slice), indices
    ], nested=True))
    return ak.any(whole_set == in_set, axis=-1)

Это решение зависит от того факта, что ваши исходные массивы имеют глубину только одного уровня (var *), хотя я думаю, что оно может быть обобщено, если вы передадите соответствующий аргумент axis в ak.cartesian, но я недостаточно об этом думал, чтобы быть уверенным.

Кроме того, #497 предлагает больше способов сделать это.