Spaces:
Sleeping
Sleeping
File size: 4,707 Bytes
f018cee | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | import pandas as pd
import numpy as np
import torch
from torch_geometric.data import Data
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.impute import SimpleImputer
import os
def load_and_preprocess_data(data_dir='data'):
# 1. Load Data
clinical = pd.read_csv(os.path.join(data_dir, 'clinical.csv'))
blood = pd.read_csv(os.path.join(data_dir, 'blood.csv'))
pathological = pd.read_csv(os.path.join(data_dir, 'pathological.csv'))
targets = pd.read_csv(os.path.join(data_dir, 'targets.csv'))
# 2. Merge Data
df = clinical.merge(blood, on='patient_id', how='inner')
df = df.merge(pathological, on='patient_id', how='inner')
df = df.merge(targets, on='patient_id', how='inner')
# 3. Target Definitions
# Task A: Recurrence (0: no, 1: yes)
df['recurrence_label'] = df['recurrence'].map({'no': 0, 'yes': 1}).fillna(0).astype(int)
# Task B: Survival Score (Normalized)
# Strategy: (days / max_days) * (1.2 if living else 0.8)
max_days = df['days_to_last_information'].max()
df['survival_score'] = (df['days_to_last_information'] / max_days)
df.loc[df['survival_status'] == 'dead', 'survival_score'] *= 0.8
df.loc[df['survival_status'] == 'living', 'survival_score'] *= 1.2
# Clip to [0, 1]
df['survival_score'] = df['survival_score'].clip(0, 1)
# 4. Preprocessing
# Identify column types
blood_cols = [c for c in blood.columns if c != 'patient_id']
pathology_cols = [c for c in pathological.columns if c != 'patient_id']
clinical_cols = [c for c in clinical.columns if c != 'patient_id']
# Impute Blood (Median)
blood_imputer = SimpleImputer(strategy='median')
df[blood_cols] = blood_imputer.fit_transform(df[blood_cols])
# Impute Pathology (Mode)
path_imputer = SimpleImputer(strategy='most_frequent')
df[pathology_cols] = path_imputer.fit_transform(df[pathology_cols])
# Feature Engineering / Encoding
# For categorical columns that are not IDs or targets
cat_cols = ['sex', 'primarily_metastasis', 'smoking_status', 'hpv_association_p16']
for col in cat_cols:
if col in df.columns:
df[col] = LabelEncoder().fit_transform(df[col].astype(str))
# Scale Numerical Features
num_cols = blood_cols + clinical_cols + ['infiltration_depth_in_mm']
num_cols = [c for c in num_cols if c in df.columns and c not in cat_cols]
scaler = StandardScaler()
df[num_cols] = scaler.fit_transform(df[num_cols])
# 5. Graph Construction (The Clinical Twin)
# Nodes: patients in df
node_features = df[num_cols + cat_cols].values
x = torch.tensor(node_features, dtype=torch.float)
# Edges: Same primary_tumor_site AND pT_stage
edge_index = []
edge_attr = []
patient_indices = df.index.tolist()
for i in range(len(patient_indices)):
for j in range(i + 1, len(patient_indices)):
p1 = df.iloc[i]
p2 = df.iloc[j]
if p1['primary_tumor_site'] == p2['primary_tumor_site'] and p1['pT_stage'] == p2['pT_stage']:
# Similarity Weight based on grading and pN_stage
# We'll use 1 / (1 + abs(diff))
grading_sim = 1.0 / (1.0 + abs(float(p1['grading']) - float(p2['grading'])))
pn_sim = 1.0 / (1.0 + abs(float(p1['pN_stage']) - float(p2['pN_stage'])))
weight = (grading_sim + pn_sim) / 2.0
edge_index.append([i, j])
edge_index.append([j, i])
edge_attr.append([weight])
edge_attr.append([weight])
if not edge_index:
# Fallback: connect everyone with a very low weight if no twins found
# Or just empty edges for now (PyG handles this)
edge_index = torch.empty((2, 0), dtype=torch.long)
edge_attr = torch.empty((0, 1), dtype=torch.float)
else:
edge_index = torch.tensor(edge_index, dtype=torch.long).t().contiguous()
edge_attr = torch.tensor(edge_attr, dtype=torch.float)
# Targets
y_recurrence = torch.tensor(df['recurrence_label'].values, dtype=torch.long)
y_survival = torch.tensor(df['survival_score'].values, dtype=torch.float).view(-1, 1)
data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr, y_recurrence=y_recurrence, y_survival=y_survival)
return data, df, num_cols + cat_cols
if __name__ == "__main__":
data, df, features = load_and_preprocess_data('../data')
print(f"Graph constructed with {data.num_nodes} nodes and {data.num_edges} edges.")
print(f"Features: {len(features)}")
print(f"Recurrence balance: {df['recurrence_label'].value_counts().to_dict()}")
|