bmi in datasets

tjiagoM · May 18, 2020 · b377ade · b377ade
1 parent b18a0bc
commit b377ade
Show file tree

Hide file tree

Showing 3 changed files with 178 additions and 92 deletions.
diff --git a/datasets.py b/datasets.py
@@ -15,8 +15,8 @@
 from torch_geometric.data import InMemoryDataset, Data
 
 from utils import Normalisation, ConnType, AnalysisType, EncodingStrategy, DatasetType
-from utils_datasets import OLD_NETMATS_PEOPLE, DESIKAN_COMPLETE_TS, DESIKAN_TRACKS, UKB_IDS_PATH, UKB_PHENOTYPE_PATH, \
-    UKB_TIMESERIES_PATH, NODE_FEATURES_NAMES, STRUCT_COLUMNS
+from utils_datasets import DESIKAN_COMPLETE_TS, DESIKAN_TRACKS, UKB_IDS_PATH, UKB_PHENOTYPE_PATH, \
+    UKB_TIMESERIES_PATH, NODE_FEATURES_NAMES, STRUCT_COLUMNS, UKB_WITHOUT_BMI
 
 PEOPLE_DEMOGRAPHICS_PATH = 'meta_data/people_demographics.csv'
 
@@ -284,7 +284,7 @@ def __init__(self, root, target_var: str, num_nodes: int, threshold: int, connec
                  encoding_strategy: EncodingStrategy = EncodingStrategy.NONE,
                  transform=None, pre_transform=None):
 
-        if target_var not in ['gender', 'age']:
+        if target_var not in ['gender', 'age', 'bmi']:
             print("UKBDataset not prepared for that target_var!")
             exit(-2)
         if connectivity_type not in [ConnType.FMRI]:
@@ -323,17 +323,23 @@ def __create_data_object(self, person: int, ts: np.ndarray, covars: pd.DataFrame
 
         if self.target_var == 'gender':
             y = torch.tensor([covars.loc[person, 'Sex']], dtype=torch.float)
+        elif self.target_var == 'bmi':
+            y = torch.tensor([covars.loc[person, 'BMI.at.scan']], dtype=torch.float)
         else:
             y = torch.tensor([covars.loc[person, 'Age.at.scan']], dtype=torch.float)
 
         data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, y=y)
         data.ukb_id = torch.tensor([person])
-        data.bmi = torch.tensor([covars.loc[person, 'BMI.at.scan']])
-        data.age = torch.tensor([covars.loc[person, 'Age.at.scan']])
+
         if self.target_var == 'gender':
             data.age = torch.tensor([covars.loc[person, 'Age.at.scan']])
+            data.bmi = torch.tensor([covars.loc[person, 'BMI.at.scan']])
+        elif self.target_var == 'bmi':
+            data.sex = torch.tensor([covars.loc[person, 'Sex']])
+            data.age = torch.tensor([covars.loc[person, 'Age.at.scan']])
         else:
             data.sex = torch.tensor([covars.loc[person, 'Sex']])
+            data.bmi = torch.tensor([covars.loc[person, 'BMI.at.scan']])
 
         return data
 
@@ -352,6 +358,8 @@ def process(self):
             if person in [1663368, 3443644]:
                 # No information in Covars file
                 continue
+            if self.target_var == 'bmi' and person in UKB_WITHOUT_BMI:
+                continue
 
             if self.connectivity_type == ConnType.FMRI:
                 ts = np.loadtxt(f'{UKB_TIMESERIES_PATH}/UKB{person}_ts_raw.txt', delimiter=',')

diff --git a/main_loop.py b/main_loop.py
@@ -705,9 +705,6 @@ def send_global_results(test_metrics: Dict[str, float]):
                                       label_scaler=scaler_labels)
         print(test_metrics)
 
-        print('{:1d}-Final: {:.7f}, Auc: {:.4f}, Acc: {:.4f}, Sens: {:.4f}, Speci: {:.4f}'
-              ''.format(outer_split_num, test_metrics['loss'], test_metrics['auc'], test_metrics['acc'],
-                        test_metrics['sensitivity'], test_metrics['specificity']))
         if scaler_labels is None:
             print('{:1d}-Final: {:.7f}, Auc: {:.4f}, Acc: {:.4f}, Sens: {:.4f}, Speci: {:.4f}'
                   ''.format(outer_split_num, test_metrics['loss'], test_metrics['auc'], test_metrics['acc'],