Source code for demos.config
import os
import toml
import orca
import pandas as pd
from pydantic import BaseModel, model_validator, Field, TypeAdapter, field_validator
from typing import Literal, Optional
from loguru import logger
from templates.calibration import CalibrationConfig, SimultaneousCalibrationConfig
from datasources import DataSourceModel
# Adapter used to coerce inline DataSourceModel dicts in module config fields
_file_adapter = TypeAdapter(DataSourceModel)
CONFIG = None
[docs]
class HHRebalancingModuleConfig(BaseModel):
"""
Configuration for Household Rebalancing module
"""
control_table: str
control_col: str
geoid_col: str
[docs]
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
[docs]
class HHReorgModuleConfig(BaseModel):
simultaneous_calibration_config: Optional[SimultaneousCalibrationConfig] = None
geoid_col: Optional[str] = None
[docs]
class MortalityModuleConfig(BaseModel):
calibration_procedure: Optional[CalibrationConfig] = None
[docs]
class BirthModuleConfig(BaseModel):
calibration_procedure: Optional[CalibrationConfig] = None
[docs]
class KidsMovingModuleConfig(BaseModel):
geoid_col: str
calibration_target_share: float = 0.12
calibration_tolerance: float = 0.001
max_iter: int = 100
[docs]
class AgingModuleConfig(BaseModel):
#: Age at which a person qualifies as senior
senior_age: int = 65
[docs]
class IncomeModuleConfig(BaseModel):
"""
Configuration for the Income module.
The income model predicts household income as nominal dollars for the year
in which it was estimated (*inflation_reference_year*). When
*inflation_reference_year* and *inflation_adjustment_table* are both
provided, predictions are automatically scaled to the current simulation
year using cumulative annual CPI adjustment factors read from that table.
Optionally, a calibration procedure can be configured to nudge aggregate
income predictions toward observed values (same mechanics as the mortality
and birth modules).
**Inflation adjustment table format** (``index_col = "year"``)::
year,adjustment
2011,0.030
2012,0.021
...
The ``adjustment`` column must contain decimal annual inflation rates
(e.g. ``0.030`` represents 3.0 %). The table must cover every year from
*inflation_reference_year* (exclusive) through the last simulation year
(inclusive). If the current simulation year equals
*inflation_reference_year*, no adjustment is applied.
"""
#: Optional calibration procedure to align aggregate predicted income with
#: observed aggregate values. Uses the same RMSE-based mechanics as the
#: mortality and birth calibration procedures. Omit (or comment out)
#: this block to skip calibration.
calibration_procedure: Optional[CalibrationConfig] = None
#: Name of the orca table containing annual CPI inflation adjustment
#: factors, or an inline DataSource definition (like a ``[[tables]]``
#: entry) that will be loaded automatically. The table must be indexed
#: by ``year`` and include an ``adjustment`` column with decimal annual
#: inflation rates. Must be set together with *inflation_reference_year*.
inflation_adjustment_table: Optional[str] = None
#: The reference year of the income model estimation. Predicted income
#: values are expressed in nominal dollars for this year and will be
#: scaled forward (or backward) to the current simulation year using the
#: cumulative product of the annual adjustment factors. Must be set
#: together with *inflation_adjustment_table*.
inflation_reference_year: Optional[int] = None
@field_validator("inflation_adjustment_table", mode="before")
@classmethod
def _coerce_inflation_table(cls, v):
"""Accept either an inline DataSourceModel dict or a plain table name."""
if isinstance(v, dict):
v = _file_adapter.validate_python(v)
if isinstance(v, BaseModel):
v.load_into_orca()
return v.table_name
return v
@model_validator(mode="after")
def _check_inflation_config_consistency(self):
has_table = self.inflation_adjustment_table is not None
has_year = self.inflation_reference_year is not None
if has_table != has_year:
raise ValueError(
"'inflation_adjustment_table' and 'inflation_reference_year' must "
"both be set together, or both omitted."
)
return self
[docs]
class DEMOSConfig(BaseModel):
"""
Global configuration for DEMOS. Individual fields in this class control the configuration of each module.
"""
random_seed: int
#: Year represented in synthetic population input
base_year: int
#: Last year of simulation
forecast_year: int = 2020
#: Path to DEMOS outputs
output_dir: str = "../data/output"
#: Name of output HDF5 file. Defaults to `demos_output_{forecast_year}.h5`.
output_fname: str = None
#: List of orca tables to include in output
output_tables: list[str] = None
#: Path to directory with calibration models
calibrated_models_dir: str = None
#: Behavior of inconsistent `persons` input table
inconsistent_persons_table_behavior: Literal["error", "fix", "ignore"] = "error"
#: Name of tables to be initialized as empty
initialize_empty_tables: list[str] = None
#: List of tables to be loaded into orca
tables: Optional[list[DataSourceModel]] = None
#: List of modules to be run
modules: Optional[list[str]] = None
# Module-specific config
aging_module_config: AgingModuleConfig = Field(default_factory=AgingModuleConfig)
employment_module_config: EmploymentModuleConfig = Field(
default_factory=EmploymentModuleConfig
)
hh_reorg_module_config: HHReorgModuleConfig = Field(
default_factory=HHReorgModuleConfig
)
mortality_module_config: MortalityModuleConfig = Field(
default_factory=MortalityModuleConfig
)
birth_module_config: BirthModuleConfig = Field(default_factory=BirthModuleConfig)
hh_rebalancing_module_config: HHRebalancingModuleConfig = Field(
default_factory=HHRebalancingModuleConfig
)
kids_moving_module_config: KidsMovingModuleConfig = Field(
default_factory=KidsMovingModuleConfig
)
income_module_config: IncomeModuleConfig = Field(default_factory=IncomeModuleConfig)
def model_post_init(self, __context) -> None:
if self.output_fname is None:
self.output_fname = (
f"{self.output_dir}/demos_output_{self.forecast_year}.h5"
)
os.makedirs(self.output_dir, exist_ok=True)
logger.info(f"Output file set to default: {self.output_fname}")
if self.output_tables is None:
self.output_tables = []
if self.initialize_empty_tables is None:
self.initialize_empty_tables = []
# Load all table datasources
for t in self.tables:
t.load_into_orca()
for n in self.initialize_empty_tables:
orca.add_table(n, pd.DataFrame())
if self.modules is None:
self.modules = [
"aging",
"laborforce_model",
"households_reorg",
"kids_moving_model",
"fatality_model",
"birth_model",
"education_model",
"household_rebalancing",
"update_income",
"normalize_table_dtypes",
]
@model_validator(mode="after")
def require_persons_and_households(self):
loaded_table_names = [t.table_name for t in self.tables]
if (
"persons" not in loaded_table_names
or "households" not in loaded_table_names
):
raise ValueError(
f"Both 'persons' and 'households' tables are required. Tables defined: {loaded_table_names}"
)
return self
@model_validator(mode="after")
def require_lcm_county_column(self):
loaded_table_names = [t.table_name for t in self.tables]
if "lcm_county_id" not in orca.get_table("households").columns:
raise ValueError(
f"`lcm_county_id` (County FIPS) is required in the households table"
)
return self
def load_config_file(dir: str) -> DEMOSConfig:
global CONFIG
CONFIG = DEMOSConfig(**toml.load(dir))
def get_config():
global CONFIG
if CONFIG is None:
CONFIG = DEMOSConfig()
return CONFIG