Add combined corr. loss function

Author: niklases

pypef/gaussian_process/composite.py

@@ -11,7 +11,8 @@
 from pypef.gaussian_process.gp_pmpnn_test import HellingerRBFKernel, get_probs_from_mutations
 from pypef.gaussian_process.gp_prosst_test import (extract_prosst_embeddings, get_prosst_models, 
                             get_structure_quantizied, read_fasta_biopython)
-from pypef.gaussian_process.metrics import spearman_soft, spearman_corr_differentiable, spearmanr2
+from pypef.plm.utils import spearman_soft, correlation_loss, hybrid_corr_mse_loss, pearson_loss
+
 
 class CombinedKernel(gpytorch.kernels.Kernel):
     """
@@ -47,30 +48,26 @@ def forward(self, 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')
+df = pd.read_csv('datasets/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
+    mutants, sequences, y, train_size=100, test_size=100, random_state=42
 )
 
-X_struct = get_probs_from_mutations(m_train)        # [N, 20]
+#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]
+pdb = 'datasets/BLAT_ECOLX/BLAT_ECOLX.pdb'
+wt_seq = list(read_fasta_biopython('datasets/BLAT_ECOLX/blat_ecolx_wt.fasta').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")
@@ -87,7 +84,7 @@ def forward(self, X):
 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 
+X_combined = torch.cat([X_seq, X_seq], dim=-1)  # Concenation is necessary as GPkernel does not accept a tuple as input 
 d_seq = X_seq.shape[1]
 
 # -----------------------------
@@ -121,10 +118,10 @@ def forward(self, X):
 # -----------------------------
 # Test
 # -----------------------------
-X_struct_test = get_probs_from_mutations(m_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)
+X_test_combined = torch.cat([X_seq_test, X_seq_test], dim=-1)
 
 model.eval()
 likelihood.eval()
@@ -143,15 +140,23 @@ def forward(self, X):
 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())
+print("Correlation loss Spearman    TRAIN:", correlation_loss(y_train, torch.from_numpy(y_pred_train), method="spearman"))
+print("Correlation hybrid MSE loss Spearman    TRAIN:", hybrid_corr_mse_loss(y_train, torch.from_numpy(y_pred_train)))
+print("Correlation loss Pearson     TRAIN:", correlation_loss(y_train, torch.from_numpy(y_pred_train), method="pearson"))
+print("Correlation loss Pearson 2   TRAIN:", pearson_loss(y_train, torch.from_numpy(y_pred_train)))
 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())
+#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())
+print("Correlation loss Spearman    TEST:", correlation_loss(y_test, torch.from_numpy(y_pred), method="spearman"))
+print("Correlation hybrid MSE loss Spearman    TEST:", hybrid_corr_mse_loss(y_test, torch.from_numpy(y_pred)))
+print("Correlation loss Pearson     TEST:", correlation_loss(y_test, torch.from_numpy(y_pred), method="pearson"))
+print("Correlation loss Pearson 2   TEST:", pearson_loss(y_test, torch.from_numpy(y_pred)))
 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())
+#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_opt.py

@@ -139,14 +139,11 @@ def extract_prosst_embeddings(
 
 
 def gp_fit():
-    pass
-
+    pass  # TODO
 
 
 def gp_fit_sklearn():
-    pass
-
-
+    pass  # TODO
 
 
 if __name__ == '__main__':

pypef/gaussian_process/gp_prosst_test.py

@@ -15,7 +15,7 @@
 
 from pypef.plm.prosst_lora_tune import get_prosst_models, get_structure_quantizied
 from pypef.utils.helpers import get_device
-from pypef.plm.utils import corr_loss, get_batches
+from pypef.plm.utils import pearson_loss, get_batches
 from pypef.plm.esm_lora_tune import get_esm_models
 
 
@@ -162,12 +162,12 @@ def sequence_log_likelihood(
                             output_hidden_states=extract_emb,
                             **model_kwargs
                         )
-
-                    token_embeddings = outputs.hidden_states[-1]  # (1, L+2, D)
-                    # Mean pool over residues (exclude CLS/EOS)
-                    seq_embedding = token_embeddings[0, 1:-1].mean(dim=0)
-                    embeddings.append(seq_embedding)
-                    continue
+                    if extract_emb:
+                        token_embeddings = outputs.hidden_states[-1]  # (1, L+2, D)
+                        # Mean pool over residues (exclude CLS/EOS)
+                        seq_embedding = token_embeddings[0, 1:-1].mean(dim=0)
+                        embeddings.append(seq_embedding)
+                        continue
 
             except TypeError as e:
                 logger.info(f"Did not find model input keyword arguments (kwargs: "
@@ -633,7 +633,7 @@ def esm_setup(wt_seq, sequences, device: str | None = None, verbose: bool = True
             'llm_optimizer': esm_optimizer,
             'llm_train_function': plm_train,
             'llm_inference_function': plm_inference,
-            'llm_loss_function': corr_loss,
+            'llm_loss_function': pearson_loss,
             'x_llm' : torch.tensor(x_esm),  # TODO: Not needed here?
             'llm_attention_mask':  torch.tensor(esm_attention_mask),  # TODO: Not needed here?
             'wt_input_ids': torch.tensor(wt_tokens),  # TODO: Not needed here?
@@ -680,7 +680,7 @@ def prosst_setup(wt_seq, pdb_file, sequences, device: str | None = None, verbose
             'llm_optimizer': prosst_optimizer,
             'llm_train_function': plm_train,
             'llm_inference_function': plm_inference,  # prosst_infer,
-            'llm_loss_function': corr_loss,
+            'llm_loss_function': pearson_loss,
             'x_llm' : x_llm_train_prosst,
             'llm_attention_mask': prosst_attention_mask,
             'llm_vocab': prosst_vocab,

pypef/gaussian_process/metrics.py

@@ -13,7 +13,70 @@
 logger = logging.getLogger('pypef.llm.utils')
 
 
-def corr_loss(y_true: torch.Tensor, y_pred: torch.Tensor):
+def hybrid_corr_mse_loss(y_true, y_pred, tau=0.1, alpha=0.5):
+    """
+    Hybrid differentiable loss combining Spearman correlation and MSE.
+    """
+    # Differentiable Spearman
+    loss_rank = correlation_loss(y_true, y_pred, method="spearman", tau=tau)
+    # MSE
+    loss_value = torch.mean((y_pred - y_true)**2)
+    # Combine
+    return alpha * loss_rank + (1 - alpha) * loss_value
+
+
+
+def correlation_loss(y_true: torch.Tensor, 
+                     y_pred: torch.Tensor, 
+                     method: str = "spearman", 
+                     tau: float = 0.1) -> torch.Tensor:
+    """
+    Differentiable correlation loss for PyTorch.
+    
+    Args:
+        y_true: Tensor of shape (..., n) or (batch, n)
+        y_pred: Tensor of same shape as y_true
+        method: "pearson" or "spearman"
+        tau: temperature for soft-rank approximation (used if method="spearman")
+        
+    Returns:
+        Scalar tensor representing the loss (to minimize)
+    """
+    if method == "spearman":
+        # Soft rank approximation
+        x = y_true
+        y = y_pred
+
+        def soft_rank(x, tau):
+            x = x.unsqueeze(-1)
+            diff = x - x.transpose(-1, -2)
+            P = torch.sigmoid(diff / tau)
+            return P.sum(dim=-1) + 0.5
+
+        rx = soft_rank(x, tau)
+        ry = soft_rank(y, tau)
+    elif method == "pearson":
+        rx = y_true
+        ry = y_pred
+    else:
+        raise ValueError(f"Unsupported method: {method}. Choose 'pearson' or 'spearman'.")
+
+    # Centering
+    rx_c = rx - rx.mean(dim=-1, keepdim=True)
+    ry_c = ry - ry.mean(dim=-1, keepdim=True)
+
+    # Normalize (like dividing by std)
+    rx_n = rx_c / (rx_c.norm(dim=-1, keepdim=True) + 1e-8)
+    ry_n = ry_c / (ry_c.norm(dim=-1, keepdim=True) + 1e-8)
+
+    # Compute correlation
+    corr = (rx_n * ry_n).sum(dim=-1)
+
+    # Return scalar loss (to minimize, so negative correlation)
+    return -corr.mean()
+
+
+def pearson_loss(y_true: torch.Tensor, y_pred: torch.Tensor):
     res_true = y_true - torch.mean(y_true)
     res_pred = y_pred - torch.mean(y_pred)
     cov = torch.mean(res_true * res_pred)
@@ -24,6 +87,41 @@ def corr_loss(y_true: torch.Tensor, y_pred: torch.Tensor):
     return - cov / (sigma_true * sigma_pred)
 
 
+def spearman_loss(y_true, y_pred, tau=0.1):
+    """Maximizing Spearman correlation"""
+    return - spearman_soft(y_true, y_pred, tau=tau).mean()
+
+
+def soft_rank_approx(x, tau=1.0):
+    """
+    A simple soft rank approximation using pairwise comparisons.
+    Args:
+        x: tensor of shape (..., n)
+        tau: temperature (larger = softer, smaller = closer to true ranks)
+    Returns:
+        approx ranks same shape as x
+    """
+    diff = x.unsqueeze(-1) - x.unsqueeze(-2)
+    # pairwise sigmoid scores
+    P = torch.sigmoid(diff / tau)
+    # sum of how many values each element is less than
+    r = P.sum(dim=-1) + 0.5  # +0.5 to approximate average rank
+    return r
+
+
+def spearman_soft(x, y, tau=0.1):
+    rx = soft_rank_approx(x, tau)
+    ry = soft_rank_approx(y, tau)
+
+    rxc = rx - rx.mean(dim=-1, keepdim=True)
+    ryc = ry - ry.mean(dim=-1, keepdim=True)
+
+    rxn = rxc / (rxc.norm(dim=-1, keepdim=True) + 1e-8)
+    ryn = ryc / (ryc.norm(dim=-1, keepdim=True) + 1e-8)
+
+    return (rxn * ryn).sum(dim=-1)
+
+
 def get_batches(a, dtype, batch_size=5,
                 keep_remaining=False, verbose: bool = False
                 ) -> list | list[np.ndarray]:
@@ -80,6 +178,7 @@ def is_model_cached(repo_id: str, cache_dir: str):
     Check if the required model and tokenizer files are cached locally.
     """
     snapshot_dir = None
+    ref_file = None
     if os.path.isdir(cache_dir):
         ref_file = os.path.join(
             cache_dir, f'models--{repo_id.replace("/", "--")}', 'refs', 'main'

pypef/plm/inference.py

@@ -20,7 +20,7 @@
 import sys  # Use local directory PyPEF files
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
 from pypef.dca.gremlin_inference import GREMLIN
-from pypef.plm.utils import get_batches, corr_loss
+from pypef.plm.utils import get_batches, pearson_loss
 from pypef.plm.esm_lora_tune import (
     get_esm_models, tokenize_sequences, 
     esm_train, esm_infer
@@ -234,7 +234,7 @@ def compute_performances(mut_data, mut_sep=':', start_i: int = 0, already_tested
                             'llm_optimizer': esm_optimizer,
                             'llm_train_function': esm_train,
                             'llm_inference_function': esm_infer,
-                            'llm_loss_function': corr_loss,
+                            'llm_loss_function': pearson_loss,
                             'x_llm' : x_llm_train_esm,
                             'llm_attention_mask':  esm_attention_mask
                         }
@@ -246,7 +246,7 @@ def compute_performances(mut_data, mut_sep=':', start_i: int = 0, already_tested
                             'llm_optimizer': prosst_optimizer,
                             'llm_train_function': prosst_train,
                             'llm_inference_function': get_logits_from_full_seqs,
-                            'llm_loss_function': corr_loss,
+                            'llm_loss_function': pearson_loss,
                             'x_llm' : x_llm_train_prosst,
                             'llm_attention_mask':  prosst_attention_mask,
                             'input_ids': input_ids,