<jupyter_start><jupyter_text>Norman 2019 Training Demo<jupyter_code>import sys
#if branch is stable, will install via pypi, else will install from source
branch = "stable"
IN_COLAB = "google.colab" in sys.modules

if IN_COLAB and branch == "stable":
    !pip install cpa-tools
elif IN_COLAB and branch != "stable":
    !pip install --quiet --upgrade jsonschema
    !pip install --quiet git+https://github.com/theislab/cpa
import os
os.chdir('/home/mohsen/projects/cpa/')
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
import cpa
import scanpy as sc
sc.settings.set_figure_params(dpi=100)
data_path = '/home/mohsen/projects/cpa/datasets/Norman2019_prep_new.h5ad'
adata = sc.read(data_path, backup_url='https://drive.google.com/u/0/uc?id=1pxT0fvXtqBBtdv1CCPVwJaMLHe9XpMHo&export=download&confirm=t')
adata
cpa.CPA.setup_anndata(adata, 
                      drug_key='condition', 
                      dose_key='dose_val',
                      control_key='control',
                      categorical_covariate_keys=['cell_type'],
                      combinatorial=True,
                     )
ae_hparams = {'adversary_depth': 2,
              'adversary_width': 256,
              'autoencoder_depth': 4,
              'autoencoder_width': 256,
              'dosers_depth': 3,
              'dosers_width': 32,
              'use_batch_norm': True,
              'use_layer_norm': False,
              'dropout_rate': 0.0,
              'variational': False,
              'seed': 31,
              }

trainer_params = {
    'n_epochs_warmup': 0,
    'adversary_lr': 0.00012219948594647382,
    'adversary_steps': 2,
    'adversary_wd': 1.4033946047401463e-05,
    'autoencoder_lr': 0.00014147035543889223,
    'autoencoder_wd': 2.2782368178455333e-08,
    'dosers_lr': 0.0007022363227033381,
    'dosers_wd': 6.249509305603141e-06,
    'penalty_adversary': 0.013702812231919399,
    'reg_adversary': 4.02272482876072,
    'step_size_lr': 45,
}
model = cpa.CPA(adata=adata, 
                n_latent=512,
                loss_ae='gauss',
                doser_type='logsigm',
                split_key='split',
                **ae_hparams,
               )
model.train(max_epochs=2000,
            use_gpu=True, 
            batch_size=256,
            early_stopping=True,
            plan_kwargs=trainer_params,
            early_stopping_patience=15,
            check_val_every_n_epoch=20,
            save_path='/home/mohsen/projects/cpa/lightning_logs/Norman2019_prep_new/',
           )
cpa.pl.plot_history(model)
cpa_api = cpa.ComPertAPI(adata, model)
cpa_plots = cpa.pl.CompertVisuals(cpa_api, fileprefix=None)
cpa_api.compute_comb_emb(thrh=30)
import pandas as pd
# construct prediction annadata
subset = adata[adata.obs['cell_type'] == 'A549'].copy()
genes_control = subset[subset.obs['condition'] == 'ctrl'].copy()

df = pd.DataFrame({'condition': ['TSC22D1+ctrl'] + ['KLF1+MAP2K6'], 
                   'dose_val': ['1+1'] + ['1+1'], 
                   'cell_type': ['A549']*2})

pred = cpa_api.predict(genes_control.X.toarray(), df)  # normally would put `sample=True` here
pred.obs['cov_drug_dose_name'] = pred.obs['cell_type'].astype(str) + '_' + pred.obs['condition'].astype(str) + '_' + pred.obs['dose_val'].astype(str)
pred.obs['cov_drug_dose_name'] = pred.obs['cov_drug_dose_name'].astype('category')
pred
adata.obs['cov_drug_dose_name'].unique()
de_dict = adata[adata.obs.split == 'train'].uns['rank_genes_groups_cov']
de_dict = {k: adata.var_names for k, v in de_dict.items()}
pred_adata = pred.concatenate(adata, batch_key='source', batch_categories=['pred', 'true'])
drug_name = 'TSC22D1+ctrl'
sc.pl.violin(
    pred_adata[pred_adata.obs.condition == drug_name],
    keys=de_dict[f'A549_{drug_name}_1+1'][:6],
    groupby='source')<jupyter_output>/home/mohsen/miniconda3/envs/cpa/lib/python3.9/site-packages/anndata/_core/anndata.py:1228: FutureWarning: The `inplace` parameter in pandas.Categorical.reorder_categories is deprecated and will be removed in a future version. Reordering categories will always return a new Categorical object.
  c.reorder_categories(natsorted(c.categories), inplace=True)
/home/mohsen/miniconda3/envs/cpa/lib/python3.9/site-packages/anndata/_core/anndata.py:1236: ImplicitModificationWarning: Initializing view as actual.
  warnings.warn(
Trying to set attribute `.obs` of view, copying.
... storing 'condition' as categorical
/home/mohsen/miniconda3/envs/cpa/lib/python3.9/site-packages/anndata/_core/anndata.py:1228: FutureWarning: The `inplace` parameter in pandas.Categorical.reorder_categories is deprecated and will be removed in a future version. Reordering categories will always return a new Categorical object.
  c.reorder_categories(natsorted(c.categories), inplace=True)
Trying to set attribute `.obs` of[...]
{'fraction_non_alphanumeric': 0.13044388609715243, 'fraction_numerical': 0.05360134003350084, 'mean_word_length': 3.4108956602031393, 'pattern_counts': {'":': 2, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '76010', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 1882, 'n_tokens_neox': 1665, 'n_words': 280, 'path': 'docs/tutorials/Norman.ipynb', 'repo_name': 'theislab/CPA', 'size': '4776'}