# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # NeuralEE on CORTEX Dataset

# `CORTEX` dataset contains 3005 mouse cortex cells and gold-standard labels for seven distinct cell types. Each cell type corresponds to a cluster to recover.

# +
import random
import numpy as np
import torch

from neuralee.embedding import NeuralEE
from neuralee.dataset import CortexDataset 
from neuralee._aux import scatter

# %matplotlib inline
# -

# Choose a GPU if a GPU available. It could be defined as follow:
# ```
# device = torch.device('cuda:0')
# device = torch.device('cuda:1')
# device = torch.device('cpu')
# ```

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# To reproduce the following results, fix the random seed.

torch.manual_seed(1234)
random.seed(1234)
np.random.seed(1234)

# First, we apply log(1 + x) transformation to each element of the cell-gene expression matrix. 
# Then, We retain top 558 genes ordered by variance as the original paper. 
# Finally, we normalize the expression of each gene by subtracting its mean and dividing its standard deviation.

cortex_dataset = CortexDataset(save_path='../')
cortex_dataset.log_shift()
cortex_dataset.subsample_genes(558)  
cortex_dataset.standardscale()

# We apply NeuralEE with different hyper-paramters.
# `N_small` takes from {1.0, 0.5, 0.25}, while `N_smalls`= 1.0 means not applied with stochastic optimization.
# `lam` takes from {1, 10}.
# `perplexity` fixs as 30.

# +
N_smalls = [1.0, 0.5, 0.25]
N_str = ["nobatch", "2batches", "4batches"]
lams = [1, 10]

cortex_dataset.affinity(perplexity=30.0)
for i in range(len(N_smalls)):
    cortex_dataset.affinity_split(N_small=N_smalls[i], perplexity=30.0)
    for lam in lams:
        NEE = NeuralEE(cortex_dataset, lam=lam, device=device)
        results_Neural = NEE.fine_tune()
        np.save('embedding/CORTEX_' + 'lam' + str(lam) + '_' + N_str[i], results_Neural['X'].numpy())
        scatter(results_Neural['X'].numpy(), NEE.labels, 
                cortex_dataset.cell_types)

StarcoderdataJupyterScriptsDedupFiltered