Merge pull request #406 from mj-will/fix-reparam-resuming

Fix reparam resuming

Author: mj-will

nessai/flowmodel/base.py

@@ -41,7 +41,6 @@ class FlowModel:
         not specified, the current working directory is used.
     """
 
-    model_config = None
     noise_scale = None
     noise_type = None
     model: BaseFlow = None

nessai/proposal/flowproposal.py

@@ -170,8 +170,8 @@ def __init__(
         super(FlowProposal, self).__init__(model)
         logger.debug("Initialising FlowProposal")
 
-        self._x_dtype = False
-        self._x_prime_dtype = False
+        self._x_dtype = None
+        self._x_prime_dtype = None
         self._draw_func = None
         self._populate_dist = None
 
@@ -193,7 +193,7 @@ def __init__(
         self._reparameterisation = None
         self.rescaling_set = False
         self.use_x_prime_prior = False
-        self.update_bounds = False
+        self.should_update_reparameterisations = False
         self.accumulate_weights = accumulate_weights
 
         self.reparameterisations = reparameterisations
@@ -279,7 +279,7 @@ def rescaled_dims(self):
     @property
     def x_dtype(self):
         """Return the dtype for the x space"""
-        if not self._x_dtype:
+        if self._x_dtype is None:
             self._x_dtype = get_dtype(
                 self.parameters, config.livepoints.default_float_dtype
             )
@@ -288,7 +288,7 @@ def x_dtype(self):
     @property
     def x_prime_dtype(self):
         """Return the dtype for the x prime space"""
-        if not self._x_prime_dtype:
+        if self._x_prime_dtype is None:
             self._x_prime_dtype = get_dtype(
                 self.prime_parameters, config.livepoints.default_float_dtype
             )
@@ -484,31 +484,37 @@ def update_flow_config(self):
         """Update the flow configuration dictionary."""
         self.flow_config["n_inputs"] = self.rescaled_dims
 
-    def initialise(self):
+    def initialise(self, resumed: bool = False) -> None:
         """
         Initialise the proposal class.
 
         This includes:
             * Setting up the rescaling
             * Verifying the rescaling is invertible
             * Initialising the FlowModel
+
+        Parameters
+        ----------
+        resumed : bool
+            Indicates if the proposal is being initialised after being resumed
+            or not. When true, the reparameterisations will not be
+            reinitialised.
         """
         if not os.path.exists(self.output):
             os.makedirs(self.output, exist_ok=True)
 
-        self._x_dtype = False
-        self._x_prime_dtype = False
-
-        self.set_rescaling()
-        self.verify_rescaling()
-        if self.expansion_fraction and self.expansion_fraction is not None:
-            logger.info("Overwriting fuzz factor with expansion fraction")
-            self.fuzz = (1 + self.expansion_fraction) ** (
-                1 / self.rescaled_dims
-            )
-            logger.info(f"New fuzz factor: {self.fuzz}")
+        # Initialise if not resuming or resuming but initialised is False
+        if not resumed or not self.initialised:
+            self.set_rescaling()
+            self.verify_rescaling()
+            if self.expansion_fraction and self.expansion_fraction is not None:
+                logger.info("Overwriting fuzz factor with expansion fraction")
+                self.fuzz = (1 + self.expansion_fraction) ** (
+                    1 / self.rescaled_dims
+                )
+                logger.info(f"New fuzz factor: {self.fuzz}")
 
-        self.configure_constant_volume()
+            self.configure_constant_volume()
         self.update_flow_config()
         self.flow = self._FlowModelClass(
             flow_config=self.flow_config,
@@ -697,10 +703,10 @@ def configure_reparameterisations(self, reparameterisations):
             r = FallbackClass(parameters=other_params, **fallback_kwargs)
             self._reparameterisation.add_reparameterisations(r)
 
-        if any(r._update_bounds for r in self._reparameterisation.values()):
-            self.update_bounds = True
+        if any(r._update for r in self._reparameterisation.values()):
+            self.should_update_reparameterisations = True
         else:
-            self.update_bounds = False
+            self.should_update_reparameterisations = False
 
         if self._reparameterisation.has_prime_prior:
             self.use_x_prime_prior = True
@@ -739,6 +745,17 @@ def rescaled_names(self):
         )
         return self.prime_parameters
 
+    @property
+    def update_bounds(self):
+        warn(
+            (
+                "`update_bounds` is deprecated, use "
+                "`should_update_reparameterisations` instead."
+            ),
+            FutureWarning,
+        )
+        return self.should_update_reparameterisations
+
     def set_rescaling(self):
         """
         Set function and parameter names for rescaling
@@ -1633,14 +1650,13 @@ def resume(self, model, flow_config, weights_file=None):
         """
         super().resume(model)
         self.flow_config = flow_config
-        self._reparameterisation = None
 
         if self.mask is not None:
             if isinstance(self.mask, list):
                 m = np.array(self.mask)
             self.flow_config["mask"] = m
 
-        self.initialise()
+        self.initialise(resumed=True)
 
         if weights_file is None:
             weights_file = self.weights_file
@@ -1652,12 +1668,6 @@ def resume(self, model, flow_config, weights_file=None):
         else:
             logger.warning("Could not reload weights for flow")
 
-        if self.update_bounds:
-            if self.training_data is not None:
-                self.check_state(self.training_data)
-            elif self.training_data is None and self.training_count:
-                raise RuntimeError("Could not resume! Missing training data!")
-
     def reset(self):
         """Reset the proposal"""
         self.indices = []
@@ -1673,6 +1683,7 @@ def reset(self):
         self.acceptance = []
         self._draw_func = None
         self._populate_dist = None
+        self._reparameterisation.reset()
 
     def __getstate__(self):
         state = self.__dict__.copy()
@@ -1696,7 +1707,6 @@ def __getstate__(self):
 
         # user provides model and config for resume
         # flow can be reconstructed from resume
-        del state["_reparameterisation"]
         del state["model"]
         del state["_flow_config"]
         del state["flow"]

nessai/reparameterisations/base.py

@@ -21,7 +21,7 @@ class Reparameterisation:
         Prior bounds for the parameter(s).
     """
 
-    _update_bounds = False
+    _update = False
     has_prior = False
     has_prime_prior = False
     requires_prime_prior = False

nessai/reparameterisations/rescale.py

@@ -294,7 +294,7 @@ def __init__(
             for p in self.boundary_inversion:
                 self.rescale_bounds[p] = [0, 1]
 
-        self._update_bounds = update_bounds if not detect_edges else True
+        self._update = update_bounds if not detect_edges else True
         self.detect_edges = detect_edges
         if self.boundary_inversion:
             self._edges = {n: None for n in self.parameters}
@@ -404,7 +404,7 @@ def configure_post_rescaling(self, post_rescaling):
             self.has_prime_prior = False
 
             if post_rescaling in ["logit", "log"]:
-                if self._update_bounds:
+                if self._update:
                     raise RuntimeError(
                         "Cannot use log or logit with update bounds"
                     )
@@ -608,7 +608,7 @@ def set_bounds(self, prior_bounds):
 
     def update_bounds(self, x):
         """Update the bounds used for the reparameterisation"""
-        if self._update_bounds:
+        if self._update:
             self.bounds = {
                 p: [
                     self.pre_rescaling(np.min(x[p]))[0] - self.offsets[p],

tests/test_deprecation_warnings.py

@@ -46,3 +46,12 @@ def test_flowproposal_rescaled_names_warning():
     proposal.prime_parameters = ["x"]
     with pytest.warns(FutureWarning, match=r"`rescaled_names` is deprecated"):
         assert FlowProposal.rescaled_names.__get__(proposal) == ["x"]
+
+
+def test_flowproposal_update_bounds_warning():
+    from nessai.proposal import FlowProposal
+
+    proposal = create_autospec(FlowProposal)
+    proposal.should_update_reparameterisations = True
+    with pytest.warns(FutureWarning, match=r"`update_bounds` is deprecated"):
+        assert FlowProposal.update_bounds.__get__(proposal) is True

tests/test_proposal/test_flowproposal/test_flowproposal_init_resume.py

@@ -30,6 +30,7 @@ def test_init_use_default_reparams(model, proposal, value, expected):
 def test_initialise(tmpdir, proposal, ef, fuzz):
     """Test the initialise method"""
     p = tmpdir.mkdir("test")
+    proposal.initialised = False
     proposal.output = os.path.join(p, "output")
     proposal.rescaled_dims = 2
     proposal.expansion_fraction = ef
@@ -44,7 +45,7 @@ def test_initialise(tmpdir, proposal, ef, fuzz):
     fm.initialise = MagicMock()
     proposal._FlowModelClass = MagicMock(new=fm)
 
-    FlowProposal.initialise(proposal)
+    FlowProposal.initialise(proposal, resumed=False)
 
     proposal.set_rescaling.assert_called_once()
     proposal.verify_rescaling.assert_called_once()
@@ -95,33 +96,6 @@ def test_resume_w_weights(osexist, proposal):
     proposal.flow.reload_weights.assert_called_once_with("weights.pt")
 
 
-@pytest.mark.parametrize("data", [[1], None])
-@pytest.mark.parametrize("count", [0, 1])
-def test_resume_w_update_bounds(proposal, data, count):
-    """Test the resume method with update bounds"""
-    proposal.initialise = MagicMock()
-    proposal.flow = MagicMock()
-    proposal.mask = None
-    proposal.update_bounds = True
-    proposal.weights_file = None
-    proposal.training_data = data
-    proposal.training_count = count
-    proposal.check_state = MagicMock()
-    model = MagicMock()
-    if count and data is None:
-        with pytest.raises(RuntimeError) as excinfo, patch(
-            "nessai.proposal.base.Proposal.resume"
-        ) as mock:
-            FlowProposal.resume(proposal, model, {})
-        assert "Could not resume" in str(excinfo.value)
-    else:
-        with patch("nessai.proposal.base.Proposal.resume") as mock:
-            FlowProposal.resume(proposal, model, {})
-        if data:
-            proposal.check_state.assert_called_once_with(data)
-    mock.assert_called_once_with(model)
-
-
 @pytest.mark.parametrize("populated", [False, True])
 @pytest.mark.parametrize("mask", [None, [1, 0]])
 def test_get_state(proposal, populated, mask):
@@ -149,7 +123,6 @@ def test_get_state(proposal, populated, mask):
     assert state["initialised"] is False
     assert state["weights_file"] == "file"
     assert state["mask"] is mask
-    assert "_reparameterisation" not in state
     assert "model" not in state
     assert "flow" not in state
     assert "_flow_config" not in state
@@ -214,6 +187,7 @@ def test_reset(proposal):
     proposal.samples = 2
     proposal.populated = True
     proposal.populated_count = 10
+    proposal._reparameterisation = MagicMock()
     FlowProposal.reset(proposal)
     assert proposal.x is None
     assert proposal.samples is None
@@ -222,6 +196,7 @@ def test_reset(proposal):
     assert proposal.r is np.nan
     assert proposal.alt_dist is None
     assert proposal._checked_population
+    proposal._reparameterisation.reset.assert_called_once()
 
 
 @pytest.mark.timeout(60)
@@ -280,6 +255,7 @@ def test_reset_integration(tmpdir, model, latent_prior):
     # attributes that should be different
     ignore = [
         "population_time",
+        "_reparameterisation",
     ]
 
     d1 = proposal.__getstate__()

tests/test_reparameterisations/test_rescale_to_bounds.py

@@ -331,7 +331,7 @@ def test_post_rescaling_with_str(reparam, rescaling):
 
     Also test the config for the logit
     """
-    reparam._update_bounds = False
+    reparam._update = False
     reparam.parameters = ["x"]
     from nessai.utils.rescaling import rescaling_functions
 
@@ -346,7 +346,7 @@ def test_post_rescaling_with_str(reparam, rescaling):
 @pytest.mark.parametrize("rescaling", ["log", "logit"])
 def test_post_rescaling_with_logit_update_bounds(reparam, rescaling):
     """Assert an error is raised if using logit and update bounds"""
-    reparam._update_bounds = True
+    reparam._update = True
     with pytest.raises(
         RuntimeError, match=r"Cannot use log or logit with update bounds"
     ):
@@ -370,7 +370,7 @@ def test_post_rescaling_invalid_input(reparam):
 def test_update_bounds_disabled(reparam, caplog):
     """Assert nothing happens in _update_bounds is False"""
     caplog.set_level("DEBUG")
-    reparam._update_bounds = False
+    reparam._update = False
     RescaleToBounds.update_bounds(reparam, [0, 1])
     assert "Update bounds not enabled" in str(caplog.text)
 

tests/test_sampling/test_standard_sampling.py

@@ -379,6 +379,58 @@ def test_resume_fallback_reparameterisation(tmpdir, model, flow_config):
     )
 
 
+@pytest.mark.slow_integration_test
+def test_resume_reparameterisation_values(tmpdir, model, flow_config):
+    """
+    Assert the scale and shift values are correct.
+    """
+    from nessai.reparameterisations.rescale import ScaleAndShift
+
+    output = str(tmpdir.mkdir("resume"))
+    fp = FlowSampler(
+        model,
+        output=output,
+        resume=True,
+        nlive=100,
+        plot=False,
+        flow_config=flow_config,
+        training_frequency=10,
+        maximum_uninformed=9,
+        reparameterisations="z-score",
+        checkpoint_on_iteration=True,
+        checkpoint_interval=5,
+        seed=1234,
+        max_iteration=11,
+        poolsize=10,
+    )
+    fp.run()
+
+    reparam = fp.ns._flow_proposal._reparameterisation
+    reparam = next(iter(fp.ns._flow_proposal._reparameterisation.values()))
+    assert isinstance(reparam, ScaleAndShift)
+    original_scale = reparam.scale
+    original_shift = reparam.shift
+    assert os.path.exists(os.path.join(output, "nested_sampler_resume.pkl"))
+
+    fp = FlowSampler(
+        model,
+        output=output,
+        resume=True,
+        flow_config=flow_config,
+    )
+    assert fp.ns.iteration == 11
+    fp.ns.max_iteration = 21
+    reparam = next(iter(fp.ns._flow_proposal._reparameterisation.values()))
+    assert isinstance(reparam, ScaleAndShift)
+    assert reparam.scale == original_scale
+    assert reparam.shift == original_shift
+    fp.run()
+    assert fp.ns.iteration == 21
+    assert os.path.exists(
+        os.path.join(output, "nested_sampler_resume.pkl.old")
+    )
+
+
 @pytest.mark.slow_integration_test
 def test_sampling_with_infinite_prior_bounds(tmpdir):
     """