如何为scipy.fft.fft2与numpy.argsort添加tqdm进度条？

一、二维FFT（scipy.fft.fft2）的进度提示 #

scipy.fft.fft2本质是对数组的两个维度依次执行一维FFT，你可以手动拆分这两步，用tqdm包裹循环跟踪每行/列的处理进度，完全复用scipy原生的fft函数，无需自行实现FFT算法。

实现思路 #

1. 先对数组所有行执行一维FFT，用tqdm显示行处理进度
2. 再对上述结果的所有列执行一维FFT，用tqdm显示列处理进度

代码示例 #

import numpy as np
from scipy.fft import fft
from tqdm import tqdm

# 创建大规模随机数据（示例为10000x10000，可替换为实际尺寸）
data = np.random.uniform(0, 1E4, (10000, 10000))

# 分步执行FFT并添加进度提示
print("Processing row-wise FFT...")
row_fft = np.empty_like(data, dtype=np.complex128)
for i in tqdm(range(data.shape[0]), total=data.shape[0]):
    row_fft[i] = fft(data[i])

print("Processing column-wise FFT...")
final_fft = np.empty_like(row_fft, dtype=np.complex128)
for j in tqdm(range(row_fft.shape[1]), total=row_fft.shape[1]):
    final_fft[:, j] = fft(row_fft[:, j])

print("FFT complete!")

注：该方式与原生fft2计算结果完全一致，仅存在微小的性能损耗（来自Python循环），但对于超大规模数据而言，进度提示的价值远超过这部分损耗；若追求极致性能，可结合numpy向量化操作或joblib并行处理，但循环方式的可读性最高。

二、numpy.argsort的进度提示 #

numpy.argsort是底层优化的排序实现（默认用快速排序，会根据数据类型自动切换算法），未暴露中间回调接口，无法直接跟踪排序的每一步进度。以下是两种可行方案：

方案1：基于历史运行时间的预估进度条 #

先运行一次同结构的小规模数据，估算单位数据量的排序耗时，再以此为基础生成大规模排序的预估总时间，最后用tqdm配合后台线程更新进度。

代码示例 #

import numpy as np
import time
import threading
from tqdm import tqdm

# 先估算排序速度
sample_data = np.random.uniform(0, 1E4, (1000, 1000)).flatten()
start = time.time()
_ = np.argsort(sample_data)
sample_time = time.time() - start

# 计算大规模数据的预估总排序时间
total_elements = data.size
estimated_total_time = (total_elements / sample_data.size) * sample_time

# 启动进度更新线程
done_event = threading.Event()
start_time = time.time()

with tqdm(total=estimated_total_time, unit="s", bar_format="{l_bar}{bar} [Elapsed: {elapsed} | Remaining: {remaining}]") as pbar:
    def update_progress():
        while not done_event.is_set():
            elapsed = time.time() - start_time
            if elapsed >= estimated_total_time:
                pbar.update(estimated_total_time - pbar.n)
            else:
                pbar.update(elapsed - pbar.n)
            time.sleep(0.1)
        pbar.update(estimated_total_time - pbar.n)

    progress_thread = threading.Thread(target=update_progress)
    progress_thread.start()

    # 主线程执行排序
    indices = np.argsort(data.flatten())

    # 标记排序完成
    done_event.set()
    progress_thread.join()

print("Sort complete!")

方案2：基于CPU负载的近似进度监控 #

如果预估时间偏差较大，可使用psutil库监控当前进程的CPU使用率，间接判断排序进度（排序过程中CPU会持续高负载，完成后负载骤降），适合对进度精度要求不高的场景。

代码示例（需先安装psutil：pip install psutil） #

import numpy as np
import time
import threading
import psutil
from tqdm import tqdm

process = psutil.Process()
done_event = threading.Event()
start_time = time.time()

# 基于样本数据估算总时间（同方案1）
sample_data = np.random.uniform(0, 1E4, (1000, 1000)).flatten()
start = time.time()
_ = np.argsort(sample_data)
sample_time = time.time() - start
total_elements = data.size
estimated_total_time = (total_elements / sample_data.size) * sample_time

with tqdm(total=100, unit="%", bar_format="{l_bar}{bar} [Elapsed: {elapsed}]") as pbar:
    def monitor_progress():
        while not done_event.is_set():
            elapsed = time.time() - start_time
            # 结合CPU负载和预估时间计算进度
            cpu_percent = process.cpu_percent()
            if cpu_percent > 90:
                progress = min(100, (elapsed / estimated_total_time) * 100)
                pbar.update(progress - pbar.n)
            time.sleep(0.5)
        pbar.update(100 - pbar.n)

    progress_thread = threading.Thread(target=monitor_progress)
    progress_thread.start()

    # 执行排序
    indices = np.argsort(data.flatten())

    done_event.set()
    progress_thread.join()

print("Sort complete!")

注：numpy.argsort的性能高度依赖底层线性代数库（如MKL、OpenBLAS）的优化，建议多次运行小样本取平均来提高预估时间的准确性。

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