Merge pull request #73 from NREL/yep/calibration-refactor

Calibration refactoring, initial configuration file and first cleanup

Author: yamilbknsu

demos/config.py

@@ -0,0 +1,75 @@
+import toml
+from pydantic import BaseModel, model_validator
+from typing import Optional
+from templates.calibration import CalibrationConfig, SimultaneousCalibrationConfig
+
+CONFIG = None
+
+class HHRebalancingModuleConfig(BaseModel):
+    control_table: str
+    control_col: str
+    geoid_col: str
+
+class EmploymentModuleConfig(BaseModel):
+    simultaneous_calibration_config: Optional[SimultaneousCalibrationConfig] = None
+    enter_model_calibration_procedure: Optional[CalibrationConfig] = None
+    exit_model_calibration_procedure: Optional[CalibrationConfig] = None
+
+    @model_validator(mode="after")
+    def check_calibration_config_exclusivity(self):
+        sim_cal = self.simultaneous_calibration_config is not None
+        enter_cal = self.enter_model_calibration_procedure is not None
+        exit_cal = self.exit_model_calibration_procedure is not None
+        if sim_cal and (enter_cal or exit_cal):
+            raise ValueError(f"Simultaneous calibration cannot be used at the same time as " + \
+                             f"individual model calibration. Simultaneous selected: {sim_cal}, " + \
+                             f"EnterModel selected: {enter_cal}, ExitModel selected: {exit_cal}")
+        return self
+
+class HHReorgModuleConfig(BaseModel):
+    simultaneous_calibration_config: Optional[SimultaneousCalibrationConfig] = None
+    geoid_col: Optional[str] = None
+
+class MortalityModuleConfig(BaseModel):
+    calibration_procedure: Optional[CalibrationConfig] = None
+
+class BirthModuleConfig(BaseModel):
+    calibration_procedure: Optional[CalibrationConfig] = None
+
+class DEMOSConfig(BaseModel):
+    # Global config
+    region_code: str
+    forecast_year: int = 2020
+    random_seed: int
+    base_year: int
+    output_fname: str = None
+    calibrated_folder: str = "custom"
+
+    # Module-specific config
+    employment_module_config: EmploymentModuleConfig
+    mortality_module_config: MortalityModuleConfig
+    birth_module_config: BirthModuleConfig
+    hh_reorg_module_config: HHReorgModuleConfig
+    hh_rebalancing_module_config: HHRebalancingModuleConfig
+    
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if self.output_fname is None:
+            self.output_fname = "data/model_data_{0}.h5".format(self.forecast_year)
+
+
+    def update(self, **kwargs):
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+def load_config_file(dir: str) -> DEMOSConfig:
+    global CONFIG
+    CONFIG = DEMOSConfig(**toml.load(dir))
+
+def set_config(params_dict: dict):
+    global CONFIG
+    CONFIG.update(**vars(params_dict))
+
+def get_config():
+    global CONFIG
+    return CONFIG