# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

clean_up=True # removes gams-related files in work-folder if true
# %run StdPackages.ipynb
os.chdir(directory['py'])
import techdata_to_tree, sys, ShockFunction
import abatement_ID_calib as abatement
os.chdir(directory['curr'])

# ## 1: Set up tree and read in data, akin to A1\_ID\_Calib

# *Read data:*

inputfile = "techdata_adv.xlsx" # file
inputoth  = 'Othdata_adv.xlsx'
output = techdata_to_tree.load_techcats(pd.read_excel(directory['data'] + "/" + inputfile, sheet_name=["inputdisp", "endofpipe", "inputprices"]))
modules = ["ID", "EOP"] # modules to use

# *Set up trees:*

nts = {m: nesting_tree.nesting_tree(name=m) for m in modules}

# #### 1.1: Input displacing tree

nts["ID"].add_tree(output["ID"]["upper_categories"], tree_name = 'ID_EC', **{"type_f":"CES_norm"}) # E to C
nts["ID"].add_tree(output["ID"]["components"], tree_name = "ID_CU", **{"type_f":"MNL"}) # C to U
nts["ID"].add_tree(output["ID"]["techs"], tree_name="ID_TU", **{'type_io': 'output', 'type_f': 'CET_norm'}) # U to T
nts["ID"].add_tree(output["ID"]["techs_inputs"], tree_name="ID_TX") # T to inputs X
nts["ID"].add_tree(output["ID"]["basetechs"], tree_name="ID_BU", **{"type_io":"output", "type_f":"linear_out"}) # baseline technologies
nts["ID"].add_tree(output["ID"]["basetech_inputs"], tree_name="ID_BX")

# #### 1.2: End of pipe tree

nts["EOP"].add_tree(output["EOP"]["components"], tree_name = "EOP_CU", **{"type_f":"MNL"}) # C to U
nts["EOP"].add_tree(output["EOP"]["techs"], tree_name="EOP_TU", **{'type_io': 'output', 'type_f': 'CET_norm'}) # T to U 
nts["EOP"].add_tree(output["EOP"]["techs_inputs"], tree_name="EOP_TX") # T to inputs X

# #### 1.3: Trees related to final goods and emissions account

# *Read in tree and data from xlsx:*

nts['ID'].add_tree(directory['data']+'\\'+inputoth,tree_name='ID_Y',**{'sheet':'Y'})
DataBase.GPM_database.merge_dbs(nts['ID'].trees['ID_Y'].database,excel2py.xl2PM.pm_from_workbook(directory['data']+'\\'+inputoth,{'Y': 'vars'}), 'first')

#
# *Add new Q2P:*

output['ID']['Q2P'] = output['ID']['Q2P'].append(excel2py.xl2PM.pm_from_workbook(directory['data']+'\\'+inputoth,{'Q2P': 'maps'}).database['Q2P'].vals)

# #### 1.3: Namespaces

# *Give modules local namespaces:*

standard_sets = ('inp','out','int','wT','map_all','kno_out','kno_inp')
for module in modules:
    namespace = {k: module + '_' + k for k in standard_sets}
    nts[module].run_all(**namespace)
    #Also replaces keys with module-specific names, i.e. changes 'inp' to 'ID_inp' in the attributes/keys.
    for std_set in standard_sets:
        if hasattr(nts[module], std_set):
            setattr(nts[module], module + "_" + std_set, getattr(nts[module], std_set))
            delattr(nts[module], std_set)

# *Add emissions data:*

DataBase.GPM_database.merge_dbs(nts['ID'].database,excel2py.xl2PM.pm_from_workbook(directory['data']+'\\'+inputoth, {'M': 'vars','M_sets': 'sets'}),'first')

# ## 2: Initialize model

# *Initialize:*

gm = abatement.abate(nt=nts['ID'],tech=output,work_folder=directory['work'],**{'data_folder':directory['gams'],'name': 'atest'})
gm.model.settings.tech_catalog = inputfile.split(".")[0]

# *Add a couple of extra aliases, and update sets with elements:*

gm.add_aliases([('n','nnnn'),('n','nnnnn')])
gm.model.database.update_all_sets(clean_up=False)

gm.initialize_variables_leontief()
gm.initialize_variables(**{'check_variables': True})

# +
db_star = DataBase.GPM_database(**{'name': 'db_star'})
sigma_star = gm.get("sigma").copy()
eta_star = gm.get("eta").copy()
mu_star = gm.get("mu").copy()

eta_star[:] = -3
sigma_star[gm.get("kno_ID_CU")] = 0.2
sigma_star[gm.get("kno_no_ID_BX")] = 0.5
sigma_star["Y"] = 0.2

mu_star[gm.get("map_ID_CU")] = 1.0

db_star["sigma"] = sigma_star
db_star["eta"] = eta_star
db_star["mu"] = mu_star
# -

name='v0'
gm.write_and_run(name=name, add_checkpoint=name)
gm.checkpoints['v1'] = gm.model_instances[name].ws.add_checkpoint()
gm.model_instances[name].solve_sneakily(db_star=db_star,from_cp = True, cp_init = gm.checkpoints[name], options_run={'checkpoint': gm.checkpoints['v1']})
gm.model.database.merge_dbs(gm.model.database,gm.model_instances[name].out_db,'second')

node_to_shock = "C_EL_1"
db_star = DataBase.GPM_database(**{'name': 'db_star'})
x_var = "sigma"
sigma_star = gm.get(x_var).copy()
sigma_star[node_to_shock] = 0.8
db_star[x_var] = sigma_star

(shock_db, shock_kwargs) = ShockFunction.sneaky_db(gm.model.database, db_star, n_steps = 160, loop_name = x_var + '_loop')
store_sol = {"PwThat":{'domains': shock_kwargs['loop_name']}, "qD":{'domains': shock_kwargs['loop_name']}, "sigma":{'domains': shock_kwargs['loop_name']}}
mi = gm.model_instances[name]
shock = mi.std_UEVAS_from_db(shock_db, loop_name=shock_kwargs['loop_name'], shock_name=shock_kwargs['shock_name'], store_sol=store_sol)

mi.execute_shock_from_cp(shock=shock, cp=gm.checkpoints['v1'])
gm.model.database.merge_dbs(gm.model.database,mi.out_db,'second')

# +

nodes_to_save = {
    "techdata_dors_2":["ID_t2", "U_ID_t2_1", "C_EL_1", "U0_ID_C_EL_1"]
}
sols = {key:mi.out_db.series["sol_" + key].vals for key in store_sol.keys()}


# -

def organize_sol_data(sols, x_var, nodes_to_save, node_to_shock, tech_catalog):
    list_of_dfs = []
    for var in sols:
        if var == x_var:
            continue
        df = sols[var]
        if len(df.index.get_level_values("n").unique()) == 1:
            only_one_series = True
        else:
            only_one_series = False
        df = df.reset_index()
        df2 = df.pivot_table(df.columns[-1], df.columns[0], df.columns[1]).reset_index()
        if gm.model.settings.tech_catalog in nodes_to_save and not only_one_series:
            df2 = df2[[x_var + "_loop"] + [c for c in nodes_to_save[tech_catalog] if c in df2.columns]]
        df3 = df2.merge(sols[x_var].iloc[sols[x_var].index.get_level_values(1) == node_to_shock].reset_index().drop(columns="n"), \
                        how="inner", on=x_var+"_loop").drop(columns=x_var+"_loop")
        df3 = df3.rename(columns={"sol_"+x_var:x_var}).set_index(x_var).sort_values(by=x_var)
        if only_one_series:
            df3.rename(columns={df3.columns[0]:var}, inplace=True)
        else:
            df3.rename(columns={key:var + "_"+ key for key in df3.columns}, inplace=True)
        list_of_dfs.append(df3)
    return list_of_dfs


out = {}
out[name] = organize_sol_data(sols=sols, x_var="sigma", nodes_to_save=nodes_to_save, node_to_shock=node_to_shock, tech_catalog=gm.model.settings.tech_catalog)

[gm.model.database.series.__delitem__(sym) for sym in ["sol_" + k for k in sols.keys()]]; # delete the symbols created from the loop
[gm.model.database.series.__delitem__(sym) for sym in gm.model.database.symbols if "loop" in sym]; # delete the symbols created from the loop
gm.export(name="medium_sigma")

# #### Second round, closer to where the error happens

db_star["sigma"].vals[node_to_shock] = 5

name = "v2"
(shock_db, shock_kwargs) = ShockFunction.sneaky_db(gm.model.database, db_star, n_steps = 240, loop_name=x_var+'_loop')
store_sol = {"PwThat":{'domains': shock_kwargs['loop_name']}, "qD":{'domains': shock_kwargs['loop_name']}, "sigma":{'domains': shock_kwargs['loop_name']}}
gm.write_and_run(name=name, add_checkpoint=name)
mi = gm.model_instances[name]
shock = mi.std_UEVAS_from_db(shock_db, loop_name=shock_kwargs['loop_name'], shock_name=shock_kwargs['shock_name'], store_sol=store_sol)

mi.execute_shock_from_cp(shock=shock, cp=gm.checkpoints[name])
sols = {key:mi.out_db.series["sol_" + key].vals for key in store_sol.keys()}
gm.model.database.merge_dbs(gm.model.database,mi.out_db,'second')

out[name] = organize_sol_data(sols=sols, x_var="sigma", nodes_to_save=nodes_to_save, node_to_shock=node_to_shock, tech_catalog=gm.model.settings.tech_catalog)

#Export as pickle
df_out = pd.concat([pd.concat(out["v0"], axis = 1), pd.concat(out["v2"], axis = 1)], axis=0)
df_out.to_pickle(x_var + "_" + gm.model.settings.tech_catalog + ".pkl")

#Export the model
[gm.model.database.series.__delitem__(sym) for sym in ["sol_" + k for k in sols.keys()]]; # delete the symbols created from the loop
[gm.model.database.series.__delitem__(sym) for sym in gm.model.database.symbols if "loop" in sym]; # delete the symbols created from the loop
gm.export(name="high_sigma")

 





StarcoderdataJupyterScriptsDedupFiltered