Python嵌套循环生成多文件及np.savez数据保存问题咨询

问题背景 #

我正在完成一项Python作业，现有代码需要遍历beta的每个取值，针对每个beta值再遍历reduction_factor的每个取值执行后续计算步骤。要求每次beta与reduction_factor的迭代结果，都保存到listofsolutions列表中对应顺序的文件内。目前有两个核心疑问：

1. 如何将嵌套循环的迭代与listofsolutions中的文件名对应关联；
2. 如何正确使用np.savez将计算数据保存到对应文件中。

修正后的完整代码 #

import numpy as np

# 假设fe是你的有限元工具模块，E、rho_tilde、le、n_dof、Ae、A_bar、n_el这些变量已定义
b = [1/4, 0]
beta = np.asarray(b)
gamma = 0.5
listofsolutions = [
    'Q2_AA_0.1','Q2_AA_0.9','Q2_AA_0.99', 'Q2_AA_1', 'Q2_AA_1.1', 'Q2_AA_2',
    'Q2_CD_0.1','Q2_CD_0.9','Q2_CD_0.99', 'Q2_CD_1', 'Q2_CD_1.1', 'Q2_CD_2'
]
consistent = True # use a consistent mass matrix

file_idx = 0  # 关键：用计数器跟踪当前要保存的文件索引

# 直接遍历beta数组，写法更直观
for bb in beta:
    c = np.sqrt(E / rho_tilde) # wave speed
    T = 0.016 # total time
    # compute the critical time-step
    # note: uncondionally stable AA scheme will return 1.0
    delta_t_crit = fe.get_delta_t_crit(le = le, gamma = gamma, beta = bb, consistent = consistent, c = c)
    # actual times-step used is a factor of the critical time-step
    reduction_factor = [0.1, 0.9, 0.99, 1, 1.1, 2]
    
    for rf in reduction_factor:
        delta_t = rf * delta_t_crit
        n_t_steps = int(np.ceil(T / delta_t)); # number of time step
        # initialise the time domain, K and M
        t = np.linspace(0, T, n_t_steps)
        K = np.zeros((n_dof, n_dof))
        M = np.zeros((n_dof, n_dof))
        # assemble K and M
        for ee in range(n_el):
            dof_index = fe.get_dof_index(ee)
            M[np.ix_(dof_index, dof_index)] += fe.get_Me(le = le, Ae = Ae, rho_tilde_e = rho_tilde, consistent = consistent)
        # damping matrix
        C = np.zeros((n_dof, n_dof))
        # assemble the system matrix A
        A_matrix = M + (gamma * delta_t) * C + (bb * delta_t**2)*K  # 使用当前循环的bb，而非全局beta数组
        # define the free dofs
        free_dof = np.arange(1,n_dof)
        # initial conditions
        d = np.zeros((n_dof, 1))
        v = np.zeros((n_dof, 1))
        F = np.zeros((n_dof, 1))
        # compute the initial acceleration
        a = np.linalg.solve(M, F - C.dot(v) - K.dot(d))
        # store the history data
        # rows -> each node
        # columns -> each time step including initial at 0
        d_his = np.zeros((n_dof, n_t_steps))
        v_his = np.zeros((n_dof, n_t_steps))
        a_his = np.zeros((n_dof, n_t_steps))
        d_his[:,0] = d[:,0]
        v_his[:,0] = v[:,0]
        a_his[:,0] = a[:,0]
        # loop over the time domain and solve the problem at each step
        for n in range(1,n_t_steps):
            # data at beginning of the time-step n
            a_n = a
            v_n = v
            d_n = d
            # applied loading
            t_current = n * delta_t # current time
            if t_current<0.001:
                F[-1] = A_bar * Ae * np.sin(1000 * t_current * np.pi)
            else:
                F[-1]=0.
            # define predictors
            d_tilde = d_n + delta_t*v_n + ((delta_t**2)/2.) * (1 - 2*bb) * a_n
            v_tilde = v_n + (1 - gamma) * delta_t * a_n
            # assemble the right-hand side from the known data
            R = F - C.dot(v_tilde) - K.dot(d_tilde)
            # impose essential boundary condition and solve A a = RHS
            A_free = A_matrix[np.ix_(free_dof, free_dof)]
            R_free = R[np.ix_(free_dof)]
            # solve for the accelerations at the free nodes
            a_free = np.linalg.solve(A_free, R_free)
            a = np.zeros((n_dof, 1))
            a[1:] = a_free
            # update displacement and vecloity predictors using the acceleration
            d = d_tilde + (bb * delta_t**2) * a
            v = v_tilde + (gamma * delta_t) * a
            # store solutions
            d_his[:,n] = d[:,0]
            v_his[:,n] = v[:,0]
            a_his[:,n] = a[:,0]
        # post-processing
        mid_node = int(np.ceil(n_dof / 2)) # mid node
        # compute the stress in each element
        # assuming constant E
        stress = (E / le) * np.diff(d_his, axis=0)
        # 保存数据到对应文件
        np.savez(f"{listofsolutions[file_idx]}.npz", time=t, mid_element_stress=stress[mid_node,:])
        file_idx += 1  # 计数器递增，对应下一个文件名

关键修改与解释 #

1. 嵌套循环与文件名的关联 #

• 新增了file_idx计数器变量，初始值为0。每次完成一组beta+reduction_factor的计算并保存文件后，计数器递增1。
• 这个逻辑完美匹配你的listofsolutions顺序：先处理第一个beta的6个reduction_factor（对应列表前6个文件名），再处理第二个beta的6个reduction_factor（对应列表后6个文件名）。
• 替换了你之前错误的for i in b and r in reduction_factor:写法，计数器方式简洁且不易出错。

2. 正确使用np.savez #

• 修正了np.savez的语法错误：补充了闭合括号，并且给保存的数组加上了关键字参数（time=t和mid_element_stress=stress[mid_node,:]）。后续读取数据时，可以通过data = np.load("filename.npz")，再用data['time']、data['mid_element_stress']直接获取对应数组，可读性更强。
• 给文件名加上了.npz后缀，这是numpy压缩文件的标准后缀，方便识别和后续读取。

3. 其他细节修正 #

• 把itertools.zip_longest(beta)改成直接遍历beta数组（for bb in beta:），因为beta是长度为2的数组，写法更直观，没必要用zip_longest。
• 代码中所有用到beta的计算步骤，都替换为当前循环的bb变量，避免误用全局beta数组导致错误。

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