policyengine 3.1.3__tar.gz → 3.1.4__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (90) hide show
  1. {policyengine-3.1.3 → policyengine-3.1.4}/CHANGELOG.md +7 -0
  2. {policyengine-3.1.3 → policyengine-3.1.4}/PKG-INFO +1 -1
  3. {policyengine-3.1.3 → policyengine-3.1.4}/changelog.yaml +5 -0
  4. {policyengine-3.1.3 → policyengine-3.1.4}/pyproject.toml +1 -1
  5. policyengine-3.1.4/src/policyengine/__pycache__/__init__.cpython-313.pyc +0 -0
  6. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/core/simulation.py +7 -0
  7. policyengine-3.1.4/src/policyengine/tax_benefit_models/uk/__init__.py +44 -0
  8. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/tax_benefit_models/uk/datasets.py +78 -7
  9. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/tax_benefit_models/uk/model.py +16 -4
  10. policyengine-3.1.4/src/policyengine/tax_benefit_models/uk.py +38 -0
  11. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/tax_benefit_models/us/__init__.py +9 -1
  12. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/tax_benefit_models/us/datasets.py +90 -7
  13. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/tax_benefit_models/us/model.py +16 -4
  14. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/tax_benefit_models/us.py +6 -0
  15. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine.egg-info/PKG-INFO +1 -1
  16. policyengine-3.1.3/src/policyengine/__pycache__/__init__.cpython-313.pyc +0 -0
  17. policyengine-3.1.3/src/policyengine/tax_benefit_models/uk/__init__.py +0 -26
  18. policyengine-3.1.3/src/policyengine/tax_benefit_models/uk.py +0 -33
  19. {policyengine-3.1.3 → policyengine-3.1.4}/.claude/policyengine-guide.md +0 -0
  20. {policyengine-3.1.3 → policyengine-3.1.4}/.claude/quick-reference.md +0 -0
  21. {policyengine-3.1.3 → policyengine-3.1.4}/.github/CONTRIBUTING.md +0 -0
  22. {policyengine-3.1.3 → policyengine-3.1.4}/.github/changelog_template.md +0 -0
  23. {policyengine-3.1.3 → policyengine-3.1.4}/.github/fetch_version.py +0 -0
  24. {policyengine-3.1.3 → policyengine-3.1.4}/.github/get-changelog-diff.sh +0 -0
  25. {policyengine-3.1.3 → policyengine-3.1.4}/.github/has-functional-changes.sh +0 -0
  26. {policyengine-3.1.3 → policyengine-3.1.4}/.github/is-version-number-acceptable.sh +0 -0
  27. {policyengine-3.1.3 → policyengine-3.1.4}/.github/publish-git-tag.sh +0 -0
  28. {policyengine-3.1.3 → policyengine-3.1.4}/.github/workflows/code_changes.yaml +0 -0
  29. {policyengine-3.1.3 → policyengine-3.1.4}/.github/workflows/docs.yml +0 -0
  30. {policyengine-3.1.3 → policyengine-3.1.4}/.github/workflows/pr_code_changes.yaml +0 -0
  31. {policyengine-3.1.3 → policyengine-3.1.4}/.github/workflows/pr_docs_changes.yaml +0 -0
  32. {policyengine-3.1.3 → policyengine-3.1.4}/.github/workflows/versioning.yaml +0 -0
  33. {policyengine-3.1.3 → policyengine-3.1.4}/.gitignore +0 -0
  34. {policyengine-3.1.3 → policyengine-3.1.4}/CLAUDE.md +0 -0
  35. {policyengine-3.1.3 → policyengine-3.1.4}/LICENSE +0 -0
  36. {policyengine-3.1.3 → policyengine-3.1.4}/Makefile +0 -0
  37. {policyengine-3.1.3 → policyengine-3.1.4}/README.md +0 -0
  38. {policyengine-3.1.3 → policyengine-3.1.4}/changelog_entry.yaml +0 -0
  39. {policyengine-3.1.3 → policyengine-3.1.4}/docs/.gitignore +0 -0
  40. {policyengine-3.1.3 → policyengine-3.1.4}/docs/core-concepts.md +0 -0
  41. {policyengine-3.1.3 → policyengine-3.1.4}/docs/country-models-uk.md +0 -0
  42. {policyengine-3.1.3 → policyengine-3.1.4}/docs/country-models-us.md +0 -0
  43. {policyengine-3.1.3 → policyengine-3.1.4}/docs/dev.md +0 -0
  44. {policyengine-3.1.3 → policyengine-3.1.4}/docs/index.md +0 -0
  45. {policyengine-3.1.3 → policyengine-3.1.4}/docs/myst.yml +0 -0
  46. {policyengine-3.1.3 → policyengine-3.1.4}/docs/visualisation.md +0 -0
  47. {policyengine-3.1.3 → policyengine-3.1.4}/examples/employment_income_variation_uk.py +0 -0
  48. {policyengine-3.1.3 → policyengine-3.1.4}/examples/employment_income_variation_us.py +0 -0
  49. {policyengine-3.1.3 → policyengine-3.1.4}/examples/income_bands_uk.py +0 -0
  50. {policyengine-3.1.3 → policyengine-3.1.4}/examples/income_distribution_us.py +0 -0
  51. {policyengine-3.1.3 → policyengine-3.1.4}/examples/policy_change_uk.py +0 -0
  52. {policyengine-3.1.3 → policyengine-3.1.4}/examples/speedtest_us_simulation.py +0 -0
  53. {policyengine-3.1.3 → policyengine-3.1.4}/setup.cfg +0 -0
  54. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/__init__.py +0 -0
  55. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/core/__init__.py +0 -0
  56. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/core/dataset.py +0 -0
  57. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/core/dataset_version.py +0 -0
  58. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/core/dynamic.py +0 -0
  59. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/core/output.py +0 -0
  60. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/core/parameter.py +0 -0
  61. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/core/parameter_value.py +0 -0
  62. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/core/policy.py +0 -0
  63. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/core/tax_benefit_model.py +0 -0
  64. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/core/tax_benefit_model_version.py +0 -0
  65. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/core/variable.py +0 -0
  66. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/outputs/__init__.py +0 -0
  67. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/outputs/aggregate.py +0 -0
  68. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/outputs/change_aggregate.py +0 -0
  69. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/outputs/decile_impact.py +0 -0
  70. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/tax_benefit_models/uk/analysis.py +0 -0
  71. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/tax_benefit_models/uk/outputs.py +0 -0
  72. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/tax_benefit_models/us/analysis.py +0 -0
  73. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/tax_benefit_models/us/outputs.py +0 -0
  74. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/utils/__init__.py +0 -0
  75. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/utils/dates.py +0 -0
  76. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/utils/parametric_reforms.py +0 -0
  77. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine/utils/plotting.py +0 -0
  78. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine.egg-info/SOURCES.txt +0 -0
  79. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine.egg-info/dependency_links.txt +0 -0
  80. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine.egg-info/requires.txt +0 -0
  81. {policyengine-3.1.3 → policyengine-3.1.4}/src/policyengine.egg-info/top_level.txt +0 -0
  82. {policyengine-3.1.3 → policyengine-3.1.4}/tests/test_aggregate.py +0 -0
  83. {policyengine-3.1.3 → policyengine-3.1.4}/tests/test_change_aggregate.py +0 -0
  84. {policyengine-3.1.3 → policyengine-3.1.4}/tests/test_entity_mapping.py +0 -0
  85. {policyengine-3.1.3 → policyengine-3.1.4}/tests/test_get_parameter_variable.py +0 -0
  86. {policyengine-3.1.3 → policyengine-3.1.4}/tests/test_uk_dataset.py +0 -0
  87. {policyengine-3.1.3 → policyengine-3.1.4}/tests/test_us_datasets.py +0 -0
  88. {policyengine-3.1.3 → policyengine-3.1.4}/tests/test_us_entity_mapping.py +0 -0
  89. {policyengine-3.1.3 → policyengine-3.1.4}/tests/test_us_simulation.py +0 -0
  90. {policyengine-3.1.3 → policyengine-3.1.4}/uv.lock +0 -0
@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
5
5
  The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
6
6
  and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
7
7
 
8
+ ## [3.1.4] - 2025-11-20 14:06:49
9
+
10
+ ### Fixed
11
+
12
+ - Minor fixes
13
+
8
14
  ## [3.1.3] - 2025-11-18 13:46:23
9
15
 
10
16
  ### Fixed
@@ -215,6 +221,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
215
221
 
216
222
 
217
223
 
224
+ [3.1.4]: https://github.com/PolicyEngine/policyengine.py/compare/3.1.3...3.1.4
218
225
  [3.1.3]: https://github.com/PolicyEngine/policyengine.py/compare/3.1.2...3.1.3
219
226
  [3.1.2]: https://github.com/PolicyEngine/policyengine.py/compare/3.1.1...3.1.2
220
227
  [3.1.1]: https://github.com/PolicyEngine/policyengine.py/compare/3.1.0...3.1.1
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: policyengine
3
- Version: 3.1.3
3
+ Version: 3.1.4
4
4
  Summary: A package to conduct policy analysis using PolicyEngine tax-benefit models.
5
5
  Author-email: PolicyEngine <hello@policyengine.org>
6
6
  License: GNU AFFERO GENERAL PUBLIC LICENSE
@@ -175,3 +175,8 @@
175
175
  fixed:
176
176
  - Entity variables moved out to an editable constant.
177
177
  date: 2025-11-18 13:46:23
178
+ - bump: patch
179
+ changes:
180
+ fixed:
181
+ - Minor fixes
182
+ date: 2025-11-20 14:06:49
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "policyengine"
7
- version = "3.1.3"
7
+ version = "3.1.4"
8
8
  description = "A package to conduct policy analysis using PolicyEngine tax-benefit models."
9
9
  readme = "README.md"
10
10
  authors = [
@@ -24,6 +24,13 @@ class Simulation(BaseModel):
24
24
  def run(self):
25
25
  self.tax_benefit_model_version.run(self)
26
26
 
27
+ def ensure(self):
28
+ try:
29
+ self.tax_benefit_model_version.load(self)
30
+ except Exception:
31
+ self.run()
32
+ self.save()
33
+
27
34
  def save(self):
28
35
  """Save the simulation's output dataset."""
29
36
  self.tax_benefit_model_version.save(self)
@@ -0,0 +1,44 @@
1
+ """PolicyEngine UK tax-benefit model."""
2
+
3
+ from importlib.util import find_spec
4
+
5
+ if find_spec("policyengine_uk") is not None:
6
+ from policyengine.core import Dataset
7
+
8
+ from .analysis import general_policy_reform_analysis
9
+ from .datasets import (
10
+ PolicyEngineUKDataset,
11
+ UKYearData,
12
+ create_datasets,
13
+ ensure_datasets,
14
+ load_datasets,
15
+ )
16
+ from .model import (
17
+ PolicyEngineUK,
18
+ PolicyEngineUKLatest,
19
+ uk_latest,
20
+ uk_model,
21
+ )
22
+ from .outputs import ProgrammeStatistics
23
+
24
+ # Rebuild Pydantic models to resolve forward references
25
+ Dataset.model_rebuild()
26
+ UKYearData.model_rebuild()
27
+ PolicyEngineUKDataset.model_rebuild()
28
+ PolicyEngineUKLatest.model_rebuild()
29
+
30
+ __all__ = [
31
+ "UKYearData",
32
+ "PolicyEngineUKDataset",
33
+ "create_datasets",
34
+ "load_datasets",
35
+ "ensure_datasets",
36
+ "PolicyEngineUK",
37
+ "PolicyEngineUKLatest",
38
+ "uk_model",
39
+ "uk_latest",
40
+ "general_policy_reform_analysis",
41
+ "ProgrammeStatistics",
42
+ ]
43
+ else:
44
+ __all__ = []
@@ -37,11 +37,7 @@ class PolicyEngineUKDataset(Dataset):
37
37
  if self.data is not None:
38
38
  self.save()
39
39
  elif self.filepath and not self.data:
40
- try:
41
- self.load()
42
- except FileNotFoundError:
43
- # File doesn't exist yet, that's OK
44
- pass
40
+ self.load()
45
41
 
46
42
  def save(self) -> None:
47
43
  """Save dataset to HDF5 file."""
@@ -85,7 +81,9 @@ def create_datasets(
85
81
  "hf://policyengine/policyengine-uk-data/enhanced_frs_2023_24.h5",
86
82
  ],
87
83
  years: list[int] = [2026, 2027, 2028, 2029, 2030],
88
- ) -> None:
84
+ data_folder: str = "./data",
85
+ ) -> dict[str, PolicyEngineUKDataset]:
86
+ result = {}
89
87
  for dataset in datasets:
90
88
  from policyengine_uk import Microsimulation
91
89
 
@@ -139,9 +137,10 @@ def create_datasets(
139
137
  )
140
138
 
141
139
  uk_dataset = PolicyEngineUKDataset(
140
+ id=f"{Path(dataset).stem}_year_{year}",
142
141
  name=f"{dataset}-year-{year}",
143
142
  description=f"UK Dataset for year {year} based on {dataset}",
144
- filepath=f"./data/{Path(dataset).stem}_year_{year}.h5",
143
+ filepath=f"{data_folder}/{Path(dataset).stem}_year_{year}.h5",
145
144
  year=year,
146
145
  data=UKYearData(
147
146
  person=MicroDataFrame(person_df, weights="person_weight"),
@@ -154,3 +153,75 @@ def create_datasets(
154
153
  ),
155
154
  )
156
155
  uk_dataset.save()
156
+
157
+ dataset_key = f"{Path(dataset).stem}_{year}"
158
+ result[dataset_key] = uk_dataset
159
+
160
+ return result
161
+
162
+
163
+ def load_datasets(
164
+ datasets: list[str] = [
165
+ "hf://policyengine/policyengine-uk-data/frs_2023_24.h5",
166
+ "hf://policyengine/policyengine-uk-data/enhanced_frs_2023_24.h5",
167
+ ],
168
+ years: list[int] = [2026, 2027, 2028, 2029, 2030],
169
+ data_folder: str = "./data",
170
+ ) -> dict[str, PolicyEngineUKDataset]:
171
+ result = {}
172
+ for dataset in datasets:
173
+ for year in years:
174
+ filepath = f"{data_folder}/{Path(dataset).stem}_year_{year}.h5"
175
+ uk_dataset = PolicyEngineUKDataset(
176
+ name=f"{dataset}-year-{year}",
177
+ description=f"UK Dataset for year {year} based on {dataset}",
178
+ filepath=filepath,
179
+ year=year,
180
+ )
181
+ uk_dataset.load()
182
+
183
+ dataset_key = f"{Path(dataset).stem}_{year}"
184
+ result[dataset_key] = uk_dataset
185
+
186
+ return result
187
+
188
+
189
+ def ensure_datasets(
190
+ datasets: list[str] = [
191
+ "hf://policyengine/policyengine-uk-data/frs_2023_24.h5",
192
+ "hf://policyengine/policyengine-uk-data/enhanced_frs_2023_24.h5",
193
+ ],
194
+ years: list[int] = [2026, 2027, 2028, 2029, 2030],
195
+ data_folder: str = "./data",
196
+ ) -> dict[str, PolicyEngineUKDataset]:
197
+ """Ensure datasets exist, loading if available or creating if not.
198
+
199
+ Args:
200
+ datasets: List of HuggingFace dataset paths
201
+ years: List of years to load/create data for
202
+ data_folder: Directory containing or to save the dataset files
203
+
204
+ Returns:
205
+ Dictionary mapping dataset keys to PolicyEngineUKDataset objects
206
+ """
207
+ # Check if all dataset files exist
208
+ all_exist = True
209
+ for dataset in datasets:
210
+ for year in years:
211
+ filepath = Path(
212
+ f"{data_folder}/{Path(dataset).stem}_year_{year}.h5"
213
+ )
214
+ if not filepath.exists():
215
+ all_exist = False
216
+ break
217
+ if not all_exist:
218
+ break
219
+
220
+ if all_exist:
221
+ return load_datasets(
222
+ datasets=datasets, years=years, data_folder=data_folder
223
+ )
224
+ else:
225
+ return create_datasets(
226
+ datasets=datasets, years=years, data_folder=data_folder
227
+ )
@@ -265,17 +265,29 @@ class PolicyEngineUKLatest(TaxBenefitModelVersion):
265
265
 
266
266
  def load(self, simulation: "Simulation"):
267
267
  """Load the simulation's output dataset."""
268
+ import os
269
+
270
+ filepath = str(
271
+ Path(simulation.dataset.filepath).parent / (simulation.id + ".h5")
272
+ )
273
+
268
274
  simulation.output_dataset = PolicyEngineUKDataset(
269
275
  id=simulation.id,
270
276
  name=simulation.dataset.name,
271
277
  description=simulation.dataset.description,
272
- filepath=str(
273
- Path(simulation.dataset.filepath).parent
274
- / (simulation.id + ".h5")
275
- ),
278
+ filepath=filepath,
276
279
  year=simulation.dataset.year,
277
280
  is_output_dataset=True,
278
281
  )
279
282
 
283
+ # Load timestamps from file system metadata
284
+ if os.path.exists(filepath):
285
+ simulation.created_at = datetime.datetime.fromtimestamp(
286
+ os.path.getctime(filepath)
287
+ )
288
+ simulation.updated_at = datetime.datetime.fromtimestamp(
289
+ os.path.getmtime(filepath)
290
+ )
291
+
280
292
 
281
293
  uk_latest = PolicyEngineUKLatest()
@@ -0,0 +1,38 @@
1
+ """PolicyEngine UK tax-benefit model - imports from uk/ module."""
2
+
3
+ from importlib.util import find_spec
4
+
5
+ if find_spec("policyengine_uk") is not None:
6
+ from .uk import (
7
+ PolicyEngineUK,
8
+ PolicyEngineUKDataset,
9
+ PolicyEngineUKLatest,
10
+ ProgrammeStatistics,
11
+ UKYearData,
12
+ create_datasets,
13
+ ensure_datasets,
14
+ general_policy_reform_analysis,
15
+ load_datasets,
16
+ uk_latest,
17
+ uk_model,
18
+ )
19
+
20
+ __all__ = [
21
+ "UKYearData",
22
+ "PolicyEngineUKDataset",
23
+ "create_datasets",
24
+ "load_datasets",
25
+ "ensure_datasets",
26
+ "PolicyEngineUK",
27
+ "PolicyEngineUKLatest",
28
+ "uk_model",
29
+ "uk_latest",
30
+ "general_policy_reform_analysis",
31
+ "ProgrammeStatistics",
32
+ ]
33
+
34
+ # Rebuild models to resolve forward references
35
+ PolicyEngineUKDataset.model_rebuild()
36
+ PolicyEngineUKLatest.model_rebuild()
37
+ else:
38
+ __all__ = []
@@ -6,7 +6,13 @@ if find_spec("policyengine_us") is not None:
6
6
  from policyengine.core import Dataset
7
7
 
8
8
  from .analysis import general_policy_reform_analysis
9
- from .datasets import PolicyEngineUSDataset, USYearData, create_datasets
9
+ from .datasets import (
10
+ PolicyEngineUSDataset,
11
+ USYearData,
12
+ create_datasets,
13
+ ensure_datasets,
14
+ load_datasets,
15
+ )
10
16
  from .model import (
11
17
  PolicyEngineUS,
12
18
  PolicyEngineUSLatest,
@@ -25,6 +31,8 @@ if find_spec("policyengine_us") is not None:
25
31
  "USYearData",
26
32
  "PolicyEngineUSDataset",
27
33
  "create_datasets",
34
+ "load_datasets",
35
+ "ensure_datasets",
28
36
  "PolicyEngineUS",
29
37
  "PolicyEngineUSLatest",
30
38
  "us_model",
@@ -44,11 +44,7 @@ class PolicyEngineUSDataset(Dataset):
44
44
  if self.data is not None:
45
45
  self.save()
46
46
  elif self.filepath and not self.data:
47
- try:
48
- self.load()
49
- except FileNotFoundError:
50
- # File doesn't exist yet, that's OK
51
- pass
47
+ self.load()
52
48
 
53
49
  def save(self) -> None:
54
50
  """Save dataset to HDF5 file."""
@@ -112,15 +108,21 @@ def create_datasets(
112
108
  "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5",
113
109
  ],
114
110
  years: list[int] = [2024, 2025, 2026, 2027, 2028],
115
- ) -> None:
111
+ data_folder: str = "./data",
112
+ ) -> dict[str, PolicyEngineUSDataset]:
116
113
  """Create PolicyEngineUSDataset instances from HuggingFace dataset paths.
117
114
 
118
115
  Args:
119
116
  datasets: List of HuggingFace dataset paths (e.g., "hf://policyengine/policyengine-us-data/cps_2024.h5")
120
117
  years: List of years to extract data for
118
+ data_folder: Directory to save the dataset files
119
+
120
+ Returns:
121
+ Dictionary mapping dataset keys (e.g., "enhanced_cps_2024") to PolicyEngineUSDataset objects
121
122
  """
122
123
  from policyengine_us import Microsimulation
123
124
 
125
+ result = {}
124
126
  for dataset in datasets:
125
127
  sim = Microsimulation(dataset=dataset)
126
128
 
@@ -263,9 +265,10 @@ def create_datasets(
263
265
  tax_unit_df = entity_df
264
266
 
265
267
  us_dataset = PolicyEngineUSDataset(
268
+ id=f"{Path(dataset).stem}_year_{year}",
266
269
  name=f"{dataset}-year-{year}",
267
270
  description=f"US Dataset for year {year} based on {dataset}",
268
- filepath=f"./data/{Path(dataset).stem}_year_{year}.h5",
271
+ filepath=f"{data_folder}/{Path(dataset).stem}_year_{year}.h5",
269
272
  year=year,
270
273
  data=USYearData(
271
274
  person=MicroDataFrame(person_df, weights="person_weight"),
@@ -285,3 +288,83 @@ def create_datasets(
285
288
  ),
286
289
  )
287
290
  us_dataset.save()
291
+
292
+ dataset_key = f"{Path(dataset).stem}_{year}"
293
+ result[dataset_key] = us_dataset
294
+
295
+ return result
296
+
297
+
298
+ def load_datasets(
299
+ datasets: list[str] = [
300
+ "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5",
301
+ ],
302
+ years: list[int] = [2024, 2025, 2026, 2027, 2028],
303
+ data_folder: str = "./data",
304
+ ) -> dict[str, PolicyEngineUSDataset]:
305
+ """Load PolicyEngineUSDataset instances from saved HDF5 files.
306
+
307
+ Args:
308
+ datasets: List of HuggingFace dataset paths (used to derive file names)
309
+ years: List of years to load data for
310
+ data_folder: Directory containing the dataset files
311
+
312
+ Returns:
313
+ Dictionary mapping dataset keys (e.g., "enhanced_cps_2024") to PolicyEngineUSDataset objects
314
+ """
315
+ result = {}
316
+ for dataset in datasets:
317
+ for year in years:
318
+ filepath = f"{data_folder}/{Path(dataset).stem}_year_{year}.h5"
319
+ us_dataset = PolicyEngineUSDataset(
320
+ name=f"{dataset}-year-{year}",
321
+ description=f"US Dataset for year {year} based on {dataset}",
322
+ filepath=filepath,
323
+ year=year,
324
+ )
325
+ us_dataset.load()
326
+
327
+ dataset_key = f"{Path(dataset).stem}_{year}"
328
+ result[dataset_key] = us_dataset
329
+
330
+ return result
331
+
332
+
333
+ def ensure_datasets(
334
+ datasets: list[str] = [
335
+ "hf://policyengine/policyengine-us-data/enhanced_cps_2024.h5",
336
+ ],
337
+ years: list[int] = [2024, 2025, 2026, 2027, 2028],
338
+ data_folder: str = "./data",
339
+ ) -> dict[str, PolicyEngineUSDataset]:
340
+ """Ensure datasets exist, loading if available or creating if not.
341
+
342
+ Args:
343
+ datasets: List of HuggingFace dataset paths
344
+ years: List of years to load/create data for
345
+ data_folder: Directory containing or to save the dataset files
346
+
347
+ Returns:
348
+ Dictionary mapping dataset keys to PolicyEngineUSDataset objects
349
+ """
350
+ # Check if all dataset files exist
351
+ all_exist = True
352
+ for dataset in datasets:
353
+ for year in years:
354
+ filepath = Path(
355
+ f"{data_folder}/{Path(dataset).stem}_year_{year}.h5"
356
+ )
357
+ if not filepath.exists():
358
+ all_exist = False
359
+ break
360
+ if not all_exist:
361
+ break
362
+
363
+ if all_exist:
364
+ return load_datasets(
365
+ datasets=datasets, years=years, data_folder=data_folder
366
+ )
367
+ else:
368
+ return create_datasets(
369
+ datasets=datasets, years=years, data_folder=data_folder
370
+ )
@@ -311,18 +311,30 @@ class PolicyEngineUSLatest(TaxBenefitModelVersion):
311
311
 
312
312
  def load(self, simulation: "Simulation"):
313
313
  """Load the simulation's output dataset."""
314
+ import os
315
+
316
+ filepath = str(
317
+ Path(simulation.dataset.filepath).parent / (simulation.id + ".h5")
318
+ )
319
+
314
320
  simulation.output_dataset = PolicyEngineUSDataset(
315
321
  id=simulation.id,
316
322
  name=simulation.dataset.name,
317
323
  description=simulation.dataset.description,
318
- filepath=str(
319
- Path(simulation.dataset.filepath).parent
320
- / (simulation.id + ".h5")
321
- ),
324
+ filepath=filepath,
322
325
  year=simulation.dataset.year,
323
326
  is_output_dataset=True,
324
327
  )
325
328
 
329
+ # Load timestamps from file system metadata
330
+ if os.path.exists(filepath):
331
+ simulation.created_at = datetime.datetime.fromtimestamp(
332
+ os.path.getctime(filepath)
333
+ )
334
+ simulation.updated_at = datetime.datetime.fromtimestamp(
335
+ os.path.getmtime(filepath)
336
+ )
337
+
326
338
  def _build_simulation_from_dataset(self, microsim, dataset, system):
327
339
  """Build a PolicyEngine Core simulation from dataset entity IDs.
328
340
 
@@ -9,7 +9,10 @@ if find_spec("policyengine_us") is not None:
9
9
  PolicyEngineUSLatest,
10
10
  ProgramStatistics,
11
11
  USYearData,
12
+ create_datasets,
13
+ ensure_datasets,
12
14
  general_policy_reform_analysis,
15
+ load_datasets,
13
16
  us_latest,
14
17
  us_model,
15
18
  )
@@ -17,6 +20,9 @@ if find_spec("policyengine_us") is not None:
17
20
  __all__ = [
18
21
  "USYearData",
19
22
  "PolicyEngineUSDataset",
23
+ "create_datasets",
24
+ "load_datasets",
25
+ "ensure_datasets",
20
26
  "PolicyEngineUS",
21
27
  "PolicyEngineUSLatest",
22
28
  "us_model",
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: policyengine
3
- Version: 3.1.3
3
+ Version: 3.1.4
4
4
  Summary: A package to conduct policy analysis using PolicyEngine tax-benefit models.
5
5
  Author-email: PolicyEngine <hello@policyengine.org>
6
6
  License: GNU AFFERO GENERAL PUBLIC LICENSE
@@ -1,26 +0,0 @@
1
- """PolicyEngine UK tax-benefit model."""
2
-
3
- from .analysis import general_policy_reform_analysis
4
- from .datasets import PolicyEngineUKDataset, UKYearData, create_datasets
5
- from .model import PolicyEngineUK, PolicyEngineUKLatest, uk_latest, uk_model
6
- from .outputs import ProgrammeStatistics
7
-
8
- __all__ = [
9
- "UKYearData",
10
- "PolicyEngineUKDataset",
11
- "create_datasets",
12
- "PolicyEngineUK",
13
- "PolicyEngineUKLatest",
14
- "uk_model",
15
- "uk_latest",
16
- "general_policy_reform_analysis",
17
- "ProgrammeStatistics",
18
- ]
19
-
20
- # Rebuild models to resolve forward references
21
- from policyengine.core import Dataset
22
-
23
- Dataset.model_rebuild()
24
- UKYearData.model_rebuild()
25
- PolicyEngineUKDataset.model_rebuild()
26
- PolicyEngineUKLatest.model_rebuild()
@@ -1,33 +0,0 @@
1
- """PolicyEngine UK tax-benefit model - imports from uk/ module."""
2
-
3
- from .uk import (
4
- PolicyEngineUK,
5
- PolicyEngineUKDataset,
6
- PolicyEngineUKLatest,
7
- ProgrammeStatistics,
8
- UKYearData,
9
- create_datasets,
10
- general_policy_reform_analysis,
11
- uk_latest,
12
- uk_model,
13
- )
14
-
15
- __all__ = [
16
- "UKYearData",
17
- "PolicyEngineUKDataset",
18
- "create_datasets",
19
- "PolicyEngineUK",
20
- "PolicyEngineUKLatest",
21
- "uk_model",
22
- "uk_latest",
23
- "general_policy_reform_analysis",
24
- "ProgrammeStatistics",
25
- ]
26
-
27
- # Rebuild models to resolve forward references
28
- from policyengine.core import Dataset
29
-
30
- Dataset.model_rebuild()
31
- UKYearData.model_rebuild()
32
- PolicyEngineUKDataset.model_rebuild()
33
- PolicyEngineUKLatest.model_rebuild()
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes