Индексируйте DataFrame с помощью строк и столбцов MultiIndex через другой DataFrame, содержащий в качестве столбцов индексы строк и столбцов.

Question

У меня есть список пар частиц, в которых каждая пара называется комбинацией индекса цепи и индекса внутри цепи обеих частиц.Я сохранил их в Dataframe (назовем его index_array), и теперь я хочу построить матрицу всех пар частиц, где я отображаю все матричные элементы, которые соответствуют одной паре в списке, одним цветом, а все остальные - другим.цвет.Моя идея заключалась в том, чтобы создать DataFrame (назовем его to_fill) с цепочечным и внутрицепочечным индексом как MultiIndex для строк и столбцов, который, таким образом, имеет две записи на пару, а затем использует index_array дляиндекс to_fill, чтобы изменить соответствующие значения, так что я могу затем построить значения to_fill через matplotlib.pyplot.pcolormesh.

Таким образом, чтобы разбить его на более или менее четко определенную проблему: у меня естьлогический DataFrame to_fill, который имеет многоиндексированные строки и столбцы (по 2 уровня в каждом), который содержит только False s.У меня также есть другой DataFrame index_array, который имеет четыре столбца, содержащих значения индекса для уровней строк и столбцов.Теперь я хочу установить все элементы, на которые указывает index_array, на True.Например, можно создать игрушечную версию с кодом ниже:

import numpy as np
import pandas as pd
lengths = pd.Series(data=[2, 4], index=[1, 2])  # Corresponds to the chains' lengths
index = pd.MultiIndex.from_tuples([(i, j) for i in lengths.index
                                          for j in np.arange(1, lengths.loc[i]+1)])
to_fill = pd.DataFrame(index=index, columns=index, dtype=np.bool)
to_fill.loc[slice(None), slice(None)] = 0
print(to_fill)
#          1             2                     
#          1      2      1      2      3      4
# 1 1  False  False  False  False  False  False
#   2  False  False  False  False  False  False
# 2 1  False  False  False  False  False  False
#   2  False  False  False  False  False  False
#   3  False  False  False  False  False  False
#   4  False  False  False  False  False  False
index_array = pd.DataFrame([[1, 1, 1, 1],
                            [1, 1, 1, 2],
                            [2, 3, 2, 3],
                            [2, 3, 2, 4]],
                           columns=["i_1", "j_1", "i_2", "j_2"])
print(index_array)
#    i_1  j_1  i_2  j_2
# 0    1    1    1    1
# 1    1    1    1    2
# 2    2    3    2    3
# 3    2    3    2    4

Теперь я хочу установить все записи в to_fill, которые соответствуют (i_1, j_1), (i_2, j_2) для строки в index_array, в True.В общем, index_array относится к записям в to_fill, которые должны быть изменены.Таким образом, ожидаемый результат будет:

print(to_fill)
#          1             2                     
#          1      2      1      2      3      4
# 1 1  True   True   False  False  False  False
#   2  False  False  False  False  False  False
# 2 1  False  False  False  False  False  False
#   2  False  False  False  False  False  False
#   3  False  False  False  False  True   True 
#   4  False  False  False  False  False  False

Но мне не удалось правильно использовать index_array в качестве индекса.Как я могу сказать to_fill обрабатывать индексные массивы i_1, j_1, i_2 и j_2 как соответствующие значения индекса для уровней строки и столбца MultiIndex соответственно?

Answer 1

Это немного лучше - хм, может быть, и не совсем:

tuples = [tuple(x) for x in index_array.values]
stacked = to_fill.stack(level=0).stack() # double stack carefully ordered
stacked.loc[tuples] = True
result = stacked.unstack(level=2).unstack().dropna(axis=1) #unstack and drop NaN cols

Answer 2

Это не очень хорошо, так как я не пытаюсь использовать iterrows (), если это может помочь.

idx = pd.IndexSlice
for row in index_array.iterrows():
    r = row[1]
    i_1= r.loc['i_1']
    j_1= r.loc['j_1']
    i_2=  r.loc['i_2']
    j_2 = r.loc['j_2']

    to_fill.loc[idx[i_1,j_1],idx[i_2,j_2]] = True