def _get_valid_data_inputs(
initialization_params: YAMLDict) -> Dict[str, Any]:
"""Helper to retrieve data_inputs for get_model_params"""
given_data_inputs = initialization_params.pop_dict("data_inputs")
if len(given_data_inputs) != 1:
raise ValueError(
f"Only one data input can be set in the yaml file, not {len(given_data_inputs)}"
)
data_inputs: Dict[str, Any] = dict()
if "big_query_inputs" in given_data_inputs.keys():
big_query_inputs_yaml_dict = given_data_inputs.pop_dict(
"big_query_inputs")
big_query_inputs_keys = big_query_inputs_yaml_dict.keys()
big_query_inputs_dict: Dict[str, str] = dict()
for k in big_query_inputs_keys:
big_query_inputs_dict[k] = big_query_inputs_yaml_dict.pop(
k, str)
data_inputs["big_query_inputs"] = big_query_inputs_dict
elif "big_query_simulation_tag" in given_data_inputs.keys():
data_inputs["big_query_simulation_tag"] = given_data_inputs.pop(
"big_query_simulation_tag", str)
else:
raise ValueError(
f"Received unexpected key in data_inputs: {given_data_inputs.keys()[0]}"
)
return data_inputs
def _get_user_inputs(
initialization_params: YAMLDict, model_params: Dict[str, Any]
) -> Dict[str, Any]:
"""Helper to retrieve user_inputs for get_model_params"""
user_inputs: Dict[str, Any] = dict()
user_inputs_yaml_dict = initialization_params.pop_dict("user_inputs")
if "big_query_simulation_tag" in model_params["data_inputs_raw"].keys():
user_inputs = {
"policy_year": user_inputs_yaml_dict.pop("policy_year", float),
"start_year": user_inputs_yaml_dict.pop("start_year", float),
"projection_years": user_inputs_yaml_dict.pop(
"projection_years", float
),
}
if "big_query_inputs" in model_params["data_inputs_raw"].keys():
user_inputs = {
"start_year": user_inputs_yaml_dict.pop("start_year", float),
"projection_years": user_inputs_yaml_dict.pop(
"projection_years", float
),
"run_date": user_inputs_yaml_dict.pop("run_date", str),
}
# Check for optional arguments
if user_inputs_yaml_dict:
user_inputs_keys = user_inputs_yaml_dict.keys()
for k in user_inputs_keys:
if k not in {"constant_admissions", "speed_run"}:
raise ValueError(f"Received unexpected key in user_inputs: {k}")
user_inputs[k] = user_inputs_yaml_dict.pop(k, bool)
return user_inputs
def _get_valid_compartments(
initialization_params: YAMLDict,
) -> Tuple[Dict[str, str], Dict[str, float]]:
"""Helper to retrieve model_architecture and compartment costs for get_model_params"""
compartments_architecture_yaml_key = "compartments_architecture"
compartments_architecture_raw = initialization_params.pop_dict(
compartments_architecture_yaml_key
)
compartments_architecture_keys = compartments_architecture_raw.keys()
compartments_architecture_dict: Dict[str, Any] = dict()
for k in compartments_architecture_keys:
compartments_architecture_dict[
k
] = compartments_architecture_raw.pop_optional(k, str)
compartment_costs_key = "per_year_costs"
compartment_costs_raw = initialization_params.pop_dict(compartment_costs_key)
compartment_costs_keys = compartment_costs_raw.keys()
compartment_costs_dict: Dict[str, float] = dict()
for k in compartment_costs_keys:
compartment_costs_dict[k] = compartment_costs_raw.pop(k, float)
# Ensure there are compartment costs for every compartment in the model architecture
model_compartments = set(
c
for c in compartments_architecture_keys
if compartments_architecture_dict[c] != "shell"
)
compartment_costs = set(compartment_costs_keys)
if compartment_costs != model_compartments:
raise ValueError(
f"Compartments do not match in the YAML '{compartment_costs_key}' "
f"and '{compartments_architecture_yaml_key}'\n"
f"Mismatched values: {compartment_costs ^ model_compartments}"
)
return compartments_architecture_dict, compartment_costs_dict