from typing import Dict, Type
from microdf import MicroDataFrame, MicroSeries
import numpy as np
from policyengine_core.data.dataset import Dataset
from policyengine_core.periods import Period
from policyengine_core.periods import period as get_period
from policyengine_core.periods.config import MONTH, YEAR
from policyengine_core.simulations.simulation import Simulation
from policyengine_core.types import ArrayLike
[docs]class Microsimulation(Simulation):
"""A `Simulation` whose entities use weights to represent larger populations."""
[docs] def get_weights(
self, variable_name: str, period: Period, map_to: str = None
) -> ArrayLike:
"""Return weights for a variable's entity from ``household_weight``."""
time_period = get_period(period)
variable = self.tax_benefit_system.get_variable(
variable_name, check_existence=True
)
entity_key = map_to or variable.entity.key
weight_variable = self.tax_benefit_system.get_variable(
"household_weight", check_existence=True
)
weight_period = None
if time_period.unit == weight_variable.definition_period:
weight_period = time_period
elif (time_period.unit == MONTH) and (
weight_variable.definition_period == YEAR
):
# Common use-case. To-do: implement others if needed.
weight_period = time_period.this_year
if weight_period is None:
return None
household_weights = self.calculate(
"household_weight",
weight_period,
use_weights=False,
)
return self._project_household_weights(household_weights, entity_key)
def _project_household_weights(
self,
household_weights: ArrayLike,
entity_key: str,
) -> ArrayLike:
household_entity = self.tax_benefit_system.get_variable(
"household_weight", check_existence=True
).entity.key
if entity_key == household_entity:
return household_weights
household_population = self.populations[household_entity]
person_weights = household_population.project(household_weights)
if entity_key == self.tax_benefit_system.person_entity.key:
return person_weights
return self.populations[entity_key].value_from_first_person(person_weights)
[docs] def calculate(
self,
variable_name: str,
period: Period = None,
map_to: str = None,
use_weights: bool = True,
decode_enums: bool = True,
) -> MicroSeries:
if period is not None and not isinstance(period, Period):
period = get_period(period)
elif period is None and self.default_calculation_period is not None:
period = get_period(self.default_calculation_period)
values = super().calculate(variable_name, period, map_to, decode_enums)
if not use_weights:
return values
weights = self.get_weights(variable_name, period, map_to)
return MicroSeries(np.array(values), weights=weights)
[docs] def calculate_add(
self,
variable_name: str,
period: Period = None,
map_to: str = None,
use_weights: bool = True,
) -> MicroSeries:
values = super().calculate_add(variable_name, period)
if map_to is not None:
source_entity = self.tax_benefit_system.get_variable(
variable_name, check_existence=True
).entity.key
values = self.map_result(np.array(values), source_entity, map_to)
if not use_weights:
return values
weights = self.get_weights(variable_name, period, map_to)
return MicroSeries(np.array(values), weights=weights)
[docs] def calculate_divide(
self,
variable_name: str,
period: Period = None,
map_to: str = None,
use_weights: bool = True,
) -> MicroSeries:
values = super().calculate_divide(variable_name, period)
if map_to is not None:
source_entity = self.tax_benefit_system.get_variable(
variable_name, check_existence=True
).entity.key
values = self.map_result(np.array(values), source_entity, map_to)
if not use_weights:
return values
weights = self.get_weights(variable_name, period, map_to)
return MicroSeries(np.array(values), weights=weights)
[docs] def calculate_dataframe(
self,
variable_names: list,
period: Period = None,
map_to: str = None,
use_weights: bool = True,
) -> MicroDataFrame:
if period is not None and not isinstance(period, Period):
period = get_period(period)
elif period is None and self.default_calculation_period is not None:
period = get_period(self.default_calculation_period)
values = super().calculate_dataframe(variable_names, period, map_to)
if not use_weights:
return values
weights = self.get_weights(variable_names[0], period, map_to)
return MicroDataFrame(values, weights=weights)