Replace a simulation's population when the builder writes a situation into it

changelog.d/builder-replaces-population.fixed.md

@@ -0,0 +1 @@
+Writing a situation into an existing simulation with `SimulationBuilder` now replaces its population: it is no longer over a dataset (`is_over_dataset` is false and `dataset` is `None`), keeps nothing loaded or calculated for the old population, and rebuilds its baseline branch on the situation; `subsample` explains that it needs a dataset instead of failing.

policyengine_core/simulations/simulation.py

@@ -1077,6 +1077,90 @@ def build_from_populations(self, populations: Dict[str, Population]) -> None:
         self.link_to_entities_instances()
         self.create_shortcuts()
 
+    def _replace_population(self) -> None:
+        """Give this simulation new, empty populations, and forget the old ones.
+
+        ``SimulationBuilder`` calls this before it writes a situation into a
+        simulation it was given, which may hold a population loaded from a
+        dataset and values calculated for it. Writing the situation into the
+        old populations kept the old population's arrays for every variable
+        the situation did not set, the fast cache and the input record kept
+        answering for it, and ``dataset`` and ``is_over_dataset`` still named
+        the dataset, so the macro cache, and country models that treat a
+        simulation over data differently, treated the situation as data.
+
+        The simulation keeps its tax-benefit system and settings. Its branches
+        keep the old populations until :meth:`_rebuild_branches`.
+
+        Raises:
+            ValueError: if the simulation is a branch, which shares its
+                parent's population and baseline.
+            RuntimeError: if a calculation is running in the simulation.
+        """
+        if getattr(self, "parent_branch", None) is not None:
+            raise ValueError(
+                f"Cannot replace the population of branch '{self.branch_name}': "
+                "it shares its parent's population and baseline. Write the "
+                "situation into the simulation it was made from, or a new one."
+            )
+        if self._calculations_in_flight:
+            raise RuntimeError(
+                "Cannot replace a simulation's population while it is calculating."
+            )
+        self.dataset = None
+        self.is_over_dataset = False
+        self.situation_input = None
+        self.has_axes = False
+        self.input_variables = []
+        self._store_history = StoreHistory()
+        self._user_input_keys = set()
+        self._requested_variables = None
+        self.invalidated_caches = set()
+        self._fast_cache = {}
+        # A calculation elsewhere that read this simulation does not keep what
+        # it calculated from the old population (``_Frame.is_stale``).
+        self._input_epoch = self._input_epoch + 1
+        if isinstance(getattr(self, "tracer", None), FullTracer):
+            # Its trees hold the old population's values.
+            self.tracer = FullTracer()
+        self.build_from_populations(self.tax_benefit_system.instantiate_entities())
+
+    def _rebuild_branches(self) -> None:
+        """Recreate this simulation's branches on its current populations.
+
+        Called after the populations are replaced in place (``subsample``, or
+        ``SimulationBuilder`` writing a situation into this simulation).
+        """
+        # Other branches still contain the old populations. Forget them so a
+        # formula's next ``get_branch`` starts from the new ones. Keep the
+        # baseline long enough to preserve its policy below.
+        self.branches = {
+            name: branch for name, branch in self.branches.items() if name == "baseline"
+        }
+
+        # Ensure the baseline branch has the new data: rebuild it from this
+        # simulation through ``get_branch`` (the same wiring ``__init__``
+        # uses), then restore the saved baseline tax-benefit system.
+        # Previously ``subsample`` assigned the saved system to
+        # ``self.branches["tax_benefit_system"]`` — a stray dict key — so the
+        # rebuilt baseline branch kept the reform system and reform-vs-baseline
+        # comparisons after ``subsample`` compared the reform against itself.
+        if "baseline" in self.branches:
+            baseline_tax_benefit_system = self.branches["baseline"].tax_benefit_system
+            del self.branches["baseline"]
+            baseline = self.get_branch("baseline")
+            # As in ``__init__``, the branch is traced in this simulation,
+            # uses the baseline system's entities and variables, and has no
+            # baseline of its own: ``get_branch`` gave it this simulation's,
+            # the branch it replaces, which kept the old population alive.
+            baseline.trace = self.trace
+            baseline.tracer = self.tracer
+            baseline.tax_benefit_system = baseline_tax_benefit_system
+            baseline._bind_to_tax_benefit_system()
+            baseline.baseline = None
+            if getattr(self, "baseline", None) is not None:
+                self.baseline = baseline
+
     def build_from_dataset(self) -> None:
         """Build a simulation from a dataset."""
         self.build_from_populations(self.tax_benefit_system.instantiate_entities())
@@ -3384,7 +3468,19 @@ def subsample(
 
         Returns:
             Simulation: The quantized simulation.
+
+        Raises:
+            ValueError: if the simulation has no dataset, as one built from a
+                situation, or whose data a situation replaced, has none.
         """
+        if getattr(self, "dataset", None) is None:
+            # The sample is rebuilt as a dataset in the dataset's period, and
+            # seeded with its name by default.
+            raise ValueError(
+                "Only a simulation with a dataset can be subsampled. This one "
+                "has none: it was built from a situation or populations, or a "
+                "situation replaced its data."
+            )
         default_calculation_period = self.default_calculation_period
         # Set default key if not provided
         if seed is None:
@@ -3567,35 +3663,8 @@ def retain_structure(name, values):
         # as "size X != Y = count" projection errors.
         self._invalidate_all_caches()
 
-        # Other branches still contain the pre-subsample populations. Forget
-        # them so a formula's next ``get_branch`` starts from the new sample.
-        # Keep the baseline long enough to preserve its policy below.
-        self.branches = {
-            name: branch for name, branch in self.branches.items() if name == "baseline"
-        }
-
-        # Ensure the baseline branch has the new data: rebuild it from the
-        # subsampled simulation through ``get_branch`` (the same wiring
-        # ``__init__`` uses), then restore the saved baseline tax-benefit
-        # system. Previously the saved system was assigned to
-        # ``self.branches["tax_benefit_system"]`` — a stray dict key — so the
-        # rebuilt baseline branch kept the reform system and reform-vs-baseline
-        # comparisons after ``subsample`` compared the reform against itself.
-        if "baseline" in self.branches:
-            baseline_tax_benefit_system = self.branches["baseline"].tax_benefit_system
-            del self.branches["baseline"]
-            baseline = self.get_branch("baseline")
-            # As in ``__init__``, the branch is traced in this simulation,
-            # uses the baseline system's entities and variables, and has no
-            # baseline of its own: ``get_branch`` gave it this simulation's,
-            # the branch it replaces, which kept the old population alive.
-            baseline.trace = self.trace
-            baseline.tracer = self.tracer
-            baseline.tax_benefit_system = baseline_tax_benefit_system
-            baseline._bind_to_tax_benefit_system()
-            baseline.baseline = None
-            if getattr(self, "baseline", None) is not None:
-                self.baseline = baseline
+        # Branches still contain the pre-subsample populations.
+        self._rebuild_branches()
 
         self.default_calculation_period = default_calculation_period
         return self

policyengine_core/simulations/simulation_builder.py

@@ -73,8 +73,18 @@ def build_from_dict(
         Build a simulation from ``input_dict``. Optionally overwrites an existing simulation.
 
         This method uses :any:`build_from_entities` if entities are fully specified, or :any:`build_from_variables` if not.
+
+        Overwriting replaces the simulation's population with the situation's
+        (see :meth:`Simulation._replace_population`): it is no longer over a
+        dataset, and holds nothing loaded or calculated for its old population.
+        It keeps its tax-benefit system and settings, and its baseline branch
+        is rebuilt on the new population. A situation that fails to parse
+        (unknown entities or variables, say) leaves the simulation as it was;
+        one that fails while it is written leaves the simulation replaced,
+        with what was written. A branch cannot be overwritten.
         """
 
+        replacing = simulation is not None
         input_dict = self.explicit_singular_entities(tax_benefit_system, input_dict)
         if any(
             key in tax_benefit_system.entities_plural() for key in input_dict.keys()
@@ -87,13 +97,32 @@ def build_from_dict(
                 tax_benefit_system, input_dict, simulation
             )
 
+        if not replacing:
+            # A replaced simulation's were recorded before its branches were
+            # rebuilt (see ``_finish_replacing``).
+            self._record_input_variables(simulation)
+
+        return simulation
+
+    def _record_input_variables(self, simulation: Simulation) -> None:
         simulation.input_variables = [
             variable.name
             for variable in simulation.tax_benefit_system.variables.values()
             if len(simulation.get_holder(variable.name).get_known_periods()) > 0
         ]
 
-        return simulation
+    def _finish_replacing(self, simulation: Simulation) -> None:
+        """Finish writing a situation into a simulation whose population it replaced.
+
+        Also when writing failed partway, so the simulation's records and
+        branches describe what was written.
+        """
+        simulation.has_axes = self.has_axes
+        self._record_input_variables(simulation)
+        # Its branches still hold the old population. Rebuilt after the
+        # inputs are recorded, the baseline records them too, as in
+        # ``Simulation.__init__``.
+        simulation._rebuild_branches()
 
     def build_from_entities(
         self,
@@ -113,6 +142,7 @@ def build_from_entities(
         """
         input_dict = copy.deepcopy(input_dict)
 
+        replacing = simulation is not None
         if simulation is None:
             simulation = Simulation(
                 tax_benefit_system=tax_benefit_system,
@@ -177,17 +207,25 @@ def build_from_entities(
             self.axes = axes
             self.expand_axes()
 
-        try:
-            self.finalize_variables_init(simulation.persons)
-        except PeriodMismatchError as e:
-            self.raise_period_mismatch(simulation.persons.entity, persons_json, e)
+        if replacing:
+            # The situation parsed: write it into new populations.
+            simulation._replace_population()
 
-        for entity_class in tax_benefit_system.group_entities:
+        try:
             try:
-                population = simulation.populations[entity_class.key]
-                self.finalize_variables_init(population)
+                self.finalize_variables_init(simulation.persons)
             except PeriodMismatchError as e:
-                self.raise_period_mismatch(population.entity, instances_json, e)
+                self.raise_period_mismatch(simulation.persons.entity, persons_json, e)
+
+            for entity_class in tax_benefit_system.group_entities:
+                try:
+                    population = simulation.populations[entity_class.key]
+                    self.finalize_variables_init(population)
+                except PeriodMismatchError as e:
+                    self.raise_period_mismatch(population.entity, instances_json, e)
+        finally:
+            if replacing:
+                self._finish_replacing(simulation)
 
         return simulation
 
@@ -209,20 +247,29 @@ def build_from_variables(
             )
         """
         count = _get_person_count(input_dict)
-        simulation = self.build_default_simulation(
+        # Check the situation before building on it, so one that fails here
+        # leaves a given simulation as it was (as ``build_from_entities``).
+        for variable, value in input_dict.items():
+            tax_benefit_system.get_variable(variable, check_existence=True)
+            if not isinstance(value, dict) and self.default_period is None:
+                raise SituationParsingError(
+                    [variable],
+                    "Can't deal with type: expected object. Input variables should be set for specific periods. For instance: {'salary': {'2017-01': 2000, '2017-02': 2500}}, or {'birth_date': {'ETERNITY': '1980-01-01'}}.",
+                )
+        replacing = simulation is not None
+        simulation = self._build_default_populations(
             tax_benefit_system, count, simulation
         )
-        for variable, value in input_dict.items():
-            if not isinstance(value, dict):
-                if self.default_period is None:
-                    raise SituationParsingError(
-                        [variable],
-                        "Can't deal with type: expected object. Input variables should be set for specific periods. For instance: {'salary': {'2017-01': 2000, '2017-02': 2500}}, or {'birth_date': {'ETERNITY': '1980-01-01'}}.",
-                    )
-                simulation.set_input(variable, self.default_period, value)
-            else:
-                for period_str, dated_value in value.items():
-                    simulation.set_input(variable, period_str, dated_value)
+        try:
+            for variable, value in input_dict.items():
+                if not isinstance(value, dict):
+                    simulation.set_input(variable, self.default_period, value)
+                else:
+                    for period_str, dated_value in value.items():
+                        simulation.set_input(variable, period_str, dated_value)
+        finally:
+            if replacing:
+                self._finish_replacing(simulation)
         return simulation
 
     def build_default_simulation(
@@ -236,13 +283,31 @@ def build_default_simulation(
             - There are ``count`` persons
             - There are ``count`` instances of each group entity, containing one person
             - Every person has, in each entity, the first role
+
+        Given a simulation, replace its population with this one, as
+        :meth:`build_from_dict` does.
         """
+        replacing = simulation is not None
+        simulation = self._build_default_populations(
+            tax_benefit_system, count, simulation
+        )
+        if replacing:
+            self._finish_replacing(simulation)
+        return simulation
 
+    def _build_default_populations(
+        self,
+        tax_benefit_system: "TaxBenefitSystem",
+        count: int,
+        simulation: Simulation = None,
+    ) -> Simulation:
         if simulation is None:
             simulation = Simulation(
                 tax_benefit_system=tax_benefit_system,
                 populations=tax_benefit_system.instantiate_entities(),
             )
+        else:
+            simulation._replace_population()
         for population in simulation.populations.values():
             population.count = count
             population.ids = np.array(range(count))