Tidy up resetting of bridging densities

Author: alexgdebeer

deep_tensor/bridging_densities/bridge.py

@@ -71,11 +71,19 @@ def update(self, us: Tensor, neglogfus_dirt: Tensor) -> Tuple[Tensor, Tensor]:
         """
         pass
 
+    @abc.abstractmethod
+    def reset(self) -> None:
+        """Resets the parameters of the bridging density to those at 
+        initialisation.
+        """
+        pass
+
     def initialise(
         self, 
         preconditioner: Preconditioner, 
         target_func: TargetFunc
     ) -> None:
+        self.reset()
         self.preconditioner = preconditioner
         self.reference = self.preconditioner.reference
         self.target_func = target_func

deep_tensor/bridging_densities/rare_event.py

@@ -67,6 +67,7 @@ def __init__(
         self.gammas, self.betas = self._parse_bridging_params(gammas, betas)
         self.num_layers = 0
         self.initialised = False
+        self.is_adaptive = False
 
         self._ratio_weight_funcs = {
             "aratio": self._compute_weights_aratio,
@@ -84,7 +85,7 @@ def _parse_bridging_params(
         gammas, 
         betas
     ) -> Tuple[Dict[int, float], Dict[int, float]]:
-        """TODO: this needs tidying up, I think..."""
+        # TODO: this could be tidied up.
 
         if isinstance(gammas, Tensor):
             gammas = gammas.tolist()
@@ -125,9 +126,7 @@ def initialise(
             msg = "Target function must be of type 'RareEventFunc'."
             raise Exception(msg)
 
-        self.preconditioner = preconditioner
-        self.reference = preconditioner.reference
-        self.target_func = target_func
+        Bridge.initialise(self, preconditioner, target_func)
         self.initialised = True
         return
 

deep_tensor/bridging_densities/tempering.py

@@ -115,9 +115,7 @@ def initialise(
         preconditioner: Preconditioner, 
         target_func: TargetFunc
     ) -> None:
-        self.preconditioner = preconditioner
-        self.reference = self.preconditioner.reference
-        self.target_func = target_func
+        Bridge.initialise(self, preconditioner, target_func)
         self.initialised = True
         return
 

deep_tensor/irt/dirt.py

@@ -48,9 +48,6 @@ def __init__(
         bridge: Bridge | None = None,
         options: DIRTOptions | None = None
     ):
-        # TODO: need to reset the bridge prior to starting. Ideally we 
-        # should be able to use the same bridge object to build 
-        # multiple DIRT objects.
 
         if not isinstance(target_func, TargetFunc):
             target_func = TargetFunc(target_func)
@@ -66,7 +63,6 @@ def __init__(
         self.domain = self.reference.domain
         self.ftt = ftt
         self.bridge = bridge
-        self.bridge.reset()
         self.bridge.initialise(preconditioner, target_func)
 
         self.ratio_type = options.ratio_type