Add updates

Author: nikhilwoodruff

src/policyengine/countries/general.py

@@ -0,0 +1,68 @@
+from policyengine_core.simulations import Simulation
+import pandas as pd
+from typing import Dict, List
+from pydantic import BaseModel, ConfigDict
+import time
+from ..database.models import ParameterMetadata, ParameterChangeMetadata
+from policyengine_core.parameters import Parameter
+
+class ModelOutput(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    table_names: List[str] = []
+    variable_whitelist: List[str] = []
+
+    def get_tables(self) -> Dict[str, pd.DataFrame]:
+        return {
+            table_name: getattr(self, table_name)
+            for table_name in self.table_names
+        }
+
+def process_simulation(simulation: Simulation, year: int, variable_whitelist: List[str] = []) -> Dict[str, pd.DataFrame]:
+    variables = list(simulation.tax_benefit_system.variables.values())
+
+    entity_tables = {}
+
+    for variable in variable_whitelist:
+        simulation.calculate(variable, year)
+
+    known_variables = [
+        variable for variable in variables
+        if len(simulation.get_holder(variable.name).get_known_periods()) > 0
+    ]
+
+    for variable in known_variables:
+        if variable.definition_period != "year":
+            continue
+        if variable.entity.key not in entity_tables:
+            entity_tables[variable.entity.key] = pd.DataFrame()
+
+        try:
+            start = time.time()
+            result = simulation.calculate(variable.name, year)
+            end = time.time()
+            if end - start > 1.0:
+                print(f"Time taken to calculate {variable.name} for {year}: {end - start} seconds")
+        except Exception as e:
+            print(f"Error calculating {variable.name} for {year}: {e}")
+            continue
+
+        entity_tables[variable.entity.key][variable.name] = result
+
+    return entity_tables
+
+
+def create_default_parameters(simulation: Simulation, country: str) -> List[ParameterMetadata]:
+    """Create default parameters for the simulation."""
+    parameter_tree = simulation.tax_benefit_system.parameters
+    parameters = []
+
+    for parameter in parameter_tree.get_descendants():
+        parameter_meta = ParameterMetadata(
+            name=parameter.name,
+            country=country,
+            #parent somehow link to the parameter.parent Parameter object
+            # rest of relevant attributes, look at policyengine_core.parameters.Parameter
+        )
+
+        # add parameterchanges as well?

src/policyengine/countries/uk.py

@@ -0,0 +1,28 @@
+from policyengine_uk import Simulation, Microsimulation
+from .general import process_simulation, ModelOutput
+from pydantic import BaseModel, ConfigDict
+import pandas as pd
+from typing import List
+from ..database.models import ParameterMetadata, ParameterChangeMetadata
+
+class UKModelOutput(ModelOutput):
+    person: pd.DataFrame
+    benunit: pd.DataFrame
+    household: pd.DataFrame
+
+    table_names: List[str] = ["person", "benunit", "household"]
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+UK_VARIABLE_WHITELIST = [
+    "household_net_income",
+]
+
+def process_uk_simulation(simulation: Simulation, year: int) -> UKModelOutput:
+    entity_tables = process_simulation(simulation, year, variable_whitelist=UK_VARIABLE_WHITELIST)
+
+    return UKModelOutput(
+        person=entity_tables.get("person"),
+        benunit=entity_tables.get("benunit"),
+        household=entity_tables.get("household"),
+    )

src/policyengine/countries/us.py

@@ -0,0 +1,33 @@
+from policyengine_us import Simulation, Microsimulation
+from .general import process_simulation, ModelOutput
+from pydantic import BaseModel, ConfigDict
+import pandas as pd
+from typing import List
+
+class USModelOutput(ModelOutput):
+    person: pd.DataFrame
+    marital_unit: pd.DataFrame
+    family: pd.DataFrame
+    tax_unit: pd.DataFrame
+    spm_unit: pd.DataFrame
+    household: pd.DataFrame
+
+    table_names: List[str] = ["person", "marital_unit", "family", "tax_unit", "spm_unit", "household"]
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+US_VARIABLE_WHITELIST = [
+    "household_net_income",
+]
+
+def process_us_simulation(simulation: Simulation, year: int) -> USModelOutput:
+    entity_tables = process_simulation(simulation, year, variable_whitelist=US_VARIABLE_WHITELIST)
+
+    return USModelOutput(
+        person=entity_tables.get("person"),
+        marital_unit=entity_tables.get("marital_unit"),
+        family=entity_tables.get("family"),
+        tax_unit=entity_tables.get("tax_unit"),
+        spm_unit=entity_tables.get("spm_unit"),
+        household=entity_tables.get("household"),
+    )

src/policyengine/database/__init__.py

@@ -1,12 +1,8 @@
 """Database abstraction for PolicyEngine with hybrid storage."""
 
 from .database import Database, DatabaseConfig
-from .simulation import save_simulation, load_simulation, list_simulations
 
 __all__ = [
     "Database",
     "DatabaseConfig",
-    "save_simulation",
-    "load_simulation",
-    "list_simulations",
 ]

src/policyengine/database/models.py

@@ -40,6 +40,7 @@ class SimulationMetadata(Base):
     # Scenario information
     dataset = Column(String, nullable=False)  # e.g., "frs_2023_24"
     scenario = Column(String, nullable=False)  # e.g., "baseline"
+    model_version = Column(String, nullable=True)  # e.g., "0.5.2"
 
     # Processing metadata
     status = Column(Enum(SimulationStatus), default=SimulationStatus.PENDING, nullable=False)
@@ -65,24 +66,83 @@ class DatasetMetadata(Base):
     # Dataset characteristics
     source = Column(String, nullable=True)  # "FRS", "CPS", etc.
     version = Column(String, nullable=True)
+    model_version = Column(String, nullable=True)  # e.g., "0.5.2"
 
     # Metadata
     description = Column(Text, nullable=True)
     created_at = Column(DateTime, default=datetime.now)
     updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now)
 
 
-class Scenario(Base):
+class ScenarioMetadata(Base):
     """Modifications made to baseline simulation behaviour."""
     __tablename__ = "scenarios"
 
     id = Column(String, primary_key=True, default=generate_uuid)
     name = Column(String, nullable=False, unique=True, index=True)
     country = Column(String, nullable=False, index=True)
+    model_version = Column(String, nullable=True)  # e.g., "0.5.2"
 
     # Metadata
     description = Column(Text, nullable=True)
 
     created_at = Column(DateTime, default=datetime.now)
     updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now)
     created_by = Column(String, nullable=True)
+    
+    # Relationships
+    parameter_changes = relationship("ParameterChange", back_populates="scenario", cascade="all, delete-orphan")
+
+
+class ParameterMetadata(Base):
+    """Registry of all parameters that can be modified."""
+    __tablename__ = "parameters"
+    
+    id = Column(String, primary_key=True, default=generate_uuid)
+    name = Column(String, nullable=False, unique=True, index=True)  # e.g., "gov.basic_rate"
+    country = Column(String, nullable=False, index=True)
+    parent_id = Column(String, ForeignKey("parameters.id"), nullable=True)
+
+    # Parameter metadata
+    label = Column(String, nullable=True)  # Human-readable name
+    description = Column(Text, nullable=True)
+    unit = Column(String, nullable=True)  # e.g., "GBP", "percent", "boolean"
+    data_type = Column(String, nullable=False)  # "float", "int", "bool", "string"
+    
+    # Metadata
+    created_at = Column(DateTime, default=datetime.now)
+    updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now)
+    
+    # Relationships
+    changes = relationship("ParameterChange", back_populates="parameter", cascade="all, delete-orphan")
+
+
+class ParameterChangeMetadata(Base):
+    """Individual parameter change within a scenario."""
+    __tablename__ = "parameter_changes"
+    
+    id = Column(String, primary_key=True, default=generate_uuid)
+    
+    # Foreign keys
+    scenario_id = Column(String, ForeignKey("scenarios.id"), nullable=False, index=True)
+    parameter_id = Column(String, ForeignKey("parameters.id"), nullable=False, index=True)
+    
+    # Time period for this change
+    start_date = Column(DateTime, nullable=False, index=True)  # When this change takes effect
+    end_date = Column(DateTime, nullable=True, index=True)  # When this change expires (null = indefinite)
+    
+    # The actual change
+    value = Column(JSON, nullable=False)  # JSON to handle different data types
+    
+    # Ordering within scenario (for applying changes in sequence)
+    order_index = Column(Integer, nullable=False, default=0)
+    
+    # Metadata
+    model_version = Column(String, nullable=True)  # e.g., "0.5.2"
+    description = Column(Text, nullable=True)  # Optional description of this specific change
+    created_at = Column(DateTime, default=datetime.now)
+    updated_at = Column(DateTime, default=datetime.now, onupdate=datetime.now)
+    
+    # Relationships
+    scenario = relationship("Scenario", back_populates="parameter_changes")
+    parameter = relationship("Parameter", back_populates="changes")