Tidy up ACA code

Author: alexgdebeer

deep_tensor/ftt/ftt.py

@@ -1,4 +1,5 @@
-from typing import Callable, Dict
+from typing import Callable, Dict, Tuple
+import warnings
 
 import torch
 from torch import linalg
@@ -421,80 +422,148 @@ def compute_fibre_submatrix_random(
 
         return fibre_matrix
 
-    def compute_fibre_submatrix_aca(self, grid: Grid, k: int) -> Tensor:
-
-        for iter in range(self.options.max_fibres):
-                
-            random_inds = grid.sample_indices(self.options.num_aca)
-            random_points = grid.indices2points(random_inds)
-
-            M_vals = self.target_func(random_points)
-            
-            if iter == 0:
+    @staticmethod
+    def _find_evaluated_points(
+        new_inds: Tensor, 
+        inds_eval: Tensor,
+        vals_eval: Tensor
+    ) -> Tuple[Tensor, Tensor]:
+        """Returns a mask elements of a set of indices that have been 
+        computed previously, as well as the computed values.
+        """
+        diffs = (new_inds[:, None, :] - inds_eval[None, ...]).abs().sum(dim=2)
+        inds_prev = diffs.argmin(dim=1)
+        mask = diffs.min(dim=1).values < EPS
+        mask_vals = vals_eval[inds_prev[mask]]
+        return mask, mask_vals
+    
+    def _generate_points_aca(self, n: int, grid: Grid) -> Tuple[Tensor, Tensor]:
+        """Returns a set of random indices and the corresponding 
+        function values.
+        """
+        inds_rand = grid.sample_indices(n)
+        random_points = grid.indices2points(inds_rand)
+        func_vals = self.target_func(random_points)
+        self.num_eval_fibres += func_vals.numel()
+        return inds_rand, func_vals
+    
+    def _initialise_index_set_aca(self, grid: Grid) -> Tuple[Tensor, Tensor]:
+        """Initialises the index set defining the current cross by 
+        sampling from the coefficient tensor at random. This is 
+        repeated multiple times in case the sampled elements are 
+        uniformly zero (see also implementation by Strossner et al.).
+        """
 
-                max_residual = M_vals.max()
-                max_residual_index = M_vals.abs().argmax()
-                max_index = random_inds[max_residual_index, :]
+        num_initialisation_batches = 5
+        num_aca = self.options.num_aca
 
-                inds = torch.atleast_2d(max_index)
+        for _ in range(num_initialisation_batches):
+            inds_rand, func_vals = self._generate_points_aca(num_aca, grid)
+            if func_vals.abs().max() > 0.0:
+                break
+        
+        if func_vals.abs().max() == 0.0:
+            msg = (
+                "ACA: None of the sampled fibre elements are nonzero. "
+                "Consider rescaling the target function. If you are "
+                "confident the target function is scaled appropriately, "
+                "consider using a refined grid, larger core ranks, an "
+                "increased number of bridging densities, or a larger "
+                "value for num_aca."
+            )
+            warnings.warn(msg)
+        
+        max_residual_index = func_vals.abs().argmax()
+        inds = torch.atleast_2d(inds_rand[max_residual_index])
+        vals = torch.atleast_1d(func_vals[max_residual_index])
+        return inds, vals
+    
+    def compute_fibre_submatrix_aca(self, grid: Grid, k: int) -> Tensor:
 
-            else:
+        num_aca = self.options.num_aca
+        inds, vals = self._initialise_index_set_aca(grid)
 
-                num_inds = inds.shape[0]
+        # Keep track of elements of the cross that have been evaluated
+        inds_eval = inds.clone()
+        vals_eval = vals.clone()
 
-                # Compute intersection matrix (NOTE: some of this 
-                # will have actually been computed at previous 
-                # iterations...)
-                inds_int = inds.repeat(num_inds, 1)
-                inds_int[:, k] = inds[:, k].repeat_interleave(num_inds, dim=0)
+        for _ in range(1, self.options.max_fibres):
 
-                inds_row = random_inds.repeat(num_inds, 1)
-                inds_row[:, k] = inds[:, k].repeat_interleave(self.options.num_aca, dim=0)
+            num_inds = inds.shape[0]
+            inds_rand, func_vals = self._generate_points_aca(num_aca, grid)
 
-                inds_col = inds.repeat(self.options.num_aca, 1)
-                inds_col[:, k] = random_inds[:, k].repeat_interleave(num_inds, dim=0)
+            inds_int = inds.repeat(num_inds, 1)
+            inds_int[:, k] = inds[:, k].repeat_interleave(num_inds, dim=0)
+            inds_row = inds_rand.repeat(num_inds, 1)
+            inds_row[:, k] = inds[:, k].repeat_interleave(num_aca, dim=0)
+            inds_col = inds.repeat(self.options.num_aca, 1)
+            inds_col[:, k] = inds_rand[:, k].repeat_interleave(num_inds, dim=0)
 
-                points_int = grid.indices2points(inds_int)
-                points_row = grid.indices2points(inds_row)
-                points_col = grid.indices2points(inds_col)
+            points_int = grid.indices2points(inds_int)
+            points_row = grid.indices2points(inds_row)
+            points_col = grid.indices2points(inds_col)
 
-                B_int = self.target_func(points_int)
-                B_int = B_int.reshape(num_inds, num_inds)
-                B_rows = self.target_func(points_row)
-                B_rows = B_rows.reshape(num_inds, self.options.num_aca)
-                B_cols = self.target_func(points_col)
-                B_cols = B_cols.reshape(self.options.num_aca, num_inds)
+            mask, mask_vals = self._find_evaluated_points(
+                inds_int, inds_eval, vals_eval
+            )
 
-                self.num_eval_fibres += (
-                    2 * num_inds * self.options.num_aca 
-                    + num_inds ** 2
-                )
+            # Form intersection submatrix (avoiding the evaluation 
+            # of function values that were previously computed)
+            B_int = torch.zeros(inds_int.shape[0])
+            B_int[mask] = mask_vals
+            if (~mask).any():
+                B_int[~mask] = self.target_func(points_int[~mask])
+            
+            B_rows = self.target_func(points_row)
+            B_cols = self.target_func(points_col)
+            
+            B_int = B_int.reshape(num_inds, num_inds)
+            B_rows = B_rows.reshape(num_inds, num_aca)
+            B_cols = B_cols.reshape(num_aca, num_inds)
+            
+            inds_eval = inds_int.clone()
+            vals_eval = B_int.flatten()
+            
+            num_eval_int = int((~mask).sum())
+            self.num_eval_fibres += (
+                num_eval_int + B_rows.numel() + B_cols.numel()
+            )
 
-                # Check for (near-)singularity of intersection matrix
-                # (also done in implementation by Strossner et al.).
-                if linalg.cond(B_int) > 1.0 / EPS:
-                    break
-
-                # Update index set with index of maximum residual
-                B_vals = B_cols @ linalg.solve(B_int, B_rows)
-                residuals = torch.diag(M_vals - B_vals).abs()
-                max_residual = residuals.max()
-                max_residual_index = residuals.abs().argmax()
-                max_index = random_inds[max_residual_index, :]
-                inds = torch.vstack((inds, max_index))
+            # Check for (near-)singularity of intersection matrix
+            # (also done in implementation by Strossner et al.).
+            # This occurs for functions where the fibre matrices 
+            # are exactly low rank.
+            if linalg.cond(B_int) > 1.0 / EPS:
+                break
 
-            if max_residual < self.options.tol_aca and iter > 1:
+            cross_vals = B_cols @ linalg.solve(B_int, B_rows)
+            residuals = torch.diag(func_vals - cross_vals).abs()
+            if residuals.max() < self.options.tol_aca:
                 break
+
+            # Update index set
+            max_index = inds_rand[residuals.argmax(), :]
+            inds = torch.vstack((inds, max_index))
 
         n_k = self.bases[k].cardinality
         num_inds = inds.shape[0]
 
         fibre_inds = inds.repeat(n_k, 1)
-        fibre_inds[:, k] = torch.arange(n_k, device=self.device).repeat_interleave(num_inds, dim=0)
-
+        ii = torch.arange(n_k, device=self.device)
+        fibre_inds[:, k] = ii.repeat_interleave(num_inds, dim=0)
         fibre_points = grid.indices2points(fibre_inds)
-        fibre_matrix = self.target_func(fibre_points).reshape(n_k, num_inds)
-        self.num_eval_fibres += n_k * num_inds
+
+        mask, mask_vals = self._find_evaluated_points(
+            fibre_inds, inds_eval, vals_eval
+        )
+
+        fibre_matrix = torch.zeros((n_k*num_inds,))
+        fibre_matrix[mask] = mask_vals
+        fibre_matrix[~mask] = self.target_func(fibre_points[~mask])
+        fibre_matrix = fibre_matrix.reshape(n_k, num_inds)
+
+        num_eval_new = int((~mask).sum())
+        self.num_eval_fibres += num_eval_new
 
         return fibre_matrix
 
@@ -580,9 +649,10 @@ def approximate(
         return
 
     def clone(self):
-        # Note: can't copy the cores and index sets over, because the 
-        # indices corresponding to the DEIM projection onto the reduced
-        # bases in each dimension can change.
+        # Note: we cannot copy the cores and index sets over, because 
+        # the indices corresponding to the DEIM projection onto the 
+        # reduced bases in each dimension can change. Instead we start 
+        # from scratch.
         tt = TT(self.tt.options, device=self.device)
         ftt = EFTT(self.bases, tt, self.options, device=self.device)
         return ftt