Move simulation parameters down to the function level

policyengine/__init__.py

@@ -1 +1 @@
-from .simulation import Simulation, SimulationOptions
+from .policyengine import PolicyEngine

policyengine/outputs/household/comparison/calculate_household_comparison.py

@@ -2,7 +2,7 @@
 
 import typing
 
-from policyengine import Simulation
+from policyengine import PolicyEngine
 
 from pydantic import BaseModel
 from policyengine.utils.calculations import get_change
@@ -27,7 +27,7 @@ class HouseholdComparison(BaseModel):
 
 
 def calculate_household_comparison(
-    simulation: Simulation,
+    engine: PolicyEngine,
 ) -> HouseholdComparison:
     """Calculate comparison statistics between two household scenarios."""
     if not simulation.is_comparison:

policyengine/outputs/household/single/calculate_single_household.py

@@ -2,7 +2,7 @@
 
 import typing
 
-from policyengine import Simulation
+from policyengine import PolicyEngine
 
 from pydantic import BaseModel
 from policyengine_core.simulations import Simulation as CountrySimulation
@@ -33,7 +33,7 @@ class SingleHousehold(BaseModel):
 
 
 def calculate_single_household(
-    simulation: Simulation,
+    engine: PolicyEngine,
 ) -> SingleHousehold:
     """Calculate household statistics for a single household scenario."""
     if simulation.is_comparison:

policyengine/outputs/macro/comparison/calculate_economy_comparison.py

@@ -2,10 +2,11 @@
 
 import typing
 
-from policyengine import Simulation
+from policyengine import PolicyEngine
 
 from pydantic import BaseModel
 from policyengine.utils.calculations import get_change
+from policyengine.constants import DEFAULT_DATASETS_BY_COUNTRY
 
 from policyengine.outputs.macro.single import (
     _calculate_government_balance,
@@ -17,7 +18,24 @@
 from .decile import calculate_decile_impacts, DecileImpacts
 
 from typing import Literal, List
-
+from policyengine.utils.types import *
+from pydantic import Field
+
+class EconomyComparisonOptions(BaseModel):
+    country: CountryType = Field(..., description="The country to simulate.")
+    data: DataType = Field(None, description="The data to simulate.")
+    time_period: TimePeriodType = Field(
+        2025, description="The time period to simulate.", ge=2024, le=2035
+    )
+    reform: PolicyType = Field(None, description="The reform to simulate.")
+    baseline: PolicyType = Field(None, description="The baseline to simulate.")
+    region: RegionType = Field(
+        None, description="The region to simulate within the country."
+    )
+    subsample: SubsampleType = Field(
+        None,
+        description="How many, if a subsample, households to randomly simulate.",
+    )
 
 class FiscalComparison(BaseModel):
     baseline: FiscalSummary
@@ -52,15 +70,37 @@ class EconomyComparison(BaseModel):
 
 
 def calculate_economy_comparison(
-    simulation: Simulation,
+    engine: PolicyEngine,
+    **options,
 ) -> EconomyComparison:
     """Calculate comparison statistics between two economic scenarios."""
-    if not simulation.is_comparison:
-        raise ValueError("Simulation must be a comparison simulation.")
 
-    baseline = simulation.baseline_simulation
-    reform = simulation.reform_simulation
-    options = simulation.options
+    options = EconomyComparisonOptions(**options)
+
+    if options.data is None:
+        options.data = DEFAULT_DATASETS_BY_COUNTRY[
+            options.country
+        ]
+
+    baseline = engine.expect_simulation(
+        name="baseline",
+        country=options.country,
+        scope="macro",
+        policy=options.baseline,
+        data=options.data,
+        time_period=options.time_period,
+        region=options.region,
+    )
+
+    reform = engine.expect_simulation(
+        name="reform",
+        country=options.country,
+        scope="macro",
+        policy=options.reform,
+        data=options.data,
+        time_period=options.time_period,
+        region=options.region,
+    )
 
     baseline_balance = _calculate_government_balance(baseline, options)
     reform_balance = _calculate_government_balance(reform, options)

policyengine/outputs/macro/comparison/charts/budget.py

@@ -2,14 +2,14 @@
 import plotly.graph_objects as go
 import typing
 
-from policyengine import Simulation
+from policyengine import PolicyEngine
 
 from pydantic import BaseModel
 from policyengine.utils.charts import *
 
 
 def create_budget_comparison_chart(
-    simulation: Simulation,
+    engine: PolicyEngine,
 ) -> go.Figure:
     """Create a budget comparison chart."""
     if not simulation.is_comparison:

policyengine/outputs/macro/comparison/charts/budget_by_program.py

@@ -2,14 +2,14 @@
 import plotly.graph_objects as go
 import typing
 
-from policyengine import Simulation
+from policyengine import PolicyEngine
 
 from pydantic import BaseModel
 from policyengine.utils.charts import *
 
 
 def create_budget_program_comparison_chart(
-    simulation: Simulation,
+    engine: PolicyEngine,
 ) -> go.Figure:
     """Create a budget comparison chart."""
     if not simulation.is_comparison:

policyengine/outputs/macro/comparison/charts/decile.py

@@ -2,15 +2,15 @@
 import plotly.graph_objects as go
 import typing
 
-from policyengine import Simulation
+from policyengine import PolicyEngine
 
 from pydantic import BaseModel
 from policyengine.utils.charts import *
 from typing import Literal
 
 
 def create_decile_chart(
-    simulation: Simulation,
+    engine: PolicyEngine,
     decile_variable: Literal["income", "wealth"],
     relative: bool,
 ) -> go.Figure:

policyengine/outputs/macro/comparison/charts/inequality.py

@@ -2,14 +2,14 @@
 import plotly.graph_objects as go
 import typing
 
-from policyengine import Simulation
+from policyengine import PolicyEngine
 
 from pydantic import BaseModel
 from policyengine.utils.charts import *
 
 
 def create_inequality_chart(
-    simulation: Simulation,
+    engine: PolicyEngine,
     relative: bool,
 ) -> go.Figure:
     """Create a budget comparison chart."""

policyengine/outputs/macro/comparison/charts/winners_losers.py

@@ -2,7 +2,7 @@
 import plotly.graph_objects as go
 import typing
 
-from policyengine import Simulation
+from policyengine import PolicyEngine
 
 from pydantic import BaseModel
 from policyengine.utils.charts import *
@@ -27,7 +27,7 @@
 
 
 def create_winners_losers_chart(
-    simulation: Simulation,
+    engine: PolicyEngine,
     decile_variable: Literal["income", "wealth"],
 ) -> go.Figure:
     """Create a budget comparison chart."""

policyengine/outputs/macro/comparison/decile.py

@@ -1,6 +1,6 @@
 import typing
 
-from policyengine import Simulation, SimulationOptions
+from policyengine import PolicyEngine
 
 from policyengine_core.simulations import Microsimulation
 

policyengine/outputs/macro/single/budget.py

@@ -1,6 +1,6 @@
 import typing
 
-from policyengine import Simulation, SimulationOptions
+from policyengine import PolicyEngine
 
 from policyengine_core.simulations import Microsimulation
 

policyengine/outputs/macro/single/calculate_average_earnings.py

@@ -1,7 +1,7 @@
-from policyengine import Simulation
+from policyengine import PolicyEngine
 
 
-def calculate_average_earnings(simulation: Simulation) -> float:
+def calculate_average_earnings(engine: PolicyEngine) -> float:
     """Calculate average earnings."""
     employment_income = simulation.baseline_simulation.calculate(
         "employment_income"

policyengine/outputs/macro/single/calculate_single_economy.py

@@ -2,7 +2,7 @@
 
 import typing
 
-from policyengine import Simulation
+from policyengine import PolicyEngine
 
 from pydantic import BaseModel
 
@@ -19,7 +19,7 @@ class SingleEconomy(BaseModel):
 
 
 def calculate_single_economy(
-    simulation: Simulation,
+    engine: PolicyEngine,
 ) -> SingleEconomy:
     """Calculate economy statistics for a single economic scenario."""
     options = simulation.options

policyengine/outputs/macro/single/inequality.py

@@ -1,6 +1,6 @@
 import typing
 
-from policyengine import Simulation, SimulationOptions
+from policyengine import PolicyEngine
 
 from policyengine_core.simulations import Microsimulation