Add dataset schema, remove State variables and add HBAI net income (#1155)

Author: nikhilwoodruff

policyengine_uk/data/dataset_schema.py

@@ -0,0 +1,57 @@
+import pandas as pd
+
+
+class UKDataset:
+    person: pd.DataFrame
+    benunit: pd.DataFrame
+    household: pd.DataFrame
+
+    def __init__(
+        self,
+        file_path: str = None,
+        person: pd.DataFrame = None,
+        benunit: pd.DataFrame = None,
+        household: pd.DataFrame = None,
+        fiscal_year: int = 2025,
+    ):
+        if file_path is not None:
+            with pd.HDFStore(file_path) as f:
+                self.person = f["person"]
+                self.benunit = f["benunit"]
+                self.household = f["household"]
+        else:
+            if person is None or benunit is None or household is None:
+                raise ValueError(
+                    "Must provide either a file path or all three DataFrames (person, benunit, household)."
+                )
+            self.person = person
+            self.benunit = benunit
+            self.household = household
+
+        self.data_format = "time_period_arrays"
+        self.time_period = fiscal_year
+
+    def save(self, file_path: str):
+        with pd.HDFStore(file_path) as f:
+            f.put("person", self.person, format="table", data_columns=True)
+            f.put("benunit", self.benunit, format="table", data_columns=True)
+            f.put(
+                "household", self.household, format="table", data_columns=True
+            )
+
+    def load(self):
+        data = {}
+        for df in (self.person, self.benunit, self.household):
+            for col in df.columns:
+                data[col] = {
+                    self.time_period: df[col].values,
+                }
+
+        return data
+
+    def copy(self):
+        return UKDataset(
+            person=self.person.copy(),
+            benunit=self.benunit.copy(),
+            household=self.household.copy(),
+        )

policyengine_uk/entities.py

@@ -3,20 +3,6 @@
 # This file defines the entities needed by our legislation.
 from policyengine_core.entities import build_entity
 
-State = build_entity(
-    key="state",
-    plural="states",
-    label="State",
-    roles=[
-        {
-            "key": "member",
-            "plural": "members",
-            "label": "Member",
-            "doc": "A person who is a citizen of a country.",
-        }
-    ],
-)
-
 Household = build_entity(
     key="household",
     plural="households",
@@ -57,4 +43,4 @@
     is_person=True,
 )
 
-entities = [State, Household, BenUnit, Person]
+entities = [Household, BenUnit, Person]

policyengine_uk/variables/gov/hmrc/expected_sdlt.py

@@ -13,6 +13,6 @@ def formula(household, period, parameters):
         if parameters(period).gov.hmrc.stamp_duty.abolish:
             return 0
         return (
-            household.state("property_sale_rate", period)
-            * household("stamp_duty_land_tax", period)
+            household("stamp_duty_land_tax", period)
+            + parameters(period).gov.hmrc.stamp_duty.property_sale_rate
         ) + household("corporate_sdlt", period)

policyengine_uk/variables/gov/revenue_scotland/expected_lbtt.py

@@ -9,7 +9,9 @@ class expected_lbtt(Variable):
     value_type = float
     unit = GBP
 
-    def formula(household, period):
-        property_sale_rate = household.state("property_sale_rate", period)
+    def formula(household, period, parameters):
+        property_sale_rate = parameters(
+            period
+        ).gov.hmrc.stamp_duty.property_sale_rate
         lbtt = household("land_and_buildings_transaction_tax", period)
         return property_sale_rate * lbtt

policyengine_uk/variables/gov/wra/expected_ltt.py

@@ -9,7 +9,9 @@ class expected_ltt(Variable):
     value_type = float
     unit = GBP
 
-    def formula(household, period):
-        property_sale_rate = household.state("property_sale_rate", period)
+    def formula(household, period, parameters):
+        property_sale_rate = parameters(
+            period
+        ).gov.hmrc.stamp_duty.property_sale_rate
         land_transaction_tax = household("land_transaction_tax", period)
         return property_sale_rate * land_transaction_tax