Update bridging densities to support case where no debugging samples are passed in

alexgdebeer · alexgdebeer · commit 192cde09d67c · 2025-10-03T10:29:13.000-05:00
diff --git a/deep_tensor/bridging_densities/bridge.py b/deep_tensor/bridging_densities/bridge.py
@@ -16,12 +16,12 @@ def is_last(self) -> bool:
         pass
 
     @property 
-    def n_layers(self) -> int:
-        return self._n_layers
+    def num_layers(self) -> int:
+        return self._num_layers
     
-    @n_layers.setter
-    def n_layers(self, value: int) -> None:
-        self._n_layers = value
+    @num_layers.setter
+    def num_layers(self, value: int) -> None:
+        self._num_layers = value
         return
     
     @property
@@ -86,6 +86,15 @@ def apply_preconditioner(self, us: Tensor) -> Tuple[Tensor, Tensor]:
         neglogdets = self.preconditioner.neglogdet_Q(us)
         return xs, neglogdets
     
+    def _eval_pullback(self, us: Tensor) -> Tensor:
+        """Evaluates the pullback of the target density under the 
+        preconditioning mapping.
+        """
+        xs, neglogdets = self.apply_preconditioner(us)
+        neglogfxs = self.target_func(xs)
+        neglogfus = neglogfxs + neglogdets
+        return neglogfus
+    
     def _reorder(
         self, 
         xs: Tensor, 
diff --git a/deep_tensor/bridging_densities/rare_event.py b/deep_tensor/bridging_densities/rare_event.py
@@ -65,7 +65,7 @@ def __init__(
         betas: Sequence | Tensor | float = 1.0
     ):
         self.gammas, self.betas = self._parse_bridging_params(gammas, betas)
-        self.n_layers = 0
+        self.num_layers = 0
         self.initialised = False
 
         self._ratio_weight_funcs = {
@@ -77,7 +77,7 @@ def __init__(
     
     @property
     def is_last(self) -> bool:
-        return self.n_layers == (len(self.betas) - 1)
+        return self.num_layers == (len(self.betas) - 1)
     
     @staticmethod
     def _parse_bridging_params(
@@ -132,14 +132,23 @@ def neglogsigmoid(self, gamma: float, responses: Tensor) -> Tensor:
         neglogsigmoids = torch.log1p(torch.exp(gamma * dzs))
         return neglogsigmoids
     
+    def _eval_pullback(self, us: Tensor) -> Tuple[Tensor, Tensor]:
+        """Evaluates the pullback of the target density under the 
+        preconditioning mapping.
+        """
+        xs, neglogdets = self.apply_preconditioner(us)
+        neglogfxs, responses = self.target_func.func(xs)
+        neglogfus = neglogfxs + neglogdets
+        return neglogfus, responses
+
     def _compute_neglogbridges(
         self, 
         neglogref_us: Tensor,
         neglogfus: Tensor,
         responses: Tensor
     ) -> Tensor:
         
-        k = self.n_layers
+        k = self.num_layers
         neglogsigmoids_p = self.neglogsigmoid(self.gammas[k-1], responses)
         neglogbridges = (
             + (1.0 - self.betas[k-1]) * neglogref_us 
@@ -159,7 +168,7 @@ def _compute_weights_aratio(
         the previous bridging density for each particle.
         """
         
-        k = self.n_layers
+        k = self.num_layers
         neglogsigmoids = self.neglogsigmoid(self.gammas[k], responses)
         neglogsigmoids_p = self.neglogsigmoid(self.gammas[k-1], responses)
         neglogsigmoids_p[neglogsigmoids_p.isinf()] = 0.0
@@ -183,7 +192,7 @@ def _compute_weights_eratio(
         each particle.
         """
 
-        k = self.n_layers
+        k = self.num_layers
         neglogsigmoids = self.neglogsigmoid(self.gammas[k], responses)
         
         neglogweights = (
@@ -243,10 +252,7 @@ def ratio_func(
         
         neglogref_rs = self.reference.eval_potential(rs)[0]
         neglogref_us = self.reference.eval_potential(us)[0]
-
-        xs, neglogdets = self.apply_preconditioner(us)
-        neglogfxs, responses = self.target_func.func(xs)
-        neglogfus = neglogfxs + neglogdets
+        neglogfus, responses = self._eval_pullback(us)
 
         neglogratios = self._compute_ratio_func(
             method,
@@ -264,10 +270,7 @@ def update(self, us: Tensor, neglogfus_dirt: Tensor) -> Tuple[Tensor, Tensor]:
             raise Exception("Need to call self.initialise().")
         
         neglogref_us = self.reference.eval_potential(us)[0]
-        
-        xs, neglogdets = self.apply_preconditioner(us)
-        neglogfxs, responses = self.target_func.func(xs)
-        neglogfus = neglogfxs + neglogdets
+        neglogfus, responses = self._eval_pullback(us)
 
         neglogbridges = self._compute_neglogbridges(
             neglogref_us,
@@ -290,14 +293,16 @@ def _get_diagnostics(
         neglogfus: Tensor,
         neglogfus_dirt: Tensor
     ) -> List[str]:
+        
+        msg = [
+            f"Gamma: {self.gammas[self.num_layers]:.4f}",
+            f"Beta: {self.betas[self.num_layers]:.4f}"
+        ]
+        
+        if None in (log_weights, neglogfus, neglogfus_dirt):
+            return msg
 
         div_h2 = compute_f_divergence(-neglogfus_dirt, -neglogfus)
         ess = estimate_ess_ratio(log_weights)
-
-        msg = [
-            f"DHell: {div_h2.sqrt():.4f}",
-            f"Gamma: {self.gammas[self.n_layers]:.4f}",
-            f"Beta: {self.betas[self.n_layers]:.4f}", 
-            f"ESS: {ess:.4f}"
-        ]
+        msg += [f"DHell: {div_h2.sqrt():.4f}", f"ESS: {ess:.4f}"]
         return msg
diff --git a/deep_tensor/bridging_densities/single_layer.py b/deep_tensor/bridging_densities/single_layer.py
@@ -14,7 +14,7 @@ class SingleLayer(Bridge):
     """
 
     def __init__(self):
-        self.n_layers = 0
+        self.num_layers = 0
         self.is_adaptive = False
         return
     
@@ -27,11 +27,7 @@ def update(
         us: Tensor, 
         neglogfus_dirt: Tensor
     ) -> Tuple[Tensor, Tensor]:
-        
-        xs, neglogdets = self.apply_preconditioner(us)
-        neglogfxs = self.target_func(xs)
-        neglogfus = neglogfxs + neglogdets
-
+        neglogfus = self._eval_pullback(us)
         log_weights = -neglogfus + neglogfus_dirt
         return log_weights, neglogfus
         
@@ -42,7 +38,4 @@ def ratio_func(
         us: Tensor,
         neglogfus_dirt: Tensor
     ) -> Tensor:
-        xs, neglogdets = self.apply_preconditioner(us)
-        neglogfxs = self.target_func(xs)
-        neglogfus = neglogfxs + neglogdets
-        return neglogfus
+        return self._eval_pullback(us)
diff --git a/deep_tensor/bridging_densities/tempering.py b/deep_tensor/bridging_densities/tempering.py
@@ -87,7 +87,7 @@ def __init__(
         self.init_beta = min_beta
         self.max_layers = max_layers
         self.is_adaptive = len(self.betas) == 1
-        self.n_layers = 0
+        self.num_layers = 0
         self.initialised = False
 
         self._ratio_weight_funcs = {
@@ -99,8 +99,8 @@ def __init__(
     
     @property 
     def is_last(self) -> bool:
-        max_layers_reached = self.n_layers == self.max_layers
-        final_beta_reached = abs(self.betas[self.n_layers-1] - 1.0) < 1e-6
+        max_layers_reached = self.num_layers == self.max_layers
+        final_beta_reached = abs(self.betas[self.num_layers-1] - 1.0) < 1e-6
         return bool(max_layers_reached or final_beta_reached)
 
     def initialise(
@@ -120,7 +120,7 @@ def _compute_neglogbridges(
         neglogfus: Tensor
     ) -> Tensor:
         
-        k = self.n_layers
+        k = self.num_layers
         neglogbridges = (
             + (1.0 - self.betas[k-1]) * neglogref_us 
             + self.betas[k-1] * neglogfus
@@ -136,7 +136,7 @@ def _compute_weights_aratio(
         """Computes the ratio between the current bridging density and 
         the previous bridging density for each particle.
         """
-        k = self.n_layers
+        k = self.num_layers
         neglogweights = (
             + (self.betas[k-1] - self.betas[k]) * neglogref_us 
             + (self.betas[k] - self.betas[k-1]) * neglogfus
@@ -153,7 +153,7 @@ def _compute_weights_eratio(
         the DIRT approximation to the previous bridging density for 
         each particle.
         """
-        k = self.n_layers
+        k = self.num_layers
         neglogweights = (
             + (1.0 - self.betas[k]) * neglogref_us 
             + self.betas[k] * neglogfus
@@ -183,7 +183,7 @@ def _compute_log_weights(
         neglogfus: Tensor,
         neglogfus_dirt: Tensor
     ) -> Tensor:
-        beta = self.betas[self.n_layers]
+        beta = self.betas[self.num_layers]
         log_weights = -beta*neglogfus - (1-beta)*neglogrefs + neglogfus_dirt
         return log_weights
     
@@ -200,10 +200,7 @@ def ratio_func(
         
         neglogref_rs = self.reference.eval_potential(rs)[0]
         neglogref_us = self.reference.eval_potential(us)[0]
-
-        xs, neglogdets = self.apply_preconditioner(us)
-        neglogfxs = self.target_func(xs)
-        neglogfus = neglogfxs + neglogdets
+        neglogfus = self._eval_pullback(us)
 
         neglogratios = self._compute_ratio_func(
             method,
@@ -221,11 +218,11 @@ def _adapt_beta(
         neglogfus_dirt: Tensor
     ):
         
-        if self.n_layers == 0:
+        if self.num_layers == 0:
             self.betas[0] = self.init_beta
             return
         
-        k = self.n_layers
+        k = self.num_layers
         self.betas[k] = self.betas[k-1] * self.beta_factor
 
         while True:
@@ -250,10 +247,7 @@ def update(
     ) -> Tuple[Tensor, Tensor]:
         
         neglogref_us = self.reference.eval_potential(us)[0]
-        
-        xs, neglogdets = self.apply_preconditioner(us)
-        neglogfxs = self.target_func(xs)
-        neglogfus = neglogfxs + neglogdets
+        neglogfus = self._eval_pullback(us)
 
         if self.is_adaptive:
             self._adapt_beta(neglogref_us, neglogfus, neglogfus_dirt)
@@ -277,13 +271,13 @@ def _get_diagnostics(
         neglogfus: Tensor,
         neglogfus_dirt: Tensor
     ) -> List[str]:
+        
+        msg = [f"Beta: {self.betas[self.num_layers]:.4f}"]
+
+        if None in (log_weights, neglogfus, neglogfus_dirt):
+            return msg
 
         div_h2 = compute_f_divergence(-neglogfus_dirt, -neglogfus)
         ess = estimate_ess_ratio(log_weights)
-
-        msg = [
-            f"DHell: {div_h2.sqrt():.4f}",
-            f"Beta: {self.betas[self.n_layers]:.4f}", 
-            f"ESS: {ess:.4f}"
-        ]
+        msg += [f"DHell: {div_h2.sqrt():.4f}", f"ESS: {ess:.4f}"]
         return msg
diff --git a/deep_tensor/irt/dirt.py b/deep_tensor/irt/dirt.py
diff --git a/deep_tensor/preconditioners/gaussian_mapping.py b/deep_tensor/preconditioners/gaussian_mapping.py
