多国家同质化方程代码重构需求：简化维护与缩减代码行数

需求背景 #

需要为求解器添加20余个国家的同质化方程，所有变量名称均带有UK、US这类国家标识后缀；当前代码存在大量重复结构，新增国家或修改方程逻辑时需重复编写相似代码，维护成本极高。

现有重复代码示例 #

注：(-1)并非数值-1，在求解器中代表上一周期的参考值；UKUS代表“从到”关系，前两位字符为目标变量标识（如RUK对应UK）；以下代码中model.add为求解器输入方法，无需关注其内部实现。

model.add('YDrUK = (YUK + RUK(-1)*BdUKUK(-1) + XRUS*RUS(-1)*BsUKUS(-1))*(1 - thetaUK) + d(XRUS)*BsUKUS(-1)')
model.add('YDhsUK = YDrUK + d(XRUS)*BsUKUS(-1)') # 12.2 : Haig-Simons disposable income in the UK
model.add('VUK - VUK(-1) = YDrUK - CONSUK')      # 12.3 : Wealth accumulation in the UK
# 12.4 : Disposable income in the US
model.add('YDrUS = (YUS + RUS(-1)*BdUSUS(-1) + XRUK*RUK(-1)*BsUSUK(-1))*(1 - thetaUS) + d(XRUK)*BsUSUK(-1)')
model.add('YDhsUS = YDrUS + d(XRUK)*BsUSUK(-1)') # 12.5 : Haig-Simons disposable income in the US
model.add('VUS - VUS(-1) = YDrUS - CONSUS')      # 12.6 : Wealth accumulation in the US
model.add('TUK = thetaUK*(YUK + RUK(-1)*BdUKUK(-1) + XRUS*RUS(-1)*BsUKUS(-1))')  # 12.7 : Taxes in the UK
model.add('TUS = thetaUS*(YUS + RUS(-1)*BdUSUS(-1) + XRUK*RUK(-1)*BsUSUK(-1))')  # 12.8 : Taxes in the US


model.add('FcbUK = RUK(-1)*BcbdUKUK(-1) + RUS(-1)*BcbsUKUS(-1)*XRUS')  # 12.11 : UK central bank profits
model.add('FcbUS = RUS(-1)*BcbdUSUS(-1)')        # 12.12 : US central bank profits
model.add('BsUK = BsUK(-1) + GUK + RUK(-1)*BsUK(-1) - TUK - FcbUK')  # 12.13 : UK Govt budget constraint
model.add('BsUS = BsUS(-1) + GUS + RUS(-1)*BsUS(-1) - TUS - FcbUS')  # 12.14 : US Govt budget constraint

核心重构诉求 #

• 用结构化配置关联各国变量/引用，避免重复编写方程结构
• 支持通过索引快速获取国家名称，满足后续图表生成、文件输出的需求
• 最大限度减少代码行数，提升可维护性，新增国家或修改方程逻辑时只需调整配置或模板

具体重构方案 #

1. 定义国家配置与方程模板 #

先整理所有涉及的国家列表，以及每个国家对应的境外关联国家，再为每一类同质化方程定义字符串模板，用占位符替换国家标识：

# 国家配置：键为当前国家标识，值为其关联的境外国家标识
country_config = {
    'UK': 'US',
    'US': 'UK'
    # 新增国家时直接添加键值对即可，如 'DE': 'FR', 'FR': 'DE'
}

# 方程模板：每个模板对应一类同质化逻辑，{self}为当前国家标识，{foreign}为关联境外国家标识
equation_templates = [
    # 可支配收入
    'YDr{self} = (Y{self} + R{self}(-1)*Bd{self}{self}(-1) + XR{foreign}*R{foreign}(-1)*Bs{self}{foreign}(-1))*(1 - theta{self}) + d(XR{foreign})*Bs{self}{foreign}(-1)',
    # Haig-Simons可支配收入
    'YDhs{self} = YDr{self} + d(XR{foreign})*Bs{self}{foreign}(-1)',
    # 财富积累
    'V{self} - V{self}(-1) = YDr{self} - CONS{self}',
    # 税收
    'T{self} = theta{self}*(Y{self} + R{self}(-1)*Bd{self}{self}(-1) + XR{foreign}*R{foreign}(-1)*Bs{self}{foreign}(-1))',
    # 政府预算约束
    'Bs{self} = Bs{self}(-1) + G{self} + R{self}(-1)*Bs{self}(-1) - T{self} - Fcb{self}'
]

2. 循环生成并添加方程 #

通过遍历国家配置和方程模板，自动生成对应国家的方程并调用model.add：

for self_country, foreign_country in country_config.items():
    for template in equation_templates:
        # 替换模板中的占位符，生成具体国家的方程
        equation = template.format(self=self_country, foreign=foreign_country)
        model.add(equation)

3. 特殊场景处理 #

对于部分国家存在特殊逻辑的情况（比如示例中US的央行利润无境外关联项），可以单独添加特殊模板或在循环中增加条件判断：

# 特殊方程模板，仅针对特定国家
special_templates = {
    'UK': ['FcbUK = RUK(-1)*BcbdUKUK(-1) + RUS(-1)*BcbsUKUS(-1)*XRUS'],
    'US': ['FcbUS = RUS(-1)*BcbdUSUS(-1)']
}

# 处理特殊方程
for country, equations in special_templates.items():
    for eq in equations:
        model.add(eq)

4. 支持国家名称索引输出 #

如果需要在图表或文件输出中调用国家名称，可扩展国家配置为包含名称的字典：

country_config = {
    'UK': {'foreign_code': 'US', 'name': 'United Kingdom'},
    'US': {'foreign_code': 'UK', 'name': 'United States'}
}

# 后续输出时通过索引获取名称
for country_code, info in country_config.items():
    print(f"生成{info['name']}的相关图表")

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