import pytorch_lightning as pl
from pytorch_lightning.loggers import WandbLogger
import wandb
from torch import nn
import torch
import torch.nn.functional as F
from torch.nn import TransformerEncoderLayer
import numpy as np
from scipy.stats import rankdata
from allennlp.modules.attention import CosineAttention, BilinearAttention, AdditiveAttention, DotProductAttention
from utils.loss_utils import MultisimilarityCriterion, _construct_labels, unique, _construct_disease_labels
from utils.train_utils import masked_mean, masked_softmax, weighted_sum, plot_degree_vs_attention, mean_reciprocal_rank, top_k_acc
class GPAligner(pl.LightningModule):
def __init__(self, hparams, embed_dim):
super().__init__()
self.hyperparameters = hparams
print('GPAligner embedding dimension: ', embed_dim)
# attention for collapsing set of phenotype embeddings
self.attn_vector = nn.Parameter(torch.zeros((1, embed_dim), dtype=torch.float), requires_grad=True)
nn.init.xavier_uniform_(self.attn_vector)
if self.hyperparameters['attention_type'] == 'bilinear':
self.attention = BilinearAttention(embed_dim, embed_dim)
elif self.hyperparameters['attention_type'] == 'additive':
self.attention = AdditiveAttention(embed_dim, embed_dim)
elif self.hyperparameters['attention_type'] == 'dotpdt':
self.attention = DotProductAttention()
if self.hyperparameters['decoder_type'] == "dotpdt":
self.decoder = DotProductAttention(normalize=False)
elif self.hyperparameters['decoder_type'] == "bilinear":
self.decoder = BilinearAttention(embed_dim, embed_dim, activation=torch.tanh, normalize=False)
else:
raise NotImplementedError
# projection layers
self.phen_project = nn.Linear(embed_dim, embed_dim)
self.gene_project = nn.Linear(embed_dim, embed_dim)
self.phen_project2 = nn.Linear(embed_dim, embed_dim)
self.gene_project2 = nn.Linear(embed_dim, embed_dim)
# optional disease projection layer
if self.hyperparameters['use_diseases']:
self.disease_project = nn.Linear(embed_dim, embed_dim)
self.disease_project2 = nn.Linear(embed_dim, embed_dim)
self.leaky_relu = nn.LeakyReLU(hparams['leaky_relu'])
self.loss = MultisimilarityCriterion(hparams['pos_weight'], hparams['neg_weight'],
hparams['margin'], hparams['thresh'],
embed_dim, hparams['only_hard_distractors'])
if 'n_transformer_layers' in hparams and hparams['n_transformer_layers'] > 0:
encoder_layer = TransformerEncoderLayer(d_model=embed_dim, nhead=hparams['n_transformer_heads'])
self.transformer_encoder = nn.TransformerEncoder(encoder_layer, num_layers=hparams['n_transformer_layers'])
def forward(self, phenotype_embeddings, candidate_gene_embeddings, disease_embeddings=None, phenotype_mask=None, gene_mask=None, disease_mask=None):
assert phenotype_mask != None
assert gene_mask != None
if self.hyperparameters['use_diseases']: assert disease_mask != None
if 'n_transformer_layers' in self.hyperparameters and self.hyperparameters['n_transformer_layers'] > 0:
phenotype_embeddings = self.transformer_encoder(phenotype_embeddings.transpose(0, 1), src_key_padding_mask=~phenotype_mask).transpose(0, 1)
# attention weighted average of phenotype embeddings
batched_attn = self.attn_vector.repeat(phenotype_embeddings.shape[0],1)
attn_weights = self.attention(batched_attn, phenotype_embeddings, phenotype_mask)
phenotype_embedding = weighted_sum(phenotype_embeddings, attn_weights)
# project embeddings
phenotype_embedding = self.phen_project2(self.leaky_relu(self.phen_project(phenotype_embedding)))
candidate_gene_embeddings = self.gene_project2(self.leaky_relu(self.gene_project(candidate_gene_embeddings)))
if self.hyperparameters['use_diseases']:
disease_embeddings = self.disease_project2(self.leaky_relu(self.disease_project(disease_embeddings)))
else:
disease_embeddings = None
disease_mask = None
return phenotype_embedding, candidate_gene_embeddings, disease_embeddings, gene_mask, phenotype_mask, disease_mask, attn_weights
def _calc_similarity(self, phenotype_embeddings, candidate_gene_embeddings, disease_embeddings, batch_cand_gene_nid, batch_corr_gene_nid, batch_disease_nid, one_hot_labels, gene_mask, phenotype_mask, disease_mask, use_candidate_list, cand_gene_to_phenotypes_spl, alpha):
# Normalize Embeddings (within each individual patient)
phenotype_embeddings = F.normalize(phenotype_embeddings, p=2, dim=1)
batch_sz = phenotype_embeddings.shape[0]
if disease_embeddings != None: disease_embeddings = F.normalize(disease_embeddings.squeeze(), p=2, dim=1)
if candidate_gene_embeddings != None:
batch_sz, n_cand_genes, embed_dim = candidate_gene_embeddings.shape
candidate_gene_embeddings = F.normalize(candidate_gene_embeddings.view(batch_sz*n_cand_genes,-1), p=2, dim=1).view(batch_sz, n_cand_genes, embed_dim)
# Only use each patient's candidate genes/diseases
if self.hyperparameters['only_hard_distractors'] or use_candidate_list:
if disease_embeddings == None: # only use genes
mask = gene_mask
one_hot_labels = one_hot_labels
raw_sims = self.decoder(phenotype_embeddings, candidate_gene_embeddings)
if cand_gene_to_phenotypes_spl != None:
sims = alpha * raw_sims + (1 - alpha) * cand_gene_to_phenotypes_spl
else: sims = raw_sims
elif candidate_gene_embeddings == None: # only use diseases
raise NotImplementedError
else:
raise NotImplementedError
# Otherwise, use entire batch as candidate genes/diseases
else:
if disease_embeddings == None: #only use genes
candidate_gene_idx, candidate_gene_embeddings, one_hot_labels = _construct_labels(candidate_gene_embeddings, batch_cand_gene_nid, batch_corr_gene_nid, gene_mask)
raw_sims = self.decoder(phenotype_embeddings, candidate_gene_embeddings.unsqueeze(0).repeat(batch_sz,1,1))
if cand_gene_to_phenotypes_spl != None:
sims = alpha * raw_sims + (1 - alpha) * cand_gene_to_phenotypes_spl
else: sims = raw_sims
mask = None
elif candidate_gene_embeddings == None: #only use diseases
candidate_embeddings, one_hot_labels = _construct_disease_labels(disease_embeddings, batch_disease_nid)
raw_sims = self.decoder(phenotype_embeddings, candidate_embeddings.unsqueeze(0).repeat(batch_sz,1,1))
if batch_disease_nid.shape[1] > 1:
raw_sims = raw_sims[batch_disease_nid[:,0].squeeze() != 0] # remove rows where the patient doesn't have
one_hot_labels = one_hot_labels[batch_disease_nid[:,0].squeeze() != 0]
else:
raw_sims = raw_sims[batch_disease_nid.squeeze() != 0] # remove rows where the patient doesn't have
one_hot_labels = one_hot_labels[batch_disease_nid.squeeze() != 0]
sims = raw_sims
mask = None
else: # use genes + diseases
raise NotImplementedError
return sims, raw_sims, mask, one_hot_labels
def _rank_genes(self, phen_gene_sims, gene_mask, one_hot_labels):
phen_gene_sims = phen_gene_sims * gene_mask
padded_phen_gene_sims = phen_gene_sims + (~gene_mask * -100000) # we want to rank the padded values last
gene_ranks = torch.tensor(np.apply_along_axis(lambda row: rankdata(row * -1, method='average'), axis=1, arr=padded_phen_gene_sims.detach().cpu().numpy()))
if one_hot_labels is None: correct_gene_ranks = None
else:
gene_ranks = gene_ranks.to(one_hot_labels.device)
correct_gene_ranks = gene_ranks[one_hot_labels == 1]
return correct_gene_ranks, padded_phen_gene_sims
def calc_loss(self, sims, mask, one_hot_labels):
return self.loss(sims, mask, one_hot_labels)
Datasets
Models
