Add check for situations where the normalising constant of the target density is very small

Author: alexgdebeer

deep_tensor/ftt/ftt.py

@@ -315,7 +315,7 @@ def approximate(
         Parameters
         ----------
         target_func: 
-            The target function, $f : [-1, 1]^{d} \rightarrow \mathbb{R}$. 
+            The target function, $f : [-1, 1]^{d} \rightarrow \mathbb{R}$.
         reference:
             The reference measure. If provided, this will be used to 
             generate the initial index sets for the underlying TT. 

deep_tensor/irt/dirt.py

@@ -152,7 +152,6 @@ def _get_new_layer(self) -> SIRT:
             self.eval_ratio_func, 
             ftt, 
             self.reference, 
-            self.domain, 
             self.defensive, 
             self.cdf_tol
         )

deep_tensor/irt/sirt.py

@@ -1,4 +1,5 @@
 from typing import Callable, Dict, Tuple
+import warnings
 
 import torch
 from torch import Tensor
@@ -7,6 +8,7 @@
 from ..ftt import ApproxBases, Direction, FTT
 from ..linalg import batch_mul, n_mode_prod, unfold_left, unfold_right
 from ..polynomials import CDF1D, construct_cdf
+from ..references import Reference
 
 
 SUBSET2DIRECTION = {
@@ -16,38 +18,33 @@
 
 
 class SIRT():
-    r"""Squared inverse Rosenblatt transport.
+    """Squared inverse Rosenblatt transport.
     
     Parameters
     ----------
     potential:
-        A function that receives an $n \times d$ matrix of samples and 
-        returns an $n$-dimensional vector containing the potential 
+        A function that receives an n * d matrix of samples and 
+        returns an n-dimensional vector containing the potential 
         function of the target density evaluated at each sample.
     ftt:
         TODO
-    defensive: 
+    reference:
         TODO
+    domain: 
+        The domain of the reference.
     defensive:
-        The defensive parameter, $\tau$, which ensures that the tails
-        of the approximation are sufficiently heavy.
+        The defensive parameter.
     cdf_tol:
-        TODO
-
-    References
-    ----------
-    Cui, T and Dolgov, S (2022). *[Deep composition of tensor-trains 
-    using squared inverse Rosenblatt transports](https://doi.org/10.1007/s10208-021-09537-5).* 
-    Foundations of Computational Mathematics, **22**, 1863--1922.
+        The tolerance used when solving the rootfinding problem to 
+        evaluate the inverse of each conditional CDF.
 
     """
 
     def __init__(
         self, 
         target_func: Callable[[Tensor], Tensor], 
         ftt: FTT,
-        reference,
-        domain: Domain,
+        reference: Reference,
         defensive: float,
         cdf_tol: float
     ):
@@ -56,7 +53,7 @@ def __init__(
         self.ftt = ftt
         self.bases = self.ftt.bases
         self.dim = self.ftt.dim
-        self.domain = domain
+        self.domain = reference.domain
         self.defensive = defensive
         self.cdfs = self.construct_cdfs(self.bases, cdf_tol)
 
@@ -175,6 +172,26 @@ def _target_func(self, ls: Tensor) -> Tensor:
         neglogwxs = self.eval_measure_potential(xs)[0]
         gs = torch.exp(-0.5 * (neglogfxs - neglogwxs))
         return gs
+    
+    @staticmethod
+    def _check_z_func(z_func) -> None:
+        
+        dtype = torch.get_default_dtype()
+        msg = (
+            "The normalising constant of the current SIRT layer is very small "
+            f"({z_func:.2e}). This may cause numerical instability. "
+        )
+        if dtype == torch.float32 and z_func < 1.0e-5:
+            msg += (
+                "Consider rescaling the potential function "
+                "or changing to double precision."
+            )
+            warnings.warn(msg)
+        elif dtype == torch.float64 and z_func < 1.0e-10:
+            msg += "Consider rescaling the potential function."
+            warnings.warn(msg)
+
+        return
 
     def _marginalise_forward(self) -> None:
         """Computes each coefficient tensor required to evaluate the 
@@ -192,6 +209,7 @@ def _marginalise_forward(self) -> None:
             self._Rs_f[k] = torch.linalg.qr(C_k, mode="reduced")[1].T
 
         self.z_func = self._Rs_f[0].square().sum()
+        self._check_z_func(self.z_func)
         return 
 
     def _marginalise_backward(self) -> None: