【教程】实测np.fromiter 和 np.array 的性能

import numpy as np
import time
import pandas as pd
import matplotlib.pyplot as plt

# 测试数据生成函数
def generate_data(size):
    return range(size)  # 使用生成器来模拟大量数据

# 测试 np.fromiter 的性能
def test_fromiter_int32(data):
    start_time = time.time()
    np.fromiter(data, dtype=np.int32)
    return time.time() - start_time

# 测试 np.array 的性能（不计算 list 开销）
def test_array_no_list_overhead(data):
    data = list(data)  # 先将生成器转换为列表
    start_time = time.time()
    np.array(data, dtype=np.int32)  # 确保 dtype 为 int32
    return time.time() - start_time

# 测试 np.array 的性能（计算 list 开销）
def test_array_with_list_overhead(data):
    start_time = time.time()
    data = list(data)  # 生成器转换为列表的时间也包含在内
    np.array(data, dtype=np.int32)  # 确保 dtype 为 int32
    return time.time() - start_time

# 数据量从 10^1 到 10^9
data_sizes = [10**i for i in range(1, 8)]  # 从 10^1 到 10^7
results_comparison = []

for size in data_sizes:
    data = generate_data(size)
    
    # np.fromiter 性能测试
    fromiter_time_int32 = test_fromiter_int32(data)
    
    # np.array 性能测试（不计算 list 开销）
    data = generate_data(size)  # 重新生成数据
    array_time_no_list_overhead = test_array_no_list_overhead(data)
    
    # np.array 性能测试（计算 list 开销）
    data = generate_data(size)  # 重新生成数据
    array_time_with_list_overhead = test_array_with_list_overhead(data)
    
    results_comparison.append((size, fromiter_time_int32, array_time_no_list_overhead, array_time_with_list_overhead))

# 将结果显示为数据框
df_results_comparison = pd.DataFrame(results_comparison, columns=[
    'Data Size', 'Fromiter Time (s)', 'Array Time without List Overhead (s)', 'Array Time with List Overhead (s)'
])

# 绘制比较性能曲线
plt.figure(figsize=(10, 6))
plt.plot(df_results_comparison['Data Size'], df_results_comparison['Fromiter Time (s)'], marker='o', label='np.fromiter (int32)')
plt.plot(df_results_comparison['Data Size'], df_results_comparison['Array Time without List Overhead (s)'], marker='o', label='np.array without list overhead (int32)')
plt.plot(df_results_comparison['Data Size'], df_results_comparison['Array Time with List Overhead (s)'], marker='o', label='np.array with list overhead (int32)')
plt.xlabel('Data Size')
plt.ylabel('Time (s)')
plt.xscale('log')  # 使用对数刻度显示更大范围的数据
plt.yscale('log')  # 使用对数刻度显示时间差异
plt.title('Performance Comparison: np.fromiter vs np.array (dtype=int32)')
plt.legend()
plt.grid(True)
plt.show()