one nc

store.py

import xarray as xr

def gen_param_names():
# spatial subscript
    state_names = ['Alabama', 'Alaska', 'Arizona', 'Arkansas', 'California', 'Colorado', 'Connecticut', 'Delaware',
                   'Florida', 'Georgia', 'Hawaii', 'Idaho', 'Illinois', 'Indiana', 'Iowa', 'Kansas', 'Kentucky',
                   'Louisiana', 'Maine', 'Maryland', 'Massachusetts', 'Michigan', 'Minnesota', 'Mississippi', 'Missouri',
                   'Montana', 'Nebraska', 'Nevada', 'New Hampshire', 'New Jersey', 'New Mexico', 'New York',
                   'North Carolina', 'North Dakota', 'Ohio', 'Oklahoma', 'Oregon', 'Pennsylvania', 'Rhode Island',
                   'South Carolina', 'South Dakota', 'Tennessee', 'Texas', 'Utah', 'Vermont', 'Virginia', 'Washington',
                   'West Virginia', 'Wisconsin', 'Wyoming', 'District of Columbia']
    variant_list = ["Delta", 'Omicron', "BA5"]
    vax_status_list = ['Vx', 'NVx']
    param_lst = []
    state_lst_1 = []
    state_lst_2 = []
    state_lst_3 = []
    state_lst_4 = []
    state_lst_5 = []
    state_lst_6 = []
    for name in state_names:
        state_lst_1.append(f"Impact of Treatment on Fatality Rate[{name}]")
        state_lst_2.append(f"Base Fatality Rate for Unit Acuity[{name}]")
        state_lst_3.append(f"Reference Force of Infection[{name}]")
        state_lst_4.append(f"Strength of Adherence Fatigue[{name}]")
        state_lst_5.append(f"Variant Intro Start Time[{name}]")
        state_lst_6.append(f"Variant Intro Start Time2[{name}]")
    t_sharing_params =  [
            state_lst_1,
            state_lst_2,
            state_lst_3,
            state_lst_4,
            state_lst_5,
            state_lst_6,
        ]
    t_sharing_sbatched_params = [param[0].replace("[Alabama]", "") for param in t_sharing_params]
    ts_sharing_params = ['Variant Impact on Immunity Loss Time[Omicron]',
                         'Variant Impact on Immunity Loss Time[Delta]',
                         'Variant Impact on Immunity Loss Time[BA5]',
                         'Extra Vaccine Impact on Responsiveness[Vx]',
                         'Extra Vaccine Impact on Responsiveness[NVx]',
                         'Extra Vaccine Impact on Responsiveness[Naive]',
                         'Immunity Loss Time[Vx]',
                         'Immunity Loss Time[NVx]',
                         'Immunity Loss Time[Naive]',
                         'Variant Accuity Multiplier[Omicron]',
                         'Variant Accuity Multiplier[Delta]',
                         'Variant Accuity Multiplier[BA5]',
                         'Variant Intro Start Time3',
                         'Variant Transmission Multiplier[Omicron]',
                         'Variant Transmission Multiplier[Delta]',
                         'Variant Transmission Multiplier[BA5]']
    return state_names, t_sharing_sbatched_params, ts_sharing_params

vax_inc_1dot1 = 1.01
vax_inc_base = 1
vax_inc_zero = 0
FINAL_TIME = 1100
TIME_PERIODS = 1071
hyperparam_time = pd.date_range(start='10/15/2019', periods=TIME_PERIODS)

state_names, t_sharing_sbatched_params, ts_sharing_params  = gen_param_names()
param_lst = t_sharing_sbatched_params + ts_sharing_params

dv = {
    "coords": {
        'value': {"dims": ("value"),
                  "data": ["marginal", "average", str(vax_inc_1dot1), str(vax_inc_base), str(vax_inc_zero)]},
        'component': {"dims": ("component"), "data": ['death', 'gdp', 'hospitalization', 'population']},
        'params': {"dims": ("params"), "data": param_lst + ["baseline"]},
        'space': {"dims": ("space"), "data": state_names},
        'time': {"dims": ("time"), "data": pd.date_range(start='10/15/2019', periods=TIME_PERIODS)},
        'hyperparam': {"dims": ("hyperparam"), "data": list(range(-50, 60, 10))}
    },
    "dims": {"params", "component", "space", "time", "value", "action"}, #action becoms hyperparmeter
    "data_vars": {  # 5+ 3
        'VV_p1_pcstv': {"dims": ("params", "component", "space", "value", "time"),
                       "data": np.zeros(shape=(len(param_lst) + 1, 4, len(state_names), 5, TIME_PERIODS))},
        'VV_p1_pctv': {"dims": ("params", "component", "value", "time"),
                      "data": np.zeros(shape=(len(param_lst) + 1, 4, 5, TIME_PERIODS))},
        'VV_p_ptv': {"dims": ("params", "value", "time"),
                     "data": np.zeros(shape=(len(param_lst) + 1, 5, TIME_PERIODS))},
        'VV_h_csth': {"dims": ("component", "space", "time", "hyperparam"),
                      "data": np.zeros(shape=(4, len(state_names), TIME_PERIODS, 11))},
        'VV_h_cth': {"dims": ("component", "time", "hyperparam"), "data": np.zeros(shape=(4, TIME_PERIODS, 11))},
        'VV_h_th': {"dims": ("time", "hyperparam"), "data": np.zeros(shape=(TIME_PERIODS, 11))},

        'epsilon_ptv': {"dims": ("params", "value", "time"),
                        "data": np.zeros(shape=(len(param_lst) + 1, 5, TIME_PERIODS))},
        'epsilon_pctv': {"dims": ("params", "component", "value", "time"),
                         "data": np.zeros(shape=(len(param_lst) + 1, 4, 5, TIME_PERIODS))},
        'epsilon_pcstv': {"dims": ("params", "component", "space", "value", "time"),
                          "data": np.zeros(shape=(len(param_lst) + 1, 4, len(state_names), 5, TIME_PERIODS))},
        # time aggregated
        'VV_csv': {"dims": ("component", "space", "value"), "data": np.zeros(shape=(3, len(state_names), 2))},
        'VV_cv': {"dims": ("component", "value"), "data": np.zeros(shape=(3, 2))},
        'VV_sv': {"dims": ("space", "value"), "data": np.zeros(shape=(len(state_names), 2))},
        'VV_v': {"dims": ("value"), "data": np.zeros(shape=(2))},

        # parameter counterfactual (for validity testing)
        'VV_p1_pcsv': {"dims": ("params", "component", "space", "value"),
                      "data": np.zeros(shape=(len(param_lst), 3, len(state_names), 2))},
        'VV_p1_pcv': {"dims": ("params", "component", "value"), "data": np.zeros(shape=(len(param_lst), 3, 2))},
        'VV_p1_pv': {"dims": ("params", "value"), "data": np.zeros(shape=(len(param_lst), 2))},
        'epsilon_pv': {"dims": ("params", "value"), "data": np.zeros(shape=(len(param_lst), 2))},
        'epsilon_pcv': {"dims": ("params", "component", "value"),
                        "data": np.zeros(shape=(len(param_lst), 3, 2))},
    },
    "attrs": {"title": "1.vaccine value disaggregated by parameter, time, space, component, 2.counterfactual by parameter for validity, 3.counterfactual by datahp for summary stat."}
    # 1 and 2 are descriptive, 3 is prescriptive
}
store_xarr = xr.Dataset.from_dict(dv)