pythonreferenceindirection

i need to restructure code for the addition of more equations with many similar use variables and keep the code succinct


I need to add 20 or so countries with part variable names like UK, US with multiple equations for each, each of which is structurally identical.

Some typical lines of code are shown ('model.add' etc is for the solver input and can be ignored): #(-1) is not -1 it has meaning in the model solver of a reference to the previous period. #UKUS has meaning 'from to'. Just think of the first two characters as the VAR of interest. RUK is of interest as UK. #I am not asking how to rewrite the code, rather how to restructure it.

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

NOTE the code is applied to a solver for which (-1) means the previous period

What i am looking for is something like [[i,j,k],[i,j,k]..] where i="a Uk variable/reference", j ="a US variable/reference" to a particular variable such as YDrUK, YDrUS in the positions i,j of [i,j,k]first cell and line "FcbUK=" is associated with a subsequent [i,,]., or better.

I thought i might be able to associate an index in an array to each variable but i don't know how to do that. Alternatively, i thought that i could write one line to select a country name corresponding to an index for the purpose of charting and file output. That is show country name and its data on plots against say trade or GNP for each. What i am looking for is ease of maintenance to add countries and change parameters, and minimal LOC overall.


Solution

  • You can write it in the following way using f-strings (I used print here instead of model.add):

    from itertools import product
    
    countries = ['UK', 'US', 'DE']
    
    
    # Example for equation with one country:
    
    for c in countries:
        print(f'V{c} - V{c}(-1) = YDr{c} - CONS{c}')
    
    
    # Example for equation with two countries:
    
    for c1, c2 in product(countries, repeat=2):
        if c1 == c2:
            continue
    
        print(f'YDhs{c1} = YDr{c1} + d(XR{c2})*Bs{c1}{c2}(-1)')
    

    Output:

    VUK - VUK(-1) = YDrUK - CONSUK
    VUS - VUS(-1) = YDrUS - CONSUS
    VDE - VDE(-1) = YDrDE - CONSDE
    YDhsUK = YDrUK + d(XRUS)*BsUKUS(-1)
    YDhsUK = YDrUK + d(XRDE)*BsUKDE(-1)
    YDhsUS = YDrUS + d(XRUK)*BsUSUK(-1)
    YDhsUS = YDrUS + d(XRDE)*BsUSDE(-1)
    YDhsDE = YDrDE + d(XRUK)*BsDEUK(-1)
    YDhsDE = YDrDE + d(XRUS)*BsDEUS(-1)