policyengine-uk 2.42.0__py3-none-any.whl → 2.43.0__py3-none-any.whl

policyengine_uk/system.py

@@ -1,21 +1,40 @@
+# Standard library imports
+import copy
 from pathlib import Path
-from policyengine_uk.entities import entities
+from typing import Any, Dict, List, Optional, Union, Type
+
+# Third-party imports
+import numpy as np
+import pandas as pd
+from microdf import MicroDataFrame, MicroSeries
+
+# PolicyEngine core imports
 from policyengine_core.data import Dataset
-from policyengine_core.taxbenefitsystems import TaxBenefitSystem
-from policyengine_core.simulations import (
-    Simulation as CoreSimulation,
-    Microsimulation as CoreMicrosimulation,
+from policyengine_core.parameters.operations.propagate_parameter_metadata import (
+    propagate_parameter_metadata,
 )
-from policyengine_uk.data.dataset_schema import (
-    UKSingleYearDataset,
-    UKMultiYearDataset,
+from policyengine_core.periods import period as period_
+from policyengine_core.parameters.operations.uprate_parameters import (
+    uprate_parameters,
 )
+from policyengine_core.parameters import Parameter
+from policyengine_core.reforms import Reform
+from policyengine_core.simulations import Simulation as CoreSimulation
+from policyengine_core.taxbenefitsystems import TaxBenefitSystem
 from policyengine_core.tools.hugging_face import download_huggingface_dataset
+from policyengine_core.tracers import FullTracer, SimpleTracer
 
-import pandas as pd
-from policyengine_uk.utils.parameters import (
-    backdate_parameters,
-    convert_to_fiscal_year_parameters,
+# PolicyEngine UK imports
+from policyengine_uk.data.dataset_schema import (
+    UKMultiYearDataset,
+    UKSingleYearDataset,
+)
+from policyengine_uk.entities import BenUnit, Household, Person
+from policyengine_uk.parameters.gov.contrib.create_private_pension_uprating import (
+    add_private_pension_uprating_factor,
+)
+from policyengine_uk.parameters.gov.dwp.state_pension.triple_lock.create_triple_lock import (
+    add_triple_lock,
 )
 from policyengine_uk.parameters.gov.economic_assumptions.create_economic_assumption_indices import (
     create_economic_assumption_indices,
@@ -26,38 +45,29 @@ from policyengine_uk.parameters.gov.economic_assumptions.lag_average_earnings im
 from policyengine_uk.parameters.gov.economic_assumptions.lag_cpi import (
     add_lagged_cpi,
 )
-from policyengine_core.reforms import Reform
-from policyengine_uk.reforms import create_structural_reforms_from_parameters
-
-from policyengine_uk.parameters.gov.contrib.create_private_pension_uprating import (
-    add_private_pension_uprating_factor,
-)
-from policyengine_uk.parameters.gov.dwp.state_pension.triple_lock.create_triple_lock import (
-    add_triple_lock,
-)
-from policyengine_core.parameters.operations.homogenize_parameters import (
-    homogenize_parameter_structures,
-)
-from policyengine_core.parameters.operations.interpolate_parameters import (
-    interpolate_parameters,
-)
-from policyengine_core.parameters.operations.propagate_parameter_metadata import (
-    propagate_parameter_metadata,
+from policyengine_uk.utils.parameters import (
+    backdate_parameters,
+    convert_to_fiscal_year_parameters,
 )
-from policyengine_core.parameters.operations.uprate_parameters import (
-    uprate_parameters,
+from policyengine_uk.utils.scenario import Scenario
+from policyengine_uk.data.economic_assumptions import (
+    apply_uprating,
+    extend_single_year_dataset,
 )
-from policyengine_core.reforms import Reform
 
+# Module constants
 COUNTRY_DIR = Path(__file__).parent
-
 ENHANCED_FRS = "hf://policyengine/policyengine-uk-data/enhanced_frs_2023_24.h5"
 
 
 class CountryTaxBenefitSystem(TaxBenefitSystem):
-    variables_dir = COUNTRY_DIR / "variables"
-    auto_carry_over_input_variables = True
-    basic_inputs = [
+    """UK-specific tax and benefit system implementation.
+
+    This class defines the UK tax-benefit system with all relevant
+    variables, parameters, and entities (Person, BenUnit, Household).
+    """
+
+    basic_inputs: List[str] = [
         "brma",
         "local_authority",
         "region",
@@ -65,193 +75,540 @@ class CountryTaxBenefitSystem(TaxBenefitSystem):
         "age",
     ]
     modelled_policies = COUNTRY_DIR / "modelled_policies.yaml"
+    auto_carry_over_input_variables: bool = True
+
+    def reset_parameters(self) -> None:
+        """Reset parameters by reloading from the parameters directory."""
+        self._parameters_at_instant_cache = {}
+        self.load_parameters(self.parameters_dir)
 
-    def process_parameters(self, reform=None):
-        if reform:
-            self.apply_reform_set(reform)
+    def process_parameters(self) -> None:
+        """Process and transform parameters with UK-specific adjustments.
+
+        Applies various parameter transformations including:
+        - Private pension uprating factors
+        - Lagged earnings and CPI indices
+        - Triple lock calculations for state pensions
+        - Economic assumption indices
+        - Parameter uprating and backdating
+        - Conversion to fiscal year parameters
+        """
+        self._parameters_at_instant_cache = {}
+        # Add various UK-specific parameter adjustments
         self.parameters = add_private_pension_uprating_factor(self.parameters)
         self.parameters = add_lagged_earnings(self.parameters)
         self.parameters = add_lagged_cpi(self.parameters)
         self.parameters = add_triple_lock(self.parameters)
         self.parameters = create_economic_assumption_indices(self.parameters)
+
+        # Create baseline parameters for reform comparisons
         self.parameters.add_child("baseline", self.parameters.clone())
-        self.parameters = homogenize_parameter_structures(
-            self.parameters, self.variables
-        )
+
+        # Apply general parameter operations
         self.parameters = propagate_parameter_metadata(self.parameters)
-        self.parameters = interpolate_parameters(self.parameters)
         self.parameters = uprate_parameters(self.parameters)
-        self.parameters = propagate_parameter_metadata(self.parameters)
-        self.add_abolition_parameters()
         self.parameters = backdate_parameters(self.parameters, "2015-01-01")
-
         self.parameters.gov = convert_to_fiscal_year_parameters(
             self.parameters.gov
         )
 
-    def __init__(self, reform=None):
-        super().__init__(entities, reform=reform)
+    def __init__(self):
+        """Initialize the UK tax-benefit system with entities and parameters."""
+        self._parameters_at_instant_cache: Dict[str, Any] = {}
+        self.variables: Dict[Any, Any] = {}
 
-        self.parameters_dir = COUNTRY_DIR / "parameters"
+        # Create copies of entity classes to avoid modifying originals
+        person, benunit, household = (
+            copy.copy(Person),
+            copy.copy(BenUnit),
+            copy.copy(Household),
+        )
 
-        self.load_parameters(self.parameters_dir)
+        # Set up entities
+        self.entities = [person, benunit, household]
+        self.person_entity = person
+        self.group_entities = [benunit, household]
+        self.group_entity_keys = [entity.key for entity in self.group_entities]
 
-        self.process_parameters(reform=reform)
+        # Link entities to this tax-benefit system
+        for entity in self.entities:
+            entity.set_tax_benefit_system(self)
 
+        self.variable_module_metadata = {}
 
-system = CountryTaxBenefitSystem()
+        # Load all variables from the variables directory
+        self.add_variables_from_directory(COUNTRY_DIR / "variables")
+
+        # Set up and process parameters
+        self.parameters_dir = COUNTRY_DIR / "parameters"
+        self.reset_parameters()
+        self.process_parameters()
 
+
+# Create system instance for module-level access
+system = CountryTaxBenefitSystem()
 parameters = system.parameters
 variables = system.variables
 
 
 class Simulation(CoreSimulation):
-    default_tax_benefit_system = CountryTaxBenefitSystem
-    default_tax_benefit_system_instance = system
-    default_calculation_period = 2023
-    default_input_period = 2023
-    default_role = "member"
-    max_spiral_loops = 10
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        reform = create_structural_reforms_from_parameters(
-            self.tax_benefit_system.parameters, "2023-01-01"
-        )
+    """UK-specific simulation class for calculating tax and benefit outcomes.
+
+    Extends the core simulation functionality with UK-specific features
+    and data handling capabilities.
+    """
+
+    default_input_period: int = 2025
+    default_calculation_period: int = 2025
+
+    def __init__(
+        self,
+        scenario: Optional[Scenario] = None,
+        situation: Optional[Dict] = None,
+        dataset: Optional[
+            Union[pd.DataFrame, str, UKSingleYearDataset, UKMultiYearDataset]
+        ] = None,
+        trace: bool = False,
+        reform: Dict | Type[Reform] = None,
+    ):
+        """Initialize a UK simulation.
+
+        Args:
+            scenario: A Scenario object defining a modification to the simulation
+            situation: A dictionary describing the situation to simulate
+            dataset: Data source - can be DataFrame, URL string, or Dataset object
+            trace: Whether to enable detailed tracing of calculations
+        """
+        # Initialize tax-benefit rules
+        self.tax_benefit_system = CountryTaxBenefitSystem()
+
+        # Migrate Reform to Scenario
+
         if reform is not None:
-            self.apply_reform(reform)
-
-        reform_dict = kwargs.get("reform")
-        if reform_dict is not None:
-            if isinstance(reform_dict, type):
-                try:
-                    reform_dict = reform_dict.parameter_values
-                except:
-                    reform_dict = None
-
-        if reform_dict is not None:
-            if any(["obr" in param for param in reform_dict]):
-                self.tax_benefit_system.load_parameters(
-                    self.tax_benefit_system.parameters_dir
-                )
-                Reform.from_dict(reform_dict).apply(self.tax_benefit_system)
-                self.tax_benefit_system.process_parameters()
+            scenario = Scenario.from_reform(reform)
+
+        # Apply parametric reforms here
+
+        if scenario is not None:
+            if scenario.parameter_changes is not None:
+                self.apply_parameter_changes(scenario.parameter_changes)
+
+        self.branch_name = "default"
+        self.invalidated_caches = set()
+        self.debug: bool = False
+        self.trace: bool = trace
+        self.tracer: SimpleTracer = (
+            SimpleTracer() if not trace else FullTracer()
+        )
+        self.opt_out_cache: bool = False
+        self.max_spiral_loops: int = 10
+        self.memory_config = None
+        self._data_storage_dir: Optional[str] = None
+
+        self.branches: Dict[str, Simulation] = {}
+
+        # Build simulation from appropriate source
+        if situation is not None:
+            self.build_from_situation(situation)
+        elif isinstance(dataset, str):
+            self.build_from_url(dataset)
+        elif isinstance(dataset, pd.DataFrame):
+            self.build_from_dataframe(dataset)
+        elif isinstance(dataset, Dataset):
+            self.build_from_dataset(dataset)
+        elif isinstance(dataset, UKSingleYearDataset):
+            self.build_from_single_year_dataset(dataset)
+        elif isinstance(dataset, UKMultiYearDataset):
+            self.build_from_multi_year_dataset(dataset)
+        elif dataset is None:
+            self.build_from_url(
+                "hf://policyengine/policyengine-uk-data/enhanced_frs_2023_24.h5"
+            )
+        else:
+            raise ValueError(f"Unsupported dataset type: {dataset.__class__}")
 
-        # Labor supply responses
+        # Handle behavioral responses for earnings and capital gains
+        self.move_values("employment_income", "employment_income_before_lsr")
+        self.move_values("capital_gains", "capital_gains_before_response")
 
-        employment_income = self.get_holder("employment_income")
-        for known_period in employment_income.get_known_periods():
-            array = employment_income.get_array(known_period)
-            self.set_input("employment_income_before_lsr", known_period, array)
-            employment_income.delete_arrays(known_period)
+        self.input_variables = self.get_known_variables()
 
-        # Capital gains responses
+        # Apply structural modifiers
 
-        cg_holder = self.get_holder("capital_gains")
-        for known_period in cg_holder.get_known_periods():
-            array = cg_holder.get_array(known_period)
-            self.set_input(
-                "capital_gains_before_response", known_period, array
+        if scenario is not None:
+            if scenario.simulation_modifier is not None:
+                scenario.simulation_modifier(self)
+
+    def get_known_variables(self):
+        variables = []
+        for variable in self.tax_benefit_system.variables:
+            if len(self.get_holder(variable).get_known_periods()) > 0:
+                variables.append(variable)
+        return variables
+
+    def apply_parameter_changes(self, changes: dict):
+        self.tax_benefit_system.reset_parameters()
+
+        for parameter in changes:
+            p: Parameter = self.tax_benefit_system.parameters.get_child(
+                parameter
             )
-            employment_income.delete_arrays(known_period)
-
-
-class Microsimulation(CoreMicrosimulation):
-    default_tax_benefit_system = CountryTaxBenefitSystem
-    default_dataset = ENHANCED_FRS
-    default_dataset_year = 2022
-    default_tax_benefit_system_instance = system
-    default_calculation_period = 2025
-    default_input_period = 2025
-    default_role = "member"
-    max_spiral_loops = 10
-
-    def __init__(self, *args, dataset=ENHANCED_FRS, **kwargs):
-        if dataset is not None:
-            if isinstance(dataset, str):
-                if "hf://" in dataset:
-                    owner, repo, filename = dataset.split("/")[-3:]
-                    if "@" in filename:
-                        version = filename.split("@")[-1]
-                        filename = filename.split("@")[0]
-                    else:
-                        version = None
-                    dataset_file_path = download_huggingface_dataset(
-                        repo=f"{owner}/{repo}",
-                        repo_filename=filename,
-                        version=version,
+            if isinstance(changes[parameter], dict):
+                # Time-period specific changes
+                for time_period in changes[parameter]:
+                    p.update(
+                        period=time_period,
+                        value=changes[parameter][time_period],
                     )
+            else:
+                p.update(period="year:2000:100", value=changes[parameter])
+
+        self.tax_benefit_system.process_parameters()
 
-                if Path(dataset_file_path).exists():
-                    if dataset_file_path.endswith(".h5"):
-                        try:
-                            UKSingleYearDataset.validate_file_path(
-                                dataset_file_path
-                            )
-                            dataset = UKSingleYearDataset(
-                                file_path=dataset_file_path
-                            )
-                        except:
-                            pass
-
-                        try:
-                            UKMultiYearDataset.validate_file_path(
-                                dataset_file_path
-                            )
-                            dataset = UKMultiYearDataset(
-                                file_path=dataset_file_path
-                            )
-                        except Exception as e:
-                            pass
-
-                        if not isinstance(
-                            dataset, (UKSingleYearDataset, UKMultiYearDataset)
-                        ):
-                            dataset = Dataset.from_file(dataset_file_path)
-
-        super().__init__(*args, dataset=dataset, **kwargs)
-
-        reform = create_structural_reforms_from_parameters(
-            self.tax_benefit_system.parameters, "2023-01-01"
+    def build_from_situation(self, situation: Dict) -> None:
+        """Build simulation from a situation dictionary.
+
+        Args:
+            situation: Dictionary describing household composition and characteristics
+        """
+        self.build_from_populations(
+            self.tax_benefit_system.instantiate_entities()
         )
-        if reform is not None:
-            self.apply_reform(reform)
-
-        reform_dict = kwargs.get("reform")
-        if reform_dict is not None:
-            if isinstance(reform_dict, type):
-                try:
-                    reform_dict = reform_dict.parameter_values
-                except:
-                    reform_dict = None
-
-        if reform_dict is not None:
-            if any(["obr" in param for param in reform_dict]):
-                self.tax_benefit_system.load_parameters(
-                    self.tax_benefit_system.parameters_dir
-                )
-                Reform.from_dict(reform_dict).apply(self.tax_benefit_system)
-                self.tax_benefit_system.process_parameters()
+        from policyengine_core.simulations.simulation_builder import (
+            SimulationBuilder,
+        )  # Import here to avoid circular dependency
+
+        builder = SimulationBuilder()
+        builder.default_period = self.default_input_period
+        builder.build_from_dict(self.tax_benefit_system, situation, self)
+        self.has_axes = builder.has_axes
+
+    def build_from_url(self, url: str) -> None:
+        """Build simulation from a HuggingFace dataset URL.
+
+        Args:
+            url: HuggingFace URL in format "hf://owner/repo/filename"
+
+        Raises:
+            ValueError: If URL is not a HuggingFace URL
+        """
+        if "hf://" not in url:
+            raise ValueError(
+                f"Non-HuggingFace URLs are currently not supported."
+            )
 
-        # Labor supply responses
+        # Parse HuggingFace URL components
+        owner, repo, filename = url.split("/")[-3:]
+        if "@" in filename:
+            version = filename.split("@")[-1]
+            filename = filename.split("@")[0]
+        else:
+            version = None
+
+        # Download dataset from HuggingFace
+        dataset = download_huggingface_dataset(
+            repo=f"{owner}/{repo}",
+            repo_filename=filename,
+            version=version,
+        )
 
-        for simulation in list(self.branches.values()) + [self]:
-            employment_income = simulation.get_holder("employment_income")
-            for known_period in employment_income.get_known_periods():
-                array = employment_income.get_array(known_period)
-                simulation.set_input(
-                    "employment_income_before_lsr", known_period, array
-                )
-                employment_income.delete_arrays(known_period)
+        # Determine dataset type and build accordingly
+        if UKMultiYearDataset.validate_file_path(dataset, False):
+            self.build_from_multi_year_dataset(UKMultiYearDataset(dataset))
+            self.dataset = dataset
+        elif UKSingleYearDataset.validate_file_path(dataset, False):
+            self.build_from_single_year_dataset(UKSingleYearDataset(dataset))
+            self.dataset = dataset
+        else:
+            dataset = Dataset.from_file(dataset, self.default_input_period)
+            self.build_from_dataset(dataset)
+
+    def build_from_dataframe(self, df: pd.DataFrame) -> None:
+        """Build simulation from a pandas DataFrame.
+
+        Args:
+            df: DataFrame with columns in format "variable_name__time_period"
+        """
+
+        def get_first_array(variable_name: str) -> pd.Series:
+            """Extract the first array for a given variable name pattern."""
+            columns = df.columns[df.columns.str.contains(variable_name + "__")]
+            return df[columns[0]]
+
+        # Extract ID columns
+        (
+            person_id,
+            person_benunit_id,
+            person_household_id,
+            benunit_id,
+            household_id,
+        ) = map(
+            get_first_array,
+            [
+                "person_id",
+                "person_benunit_id",
+                "person_household_id",
+                "benunit_id",
+                "household_id",
+            ],
+        )
 
-        # Capital gains responses
+        # Build entity structure
+        self.build_from_ids(
+            person_id,
+            person_benunit_id,
+            person_household_id,
+            benunit_id,
+            household_id,
+        )
 
-        for simulation in list(self.branches.values()) + [self]:
-            cg_holder = self.get_holder("capital_gains")
-            for known_period in cg_holder.get_known_periods():
-                array = cg_holder.get_array(known_period)
+        # Set input values for each variable and time period
+        for column in df:
+            variable, time_period = column.split("__")
+            if variable not in self.tax_benefit_system.variables:
+                continue
+            self.set_input(variable, time_period, df[column])
+
+    def build_from_dataset(self, dataset: Dataset) -> None:
+        """Build simulation from a Dataset object.
+
+        Args:
+            dataset: PolicyEngine Dataset object containing simulation data
+        """
+        data: Dict[str, Dict[str, Union[float, int, str]]] = (
+            dataset.load_dataset()
+        )
+
+        first_variable = data[list(data.keys())[0]]
+        first_time_period = list(first_variable.keys())[0]
+
+        def get_first_array(variable_name: str) -> np.ndarray:
+            """Get the first time period's values for a variable."""
+            time_period_values = data[variable_name]
+            return time_period_values[first_time_period]
+
+        # Build entity structure from IDs
+        self.build_from_ids(
+            *map(
+                get_first_array,
+                [
+                    "person_id",
+                    "person_benunit_id",
+                    "person_household_id",
+                    "benunit_id",
+                    "household_id",
+                ],
+            )
+        )
+
+        # Load all variable values
+        for variable in data:
+            for time_period in data[variable]:
+                if variable not in self.tax_benefit_system.variables:
+                    continue
                 self.set_input(
-                    "capital_gains_before_response", known_period, array
+                    variable, time_period, data[variable][time_period]
                 )
-                employment_income.delete_arrays(known_period)
+
+        # Now convert to the new UKSingleYearDataset
+        self.input_variables = self.get_known_variables()
+        self.dataset = dataset
+        dataset = UKSingleYearDataset.from_simulation(
+            self, fiscal_year=first_time_period
+        )
+        multi_year_dataset = extend_single_year_dataset(dataset)
+
+        self.build_from_multi_year_dataset(multi_year_dataset)
+        self.dataset = multi_year_dataset
+
+    def build_from_single_year_dataset(
+        self, dataset: UKSingleYearDataset
+    ) -> None:
+        """Build simulation from a single-year UK dataset.
+
+        Args:
+            dataset: UKSingleYearDataset containing one year of data
+        """
+
+        dataset = extend_single_year_dataset(dataset)
+        self.build_from_multi_year_dataset(dataset)
+
+    def build_from_multi_year_dataset(
+        self, dataset: UKMultiYearDataset
+    ) -> None:
+        """Build simulation from a multi-year UK dataset.
+
+        Args:
+            dataset: UKMultiYearDataset containing multiple years of data
+        """
+        # Use first year to establish entity structure
+        first_year = dataset[dataset.years[0]]
+        self.build_from_ids(
+            first_year.person.person_id,
+            first_year.person.person_benunit_id,
+            first_year.person.person_household_id,
+            first_year.benunit.benunit_id,
+            first_year.household.household_id,
+        )
+
+        # Load variable values for all years
+        for year in dataset.years:
+            for table in dataset[year].tables:
+                for variable in table.columns:
+                    if variable not in self.tax_benefit_system.variables:
+                        continue
+                    self.set_input(variable, year, table[variable])
+
+    def build_from_ids(
+        self,
+        person_id: np.ndarray,
+        person_benunit_id: np.ndarray,
+        person_household_id: np.ndarray,
+        benunit_id: np.ndarray,
+        household_id: np.ndarray,
+    ) -> None:
+        """Build simulation entities from ID arrays.
+
+        Args:
+            person_id: Array of person IDs
+            person_benunit_id: Array mapping persons to benefit units
+            person_household_id: Array mapping persons to households
+            benunit_id: Array of benefit unit IDs
+            household_id: Array of household IDs
+        """
+        from policyengine_core.simulations.simulation_builder import (
+            SimulationBuilder,
+        )  # Import here to avoid circular dependency
+
+        builder = SimulationBuilder()
+        builder.populations = self.tax_benefit_system.instantiate_entities()
+
+        # Declare entities
+        builder.declare_person_entity("person", person_id)
+        builder.declare_entity("benunit", np.unique(benunit_id))
+        builder.declare_entity("household", np.unique(household_id))
+
+        # Link persons to benefit units and households
+        builder.join_with_persons(
+            builder.populations["benunit"],
+            person_benunit_id,
+            np.array(["member"] * len(person_benunit_id)),
+        )
+        builder.join_with_persons(
+            builder.populations["household"],
+            person_household_id,
+            np.array(["member"] * len(person_household_id)),
+        )
+
+        self.build_from_populations(builder.populations)
+
+    def move_values(self, variable_donor: str, variable_target: str) -> None:
+        """Move values from one variable to another across all branches.
+
+        Used for behavioral response modeling where original values need
+        to be preserved.
+
+        Args:
+            variable_donor: Variable to move values from
+            variable_target: Variable to move values to
+        """
+        for simulation in list(self.branches.values()) + [self]:
+            holder = simulation.get_holder(variable_donor)
+            for known_period in holder.get_known_periods():
+                array = holder.get_array(known_period)
+                simulation.set_input(variable_target, known_period, array)
+                holder.delete_arrays(known_period)
+
+    def calculate(
+        self,
+        variable_name: str,
+        period: str = None,
+        map_to: str = None,
+        decode_enums: bool = False,
+    ):
+        tracer: SimpleTracer = self.tracer
+        if len(tracer.stack) == 0:
+            # Only decode enums to string values when we're not within
+            # the simulation tree.
+            decode_enums = True
+
+        if period is None:
+            period = self.default_calculation_period
+
+        period = period_(period)
+
+        return super().calculate(
+            variable_name, period, map_to=map_to, decode_enums=decode_enums
+        )
+
+
+class Microsimulation(Simulation):
+    """Extended simulation class with weighting support for microsimulation.
+
+    Provides weighted calculations using survey weights for population-level
+    estimates and statistics.
+    """
+
+    def get_weights(
+        self, variable_name: str, period: str, map_to: Optional[str] = None
+    ) -> np.ndarray:
+        """Get weights for the specified variable's entity.
+
+        Args:
+            variable_name: Name of the variable to get weights for
+            period: Time period for the weights
+            map_to: Optional entity key to map weights to
+
+        Returns:
+            Array of weights for the entity
+        """
+        variable = self.tax_benefit_system.get_variable(variable_name)
+        entity_key = map_to or variable.entity.key
+        weight_variable_name = f"{entity_key}_weight"
+        return self.calculate(
+            weight_variable_name, period, map_to=map_to, unweighted=True
+        )
+
+    def calculate(
+        self,
+        variable_name: str,
+        period: str = None,
+        map_to: str = None,
+        decode_enums: bool = False,
+        unweighted: bool = False,
+    ):
+        tracer: SimpleTracer = self.tracer
+
+        result = super().calculate(
+            variable_name, period, map_to=map_to, decode_enums=decode_enums
+        )
+
+        if not unweighted and len(tracer.stack) == 0:
+            weights = self.get_weights(variable_name, period, map_to=map_to)
+            return MicroSeries(result, weights=weights)
+
+        return result
+
+    def calculate_dataframe(
+        self,
+        variable_names: List[str],
+        period: Optional[str] = None,
+        map_to: Optional[str] = None,
+        use_weights: bool = True,
+    ) -> MicroDataFrame:
+        """Calculate multiple variables as a weighted DataFrame.
+
+        Args:
+            variable_names: List of variable names to calculate
+            period: Time period for calculation
+            map_to: Optional entity key to map results to
+            use_weights: Whether to apply survey weights
+
+        Returns:
+            MicroDataFrame with calculated values and weights
+        """
+        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=map_to)
+        return MicroDataFrame(values, weights=weights)