如何避免重复编写120次Python降雨模型月度雨量统计代码？

核心优化方向 #

1. 自动生成日期标签（替代手动拼接列表） #

原代码手动创建月、年列表再拼接的方式繁琐且易出错，可直接用pd.date_range生成连续10年的日期序列，一键提取月份和年份标签：

# 生成10年日期序列（3650天，忽略闰年）
dates = pd.date_range(start='2020-01-01', periods=3650, freq='D')
rainfall['Month Tag'] = dates.month
rainfall['Year Tag'] = dates.year - 2019  # 让年份从1开始计数，匹配原逻辑

2. 分组统计一键完成（替代120次重复代码） #

利用Pandas的groupby功能，按年份+月份分组后直接计算总和/平均值，完全避免手动循环：

# 计算各年份各月份的降雨量总和
monthly_rainfall = rainfall.groupby(['Year Tag', 'Month Tag'])['Amount (mm)'].sum().reset_index()

# 计算10年间各月份的平均降雨量（所有年份同月份的均值）
monthly_avg = rainfall.groupby('Month Tag')['Amount (mm)'].mean().reset_index()

3. 向量操作替代循环（提升代码效率） #

将原代码中逐行循环的状态生成、降雨量计算改为Pandas向量操作，简化代码同时提升运行效率：

• 状态生成通过shift()获取前一行状态，结合条件判断实现向量化更新
• 降雨量计算直接用列运算替代逐行遍历

完整优化后的代码 #

import numpy as np
import pandas as pd

# 参数设置
pww = 0.7
pdd = 0.6
pwd = 1 - pww
pdw = 1 - pdd
onelam = 3.1
lda = 1 / onelam

# 初始化数据框
rainfall = pd.DataFrame({
    "Day": range(1, 3651),
    "Random 1": np.random.rand(3650),
    "Random 2": np.random.rand(3650)
})

# 自动生成月份、年份标签
dates = pd.date_range(start='2020-01-01', periods=3650, freq='D')
rainfall['Month Tag'] = dates.month
rainfall['Year Tag'] = dates.year - 2019

# 向量化生成状态序列
rainfall['State'] = 1
for i in range(1, len(rainfall)):
    prev_state = rainfall.loc[i-1, 'State']
    random_val = rainfall.loc[i, 'Random 2']
    rainfall.loc[i, 'State'] = 1 if (prev_state == 1 and random_val < pww) or (prev_state == 0 and random_val < pdw) else 0

# 向量化计算降雨量
rainfall['Amount (mm)'] = np.where(
    rainfall['State'] == 1,
    (-np.log(1 - rainfall['Random 1'])) / lda,
    0
)

# 输出统计结果
print("各年份各月份降雨量总和：")
print(rainfall.groupby(['Year Tag', 'Month Tag'])['Amount (mm)'].sum().reset_index())

print("\n10年间各月份平均降雨量：")
print(rainfall.groupby('Month Tag')['Amount (mm)'].mean().reset_index())

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