implemented nplets_measures for automatic hot encoding if multi order

nplets

Author: Laouen

tests/test_multi_order_measures.py

@@ -9,7 +9,7 @@
 from thoi.measures.gaussian_copula_hot_encoded import multi_order_measures_hot_encoded
 from thoi.commons import gaussian_copula_covmat
 
-# TODO: make this test for all combinations of use_cpu in [True, False] use_cpu_dataset in [True, False] and dataset_device in ['cpu', 'gpu']
+# TODO: make this test for all combinations of device in [cpu, cuda] and different input types
 class TestMultiOrderMeasures(unittest.TestCase):
 
     # make the constructor
@@ -54,37 +54,37 @@ def _validate_same_results_for_repeated_datasets(self, df_res, rtol, atol):
             self._compare_with_ground_truth(df_dataset, rtol, atol)
 
     def test_multiorder_measures_timeseries(self):
-        df_res = multi_order_measures(self.X, use_cpu=True)
+        df_res = multi_order_measures(self.X)
         self._compare_with_ground_truth(df_res, rtol=1e-16, atol=1e-12)
 
     def test_multiorder_measures_precomputed_covmat(self):
-        df_res = multi_order_measures(self.covmat, covmat_precomputed=True, T=self.X.shape[0], use_cpu=True)
+        df_res = multi_order_measures(self.covmat, covmat_precomputed=True, T=self.X.shape[0])
         self._compare_with_ground_truth(df_res, rtol=1e-16, atol=1e-12)
 
     def test_multiorder_measures_timeseries_hot_encoded(self):
-        df_res = multi_order_measures_hot_encoded(self.X, batch_size=200000, use_cpu=True)
+        df_res = multi_order_measures_hot_encoded(self.X, batch_size=200000)
         self._compare_with_ground_truth(df_res, rtol=1e-8, atol=1e-4)
 
     def test_multiorder_measures_precomputed_hot_encoded(self):
-        df_res = multi_order_measures_hot_encoded(self.covmat, covmat_precomputed=True, T=self.X.shape[0], use_cpu=True)
+        df_res = multi_order_measures_hot_encoded(self.covmat, covmat_precomputed=True, T=self.X.shape[0])
         self._compare_with_ground_truth(df_res, rtol=1e-8, atol=1e-4)
 
     def test_multiple_times_same_datasets_timeseries(self):
-        df_res = multi_order_measures([self.X, self.X], use_cpu=True)
+        df_res = multi_order_measures([self.X, self.X])
         self._validate_same_results_for_repeated_datasets(df_res, rtol=1e-16, atol=1e-7)
 
     def test_multiple_times_same_datasets_precomputed(self):
         covmats = [self.covmat, self.covmat]
-        df_res = multi_order_measures(covmats, covmat_precomputed=True, T=self.X.shape[0], use_cpu=True)
+        df_res = multi_order_measures(covmats, covmat_precomputed=True, T=self.X.shape[0])
         self._validate_same_results_for_repeated_datasets(df_res, rtol=1e-16, atol=1e-7)
 
     def test_multiple_times_same_datasets_timeseries_hot_encoded(self):
-        df_res = multi_order_measures_hot_encoded([self.X, self.X], use_cpu=True)
+        df_res = multi_order_measures_hot_encoded([self.X, self.X])
         self._validate_same_results_for_repeated_datasets(df_res, rtol=1e-8, atol=1e-4)
 
     def test_multiple_times_same_datasets_precomputed_hot_encoded(self):
         covmats = [self.covmat, self.covmat]
-        df_res = multi_order_measures_hot_encoded(covmats, covmat_precomputed=True, T=self.X.shape[0], use_cpu=True)
+        df_res = multi_order_measures_hot_encoded(covmats, covmat_precomputed=True, T=self.X.shape[0])
         self._validate_same_results_for_repeated_datasets(df_res, rtol=1e-8, atol=1e-4)
 
 

tests/test_nplet_measures.py

@@ -13,7 +13,7 @@
 from thoi.measures.gaussian_copula_hot_encoded import nplets_measures_hot_encoded
 from thoi.commons import gaussian_copula_covmat
 
-# TODO: make this test for all combinations of use_cpu in [True, False] use_cpu_dataset in [True, False] and dataset_device in ['cpu', 'gpu']
+# TODO: make this test for all combinations of devices in [cpu, gpu] and different input types
 class TestNpletsMeasures(unittest.TestCase):
 
     # make the constructor
@@ -67,31 +67,31 @@ def test_nplets_measures_timeseries(self):
         for order in range(3,11):
             with self.subTest(order=order):
                 nplets = torch.tensor(list(combinations(full_nplet, order)))
-                res = nplets_measures(self.X, nplets, use_cpu=True)
+                res = nplets_measures(self.X, nplets)
                 self._compare_with_ground_truth(res, nplets, rtol=1e-16, atol=1e-12)
 
     def test_nplets_measures_precomputed(self):
         full_nplet = range(self.X.shape[1])
         for order in range(3,11):
             with self.subTest(order=order):
                 nplets = torch.tensor(list(combinations(full_nplet, order)))
-                res = nplets_measures(self.covmat, nplets, covmat_precomputed=True, T=self.X.shape[0], use_cpu=True)
+                res = nplets_measures(self.covmat, nplets, covmat_precomputed=True, T=self.X.shape[0])
                 self._compare_with_ground_truth(res, nplets, rtol=1e-16, atol=1e-12)
 
     def test_multiple_times_same_datasets_timeseries(self):
         full_nplet = range(self.X.shape[1])
         for order in range(3,11):
             with self.subTest(order=order):
                 nplets = torch.tensor(list(combinations(full_nplet, order)))
-                res = nplets_measures([self.X, self.X], nplets, use_cpu=True)
+                res = nplets_measures([self.X, self.X], nplets)
                 self._validate_same_results_for_repeated_datasets(res, nplets, rtol=1e-16, atol=1e-7)
 
     def test_multiple_times_same_datasets_precomputed(self):
         full_nplet = range(self.X.shape[1])
         for order in range(3,11):
             with self.subTest(order=order):
                 nplets = torch.tensor(list(combinations(full_nplet, order)))
-                res = nplets_measures([self.covmat, self.covmat], nplets, covmat_precomputed=True, T=self.X.shape[0], use_cpu=True)
+                res = nplets_measures([self.covmat, self.covmat], nplets, covmat_precomputed=True, T=self.X.shape[0])
                 self._validate_same_results_for_repeated_datasets(res, nplets, rtol=1e-16, atol=1e-7)
 
     def test_nplets_measures_timeseries_hot_encoded(self):
@@ -103,7 +103,7 @@ def test_nplets_measures_timeseries_hot_encoded(self):
                 batch_size = nplets.shape[0]
                 nplets_hot_encoded = torch.zeros((batch_size, N), dtype=torch.int)
                 nplets_hot_encoded[torch.arange(0,batch_size, dtype=int).view(-1,1), nplets] = 1
-                res = nplets_measures_hot_encoded(self.X, nplets_hot_encoded, use_cpu=True)
+                res = nplets_measures_hot_encoded(self.X, nplets_hot_encoded)
                 self._compare_with_ground_truth(res, nplets, rtol=1e-8, atol=1e-4)
 
     def test_nplets_measures_precomputed_hot_encoded(self):
@@ -115,7 +115,7 @@ def test_nplets_measures_precomputed_hot_encoded(self):
                 batch_size = nplets.shape[0]
                 nplets_hot_encoded = torch.zeros((batch_size, N), dtype=torch.int)
                 nplets_hot_encoded[torch.arange(0,batch_size, dtype=int).view(-1,1), nplets] = 1
-                res = nplets_measures_hot_encoded(self.covmat, nplets_hot_encoded, covmat_precomputed=True, T=self.X.shape[0], use_cpu=True)
+                res = nplets_measures_hot_encoded(self.covmat, nplets_hot_encoded, covmat_precomputed=True, T=self.X.shape[0])
                 self._compare_with_ground_truth(res, nplets, rtol=1e-8, atol=1e-4)
 
     def test_multiple_times_same_dataset_timeseries_hot_encoded(self):
@@ -127,7 +127,7 @@ def test_multiple_times_same_dataset_timeseries_hot_encoded(self):
                 batch_size = nplets.shape[0]
                 nplets_hot_encoded = torch.zeros((batch_size, N), dtype=torch.int)
                 nplets_hot_encoded[torch.arange(0,batch_size, dtype=int).view(-1,1), nplets] = 1
-                res = nplets_measures_hot_encoded([self.X, self.X], nplets_hot_encoded, use_cpu=True)
+                res = nplets_measures_hot_encoded([self.X, self.X], nplets_hot_encoded)
                 self._validate_same_results_for_repeated_datasets(res, nplets, rtol=1e-8, atol=1e-4)
 
     def test_multiple_times_same_dataset_precomputed_hot_encoded(self):
@@ -139,7 +139,7 @@ def test_multiple_times_same_dataset_precomputed_hot_encoded(self):
                 batch_size = nplets.shape[0]
                 nplets_hot_encoded = torch.zeros((batch_size, N), dtype=torch.int)
                 nplets_hot_encoded[torch.arange(0,batch_size, dtype=int).view(-1,1), nplets] = 1
-                res = nplets_measures_hot_encoded([self.covmat, self.covmat], nplets_hot_encoded, covmat_precomputed=True, T=self.X.shape[0], use_cpu=True)
+                res = nplets_measures_hot_encoded([self.covmat, self.covmat], nplets_hot_encoded, covmat_precomputed=True, T=self.X.shape[0])
                 self._validate_same_results_for_repeated_datasets(res, nplets, rtol=1e-8, atol=1e-4)
 
 

thoi/commons.py

@@ -73,20 +73,17 @@ def _get_device(use_cpu:bool=False):
 def _normalize_input_data(X: TensorLikeArray,
                          covmat_precomputed: bool=False,
                          T: Optional[Union[int, List[int]]]=None,
-                         use_cpu: bool=False):
+                         device: torch.device=torch.device('cpu')):
     '''
     brief: Normalize the input data to be a list of covariance matrices with shape (D, N, N) where D is the lenght of the list and N is the number of variables in the system.
 
     Parameters:
     - X: A list of 2D numpy arrays or tensors of shape: 1. (T, N) where T is the number of samples if X are multivariate series. 2. a list of 2D covariance matrices with shape (N, N).
     - covmat_precomputed: A boolean flag to indicate if the input data is a list of covariance matrices or multivariate series.
     - T (optional): A list of integers indicating the number of samples for each multivariate series.
-    - use_cpu: A boolean flag to indicate if the computation should be done on the CPU.
+    - device: The device to use for the computation. Default is 'cpu'.
     '''
 
-    # Get device to use
-    device = _get_device(use_cpu)
-
     # Handle different options for X parameter. Accept multivariate data or covariance matrix
     if covmat_precomputed:
         covmats = torch.as_tensor(X)
@@ -116,4 +113,4 @@ def _normalize_input_data(X: TensorLikeArray,
     # Send covmat to device  
     covmats = covmats.to(device).contiguous()
 
-    return covmats, D, N, T, device
+    return covmats, D, N, T

thoi/heuristics/greedy.py

@@ -5,21 +5,22 @@
 import torch
 from functools import partial
 
+from thoi.typing import TensorLikeArray
 from thoi.measures.gaussian_copula import multi_order_measures
 from thoi.collectors import batch_to_tensor, concat_batched_tensors
 from thoi.heuristics.commons import _get_valid_candidates
 from thoi.heuristics.scoring import _evaluate_nplets
 from thoi.commons import _normalize_input_data
 
 @torch.no_grad()
-def greedy(X: Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]],
+def greedy(X: TensorLikeArray,
            initial_order: int=3,
            order: Optional[int]=None,
            *,
            covmat_precomputed: bool=False,
            T: Optional[Union[int, List[int]]]=None,
            repeat: int=10,
-           use_cpu: bool=False,
+           device: torch.device=torch.device('cpu'),
            batch_size: int=1000000,
            metric: Union[str,Callable]='o',
            largest: bool=False):
@@ -28,12 +29,23 @@ def greedy(X: Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tenso
     Brief: Greedy algorithm to find the best order of nplets to maximize the metric for a given multivariate series or covariance matrices
     
     Parameters:
-    - X (Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tensor]]): The input data to compute the nplets. It can be a list of 2D numpy arrays or tensors of shape: 1. (T, N) where T is the number of samples if X are multivariate series. 2. a list of 2D covariance matrices with shape (N, N).
-    - covmat_precomputed (bool): A boolean flag to indicate if the input data is a list of covariance matrices or multivariate series.
-    - T (Optional[Union[int, List[int]]]): A list of integers indicating the number of samples for each multivariate series.
+        - X (TensorLikeArray): The input data to compute the nplets. It can be a list of 2D numpy arrays or tensors of shape: 1. (T, N) where T is the number of samples if X are multivariate series. 2. a list of 2D covariance matrices with shape (N, N).
+        - initial_order (int): The initial order to start the greedy algorithm. Default is 3.
+        - order (Optional[int]): The final order to stop the greedy algorithm. If None, it will be set to N.
+        - covmat_precomputed (bool): A boolean flag to indicate if the input data is a list of covariance matrices or multivariate series.
+        - T (Optional[Union[int, List[int]]]): A list of integers indicating the number of samples for each multivariate series.
+        - repeat (int): The number of repetition to do to obtain different solutions starting from les optimal initial solutions.
+        - device (torch.device): The device to use for the computation. Default is 'cpu'.
+        - batch_size (int): The batch size to use for the computation. Default is 1000000.
+        - metric (Union[str,Callable]): The metric to evaluate. One of tc, dtc, o, s or a callable function
+        - largest (bool): A flag to indicate if the metric is to be maximized or minimized
+        
+    Returns:
+        - best_nplets (torch.Tensor): The nplets with the best score found with shape (repeat, order)
+        - best_scores (torch.Tensor): The best scores for the best nplets with shape (repeat,)
     '''
 
-    covmats, D, N, T, device = _normalize_input_data(X, covmat_precomputed, T, use_cpu)
+    covmats, D, N, T = _normalize_input_data(X, covmat_precomputed, T, device)
 
     # Compute initial solutions
     batch_data_collector = partial(batch_to_tensor, top_k=repeat, metric=metric, largest=largest)
@@ -46,7 +58,7 @@ def greedy(X: Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tenso
                                                                min_order=initial_order,
                                                                max_order=initial_order,
                                                                batch_size=batch_size,
-                                                               use_cpu=use_cpu,
+                                                               device=device,
                                                                batch_data_collector=batch_data_collector,
                                                                batch_aggregation=batch_aggregation)
 
@@ -62,7 +74,7 @@ def greedy(X: Union[np.ndarray, torch.Tensor, List[np.ndarray], List[torch.Tenso
         best_candidate, best_score = _next_order_greedy(covmats, T, current_solution,
                                                        metric=metric,
                                                        largest=largest,
-                                                       use_cpu=use_cpu)
+                                                       device=device)
         best_scores.append(best_score)
 
         current_solution = torch.cat((current_solution, best_candidate.unsqueeze(1)) , dim=1)
@@ -101,9 +113,9 @@ def _create_all_solutions(initial_solution: torch.Tensor, valid_candidates: torc
 def _next_order_greedy(covmats: torch.Tensor,
                       T: Optional[List[int]],
                       initial_solution: torch.Tensor,
-                      metric:Union[str,Callable],
-                      largest:bool,
-                      use_cpu:bool):
+                      metric: Union[str,Callable],
+                      largest: bool,
+                      device: torch.device=torch.device('cpu')):
 
     '''
     Brief: Greedy algorithm to find the best candidate to add to the current solution
@@ -114,15 +126,14 @@ def _next_order_greedy(covmats: torch.Tensor,
     - initial_solution (torch.Tensor): The initial solution with shape (batch_size, order)
     - metric (Union[str,Callable]): The metric to evaluate. One of tc, dtc, o, s or a callable function
     - largest (bool): A flag to indicate if the metric is to be maximized or minimized
-    - use_cpu (bool): A flag to indicate if the computation should be done on the CPU
+    - device (torch.device): The device to use for the computation. Default is 'cpu'
     
     Returns:
     - best_candidates (torch.Tensor): The best candidates to add to the current solution with shape (batch_size)
     - best_score (torch.Tensor): The best score for the best candidates with shape (batch_size)
     '''
 
     # Get parameters attributes
-    device = covmats.device
     N = covmats.shape[1]
     batch_size, order = initial_solution.shape
 
@@ -137,7 +148,7 @@ def _next_order_greedy(covmats: torch.Tensor,
     all_solutions = all_solutions.view(batch_size*(N-order), order+1)
 
     # |batch_size x N-order|
-    best_score = _evaluate_nplets(covmats, T, all_solutions, metric, use_cpu=use_cpu)
+    best_score = _evaluate_nplets(covmats, T, all_solutions, metric, device=device)
 
     # |batch_size| x |N-order|
     best_score = best_score.view(batch_size, N-order)

thoi/heuristics/scoring.py

@@ -12,12 +12,13 @@ def _evaluate_nplets(covmats: torch.Tensor,
                      T: Optional[List[int]],
                      batched_nplets: torch.Tensor,
                      metric: Union[str, Callable],
-                     use_cpu:bool):
+                     device: torch.device):
     """
-        covmats (torch.Tensor): The covariance matrix or matrixes with shape (N, N) or (D, N, N)
-        T (Optional[List[int]]): The number of samples for each multivariate series or None
-        batched_nplets (torch.Tensor): The nplets to calculate the inverse of the oinformation with shape (batch_size, order)
-        metric (str): The metric to evaluate. One of tc, dtc, o, s or Callable
+        - covmats (torch.Tensor): The covariance matrix or matrixes with shape (N, N) or (D, N, N)
+        - T (Optional[List[int]]): The number of samples for each multivariate series or None
+        - batched_nplets (torch.Tensor): The nplets to calculate the inverse of the oinformation with shape (batch_size, order)
+        - metric (str): The metric to evaluate. One of tc, dtc, o, s or Callable
+        - device (torch.device): The device to use
     """
 
     if len(covmats.shape) == 2:
@@ -30,17 +31,17 @@ def _evaluate_nplets(covmats: torch.Tensor,
                                        nplets=batched_nplets,
                                        T=T,
                                        covmat_precomputed=True,
-                                       use_cpu=use_cpu)
+                                       device=device)
 
     # |batch_size|
-    return metric_func(batched_measures).to(covmats.device)
+    return metric_func(batched_measures).to(device)
 
 
 def _evaluate_nplet_hot_encoded(covmats: torch.Tensor,
-                                T:int,
+                                T: int,
                                 batched_nplets: torch.Tensor,
-                                metric:str,
-                                use_cpu:bool):
+                                metric: str,
+                                device: torch.device):
 
     """
         covmats (torch.Tensor): The covariance matrix or matrixes with shape (N, N) or (D, N, N)
@@ -59,7 +60,7 @@ def _evaluate_nplet_hot_encoded(covmats: torch.Tensor,
                                                    nplets=batched_nplets,
                                                    T=T,
                                                    covmat_precomputed=True,
-                                                   use_cpu=use_cpu)
+                                                   device=device)
 
     # |batch_size|
-    return metric_func(batched_measures).to(covmats.device)
+    return metric_func(batched_measures).to(device)