Нахождение индексов значения в матрице расстояний?

Question

Из списка 461 координат X, Y и Z я создал матрицу расстояний с помощью SciPy, которая выглядит следующим образом

[[ 0.          3.78112691  6.55159315 ... 63.40661118 62.2923149
  64.71125443]
 [ 3.78112691  0.          3.76986434 ... 60.79913069 59.55251531
  61.87364432]
 [ 6.55159315  3.76986434  0.         ... 61.12392086 59.65959803
  61.94572052]
 ...
 [63.40661118 60.79913069 61.12392086 ...  0.          3.8003808
   5.63044026]
 [62.2923149  59.55251531 59.65959803 ...  3.8003808   0.
   3.82889361]
 [64.71125443 61.87364432 61.94572052 ...  5.63044026  3.82889361
   0.        ]]

Я также написал код, позволяющий пользователю извлекать значения из расстоянияматрицы, которые меньше или равны определенному значению, а затем записывают его в текстовый файл

radius = int(input("Enter the radius threshold for contact (Angstrom): "))


###########This code will extract every pair <= radius##################
f = open('contact.txt', 'w')

for i in distance_matrix:
    for j in i:
        if j <= radius:
            if j != 0:
                f.write(str(j) + '\n')

print("File contact.txt written")
#########################################################################

А текстовый файл (contact.txt) представляет собой просто длинный список значений <= для пользователяуказанное значение.Есть ли способ также написать, из какой пары индексов эти значения происходят?Например, если значение «a» найдено на пересечении индекса b на оси x и значения c на оси y, есть ли способ добавить это в выходной текстовый файл?</p>

Answer 1

Вы можете использовать np.where

indices = np.where((distance_matrix <= radius) & (distance_matrix != 0))
row_col_vals = list(zip(*indices, distance_matrix[indices]))  # gives you a list or (row, col, value) tuples

Пример, если radius = 5 и

distance_matrix =  np.array([[2, 6], [0, 3]])

, то

row_col_vals = [(0, 0, 2), (1, 1, 3)]

Для записи в файл:

f = open('contact.txt', 'w')
for row, col, val in row_col_vals:
    print(row, col, val, file=f)
f.close()

HTH.

Answer 2

В соответствии с комментарием вы можете использовать перечисление, которое эквивалентно

zip( range( len(vec)), vec)

, чтобы сделать это:

f = open('contact.txt', 'w')

for n1, i in enumerate(distance_matrix):
    for n2 ,j in enumerate(i):
        if j <= radius:
            if j != 0:
                f.write(str(j) + '\n')

print("File contact.txt written")

, где n1 и n2 - ваши индексы