迭代物理计算中函数可视化的代码优化方法问询

我正在求解一个迭代式物理问题，需递归计算函数f，当函数满足某种收敛条件时即为解。以下是示意性最小可复现示例（MWE）：

from scipy.interpolate import CubicSpline
import numpy as np  # 补充原代码缺失的导入

class Experiment:
    def __init__(self):
        self.x = np.linspace(0, 1, 100)
        self.f = lambda x: 0  # 初始函数

    def update_f(self):
        # 更新f的示意逻辑（修正原代码闭包引用问题）
        current_f = self.f
        self.f = lambda x: current_f(x)**2 + 1

    def convergence(self):
        # 示意收敛判断，实际需根据问题定义
        return False

    def main_calculation(self):
        # 迭代直到收敛
        while not self.convergence():
             self.update_f()

experiment = Experiment()
experiment.main_calculation()

每次想要查看计算过程中变量的变化（例如绘图）时，都需要修改主代码内部的方法，既繁琐又容易让代码杂乱。比如想查看f更新前后的形态，就得修改main_calculation方法，但这样会产生大量重复代码，还可能引入不必要的依赖。

优化方案 #

1. 引入回调函数机制 #

让main_calculation支持传入回调函数，在每次迭代的关键节点（比如更新前后）触发回调，无需修改核心代码就能添加自定义的监控、绘图逻辑。

修改后的核心代码：

class Experiment:
    def __init__(self):
        self.x = np.linspace(0, 1, 100)
        self.f = lambda x: 0

    def update_f(self):
        current_f = self.f
        self.f = lambda x: current_f(x)**2 + 1

    def convergence(self):
        # 示例：假设迭代10次后收敛
        return getattr(self, 'iter_count', 0) >= 10

    def main_calculation(self, before_update=None, after_update=None):
        self.iter_count = 0
        while not self.convergence():
            if before_update is not None:
                before_update(self)  # 传入当前实例供回调使用
            self.update_f()
            if after_update is not None:
                after_update(self)
            self.iter_count += 1

单独编写绘图回调，完全与核心计算逻辑分离：

import matplotlib.pyplot as plt

def plot_before_update(exp):
    plt.plot(exp.x, exp.f(exp.x), alpha=0.5, label=f'Iter {exp.iter_count} (before)')

def plot_after_update(exp):
    plt.plot(exp.x, exp.f(exp.x), alpha=0.8, label=f'Iter {exp.iter_count} (after)')

# 使用方式
fig, ax = plt.subplots()
experiment = Experiment()
experiment.main_calculation(before_update=plot_before_update, after_update=plot_after_update)
ax.legend()
plt.show()

2. 记录迭代历史 #

如果需要事后分析或批量绘图，可以在Experiment类中添加历史记录功能，自动保存每次迭代的f值，之后再统一处理：

class Experiment:
    def __init__(self):
        self.x = np.linspace(0, 1, 100)
        self.f = lambda x: 0
        self.history = []  # 保存每次迭代的f(x)结果

    def update_f(self):
        current_f = self.f
        self.f = lambda x: current_f(x)**2 + 1

    def convergence(self):
        return len(self.history) >= 10

    def main_calculation(self):
        # 先记录初始状态
        self.history.append(self.f(self.x))
        while not self.convergence():
            self.update_f()
            self.history.append(self.f(self.x))

# 计算完成后统一绘图
experiment = Experiment()
experiment.main_calculation()

fig, ax = plt.subplots()
for i, fx in enumerate(experiment.history):
    ax.plot(experiment.x, fx, label=f'Iter {i}')
ax.legend()
plt.show()

3. 观察者模式（进阶） #

如果需要多个监控逻辑同时运行，可以用观察者模式，让Experiment维护一个观察者列表，迭代时通知所有观察者：

class Experiment:
    def __init__(self):
        self.x = np.linspace(0, 1, 100)
        self.f = lambda x: 0
        self.observers = []

    def add_observer(self, observer):
        self.observers.append(observer)

    def notify_observers(self, event_type):
        for observer in self.observers:
            observer.update(event_type, self)

    def update_f(self):
        current_f = self.f
        self.f = lambda x: current_f(x)**2 + 1

    def convergence(self):
        return getattr(self, 'iter_count', 0) >= 10

    def main_calculation(self):
        self.iter_count = 0
        while not self.convergence():
            self.notify_observers('before_update')
            self.update_f()
            self.notify_observers('after_update')
            self.iter_count += 1

# 定义观察者类
class PlotObserver:
    def __init__(self):
        self.fig, self.ax = plt.subplots()

    def update(self, event_type, exp):
        if event_type == 'after_update':
            self.ax.plot(exp.x, exp.f(exp.x), alpha=0.6, label=f'Iter {exp.iter_count}')

class LogObserver:
    def update(self, event_type, exp):
        if event_type == 'after_update':
            print(f"Iteration {exp.iter_count}: max(f(x)) = {np.max(exp.f(exp.x)):.4f}")

# 使用方式
experiment = Experiment()
experiment.add_observer(PlotObserver())
experiment.add_observer(LogObserver())
experiment.main_calculation()
plt.legend()
plt.show()

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