gpbench 1.0.0__py3-none-any.whl

gpbench/method_class/G2PDeep/base_G2PDeep_class.py

@@ -0,0 +1,209 @@
+import torch
+import torch.nn as nn
+from typing import List, Optional
+import numpy as np
+
+
+class ModelHyperparams:
+    def __init__(self,
+                 left_tower_filters_list: Optional[List[int]] = None,
+                 left_tower_kernel_size_list: Optional[List[int]] = None,
+                 right_tower_filters_list: Optional[List[int]] = None,
+                 right_tower_kernel_size_list: Optional[List[int]] = None,
+                 central_tower_filters_list: Optional[List[int]] = None,
+                 central_tower_kernel_size_list: Optional[List[int]] = None,
+                 dnn_size_list: Optional[List[int]] = None,
+                 activation: str = "linear",
+                 dropout_rate: float = 0.75):
+        self.left_tower_filters_list = left_tower_filters_list or [4, 4]
+        self.left_tower_kernel_size_list = left_tower_kernel_size_list or [3, 5]
+        self.right_tower_filters_list = right_tower_filters_list or [4]
+        self.right_tower_kernel_size_list = right_tower_kernel_size_list or [3]
+        self.central_tower_filters_list = central_tower_filters_list or [4]
+        self.central_tower_kernel_size_list = central_tower_kernel_size_list or [3]
+        self.dnn_size_list = dnn_size_list or [32] 
+        self.activation = activation
+        self.dropout_rate = dropout_rate
+
+def get_activation(name: str):
+    if name.lower() == "relu":
+        return nn.ReLU()
+    elif name.lower() == "linear":
+        return nn.Identity()
+    else:
+        raise ValueError(f"Unsupported activation: {name}")
+    
+
+class G2PDeep(nn.Module):
+    def __init__(self, nsnp: int, num_classes: int, hyperparams: ModelHyperparams):
+        super().__init__()
+        self.nsnp = nsnp
+        self.num_classes = num_classes
+        hp = hyperparams
+
+        # --- Left Tower ---
+        self.left_convs = nn.ModuleList()
+        in_ch = 4
+        for filt, k in zip(hp.left_tower_filters_list, hp.left_tower_kernel_size_list):
+            self.left_convs.append(nn.Conv1d(in_ch, filt, k, padding="same"))
+            in_ch = filt
+
+        # --- Right Tower ---
+        self.right_convs = nn.ModuleList()
+        in_ch = 4
+        for filt, k in zip(hp.right_tower_filters_list, hp.right_tower_kernel_size_list):
+            self.right_convs.append(nn.Conv1d(in_ch, filt, k, padding="same"))
+            in_ch = filt
+
+        # --- Channel alignment ---
+        left_out_ch = hp.left_tower_filters_list[-1]
+        right_out_ch = hp.right_tower_filters_list[-1]
+        self.merged_ch = max(left_out_ch, right_out_ch)
+
+        self.left_proj = nn.Conv1d(left_out_ch, self.merged_ch, 1) \
+            if left_out_ch != self.merged_ch else nn.Identity()
+        self.right_proj = nn.Conv1d(right_out_ch, self.merged_ch, 1) \
+            if right_out_ch != self.merged_ch else nn.Identity()
+
+        # --- Central Tower ---
+        self.central_convs = nn.ModuleList()
+        in_ch = self.merged_ch
+        for filt, k in zip(hp.central_tower_filters_list, hp.central_tower_kernel_size_list):
+            self.central_convs.append(nn.Conv1d(in_ch, filt, k, padding="same"))
+            in_ch = filt
+
+        # --DNN ---
+        self.dropout = nn.Dropout(p=hp.dropout_rate)
+        final_conv_ch = hp.central_tower_filters_list[-1]
+        flattened_dim = final_conv_ch * nsnp
+
+        dnn_layers = []
+        prev = flattened_dim
+        for out_sz in hp.dnn_size_list:
+            dnn_layers.append(nn.Linear(prev, out_sz))
+            dnn_layers.append(get_activation(hp.activation))
+            dnn_layers.append(nn.Dropout(hp.dropout_rate))
+            prev = out_sz
+
+        dnn_layers.append(nn.Linear(prev, num_classes))
+        self.dnn = nn.Sequential(*dnn_layers)
+
+        self.activation = get_activation(hp.activation)
+
+    def forward(self, x):
+        # (B, Seq, 4) -> (B, 4, Seq)
+        if x.shape[-1] != 4:
+            raise ValueError(f"Expected input with 4 channels, got {x.shape}")
+
+        x = x.transpose(1, 2)
+
+        # Left tower
+        left = x
+        for conv in self.left_convs:
+            left = self.activation(conv(left))
+
+        # Right tower
+        right = x
+        for conv in self.right_convs:
+            right = self.activation(conv(right))
+
+        merged = self.left_proj(left) + self.right_proj(right)
+
+        # Central tower
+        for conv in self.central_convs:
+            merged = self.activation(conv(merged))
+
+        x_flat = torch.flatten(merged, 1)
+        x_flat = self.dropout(x_flat)
+        return self.dnn(x_flat)  # (B, num_classes) logits
+
+    def train_model(self, train_loader, valid_loader, num_epochs, learning_rate, patience, device):
+        optimizer = torch.optim.Adam(self.parameters(), lr=learning_rate, weight_decay=1e-4)
+        criterion = nn.CrossEntropyLoss()
+        self.to(device)
+        
+        use_amp = device.type == 'cuda'
+        scaler = torch.amp.GradScaler('cuda') if use_amp else None
+
+        best_loss = float('inf')
+        best_state = None
+        trigger_times = 0
+
+        for epoch in range(num_epochs):
+            self.train()
+            train_loss = 0.0
+            for inputs, labels in train_loader:
+                inputs = inputs.to(device, non_blocking=True)
+                labels = labels.to(device, non_blocking=True)  # (B,) long tensor
+                
+                optimizer.zero_grad()
+                
+                if use_amp:
+                    with torch.amp.autocast('cuda'):
+                        outputs = self(inputs)  # (B, num_classes)
+                        loss = criterion(outputs, labels)
+                    scaler.scale(loss).backward()
+                    scaler.step(optimizer)
+                    scaler.update()
+                else:
+                    outputs = self(inputs)
+                    loss = criterion(outputs, labels)
+                    loss.backward()
+                    optimizer.step()
+                
+                train_loss += loss.item() * inputs.size(0)
+            train_loss /= len(train_loader.dataset)
+
+            # 验证
+            self.eval()
+            valid_loss = 0.0
+            with torch.no_grad():
+                for inputs, labels in valid_loader:
+                    inputs = inputs.to(device, non_blocking=True)
+                    labels = labels.to(device, non_blocking=True)
+                    
+                    if use_amp:
+                        with torch.amp.autocast('cuda'):
+                            outputs = self(inputs)
+                            loss = criterion(outputs, labels)
+                    else:
+                        outputs = self(inputs)
+                        loss = criterion(outputs, labels)
+                    
+                    valid_loss += loss.item() * inputs.size(0)
+            valid_loss /= len(valid_loader.dataset)
+
+            # Early stopping
+            if valid_loss < best_loss:
+                best_loss = valid_loss
+                best_state = {k: v.cpu().clone() for k, v in self.state_dict().items()}
+                trigger_times = 0
+            else:
+                trigger_times += 1
+                if trigger_times >= patience:
+                    print(f"Early stopping at epoch {epoch+1}")
+                    break
+
+        if best_state is not None:
+            cur_device = next(self.parameters()).device
+            best_state = {k: v.to(cur_device) for k, v in best_state.items()}
+            self.load_state_dict(best_state)
+        return best_loss
+
+    def predict(self, test_loader, device):
+        self.eval()
+        self.to(device)
+        y_pred_list = []
+        use_amp = device.type == 'cuda'
+        with torch.no_grad():
+            for inputs, _ in test_loader:
+                inputs = inputs.to(device, non_blocking=True)
+                if use_amp:
+                    with torch.amp.autocast('cuda'):
+                        outputs = self(inputs)  # (B, num_classes)
+                else:
+                    outputs = self(inputs)
+                preds = torch.argmax(outputs, dim=1)  # (B,)
+                y_pred_list.append(preds.cpu())
+        y_pred = torch.cat(y_pred_list, dim=0).numpy()
+        return y_pred

gpbench/method_class/GBLUP/GBLUP_class.py

@@ -0,0 +1,183 @@
+import os
+import time
+import psutil
+import swanlab
+import argparse
+import random
+import torch
+import numpy as np
+import pandas as pd
+from sklearn.model_selection import StratifiedKFold
+from sklearn.preprocessing import LabelEncoder
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support, confusion_matrix
+
+import rpy2.robjects as ro
+from rpy2.robjects import numpy2ri
+numpy2ri.activate()
+ro.r('library(rrBLUP)')
+
+
+def gblup_classification(X_train, y_train_bin, X_test):
+
+    # Pass data to R
+    ro.globalenv['X_train'] = X_train
+    ro.globalenv['y_train_bin'] = y_train_bin
+    ro.globalenv['X_test'] = X_test
+
+    r_code = """
+    library(rrBLUP)
+    
+    n_train <- nrow(X_train)
+    m <- ncol(X_train)
+
+    # Step1: allele frequencies
+    p <- colMeans(X_train) / 2
+    p <- pmax(pmin(p, 0.99), 0.01)
+
+    # Step2: VanRaden standardized genotype
+    Z_train <- sweep(X_train, 2, 2*p, "-") / sqrt(2*p*(1-p))
+    Z_train[is.na(Z_train)] <- 0
+
+    Z_test <- sweep(X_test, 2, 2*p, "-") / sqrt(2*p*(1-p))
+    Z_test[is.na(Z_test)] <- 0
+
+    # Step3: Genomic relationship matrix (VanRaden method 2)
+    denom <- sum(2*p*(1-p))
+    G <- Z_train %*% t(Z_train) / denom
+    G <- G + diag(1e-6, n_train)  # stability
+
+    # Step4: REML GBLUP
+    fit <- mixed.solve(y = y_train_bin, K = G, SE = FALSE)
+
+    # Extract variance components and fixed effect
+    Vu <- fit$Vu
+    Ve <- fit$Ve
+    mu <- as.numeric(fit$beta)  # intercept
+    h2 <- Vu / (Vu + Ve)
+
+    # Step5: GBLUP prediction for test set
+    y_centered <- y_train_bin - mu
+    A <- G + (Ve / Vu) * diag(n_train)  # G + λ I
+
+    G_test_train <- Z_test %*% t(Z_train) / denom
+    u_test <- G_test_train %*% solve(A, y_centered)  # strictly correct formula
+
+    y_pred_score <- mu + u_test
+    y_pred_score
+    """
+
+    y_pred_score = np.array(ro.r(r_code)).flatten()
+    return y_pred_score
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Argument parser")
+    parser.add_argument('--methods', type=str, default='GBLUP/', help='Method name') 
+    parser.add_argument('--species', type=str, default='')
+    parser.add_argument('--phe', type=str, default='', help='Dataset name')
+    parser.add_argument('--data_dir', type=str, default='../../data/')
+    parser.add_argument('--result_dir', type=str, default='result/')
+    return parser.parse_args()
+
+
+def load_data(args):
+    xData = np.load(os.path.join(args.data_dir, args.species, 'genotype.npz'))["arr_0"]
+    yData = np.load(os.path.join(args.data_dir, args.species, 'phenotype.npz'))["arr_0"]
+    names = np.load(os.path.join(args.data_dir, args.species, 'phenotype.npz'))["arr_1"]
+    nsample = xData.shape[0]
+    nsnp = xData.shape[1]
+    print("Number of samples: ", nsample)
+    print("Number of SNPs: ", nsnp)
+    return xData, yData, nsample, nsnp, names
+
+
+def set_seed(seed=42):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def run_nested_cv(args, data, label, process):
+    result_dir = os.path.join(args.result_dir, args.methods + args.species)
+    os.makedirs(result_dir, exist_ok=True)
+    print("Starting 10-fold cross-validation (GBLUP Classification with R)...")
+
+    skf = StratifiedKFold(n_splits=10, shuffle=True, random_state=42)
+    le = LabelEncoder()
+    label_all = le.fit_transform(label)
+    np.save(os.path.join(result_dir, 'label_mapping.npy'), le.classes_)
+
+    all_acc, all_prec, all_rec, all_f1 = [], [], [], []
+
+    for fold, (train_idx, test_idx) in enumerate(skf.split(data, label_all)):
+        fold_start = time.time()
+        print(f"===== Fold {fold} =====")
+        X_train, X_test = data[train_idx], data[test_idx]
+        Y_train, Y_test = label_all[train_idx], label_all[test_idx]
+
+        if torch.cuda.is_available():
+            torch.cuda.reset_peak_memory_stats()
+
+        classes = np.unique(Y_train)
+        scores = np.zeros((len(classes), X_test.shape[0]))
+        for idx, cls in enumerate(classes):
+            y_train_bin = (Y_train == cls).astype(float)
+            scores[idx, :] = gblup_classification(X_train, y_train_bin, X_test)
+
+        Y_pred = np.argmax(scores, axis=0)
+
+        acc = accuracy_score(Y_test, Y_pred)
+        prec, rec, f1, _ = precision_recall_fscore_support(Y_test, Y_pred, average='macro', zero_division=0)
+        all_acc.append(acc)
+        all_prec.append(prec)
+        all_rec.append(rec)
+        all_f1.append(f1)
+
+        fold_time = time.time() - fold_start
+        fold_gpu_mem = torch.cuda.max_memory_allocated() / 1024**2 if torch.cuda.is_available() else 0
+        fold_cpu_mem = process.memory_info().rss / 1024**2
+        print(f'Fold {fold}: ACC={acc:.4f}, PREC={prec:.4f}, REC={rec:.4f}, F1={f1:.4f}, '
+              f'Time={fold_time:.2f}s, GPU={fold_gpu_mem:.2f}MB, CPU={fold_cpu_mem:.2f}MB')
+
+        Y_test_orig = le.inverse_transform(Y_test)
+        Y_pred_orig = le.inverse_transform(Y_pred)
+        results_df = pd.DataFrame({'Y_test': Y_test_orig, 'Y_pred': Y_pred_orig})
+        results_df.to_csv(os.path.join(result_dir, f"fold{fold}.csv"), index=False)
+
+    print("\n===== Cross-validation summary =====")
+    print(f"Average ACC: {np.mean(all_acc):.4f} ± {np.std(all_acc):.4f}")
+    print(f"Average PREC: {np.mean(all_prec):.4f} ± {np.std(all_prec):.4f}")
+    print(f"Average REC: {np.mean(all_rec):.4f} ± {np.std(all_rec):.4f}")
+    print(f"Average F1 : {np.mean(all_f1):.4f} ± {np.std(all_f1):.4f}")
+
+
+def GBLUP_class():
+    set_seed(42)
+    torch.cuda.empty_cache()
+    args = parse_args()
+    process = psutil.Process(os.getpid())
+
+    all_species = ["Human/Sim/"]
+    for sp in all_species:
+        args.species = sp
+        X, Y, nsamples, nsnp, names = load_data(args)
+        print("Starting run " + args.methods + args.species)
+        label = Y[:, 0]
+        s = pd.Series(label)
+        fill_val = s.mode().iloc[0] if not s.dropna().empty else 0
+        label = np.nan_to_num(label, nan=fill_val)
+
+        start_time = time.time()
+        torch.cuda.reset_peak_memory_stats()
+        run_nested_cv(args, data=X, label=label, process=process)
+
+        elapsed_time = time.time() - start_time
+        print(f"Total running time: {elapsed_time:.2f} s")
+        print("Successfully finished!")
+
+
+if __name__ == "__main__":
+    GBLUP_class()

gpbench/method_class/GBLUP/__init__.py

@@ -0,0 +1,5 @@
+from .GBLUP_class import GBLUP_class
+
+GBLUP = GBLUP_class
+
+__all__ = ["GBLUP","GBLUP_class"]

gpbench/method_class/GEFormer/GEFormer_class.py

@@ -0,0 +1,169 @@
+import os
+import torch
+import swanlab
+import argparse
+import psutil
+import time
+import random
+import numpy as np
+import pandas as pd
+import pynvml
+from . import GEFormer_he_class
+
+from .gMLP_class import GEFormer
+from sklearn.metrics import accuracy_score, precision_recall_fscore_support
+from sklearn.preprocessing import LabelEncoder
+from torch.utils.data import DataLoader, TensorDataset
+from sklearn.model_selection import StratifiedKFold, train_test_split
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--methods', type=str, default='GEFormer/')
+    parser.add_argument('--species', type=str, default='')
+    parser.add_argument('--phe', type=str, default='')
+    parser.add_argument('--data_dir', type=str, default='../../data/')
+    parser.add_argument('--result_dir', type=str, default='result/')
+    
+    parser.add_argument('--epoch', type=int, default=1000)
+    parser.add_argument('--batch_size', type=int, default=64)
+    parser.add_argument('--learning_rate', type=float, default=0.01)
+    parser.add_argument('--patience', type=int, default=10)
+    return parser.parse_args()
+
+def set_seed(seed=42):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+
+def load_data(args):
+    x = np.load(os.path.join(args.data_dir, args.species, 'genotype.npz'))["arr_0"]
+    y = np.load(os.path.join(args.data_dir, args.species, 'phenotype.npz'))["arr_0"]
+    names = np.load(os.path.join(args.data_dir, args.species, 'phenotype.npz'))["arr_1"]
+    return x, y, x.shape[0], x.shape[1], names
+
+
+def get_gpu_mem_by_pid(pid):
+    procs = pynvml.nvmlDeviceGetComputeRunningProcesses(handle)
+    for p in procs:
+        if p.pid == pid:
+            return p.usedGpuMemory / 1024**2
+    return 0.0
+
+def run_nested_cv(args, data, label, nsnp, device):
+    result_dir = os.path.join(args.result_dir, args.methods + args.species + args.phe)
+    os.makedirs(result_dir, exist_ok=True)
+
+    kf = StratifiedKFold(n_splits=10, shuffle=True, random_state=42)
+
+    all_acc, all_prec, all_rec, all_f1 = [], [], [], []
+    total_start = time.time()
+
+    num_classes = len(np.unique(label))
+
+    for fold, (train_idx, test_idx) in enumerate(kf.split(data, label)):
+        fold_start = time.time()
+        process = psutil.Process(os.getpid())
+
+        X_train, X_test = data[train_idx], data[test_idx]
+        y_train, y_test = label[train_idx], label[test_idx]
+
+        X_train_sub, X_valid, y_train_sub, y_valid = train_test_split(
+            X_train, y_train, test_size=0.1, stratify=y_train, random_state=42
+        )
+
+        x_train = torch.from_numpy(X_train_sub).float().to(device)
+        y_train = torch.from_numpy(y_train_sub).long().to(device)
+        x_valid = torch.from_numpy(X_valid).float().to(device)
+        y_valid = torch.from_numpy(y_valid).long().to(device)
+        x_test  = torch.from_numpy(X_test).float().to(device)
+        y_test_tensor = torch.from_numpy(y_test).long().to(device)
+
+        train_loader = DataLoader(TensorDataset(x_train, y_train), args.batch_size, shuffle=True)
+        valid_loader = DataLoader(TensorDataset(x_valid, y_valid), args.batch_size, shuffle=False)
+        test_loader  = DataLoader(TensorDataset(x_test, y_test_tensor), args.batch_size, shuffle=False)
+
+        model = GEFormer(nsnp=nsnp, num_classes=num_classes).to(device)
+        model.train_model(
+            train_loader, valid_loader,
+            args.epoch, args.learning_rate, args.patience, device
+        )
+
+        logits = model.predict(test_loader)
+        y_pred = np.argmax(logits, axis=1)
+
+        acc = accuracy_score(y_test, y_pred)
+        prec, rec, f1, _ = precision_recall_fscore_support(
+            y_test, y_pred, average='macro', zero_division=0
+        )
+
+        all_acc.append(acc)
+        all_prec.append(prec)
+        all_rec.append(rec)
+        all_f1.append(f1)
+
+        fold_time = time.time() - fold_start
+        gpu_mem = get_gpu_mem_by_pid(os.getpid())
+        cpu_mem = process.memory_info().rss / 1024**2
+        
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.reset_peak_memory_stats()
+
+        print(
+            f"Fold {fold}: ACC={acc:.4f}, PREC={prec:.4f}, REC={rec:.4f}, "
+            f"F1={f1:.4f}, Time={fold_time:.2f}s, "
+            f"GPU={gpu_mem:.2f}MB, CPU={cpu_mem:.2f}MB"
+        )
+
+        pd.DataFrame({"Y_test": y_test, "Y_pred": y_pred}).to_csv(
+            os.path.join(result_dir, f"fold{fold}.csv"), index=False
+        )
+
+    total_time = time.time() - total_start
+    print("\n===== CV Summary =====")
+    print(f"ACC : {np.mean(all_acc):.4f} ± {np.std(all_acc):.4f}")
+    print(f"PREC: {np.mean(all_prec):.4f} ± {np.std(all_prec):.4f}")
+    print(f"REC : {np.mean(all_rec):.4f} ± {np.std(all_rec):.4f}")
+    print(f"F1  : {np.mean(all_f1):.4f} ± {np.std(all_f1):.4f}")
+    print(f"Time: {total_time:.2f}s")
+
+
+def GEFormer_class():
+    set_seed(42)
+    pynvml.nvmlInit()
+    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+
+    args = parse_args()
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+    all_species =  ["Human/Sim/"]
+
+    for species in all_species:
+        args.species = species
+        X, Y, _, nsnp, _ = load_data(args)
+        label_raw = np.nan_to_num(Y[:, 0])
+        le = LabelEncoder()
+        label = le.fit_transform(label_raw)
+        num_classes = len(le.classes_)
+
+        best_params = GEFormer_he_class.Hyperparameter(X, label, nsnp)
+        args.learning_rate = best_params['learning_rate']
+        args.batch_size = best_params['batch_size']
+        args.patience = best_params['patience']
+        start_time = time.time() 
+        if torch.cuda.is_available():
+            torch.cuda.reset_peak_memory_stats()
+        process = psutil.Process(os.getpid())
+        run_nested_cv(args, X, label, nsnp, device)
+
+        elapsed_time = time.time() - start_time
+        print(f"Running time: {elapsed_time:.2f}s")
+        print("Successfully finished:", args.species, args.phe)
+
+
+if __name__ == "__main__":
+    GEFormer_class()