vae_model.py

import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import json
import os
import torch
import torch.nn as nn
from torch.utils.data import Dataset
from sklearn.model_selection import train_test_split
from tqdm.notebook import tqdm


def gaussian_reparameterization(mu, logvar):
    """Sample from the posterior distribution.

    Args:
        mu (Tensor): Mean of the latent Gaussian.
        logvar (Tensor): Logvar of the latent Gaussian.

    Returns:
        Tensor: random sample ~ N(mu, std)
    """

    std = torch.exp(0.5 * logvar)
    eps = torch.randn_like(std)
    return eps * std + mu


class CancerSamplesDataset(Dataset):
    def __init__(self, sample_encodings_fname, sorted_mutations_fname, mutations_mapping_fname, train=True, val_split=0.1):
        super().__init__()
        sample_encodings = pd.read_csv(sample_encodings_fname)
        self.sample_ids = sample_encodings["ID_sample"].astype(str).values
        self.encodings = sample_encodings.values[:,1:]
        with open(sorted_mutations_fname, "r") as f:
            self.mutations = json.load(f)
            self.mutations = {m: i for i, m in enumerate(self.mutations)}
        with open(mutations_mapping_fname, "r") as f:
            self.mutations_mapping = json.load(f)
        self.n_mutations = len(self.mutations)
        if train is not None:
            idxs = list(range(len(self.sample_ids)))
            idxs_train, idxs_val = train_test_split(
                idxs, test_size=val_split, random_state=42)
            idxs = idxs_train if train else idxs_val
            self.sample_ids = [self.sample_ids[i] for i in idxs]
            self.encodings = self.encodings[idxs, :]
            self.mutations_mapping = {
                k: v for k, v in self.mutations_mapping.items() if k in self.sample_ids}
        self.deleterious_mutations = {k: v[0]
                                      for k, v in self.mutations_mapping.items()}
        self.non_deleterious_mutations = {k: v[1]
                                          for k, v in self.mutations_mapping.items()}

    def __len__(self):
        return len(self.sample_ids)

    def __getitem__(self, idx):
        sample_id = self.sample_ids[idx]
        del_mutations = self.deleterious_mutations[sample_id]
        del_mutations_idxs = list(
            set([self.mutations[m] for m in del_mutations]))
        nd_mutations = self.non_deleterious_mutations[sample_id]
        nd_mutations_idxs = list(
            set([self.mutations[m] for m in nd_mutations]))
        del_encoding = torch.zeros(self.n_mutations)
        del_encoding[del_mutations_idxs] = 1
        nd_encoding = torch.zeros(self.n_mutations)
        nd_encoding[nd_mutations_idxs] = 1
        subtype_encoding = torch.tensor(self.encodings[idx])
        return del_encoding, nd_encoding, subtype_encoding


class BimodalVAE(nn.Module):
    def __init__(self, embedding_dim=50, dropout_p=0.4):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.dropout_p = dropout_p

        # deleterious encoder
        self.del_encoder = nn.Sequential(
            nn.Linear(12449, 400),
            nn.Dropout(p=dropout_p),
            nn.ReLU(),
            nn.Linear(400, 50),
            nn.Dropout(p=dropout_p),
            nn.ReLU()
        )

        # non deleterious encoder
        self.nd_encoder = nn.Sequential(
            nn.Linear(12449, 400),
            nn.Dropout(p=dropout_p),
            nn.ReLU(),
            nn.Linear(400, 50),
            nn.Dropout(p=dropout_p),
            nn.ReLU()
        )

        self.mu = nn.Linear(100, embedding_dim)
        self.logvar = nn.Linear(100, embedding_dim)

        self.del_decoder = nn.Sequential(
            nn.Linear(embedding_dim, 50),
            nn.Dropout(p=dropout_p),
            nn.ReLU(),
            nn.Linear(50, 400),
            nn.Dropout(p=dropout_p),
            nn.ReLU(),
            nn.Linear(400, 12449),
            nn.Sigmoid()
        )

        self.nd_decoder = nn.Sequential(
            nn.Linear(embedding_dim, 50),
            nn.Dropout(p=dropout_p),
            nn.ReLU(),
            nn.Linear(50, 400),
            nn.Dropout(p=dropout_p),
            nn.ReLU(),
            nn.Linear(400, 12449),
            nn.Sigmoid()
        )

    def encode(self, X_del, X_nd):
        h_del = self.del_encoder(X_del)
        h_nd = self.nd_encoder(X_nd)
        merged_tensor = torch.cat((h_del, h_nd), -1)
        mu = self.mu(merged_tensor)
        logvar = self.logvar(merged_tensor)
        return mu, logvar

    def decode(self, z):
        y_del = self.del_decoder(z)
        y_nd = self.nd_decoder(z)
        return y_del, y_nd

    def forward(self, X_del, X_nd):
        mu, logvar = self.encode(X_del, X_nd)
        z = gaussian_reparameterization(mu, logvar)
        y_del, y_nd = self.decode(z)
        return y_del, y_nd, mu, logvar