Revert "Fix duplication errors by slicing datasets by position instead of site_id"

cvasl/dataset.py

@@ -10,11 +10,7 @@ def encode_cat_features(dff,cat_features_to_encode):
 
     feature_mappings = {}
     reverse_mappings = {}
-
-    # Track the original lengths to split data back correctly
-    dataset_lengths = [len(_d.data) for _d in dff]
-
-    data = pd.concat([_d.data for _d in dff], ignore_index=True)
+    data = pd.concat([_d.data for _d in dff])
 
     for feature in cat_features_to_encode:
         if feature in data.columns:
@@ -24,15 +20,11 @@ def encode_cat_features(dff,cat_features_to_encode):
             reverse_mappings[feature] = {v: k for k, v in mapping.items()}
             data[feature] = data[feature].map(mapping)
 
-    # Split data back using original dataset lengths
-    start_idx = 0
-    for i, _d in enumerate(dff):
-        end_idx = start_idx + dataset_lengths[i]
-        _d.data = data.iloc[start_idx:end_idx].reset_index(drop=True)
+    for _d in dff:
+        _d.data = data[data['site'] == _d.site_id]
         _d.feature_mappings = feature_mappings
         _d.reverse_mappings = reverse_mappings
         _d.cat_features_to_encode = cat_features_to_encode
-        start_idx = end_idx
     return dff
 
 class MRIdataset:

cvasl/harmonizers.py

@@ -236,13 +236,15 @@ def _reintegrate_harmonized_data(self, mri_datasets, harmonized_data, original_d
         harmonized_df = harmonized_df[original_order]
 
 
-        # Split data back using positional slicing based on original dataset lengths
-        start_idx = 0
         for dataset in mri_datasets:
-            end_idx = start_idx + len(dataset.data)
-            adjusted_data = harmonized_df.iloc[start_idx:end_idx].copy()
-            dataset.data = adjusted_data.reset_index(drop=True)
-            start_idx = end_idx
+            site_value = dataset.site_id # Assuming dataset.site_id is a single value
+            # Assuming site_indicator is a list of columns, and we use the first one for filtering for now - THIS MIGHT NEED ADJUSTMENT BASED ON HOW SITE_ID and site_indicator are related.
+            site_column_to_filter = self.site_indicator[0] if self.site_indicator else None # Use the first site indicator column for filtering
+            if site_column_to_filter and site_column_to_filter in harmonized_df.columns:
+                adjusted_data = harmonized_df[harmonized_df[site_column_to_filter] == dataset.site_id].copy()
+                dataset.data = adjusted_data.reset_index(drop=True)
+            else:
+                dataset.data = harmonized_df.copy() # If no site_indicator or column not found, assign the entire harmonized data (check if this is the desired fallback)
         return mri_datasets
 
 
@@ -469,13 +471,6 @@ def _reintegrate_harmonized_data(self, mri_datasets, harmonized_data, semi_featu
         """
         for i, dataset in enumerate(mri_datasets):
             site_value = dataset.site_id
-            # data_length = mri_datasets[i].data.shape[0]
-            # if i == 0:
-            #     adjusted_data = harmonized_data[:data_length]
-            #     previous_length = data_length
-            # else:
-            #     adjusted_data = harmonized_data[previous_length:previous_length + data_length]
-            #     previous_length = adjusted_data.shape[0]
             adjusted_data = harmonized_data[harmonized_data[self.site_indicator] == site_value].copy() # copy to avoid set on copy
 
             # Drop overlapping columns from semi_features to avoid _x/_y suffixes during merge
@@ -708,11 +703,11 @@ def  _reintegrate_harmonized_data(self, mri_datasets, data_combat, bt, feature_d
             harmonized_datasets.append(harmonized_data)
             start = end
 
-        # Assign harmonized data back to each dataset directly (already split by length above)
+        harmonized_data_concat = pd.concat([_d for _d in harmonized_datasets])
         for i, dataset in enumerate(mri_datasets):
-            adjusted_data = harmonized_datasets[i].copy()
-            # Merge with semi_features to add back any missing columns
-            adjusted_data = pd.merge(adjusted_data, semi_features[i].drop(self.discrete_covariates + self.continuous_covariates + ['index'], errors='ignore'), on=self.patient_identifier, how='left') # Explicit left merge, errors='ignore'
+            site_value = dataset.site_id
+            adjusted_data = harmonized_data_concat[harmonized_data_concat[self.site_indicator] == site_value].copy() # copy to avoid set on copy
+            adjusted_data = pd.merge(adjusted_data, semi_features[i].drop(self.discrete_covariates + self.continuous_covariates + ['index'],axis = 1, errors='ignore'), on=self.patient_identifier, how='left') # Explicit left merge, errors='ignore'
             for _c in ocols:
                 if _c + '_y' in adjusted_data.columns and _c + '_x' in adjusted_data.columns:
                     adjusted_data.drop(columns=[_c+'_y'], inplace=True)
@@ -874,13 +869,9 @@ def _reintegrate_harmonized_data(self, mri_datasets, harmonized_data, covariates
             [harmonized_df, all_data[non_harmonized].reset_index(drop=True)], axis=1,
         )
 
-        # Split data back using positional slicing based on original dataset lengths
-        start_idx = 0
         for mri_dataset in mri_datasets:
-            end_idx = start_idx + len(mri_dataset.data)
-            mri_dataset.data = harmonized_df.iloc[start_idx:end_idx].copy()
+            mri_dataset.data = harmonized_df[harmonized_df["SITE"] == mri_dataset.site_id]
             mri_dataset.data = mri_dataset.data.drop(columns=["SITE", "index"], errors='ignore')
-            start_idx = end_idx
         return mri_datasets
 
     def harmonize(self, mri_datasets):
@@ -1082,13 +1073,10 @@ def _reintegrate_harmonized_data(self, mri_datasets, harmonized_data, original_d
         original_order = list(original_data.columns)
         harmonized_df = harmonized_df[original_order]
 
-        # Split data back using positional slicing based on original dataset lengths
-        start_idx = 0
         for dataset in mri_datasets:
-            end_idx = start_idx + len(dataset.data)
-            adjusted_data = harmonized_df.iloc[start_idx:end_idx].copy()
+            site_value = dataset.site_id
+            adjusted_data = harmonized_df[harmonized_df[self.site_indicator[0]] == site_value].copy()
             dataset.data = adjusted_data.reset_index(drop=True)
-            start_idx = end_idx
         return mri_datasets
 
     def harmonize(self, mri_datasets):

tests/test_basics.py

@@ -21,7 +21,6 @@ def load_datasets(shared_datadir):
         os.path.realpath(shared_datadir / "TrainingData_Site1_fake.csv")
     ]
     # Using unique site_ids to avoid singular matrix issues in neuroharmonize
-    # (the third file is also from site 1 but we assign it site 3 to ensure uniqueness)
     input_sites = [1, 2, 3]
 
     mri_datasets = [
@@ -99,6 +98,7 @@ def test_neuroharmonize(shared_datadir):
             assert feature in dataset.data.columns
 
 
+@pytest.mark.skip(reason="Reverting breaking changes in CovBat harmonizer for now")
 def test_covbat(shared_datadir):
     """Test whether the CovBat harmonizer runs."""
     datasets = load_datasets(shared_datadir)