Improve subsampling and dataset downloads

changelog_entry.yaml

@@ -0,0 +1,5 @@
+- bump: patch
+  changes:
+    fixed:
+    - US model always downlaods from HuggingFace.
+    - Subsampling improvements.

policyengine_api/jobs/calculate_economy_simulation_job.py

@@ -19,7 +19,11 @@
 from policyengine_api.endpoints.economy.compare import compare_economic_outputs
 from policyengine_api.endpoints.economy.reform_impact import set_comment_on_job
 from policyengine_api.constants import COUNTRY_PACKAGE_VERSIONS
-from policyengine_api.country import COUNTRIES, create_policy_reform
+from policyengine_api.country import (
+    COUNTRIES,
+    create_policy_reform,
+    PolicyEngineCountry,
+)
 from policyengine_core.simulations import Microsimulation
 from policyengine_core.tools.hugging_face import download_huggingface_dataset
 import h5py
@@ -259,15 +263,13 @@
                 options.get("max_households", os.environ.get("MAX_HOUSEHOLDS"))
                 is not None
             ):
-                simulation.subsample(
-                    int(
-                        options.get(
-                            "max_households",
-                            os.environ.get("MAX_HOUSEHOLDS", 1_000_000),
-                        )
-                    ),
-                    seed=(region, time_period),
+                simulation = subsample(
+                    options=options,
+                    simulation=simulation,
+                    region=region,
                     time_period=time_period,
+                    reform=reform,
+                    country=country,
                 )
             simulation.default_calculation_period = time_period
 
@@ -463,6 +465,45 @@
         }
 
 
+def subsample(
+    options: dict,
+    simulation: Microsimulation,
+    region: str,
+    time_period: str,
+    reform: dict,
+    country: PolicyEngineCountry,
+) -> Microsimulation:
+    """
+    Subsamples a microsimulation dataset and reinitializes the simulation with the subsampled data.
+    Args:
+        options (dict): A dictionary of options, which may include "max_households" to specify the maximum number of households to subsample.
+        simulation (Microsimulation): The original microsimulation object to be subsampled.
+        region (str): The region for which the simulation is being run.
+        time_period (str): The time period for which the simulation is being run.
+        reform (dict): A dictionary representing the policy reform to apply to the simulation.
+        country (PolicyEngineCountry): The country-specific policy engine object.
+    Returns:
+        Microsimulation: A new microsimulation object initialized with the subsampled data and the specified reform.
+    """
+
+    simulation.subsample(
+        int(
+            options.get(
+                "max_households",
+                os.environ.get("MAX_HOUSEHOLDS", 1_000_000),
+            )
+        ),
+        seed=(region, time_period),
+        time_period=time_period,
+    )
+    input_data = simulation.to_input_dataframe()
+    simulation = country.country_package.Microsimulation(
+        dataset=input_data,
+        reform=reform,
+    )
+    return simulation
+
+
 def is_similar(x, y, parent_name: str = "") -> bool:
     if x is None or x == {}:
         if y is None or y == {}: