Как обобщить этот расчет с пандой DataFrame на любое количество столбцов?

Question

У меня есть файл с некоторыми данными, который выглядит как

1 2 3 4
2 3 4 5
3 4 5 6
4 5 6 7

Я могу обработать эти данные и просто выполнить математические расчеты:

import sys
import numpy as np
import pandas as pd

def main():
    if(len(sys.argv) != 2):
        print "Takes one filename as argument"
        sys.exit()

    file_name = sys.argv[1]

    data = pd.read_csv(file_name, sep=" ", header=None)
    data.columns = ["timestep", "mux", "muy", "muz"]

    t = data["timestep"].count()

    c = np.zeros(t)

    for i in range(0,t):
        for j in range(0,i+1):
            c[i-j] += data["mux"][i-j] * data["mux"][i]
            c[i-j] += data["muy"][i-j] * data["muy"][i] 
            c[i-j] += data["muz"][i-j] * data["muz"][i]

    for i in range(t):
        print c[i]/(t-i)

Ожидаемый результат для моего образцавходные данные

42.5
62.0
84.5
110.0

Эта математика находит для моих данных корреляционную функцию времени , которая является средним временем всех перестановок пар продуктов в каждом столбце.Я хотел бы обобщить эту программу, чтобы

• работал на n количестве столбцов (например, в цикле i / j), и
• мог читать в столбцеимена из файла, чтобы они не были жестко запрограммированы в

Какие методы numpy или pandas можно использовать для достижения этой цели?

Answer 1

Мы можем сократить его до одного цикла, так как мы будем использовать нарезку массива и использовать sum ufunc для работы со строками кадра данных и, таким образом, в этом процессе сделать его универсальным для охвата любого количества столбцов, напримерпоэтому -

a = data.values
t = data["timestep"].count()
c = np.zeros(t)
for i in range(t):
    c[:i+1] += (a[:i+1,1:]*a[i,1:]).sum(axis=1)

Пояснение

1) a[:i+1,1:] - это срез всех строк до i+1 -ой строки и всех столбцов, начиная со второгостолбец, т.е. mux, muy и т. д.

2) Аналогично, для [i,1:] это строка i-th и все столбцы, начиная со второго столбца и далее.

Коставьте это "панда", просто замените a[ на data.iloc[.