From 559668020ba2f491cc9a283d1db9116fef8dbece Mon Sep 17 00:00:00 2001
From: arnaudbore <arnaud.bore@gmail.com>
Date: Thu, 26 Sep 2024 12:50:56 -0400
Subject: [PATCH] answer em comments

---
 dwi_ml/data/hdf5/hdf5_creation.py | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/dwi_ml/data/hdf5/hdf5_creation.py b/dwi_ml/data/hdf5/hdf5_creation.py
index e8f3c663..cad6b6bf 100644
--- a/dwi_ml/data/hdf5/hdf5_creation.py
+++ b/dwi_ml/data/hdf5/hdf5_creation.py
@@ -145,6 +145,8 @@ def __init__(self, root_folder: Path, out_hdf_filename: Path,
                  save_intermediate: bool = False,
                  intermediate_folder: Path = None):
         """
+        Params step_size, nb_points and compress are mutually exclusive.
+
         Params
         ------
         root_folder: Path
@@ -159,7 +161,7 @@ def __init__(self, root_folder: Path, out_hdf_filename: Path,
         groups_config: dict
             Information from json file loaded as a dict.
         dps_keys: List[str]
-            List of keys to keep in data_per_streamline. Default: None.
+            List of keys to keep in data_per_streamline. Default: [].
         step_size: float
             Step size to resample streamlines. Default: None.
         nb_points: int
@@ -201,7 +203,7 @@ def __init__(self, root_folder: Path, out_hdf_filename: Path,
             self._analyse_config_file()
 
         # -------- Performing checks
-
+        self._check_streamlines_operations()
         # Check that all subjects exist.
         logging.debug("Preparing hdf5 creator for \n"
                       "  training subjs {}, \n"
@@ -353,6 +355,19 @@ def flatten_list(a_list):
                                 self.enforce_files_presence,
                                 folder=subj_input_dir)
 
+    def _check_streamlines_operations(self):
+        valid = True
+        if self.step_size and self.nb_points:
+            valid = False
+        elif self.step_size and self.compress:
+            valid = False
+        elif self.nb_points and self.compress:
+            valid = False
+        if not valid:
+            raise ValueError(
+                "Only one option can be chosen: either resampling to "
+                "step_size, nb_points or compressing, not both.")
+
     def create_database(self):
         """
         Generates a hdf5 dataset from a group of subjects. Hdf5 dataset will
@@ -746,11 +761,11 @@ def _load_and_process_sft(self, tractogram_file, header):
                 "We do not support file's type: {}. We only support .trk "
                 "and .tck files.".format(tractogram_file))
         if file_extension == '.trk':
-            if header:
-                if not is_header_compatible(str(tractogram_file), header):
-                    raise ValueError("Streamlines group is not compatible "
-                                     "with volume groups\n ({})"
-                                     .format(tractogram_file))
+            if header and not is_header_compatible(str(tractogram_file),
+                                                   header):
+                raise ValueError("Streamlines group is not compatible "
+                                 "with volume groups\n ({})"
+                                 .format(tractogram_file))
             # overriding given header.
             header = 'same'