Python如何从函数返回数组？多文件读取生成DataFrame子集集合

问题描述 #

我写了一个read_files函数，接收paths和scalings数组，用来读取多个路径下的文件并生成pandas DataFrame，然后提取E列（命名为energy_）和I列（命名为intensity_）。但现在函数只能返回最后一组数据，我希望它返回包含所有路径对应数据的数组，比如energy = [energy_0, energy_1, energy_2, energy_3]、intensity = [intensity_0, intensity_1, intensity_2, intensity_3]，这样就能传给plotfunction绘图了。另外，我还想知道怎么单独调用paths里path3对应的energy和intensity数据。

原函数代码：

def read_files(paths:np.array,scalings:np.array):
    names = ['E','I']
    for p,s in zip(paths,scalings):
        df = pd.read_csv(p,engine = 'python', sep ='\s+', names=names)
        energy_ = df['E']
        intensity_ = df['I']
    return energy_, intensity_

修改后的函数 #

要返回所有路径的数据，只需在循环前初始化两个空列表，每次循环把读取到的列数据追加到列表中即可。如果需要应用scalings参数对强度做缩放，也可以在这一步处理：

import numpy as np
import pandas as pd

def read_files(paths: np.array, scalings: np.array):
    names = ['E', 'I']
    energy_list = []
    intensity_list = []
    for p, s in zip(paths, scalings):
        df = pd.read_csv(p, engine='python', sep='\s+', names=names)
        # 若需缩放强度，取消下方注释即可
        # scaled_intensity = df['I'] * s
        # intensity_list.append(scaled_intensity)
        energy_list.append(df['E'])
        intensity_list.append(df['I'])
    return energy_list, intensity_list

调用方式 #

调用函数后直接将结果传给plotfunction：

energy, intensity = read_files(paths_array, scalings_array)
fig = plotfunction(energy, intensity, etc)

单独提取path3对应的数据 #

按索引提取 #

如果path3是paths数组中索引为3的元素（Python索引从0开始），直接通过列表索引获取：

energy, intensity = read_files(paths_array, scalings_array)
path3_energy = energy[3]
path3_intensity = intensity[3]

按路径字符串提取 #

如果已知path3的具体路径字符串，先找到它在数组中的索引再提取：

target_path = "path3的具体路径"
idx = np.where(paths_array == target_path)[0][0]
path3_energy = energy[idx]
path3_intensity = intensity[idx]

内容的提问来源于stack exchange，提问作者werghuiop