dev: add first GP scripts

Author: niklases

pypef/gaussian_process/composite.py

@@ -0,0 +1,157 @@
+from typing import Literal, Tuple
+from sklearn.model_selection import train_test_split
+import torch
+
+from gpytorch.kernels import ScaleKernel
+import gpytorch
+import pandas as pd
+from tqdm import tqdm
+
+from gp_esm2_test import extract_esm_embeddings
+from gp_pmpnn_test import HellingerRBFKernel, get_probs_from_mutations
+from gp_prosst_test import (extract_prosst_embeddings, get_prosst_models, 
+                            get_structure_quantizied, read_fasta_biopython)
+from metrics import spearman_soft, spearman_corr_differentiable, spearmanr2
+
+class CombinedKernel(gpytorch.kernels.Kernel):
+    """
+    Combine two kernels: K_seq + K_struct
+    Input X is a single concatenated tensor: [seq | struct]
+    """
+
+    def __init__(self, kernel_seq, kernel_struct, d_seq):
+        super().__init__()
+        self.kernel_seq = kernel_seq
+        self.kernel_struct = kernel_struct
+        self.d_seq = d_seq  # number of sequence dimensions
+
+    def forward(self, X1, X2, **params):
+        X1_seq, X1_struct = X1[:, :self.d_seq], X1[:, self.d_seq:]
+        X2_seq, X2_struct = X2[:, :self.d_seq], X2[:, self.d_seq:]
+
+        K_seq = self.kernel_seq(X1_seq, X2_seq)
+        K_struct = self.kernel_struct(X1_struct, X2_struct)
+
+        return K_seq + K_struct  # could also use product or weighted sum
+
+
+class MultiInputGP(gpytorch.models.ExactGP):
+    def __init__(self, train_x, train_y, likelihood, kernel):
+        super().__init__(train_x, train_y, likelihood)
+        self.mean_module = gpytorch.means.ZeroMean()
+        self.covar_module = kernel
+
+    def forward(self, X):
+        mean_x = self.mean_module(X)
+        covar_x = self.covar_module(X, X)
+        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
+
+
+
+
+
+
+# -----------------------------
+# Load and preprocess data
+# -----------------------------
+df = pd.read_csv('example_data/blat_ecolx/BLAT_ECOLX_Stiffler_2015.csv')
+
+print(df.columns)
+mutants = df['mutant'].to_list()
+sequences = df['mutated_sequence'].to_list()
+y = df['DMS_score'].to_list()
+
+m_train, m_test, s_train, s_test, y_train, y_test = train_test_split(
+    mutants, sequences, y, test_size=0.33, random_state=42
+)
+
+X_struct = get_probs_from_mutations(m_train)        # [N, 20]
+
+
+print("Getting ProSST models")
+pdb = 'example_data/blat_ecolx/BLAT_ECOLX.pdb'
+wt_seq = list(read_fasta_biopython('example_data/blat_ecolx/blat_ecolx_wt_seq.fa').values())[0]
+prosst_base_model, prosst_lora_model, prosst_tokenizer, prosst_optimizer = get_prosst_models()
+prosst_vocab = prosst_tokenizer.get_vocab()
+prosst_base_model = prosst_base_model.to("cuda")
+
+input_ids, prosst_attention_mask, structure_input_ids = get_structure_quantizied(
+        pdb, prosst_tokenizer, wt_seq, verbose=True
+)
+wt_structure_input_ids = structure_input_ids[0, 1:-1].tolist()  # Remove CLS/EOS
+#X_seq = torch.tensor(extract_esm_embeddings(s_train)).float()  # [N, d_seq]
+X_seq = torch.tensor(extract_prosst_embeddings(
+    prosst_base_model, prosst_tokenizer, s_train, wt_structure_input_ids
+))
+y_train = torch.tensor(y_train).float()
+y_test = torch.tensor(y_test).float()
+
+# Concatenate features
+X_combined = torch.cat([X_seq, X_struct], dim=-1)  # Concenation is necessary as GPkernel does not accept a tuple as input 
+d_seq = X_seq.shape[1]
+
+# -----------------------------
+# Define kernels and model
+# -----------------------------
+seq_kernel = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())
+struct_kernel = HellingerRBFKernel()
+combined_kernel = CombinedKernel(seq_kernel, struct_kernel, d_seq=d_seq)
+
+likelihood = gpytorch.likelihoods.GaussianLikelihood()
+model = MultiInputGP(X_combined, y_train, likelihood, combined_kernel)
+
+# -----------------------------
+# Train
+# -----------------------------
+model.train()
+likelihood.train()
+
+optimizer = torch.optim.Adam(model.parameters(), lr=0.05)
+mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)
+
+pbar = tqdm(range(100), desc='Training')
+for i in pbar:
+    optimizer.zero_grad()
+    output = model(X_combined)
+    loss = -mll(output, y_train)
+    loss.backward()
+    optimizer.step()
+    pbar.set_description(f"Training (loss: {loss:.4f})")
+
+# -----------------------------
+# Test
+# -----------------------------
+X_struct_test = get_probs_from_mutations(m_test)
+#X_seq_test = torch.tensor(extract_esm_embeddings(s_test)).float()
+X_seq_test = torch.tensor(extract_prosst_embeddings(prosst_base_model, prosst_tokenizer, s_test, wt_structure_input_ids))
+X_test_combined = torch.cat([X_seq_test, X_struct_test], dim=-1)
+
+model.eval()
+likelihood.eval()
+
+
+with torch.no_grad(), gpytorch.settings.fast_pred_var():
+    pred_train = likelihood(model(X_combined))
+    y_pred_train = pred_train.mean.cpu().numpy()
+
+    pred = likelihood(model(X_test_combined))
+    y_pred = pred.mean.cpu().numpy()
+
+
+from scipy.stats import spearmanr
+
+rho, p = spearmanr(y_train, y_pred_train)
+print("Spearman rho SciPy           TRAIN:", rho)
+print("Spearman soft                TRAIN:", spearman_soft(y_train, torch.from_numpy(y_pred_train)).item())
+y_train_t  = y_train.float().unsqueeze(0)       # shape (1, n)
+y_pred_train_t  = torch.from_numpy(y_pred_train).float().unsqueeze(0)    # shape (1, n)
+print("Spearman corr diff (ChatGPT) TRAIN:", spearman_corr_differentiable(y_train_t, y_pred_train_t).item())
+print("Spearman2 torchsort          TRAIN:", spearmanr2(y_train_t, y_pred_train_t).item())
+
+rho, p = spearmanr(y_test, y_pred)
+print("Spearman rho SciPy           TEST:", rho)
+print("Spearman soft                TEST:", spearman_soft(y_test, torch.from_numpy(y_pred)).item())
+y_test_t  = y_test.float().unsqueeze(0)       # shape (1, n)
+y_pred_t  = torch.from_numpy(y_pred).float().unsqueeze(0)    # shape (1, n)
+print("Spearman corr diff (ChatGPT) TEST:", spearman_corr_differentiable(y_test_t, y_pred_t).item())
+print("Spearman2 torchsort          TEST:", spearmanr2(y_test_t, y_pred_t).item())

pypef/gaussian_process/gp_esm2_test.py

@@ -0,0 +1,149 @@
+import torch
+import numpy as np
+from sklearn.gaussian_process import GaussianProcessRegressor
+from sklearn.gaussian_process.kernels import RBF, WhiteKernel
+from tqdm import tqdm
+import pandas as pd
+from sklearn.model_selection import train_test_split
+from scipy.stats import spearmanr
+import gpytorch
+
+# --- Step 1: Load a pretrained ESM plm_model ---
+from esm import pretrained  # pip install fair-esm
+
+
+"""
+git clone https://github.com/facebookresearch/esm.git
+cd esm
+pip install .
+"""
+
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+USE_SCIKIT_LEARN = False
+
+
+class ExactGPModel(gpytorch.models.ExactGP):
+    def __init__(self, train_x, train_y, likelihood):
+        super().__init__(train_x, train_y, likelihood)
+        self.mean_module = gpytorch.means.ConstantMean()
+        self.covar_module = gpytorch.kernels.ScaleKernel(
+            gpytorch.kernels.RBFKernel()
+        )
+    
+    def forward(self, x):
+        mean_x = self.mean_module(x)
+        covar_x = self.covar_module(x)
+        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
+
+
+def extract_esm_embeddings(sequences):
+    embeddings = []
+
+    for seq in tqdm(sequences, 'Embedding (ESM)'):
+        data = [("protein", seq)]
+        batch_labels, batch_strs, batch_tokens = batch_converter(data)
+        batch_tokens = batch_tokens.to(device)
+        with torch.no_grad():
+            results = plm_model(batch_tokens, repr_layers=[33], return_contacts=False)
+        token_representations = results["representations"][33]
+        # Mean-pool per-residue representations (excluding special tokens)
+        seq_embedding = token_representations[0, 1:len(seq)+1].mean(0)
+        embeddings.append(seq_embedding.cpu().numpy())
+
+    X = np.vstack(embeddings)
+
+    return X
+
+
+plm_model, alphabet = pretrained.esm2_t33_650M_UR50D()
+plm_model = plm_model.to(device)
+batch_converter = alphabet.get_batch_converter()
+plm_model.eval()  # disable dropout
+
+if __name__ == '__main__':
+    # Load ESM-2 (you can choose different sizes: 35M, 150M, 650M, 3B)
+
+    # --- Example dataset ---
+    # sequences: list of amino acid strings
+    # y: list/array of experimental fitness values
+
+    df = pd.read_csv('example_data/blat_ecolx/BLAT_ECOLX_Stiffler_2015.csv')
+    sequences = df['mutated_sequence'].to_list()
+    y = df['DMS_score'].to_list()
+
+    s_train, s_test, y_train, y_test = train_test_split(
+        sequences, y, test_size=0.33, random_state=42)  # train_size=100, test_size=200,
+
+    # --- Step 2: Extract ESM embeddings ---
+    X = extract_esm_embeddings(s_train)
+    print("Embedding extraction done")
+    print(np.shape(X))
+
+    # --- Step 3: Build and fit a Gaussian Process ---
+
+    if USE_SCIKIT_LEARN:
+        # RBF kernel + WhiteKernel (noise term)
+        kernel = 1.0 * RBF(length_scale=1.0) + WhiteKernel(noise_level=0.1)
+        gpr = GaussianProcessRegressor(kernel=kernel, n_restarts_optimizer=10, normalize_y=True)
+        gpr.fit(X, y_train)
+
+    else: # GPYTORCH
+        # Likelihood
+        # Suppose X: [num_sequences, embedding_dim], y: [num_sequences]
+        X = torch.tensor(X, dtype=torch.float32).to(device)
+        y_train = torch.tensor(y_train, dtype=torch.float32).to(device)
+
+        likelihood = gpytorch.likelihoods.GaussianLikelihood().to(device)
+        gp_model = ExactGPModel(X, y_train, likelihood).to(device)
+
+        gp_model.train()
+        likelihood.train()
+
+        optimizer = torch.optim.Adam(gp_model.parameters(), lr=0.1)
+        mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, gp_model)
+
+        training_iter = 100
+        for i in range(training_iter):
+            optimizer.zero_grad()
+            output = gp_model(X)
+            loss = -mll(output, y_train)
+            loss.backward()
+            print(f"Iter {i+1}/{training_iter} - Loss: {loss.item():.3f}")
+            optimizer.step()
+
+    # --- Step 4: Predict on new sequences ---
+    test_embeddings = []
+
+    for seq in tqdm(s_test):
+        data = [("protein", seq)]
+        batch_labels, batch_strs, batch_tokens = batch_converter(data)
+        batch_tokens = batch_tokens.to(device)
+        with torch.no_grad():
+            results = plm_model(batch_tokens, repr_layers=[33], return_contacts=False)
+        seq_embedding = results["representations"][33][0, 1:len(seq)+1].mean(0)
+        test_embeddings.append(seq_embedding.cpu().numpy())
+
+    X_test = np.array(test_embeddings)  # or np.vstack
+    print("Test embeddings shape:", X_test.shape)
+
+    if USE_SCIKIT_LEARN:
+        y_mean, y_std = gpr.predict(X_test, return_std=True)
+    else:  # GPYTORCH
+        X_test = torch.tensor(X_test, dtype=torch.float32).to(device)
+        gp_model.eval()
+        likelihood.eval()
+
+        with torch.no_grad(), gpytorch.settings.fast_pred_var():
+            # Suppose X_test: [num_test, embedding_dim]
+            X_test = torch.tensor(test_embeddings, dtype=torch.float32).to(device)
+            pred = likelihood(gp_model(X_test))
+            y_mean = pred.mean  # predicted mean
+            lower, upper = pred.confidence_region()  # 95% confidence interval
+            y_mean = y_mean.cpu().numpy()
+
+    print("Predicted fitness:", y_mean)
+    #print("Uncertainty (std):", y_std)
+
+    print(spearmanr(y_test, y_mean))