[WIP] feat: create and train GNN for gene expression prediction using gene regulatory network

bactgraph/modeling/data_reader.py

@@ -0,0 +1,107 @@
+import os
+import random
+from collections.abc import Callable
+from typing import Any
+
+import numpy as np
+import pandas as pd
+from torch.utils.data import DataLoader
+
+from bactgraph.modeling.dataset import BactGraphDataset
+
+BACTMAP_PROTEINS_FILE_NAME = "bactmap_proteins_prot_embeds.parquet"
+NORMALISED_EXPRESSION_FILE_NAME = "norm_dat_pao1.tsv"
+PERTURB_NETWORK_FILE_NAME = "llcb_perturb_hits_adj_matrix.tsv"
+
+
+def preprocess_data_for_training(
+    input_dir: str,
+    transform_norm_expression_fn: Callable = np.log10,
+    train_size: float = 0.7,
+    test_size: float = 0.2,
+    batch_size: int = 32,
+    num_workers: int = 4,
+    random_seed: int = 42,
+) -> dict[str, Any]:
+    """Preprocess the data for training the BactGraph model."""
+    # read the data
+    protein_embeddings = pd.read_parquet(os.path.join(input_dir, BACTMAP_PROTEINS_FILE_NAME))
+    expression_df = pd.read_csv(os.path.join(input_dir, NORMALISED_EXPRESSION_FILE_NAME), sep="\t").set_index(
+        "feature_id"
+    )
+    perturb_network = pd.read_csv(os.path.join(input_dir, PERTURB_NETWORK_FILE_NAME), sep="\t").set_index("gene_id")
+
+    # keep only genes which are in all files
+    prot_emb_genes = set(protein_embeddings.columns.tolist())
+    expression_genes = set(expression_df.index.tolist())
+    perturb_network_genes = set(perturb_network.index.tolist() + perturb_network.columns.tolist())
+
+    genes_of_interest = list(prot_emb_genes.intersection(expression_genes).intersection(perturb_network_genes))
+    print(f"Total nr of genes available: {len(genes_of_interest)}")
+
+    # subset the genes of interest
+    protein_embeddings = protein_embeddings[genes_of_interest]
+    expression_df = expression_df[expression_df.index.isin(genes_of_interest)]
+    perturb_network = perturb_network[[g for g in genes_of_interest if g in perturb_network.columns]]
+    perturb_network = perturb_network[perturb_network.index.isin(genes_of_interest)]
+
+    # subset to the strains with expression data
+    strains_w_expression = expression_df.columns.tolist()
+    strains_w_prot_emb = protein_embeddings.index.tolist()
+    strains_of_interest = list(set(strains_w_expression).intersection(strains_w_prot_emb))
+    expression_df = expression_df[strains_of_interest]
+    protein_embeddings = protein_embeddings.loc[strains_of_interest]
+
+    # split the data
+    random.seed(random_seed)
+    random.shuffle(strains_of_interest)
+    train_size = int(len(strains_of_interest) * train_size)
+    test_size = int(len(strains_of_interest) * test_size)
+    train_strains = strains_of_interest[:train_size]
+    test_strains = strains_of_interest[train_size : train_size + test_size]
+    val_strains = strains_of_interest[train_size + test_size :]
+
+    gene2idx = {gene: idx for idx, gene in enumerate(protein_embeddings.columns)}
+
+    # create datasets
+    train_dataset = BactGraphDataset(
+        protein_embeddings=protein_embeddings.loc[train_strains],
+        expression_df=expression_df[train_strains],
+        gene2idx=gene2idx,
+        perturb_network=perturb_network,
+        transform_norm_expression_fn=transform_norm_expression_fn,
+        random_seed=random_seed,
+    )
+    val_dataset = BactGraphDataset(
+        protein_embeddings=protein_embeddings.loc[val_strains],
+        expression_df=expression_df[val_strains],
+        gene2idx=gene2idx,
+        perturb_network=perturb_network,
+        transform_norm_expression_fn=transform_norm_expression_fn,
+        random_seed=random_seed,
+    )
+    test_dataset = BactGraphDataset(
+        protein_embeddings=protein_embeddings.loc[test_strains],
+        expression_df=expression_df[test_strains],
+        gene2idx=gene2idx,
+        perturb_network=perturb_network,
+        transform_norm_expression_fn=transform_norm_expression_fn,
+        random_seed=random_seed,
+    )
+
+    # create dataloaders
+    train_dataloader = DataLoader(
+        train_dataset, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=num_workers
+    )
+    val_dataloader = DataLoader(
+        val_dataset, batch_size=batch_size, shuffle=False, pin_memory=True, num_workers=num_workers
+    )
+    test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False, pin_memory=True, num_workers=num_workers)
+
+    return dict(  # noqa
+        train_dataloader=train_dataloader,
+        val_dataloader=val_dataloader,
+        test_dataloader=test_dataloader,
+        n_train_size=len(train_strains),
+        gene2idx=gene2idx,
+    )

bactgraph/modeling/dataset.py

@@ -0,0 +1,92 @@
+from collections.abc import Callable
+
+import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import Dataset
+
+BACTMAP_PROTEINS_FILE_NAME = "bactmap_proteins_prot_embeds.parquet"
+NORMALISED_EXPRESSION_FILE_NAME = "norm_dat_pao1.tsv"
+PERTURB_NETWORK_FILE_NAME = "bactmap_proteins_prot_embeds.parquet"
+
+
+def perturb_mtx_to_triples(df: pd.DataFrame, gene2idx: dict[str, int]) -> torch.Tensor:
+    """Conver perturbation dataframe to triples with non-zero values for training."""
+    # 1. "Stack" the DataFrame so that rows become part of a MultiIndex
+    nonzero_stacked = df.stack()  # This will convert the DataFrame into a Series
+
+    # 2. Filter out zero values
+    nonzero_stacked = nonzero_stacked[nonzero_stacked != 0]
+
+    # 3. Convert to a list of (index_name, column_name, value) tuples
+    triples = list(
+        zip(
+            nonzero_stacked.index.get_level_values(0),  # index name
+            nonzero_stacked.index.get_level_values(1),  # column name
+            nonzero_stacked.values,
+            strict=False,  # value
+        )
+    )
+
+    triples = torch.tensor(
+        [
+            [gene2idx[gene1] for gene1, _, _ in triples],
+            [gene2idx[gene2] for _, gene2, _ in triples],
+            [val for _, _, val in triples],
+        ],
+        dtype=torch.float32,
+    )
+    return triples
+
+
+class BactGraphDataset(Dataset):
+    """Dataset of gene networks in bacteria for BactGraph project."""
+
+    def __init__(
+        self,
+        protein_embeddings: pd.DataFrame,
+        expression_df: pd.DataFrame,
+        gene2idx: dict[str, int],
+        perturb_network: pd.DataFrame,
+        transform_norm_expression_fn: Callable = np.log10,
+        random_seed: int = 42,
+    ):
+        self.protein_embeddings = protein_embeddings
+        self.expression_df = expression_df
+        self.gene2idx = gene2idx
+
+        # get triples
+        self.triples = perturb_mtx_to_triples(perturb_network, self.gene2idx)[:2, :]
+        # reverse the direction
+        # self.triples = self.triples[:2, :].flip(0)
+        # randomize the network experiment
+        # print("Randomizing the network experiment by randomly sampling edges.")
+        # torch.manual_seed(random_seed)
+        # self.triples = torch.randint(0, len(self.gene2idx), self.triples.shape)
+        # fully connected network
+        # self.triples = torch.stack(
+        #     [torch.arange(len(self.gene2idx)), torch.arange(len(self.gene2idx)), torch.ones(len(self.gene2idx))],
+        #     dim=0,
+        # )
+
+        # normalise the expression data
+        # revert previous log2 transformation (the data was provided like this)
+        self.expression_df = self.expression_df.apply(np.exp2)
+        # transform the data with the provided function
+        self.expression_df = self.expression_df.apply(transform_norm_expression_fn).fillna(-100.0)
+
+        self.strains = self.expression_df.columns.tolist()
+
+    def __len__(self):
+        return len(self.expression_df.columns)
+
+    def __getitem__(self, idx) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        # get the expression data for the idx-th strain
+        strain = self.strains[idx]
+        # get protein embeddings
+        prot_emb = torch.tensor(np.stack(self.protein_embeddings.loc[strain].values), dtype=torch.float32)
+        expr_values = torch.tensor(
+            [self.expression_df.loc[gene, strain] for gene in self.protein_embeddings.columns], dtype=torch.float32
+        )
+        gene_idx = torch.arange(len(self.protein_embeddings.columns), dtype=torch.long)
+        return prot_emb, self.triples, expr_values, gene_idx