multiple bug fixes and added repeat_batch_size to greedy

Author: Laouen

thoi/heuristics/greedy.py

@@ -20,6 +20,7 @@ def greedy(X: TensorLikeArray,
            T: Optional[Union[int, List[int]]]=None,
            repeat: int=10,
            batch_size: int=1000000,
+           repeat_batch_size: int=1000000,
            device: torch.device=torch.device('cpu'),
            metric: Union[str,Callable]='o',
            largest: bool=False):
@@ -50,7 +51,7 @@ def greedy(X: TensorLikeArray,
     batch_data_collector = partial(batch_to_tensor, top_k=repeat, metric=metric, largest=largest)
     batch_aggregation = partial(concat_batched_tensors, top_k=repeat, metric=None, largest=largest)
 
-    # |repeat| x |initial_order|
+    # |repeat| x |initial_order|, |repeat|
     _, current_solution, current_scores = multi_order_measures(covmats,
                                                                covmat_precomputed=True,
                                                                T=T,
@@ -60,7 +61,7 @@ def greedy(X: TensorLikeArray,
                                                                device=device,
                                                                batch_data_collector=batch_data_collector,
                                                                batch_aggregation=batch_aggregation)
-    
+
     # send current solution to the device
     current_solution = current_solution.to(device).contiguous()
 
@@ -70,15 +71,21 @@ def greedy(X: TensorLikeArray,
     # Iterate over the remaining orders to get the best solution for each order
     best_scores = [current_scores]
     for _ in trange(initial_order, order, leave=False, desc='Order'):
+        
+        # |repeat|, |repeat|
         best_candidate, best_score = _next_order_greedy(covmats, T, current_solution,
                                                        metric=metric,
                                                        largest=largest,
                                                        batch_size=batch_size,
+                                                       repeat_batch_size=repeat_batch_size,
                                                        device=device)
+        # |order - initial_order| x |repeat|
         best_scores.append(best_score)
-
+        
+        # |repeat| x |order|
         current_solution = torch.cat((current_solution, best_candidate.unsqueeze(1)) , dim=1)
 
+    # |repeat| x |order|, |repeat| x |order - initial_order|
     return current_solution, torch.stack(best_scores).T
 
 
@@ -116,6 +123,7 @@ def _next_order_greedy(covmats: torch.Tensor,
                       metric: Union[str,Callable],
                       largest: bool,
                       batch_size: int=1000000,
+                      repeat_batch_size: int=1000000,
                       device: torch.device=torch.device('cpu')):
 
     '''
@@ -146,33 +154,33 @@ def _next_order_greedy(covmats: torch.Tensor,
     best_candidates = []
     best_scores = []
 
-    for start in range(0, total_size, batch_size):
-        end = min(start + batch_size, total_size)
+    for start in trange(0, total_size, repeat_batch_size, desc='Batch repeat', leave=False):
+        end = min(start + repeat_batch_size, total_size)
         batch_initial_solution = initial_solution[start:end]
         batch_valid_candidates = valid_candidates[start:end]
 
-        # |batch_size| x |N-order| x |order+1|
+        # |repeat_batch_size| x |N-order| x |order+1|
         all_solutions = _create_all_solutions(batch_initial_solution, batch_valid_candidates)
 
-        # |batch_size x N-order| x |order+1|
+        # |repeat_batch_size x N-order| x |order+1|
         all_solutions = all_solutions.view(-1, order+1)
 
-        # |batch_size x N-order|
+        # |repeat_batch_size x N-order|
         batch_best_score = _evaluate_nplets(covmats, T,
                                             all_solutions,
                                             metric,
                                             batch_size=batch_size,
                                             device=device)
 
-        # |batch_size| x |N-order|
+        # |repeat_batch_size| x |N-order|
         batch_best_score = batch_best_score.view(end - start, N - order)
 
         if not largest:
             batch_best_score = -batch_best_score
 
         # get for each batch item the best score over the second dimension
 
-        # |batch_size|
+        # |repeat_batch_size|
         max_idxs = torch.argmax(batch_best_score, dim=1)
         batch_best_candidates = batch_valid_candidates[torch.arange(end - start), max_idxs]
         batch_best_score = batch_best_score[torch.arange(end - start), max_idxs]

thoi/heuristics/simulated_annealing.py

@@ -44,13 +44,13 @@ def simulated_annealing(X: Union[np.ndarray, torch.Tensor, List[np.ndarray], Lis
     covmats, D, N, T = _normalize_input_data(X, covmat_precomputed, T, device)
 
     # Compute current solution
-    # |batch_size| x |order|
+    # |repeat| x |order|
     if initial_solution is None:
         current_solution = random_sampler(N, order, repeat, device)
     else:
         current_solution = initial_solution.to(device).contiguous()
 
-    # |batch_size|
+    # |repeat|
     current_energy = _evaluate_nplets(covmats, T,
                                       current_solution,
                                       metric,
@@ -61,21 +61,21 @@ def simulated_annealing(X: Union[np.ndarray, torch.Tensor, List[np.ndarray], Lis
         current_energy = -current_energy
 
     # Initial valid candidates
-    # |batch_size| x |N-order|
+    # |repeat| x |N-order|
     valid_candidates = _get_valid_candidates(current_solution, N, device)
 
     # Set initial temperature
     temp = initial_temp
 
     # Best solution found
-    # |batch_size| x |order|
+    # |repeat| x |order|
     best_solution = current_solution.clone()
-    # |batch_size|
+    # |repeat|
     best_energy = current_energy.clone()
 
     # Repeat tensor for indexing the current_solution
     # |repeat| x |1|
-    i_repeat = torch.arange(repeat)
+    i_repeat = torch.arange(repeat, device=device)
 
     no_progress_count = 0
     pbar = trange(max_iterations, leave=False)
@@ -88,7 +88,7 @@ def simulated_annealing(X: Union[np.ndarray, torch.Tensor, List[np.ndarray], Lis
 
         # Generate new solution by modifying the current solution.
         # Generate the random indexes to change.
-        # |batch_size| x |order|
+        # |repeat| x |order|
         i_sol = torch.randint(0, current_solution.shape[1], (repeat,), device=device)
         i_cand = torch.randint(0, valid_candidates.shape[1], (repeat,), device=device)
 
@@ -99,7 +99,7 @@ def simulated_annealing(X: Union[np.ndarray, torch.Tensor, List[np.ndarray], Lis
         current_solution[i_repeat, i_sol] = new_candidates
 
         # Calculate energy of new solution
-        # |batch_size|
+        # |repeat|
         new_energy = _evaluate_nplets(covmats, T,
                                       current_solution,
                                       metric,
@@ -111,33 +111,33 @@ def simulated_annealing(X: Union[np.ndarray, torch.Tensor, List[np.ndarray], Lis
 
         # Calculate change in energy
         # delca_energy > 0 means new_energy is bigger (more optimal) than current_energy
-        # |batch_size|
+        # |repeat|
         delta_energy = new_energy - current_energy
 
         # Determine if we should accept the new solution
-        # |batch_size|
+        # |repeat|
         temp_probas = torch.rand(repeat, device=device) < torch.exp(delta_energy / temp)
         improves = delta_energy > 0
         accept_new_solution = torch.logical_or(improves, temp_probas)
 
         # Restore original values for rejected candidates
-        # |batch_size| x |order|
+        # |repeat| x |order|
         current_solution[i_repeat[~accept_new_solution], i_sol[~accept_new_solution]] = current_candidates[~accept_new_solution]
 
         # Update valid_candidate for the accepted answers as they are not longer valid candidates
-        # |batch_size| x |N-order|
+        # |repeat| x |N-order|
         valid_candidates[i_repeat[accept_new_solution], i_cand[accept_new_solution]] = current_candidates[accept_new_solution]
 
         # Update current energy for the accepted solutions
-        # |batch_size|
+        # |repeat|
         current_energy[accept_new_solution] = new_energy[accept_new_solution]
 
         new_global_maxima = (new_energy > best_energy)
 
-        # |batch_size| x |order|
+        # |repeat| x |order|
         best_solution[new_global_maxima] = current_solution[new_global_maxima]
 
-        # |batch_size|
+        # |repeat|
         best_energy[new_global_maxima] = new_energy[new_global_maxima]
 
         # Cool down

thoi/measures/gaussian_copula.py

@@ -502,7 +502,7 @@ def multi_order_measures(X: TensorLikeArray,
     assert min_order <= max_order, f"min_order must be lower or equal than max_order. {min_order} > {max_order}"
 
     # Ensure that final batch_size is smaller than the original batch_size 
-    batch_size = batch_size // D
+    batch_size = max(batch_size // D, 1)
 
     # To compute using pytorch, we need to compute each order separately
     batched_data = []