run black

Author: yamilbknsu

demos/__init__.py

@@ -2,4 +2,4 @@
 DEMOS - Demographic Micro-Simulator
 
 A microsimulation framework for demographic and economic modeling.
-"""
+"""

demos/config.py

@@ -10,14 +10,17 @@
 
 CONFIG = None
 
+
 class HHRebalancingModuleConfig(BaseModel):
     """
     Configuration for Household Rebalancing module
     """
+
     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
@@ -29,24 +32,31 @@ def check_calibration_config_exclusivity(self):
         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}")
+            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 KidsMovingModuleConfig(BaseModel):
     geoid_col: str
 
+
 class AgingModuleConfig(BaseModel):
     #: Age at which a person qualifies as senior
     senior_age: int = 65
@@ -56,6 +66,7 @@ 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
@@ -82,32 +93,44 @@ class DEMOSConfig(BaseModel):
 
     # 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)
+    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)
-    
+    hh_rebalancing_module_config: HHRebalancingModuleConfig = Field(
+        default_factory=HHRebalancingModuleConfig
+    )
+    kids_moving_module_config: KidsMovingModuleConfig = Field(
+        default_factory=KidsMovingModuleConfig
+    )
+
     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"
+            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",
@@ -118,21 +141,27 @@ def model_post_init(self, __context) -> None:
                 "birth_model",
                 "education_model",
                 "household_rebalancing",
-                "update_income"
+                "update_income",
             ]
-        
 
-    @model_validator(mode='after')
+    @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}")
+        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
 
+
 def load_config_file(dir: str) -> DEMOSConfig:
     global CONFIG
     CONFIG = DEMOSConfig(**toml.load(dir))
 
+
 def get_config():
     global CONFIG
     if CONFIG is None:

demos/datasources.py

@@ -7,7 +7,8 @@
 
 class CSVTableSource(BaseModel):
     """"""
-    file_type: Literal['csv']
+
+    file_type: Literal["csv"]
     #: Path to source file
     filepath: str
     #: Column in the file to be used as index (e.g. `person_id`)
@@ -21,24 +22,31 @@ class CSVTableSource(BaseModel):
 
     def load_into_orca(self):
         logger.info(f"Loading CSV '{self.table_name}' table from {self.filepath}")
-        df = pd.read_csv(self.filepath, delimiter=self.delimiter,
-            dtype=self.custom_dtype_casting).set_index(self.index_col)
+        df = pd.read_csv(
+            self.filepath, delimiter=self.delimiter, dtype=self.custom_dtype_casting
+        ).set_index(self.index_col)
         orca.add_table(self.table_name, df)
 
 
 class H5TableSource(BaseModel):
     """"""
-    file_type: Literal['h5']
+
+    file_type: Literal["h5"]
     #: Path to source file
     filepath: str
     #: key in the source HDF5 to be loaded
     h5_key: str
     #: Identifier of the table in orca
     table_name: str
-    
+
     def load_into_orca(self):
-        logger.info(f"Loading HDF5 '{self.table_name}' table from {self.filepath}/{self.h5_key}")
+        logger.info(
+            f"Loading HDF5 '{self.table_name}' table from {self.filepath}/{self.h5_key}"
+        )
         df = pd.read_hdf(self.filepath, key=self.h5_key)
         orca.add_table(self.table_name, df)
 
-DataSourceModel = Annotated[H5TableSource | CSVTableSource, Field(discriminator="file_type")]
+
+DataSourceModel = Annotated[
+    H5TableSource | CSVTableSource, Field(discriminator="file_type")
+]

demos/logging_logic.py

@@ -25,10 +25,16 @@ def flush(self):
             logger.opt(depth=1).log(self.level, f"{self.prefix}{self._buf.rstrip()}")
         self._buf = ""
 
+
 def log_execution_time(start_time, year, module_name):
     now = time.time()
-    run_table = orca.get_table('run_times')
-    run_table.local = pd.concat([run_table.local,
-                                pd.DataFrame([[year, module_name, now - start_time]],
-                                            columns=["year", "module", "walltime"])
-                                ])
+    run_table = orca.get_table("run_times")
+    run_table.local = pd.concat(
+        [
+            run_table.local,
+            pd.DataFrame(
+                [[year, module_name, now - start_time]],
+                columns=["year", "module", "walltime"],
+            ),
+        ]
+    )

demos/models/__init__.py

@@ -10,4 +10,4 @@
 from .rebalancing import *
 from .income_adjustment import *
 from .export import *
-from .main import *
+from .main import *

demos/models/aging.py

@@ -5,9 +5,8 @@
 from config import DEMOSConfig, AgingModuleConfig, get_config
 
 STEP_NAME = "aging"
-REQUIRED_COLUMNS = [
-    "persons.age"
-]
+REQUIRED_COLUMNS = ["persons.age"]
+
 
 @orca.step(STEP_NAME)
 def aging(persons):
@@ -78,6 +77,7 @@ def senior(data="persons.age"):
 
     return (data >= aging_config.senior_age).astype(int)
 
+
 @orca.column(table_name="persons")
 def age_group(data="persons.age"):
     """
@@ -96,5 +96,7 @@ def age_group(data="persons.age"):
         Categorical age group labels as strings.
     """
     age_intervals = [0, 20, 30, 40, 50, 65, 900]
-    age_labels = ['lte19', '20-29', '30-39', '40-49', '50-64', 'gte65']
-    return pd.cut(data, bins=age_intervals, labels=age_labels, include_lowest=True).astype(str)
+    age_labels = ["lte19", "20-29", "30-39", "40-49", "50-64", "gte65"]
+    return pd.cut(
+        data, bins=age_intervals, labels=age_labels, include_lowest=True
+    ).astype(str)

demos/models/birth.py

@@ -7,16 +7,23 @@
 from logging_logic import log_execution_time
 from config import DEMOSConfig, get_config
 
+
 @orca.injectable(autocall=False)
 def get_new_person_id(n):
     persons = orca.get_table("persons")
     graveyard = orca.get_table("graveyard")
     rebalanced_persons = orca.get_table("rebalanced_persons")
 
-    current_max = max([persons.local.index.max(), graveyard.local.index.max(), rebalanced_persons.local.index.max()])
+    current_max = max(
+        [
+            persons.local.index.max(),
+            graveyard.local.index.max(),
+            rebalanced_persons.local.index.max(),
+        ]
+    )
     return (
-        np.arange(n)    # = [0, 1, 2 ...] up to the number of people
-        + current_max   # = [max_person_id, max_person_id + 1, ...]
+        np.arange(n)  # = [0, 1, 2 ...] up to the number of people
+        + current_max  # = [max_person_id, max_person_id + 1, ...]
         + 1
     )
 
@@ -61,19 +68,25 @@ def birth_model(persons, households, observed_births_data, get_new_person_id, ye
 
     # Set race of babies
     # TODO: There is duplication of information between `race_id` and `race`
-    hh_races = (persons.local.groupby("household_id")
-                             .agg(num_races=("race_id", "nunique"))
-                             .reset_index()
-                             .merge(
-                                 households.to_frame(["hh_race_of_head", "hh_race_id_of_head", "household_id"])
-                                 .reset_index(),
-                               on="household_id")).set_index("household_id")
+    hh_races = (
+        persons.local.groupby("household_id")
+        .agg(num_races=("race_id", "nunique"))
+        .reset_index()
+        .merge(
+            households.to_frame(
+                ["hh_race_of_head", "hh_race_id_of_head", "household_id"]
+            ).reset_index(),
+            on="household_id",
+        )
+    ).set_index("household_id")
     one_race_hh_filter = (hh_races.loc[babies.household_id]["num_races"] == 1).values
     babies["race_id"] = 9
-    babies.loc[one_race_hh_filter, "race_id"] = hh_races.loc[babies.loc[one_race_hh_filter, "household_id"], "hh_race_id_of_head"].values
+    babies.loc[one_race_hh_filter, "race_id"] = hh_races.loc[
+        babies.loc[one_race_hh_filter, "household_id"], "hh_race_id_of_head"
+    ].values
     babies["race"] = babies["race_id"].map({1: "white", 2: "black"})
     babies["race"].fillna("other", inplace=True)
-    
+
     # Finally add babies to persons table
     persons.local = pd.concat([persons.local, babies])
 
@@ -87,13 +100,15 @@ def run_and_calibrate_birth_model(persons, households):
     # Load calibration config
     demos_config: DEMOSConfig = get_config()
     calibration_procedure = demos_config.birth_module_config.calibration_procedure
-    
+
     # Get model data
     birth_model = mm.get_step("birth")
     birth_model_variables = columns_in_formula(birth_model.model_expression)
     birth_model_data = households.to_frame(birth_model_variables).loc[ELIGIBLE_HH]
-    
+
     # Calibrate if needed
     if calibration_procedure is not None:
-        return calibration_procedure.calibrate_and_run_model(birth_model, birth_model_data)
+        return calibration_procedure.calibrate_and_run_model(
+            birth_model, birth_model_data
+        )
     return birth_model.predict(birth_model_data)