Quick Start: Train MintFlow on a Single Tissue Section
Tutorial for basic training on a single tissue section
Creator: Amir Akbarnejad (aa36@sanger.ac.uk)
Affiliation: Wellcome Sanger Institute and University of Cambridge
Date of Creation: 23.06.2025
Date of Last Modificaion: 23.06.2025
To be able to run the notebook, the parts that you need to modify are specified by TODO:MODIFY:. The rest can be left untouched, as far as the goal is to run the notebook.
This notebook demonstrates how to train MintFlow on a single tissue section. This notebook is only for demonstration, and to get biologically meaningful results you may need longer training and/or different hyper-parameter settings.
1. Download the anndata object
Download this .h5ad file from google drive (link to the file)
and place it in a directory of you choice. Thereafter, set the variable path_anndata below to the path where you placed the.h5ad file.
Dataset source declaration: This anndata object was originally obtained from the following source: https://www.10xgenomics.com/datasets/xenium-prime-ffpe-human-skin
path_anndata = './NonGit/data_train_single_section.h5ad'
# TODO:MODIFY: set to the path where you've put the `.h5ad` file that you downloaded.
import os, sys
import yaml
import mintflow
import pickle
from tqdm.autonotebook import tqdm
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import torch
import pandas as pd
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
1. Read the defualt configurations
In this section 4 default configuration objects are loaded, which are later on customised. You only need to specify
num_tissue_sections_training: Number of tissue sections to be used for training.num_tissue_sections_evaluation: Number of tissue sections to be used for evaluation.
Same tissue sections can be used for both training and evaluation, in which case these two numbers are the same.
config_data_train, config_data_evaluation, config_model, config_training = mintflow.get_default_configurations(
num_tissue_sections_training=1,
num_tissue_sections_evaluation=1
)
2. Customise the 4 configurations
In this section we customise the four configurations returned by mintflow.get_default_configurations above.
2.1. Costomise config_data_train
MintFlow requires that each tissue section is saved in a separate anndata file on disk (i.e. one anndata object for each tissue section).
The .X field of each anndata object is required to have raw counts, in integer data type and “without” row-sum normalisation or log1p transformation.
The .obs field of each anndata object is required to have
A column that specifies cell type labels
A column that specifies a unique tissue section (i.e. slice) identifier. For each anndata object you can add a column to its
.obsfield that contains, e.g., the index or barcode of each tissue section that you’ve assiened to each tissue section.A column that specifies batch identifier to correct for batch effect (biological, technological, between-patient, etc.).
In this notebook we have a single tissue section, so no batch correction is needed here.
# configure tissue section 1 =========
config_data_train['list_tissue']['anndata1']['file'] = path_anndata
# the absolute path to anndata object of tissue section 1 on disk.
config_data_train['list_tissue']['anndata1']['obskey_cell_type'] = 'broad_celltypes'
# meaning that for the 1st tissue section, cell type labels are provided in `broad_celltypes` column of `adata.obs`.
config_data_train['list_tissue']['anndata1']['obskey_sliceid_to_checkUnique'] = 'info_id'
# meaning that for the 1st tissue section, tissue section ID (i.e. slice ID) is provided in `info_id` column of `adata.obs`
config_data_train['list_tissue']['anndata1']['obskey_x'] = 'x_centroid'
# meaning that for the 1st tissue section, spatial x coordinates are provided in `x_centroid` column of `adata.obs`
config_data_train['list_tissue']['anndata1']['obskey_y'] = 'y_centroid'
# meaning that for the 1st tissue section, spatial y coordinates are provided in `y_centroid` column of `adata.obs`
config_data_train['list_tissue']['anndata1']['obskey_biological_batch_key'] = 'info_id'
# meaning that for the 1st tissue section, batch identifier is provided in `info_id` column of `adata.obs`
config_data_train['list_tissue']['anndata1']['config_dataloader_train']['width_window'] = 100
# For tissue section one, the crop size of the customised dataloader desribed in Supplementary Fig. 16 of paper.
# The larger this number, the larger the tissue crops, and the bigger the subset of cells in each training iteration.
# This implies that more GPU memory would be required during training.
# In this notebook after calling `mintflow.setup_data` in Sec 4 the crop(s) are shown on tissue,
# with some information on image title which can help you tune this parameter.
# In the manuscript we used `width_window` values between 300 and 800 depending on dataset.
config_data_train['list_tissue']['anndata1']['config_neighbourhood_graph'] = {
'n_neighs': 5,
'set_diag': 'False',
'delaunay': 'False',
}
# The parameters for creating the neighbourhood graph for training tissue section 1
2.2. Costomise config_data_evaluation
The set of tissue sections for evaluation can be the same, in which case the same values can be used at the following.
Note that in the following cell instead of ['config_dataloader_train']['width_window'] we have ['config_dataloader_test']['width_window'].
# configure tissue section 1 =======================
config_data_evaluation['list_tissue']['anndata1']['file'] = path_anndata
# the absolute path to anndata object of tissue section 1 on disk.
config_data_evaluation['list_tissue']['anndata1']['obskey_cell_type'] = 'broad_celltypes'
# meaning that for the 1st tissue section, cell type labels are provided in `broad_celltypes` column of `adata.obs`
config_data_evaluation['list_tissue']['anndata1']['obskey_sliceid_to_checkUnique'] = 'info_id'
# meaning that for the 1st tissue section, tissue section ID (i.e. slice ID) is provided in `info_id` column of `adata.obs`
config_data_evaluation['list_tissue']['anndata1']['obskey_x'] = 'x_centroid'
# meaning that for the 1st tissue section, spatial x coordinates are provided in `x_centroid` column of `adata.obs`
config_data_evaluation['list_tissue']['anndata1']['obskey_y'] = 'y_centroid'
# meaning that for the 1st tissue section, spatial y coordinates are provided in `y_centroid` column of `adata.obs`
config_data_evaluation['list_tissue']['anndata1']['obskey_biological_batch_key'] = 'info_id'
# meaning that for the 1st tissue section, batch identifier is provided in `info_id` column of `adata.obs`
config_data_evaluation['list_tissue']['anndata1']['config_dataloader_test']['width_window'] = 100
# For tissue section one, the crop size of the customised dataloader desribed in Supplementary Fig. 16 of paper.
# The larger this number, the larger the tissue crops, and the bigger the subset of cells in each training iteration.
# This implies that more GPU memory would be required during training.
# In this notebook after calling `mintflow.setup_data` in Sec 4 the crop(s) are shown on tissue,
# with some information on image title which can help you tune this parameter.
# In the manuscript we used `width_window` values between 300 and 800 depending on dataset.
config_data_evaluation['list_tissue']['anndata1']['config_neighbourhood_graph'] = {
'n_neighs': 5,
'set_diag': 'False',
'delaunay': 'False',
}
# The parameters for creating the neighbourhood graph for evaluation tissue section 1
2.3. Customise config_model
None of model configuration are essential to tune “in this tutorial notebook”. So in this tutorial we leave config_model untouched. Please refer to our documentation for changes that you can make to config_model.
2.4. Customise config_training
A note about wandb: before proceeding, it is highligy recommended (though optional) to setup wandb and track/log different values during training.
To enable wandb: Go to https://wandb.ai/ and create an account
To disable wandb: set
config_training['flag_enable_wandb']in the below cell to ‘False’.
config_training['num_training_epochs'] = 20
# number of training epochs, i.e. the number of times the model sees the dataset during training.
config_training['flag_use_GPU'] = 'True'
# whether GPU is used.
config_training['flag_enable_wandb'] = 'True'
# if set to True, during training different loss terms are logged to wandb.
# It's highly recommended to enable wandb. Please refer to wandb website for more info: `wandb.ai`
config_training['wandb_project_name'] = 'MintFlow'
# wandb project name (ignored if `config_training['flag_enable_wandb']` is set to False)
config_training['wandb_run_name'] = 'Mintflow_Tutorial_Train_Single_Tissue_Section'
# wandb run name (ignored if `config_training['flag_enable_wandb']` is set to False)
3. Verify and post-process the four configurations
In this section we verify and postprocess the four configurations.
config_data_train = mintflow.verify_and_postprocess_config_data_train(config_data_train)
config_data_evaluation = mintflow.verify_and_postprocess_config_data_evaluation(config_data_evaluation)
config_model = mintflow.verify_and_postprocess_config_model(config_model, num_tissue_sections=len(config_data_train))
config_training = mintflow.verify_and_postprocess_config_training(config_training)
4. Setup the Data/Model/Trainer
Having created and verified the 4 configurations, in this section we create the variables data_mintflow, model, and trainer.
dict_all4_configs = {
'config_data_train':config_data_train,
'config_data_evaluation':config_data_evaluation,
'config_model':config_model,
'config_training':config_training
}
data_mintflow = mintflow.setup_data(dict_all4_configs=dict_all4_configs)
model = mintflow.setup_model(
dict_all4_configs=dict_all4_configs,
data_mintflow=data_mintflow
)
trainer = mintflow.Trainer(
dict_all4_configs=dict_all4_configs,
model=model,
data_mintflow=data_mintflow
)
5. Train the Model
Set the variable path_ouptput_files below to the path where you want the training files (checkpoints etc) to be saved.
path_ouptput_files = "./NonGit/Outputs_TutorialNoboteok1"
# TODO:MODIFY: the path where checkpoints and other files are saved during training.
for index_epoch in tqdm(range(config_training['num_training_epochs']), desc='Training epoch'):
'''
IMPORTANT NOTE: To change the number of epochs, set `config_training['num_training_epochs']` in previous cells of this notebook
and please refrain from changing the for loop here to, e.g., `for index_epoch in tqdm(range(10), ...)`.
Because MintFlow's annealing module presumes that the number of epochs equals `config_training['num_training_epochs']`.
'''
# train for one epoch
trainer.train_one_epoch()
# get/save the predictions
predictions = mintflow.predict(
device=device,
dict_all4_configs=dict_all4_configs,
data_mintflow=data_mintflow,
model=model,
evalulate_on_sections="all",
)
with open(os.path.join(path_ouptput_files, "predictions_epoch_{}.pkl".format(index_epoch)), 'wb') as f:
pickle.dump(
predictions,
f
)
# evaluate the model and save the evaluation result for this checkpoint
df_evaluation_result = mintflow.evaluate_by_known_signalling_genes(
device=device,
dict_all4_configs=dict_all4_configs,
data_mintflow=data_mintflow,
model=model,
evalulate_on_sections='all',
optional_list_colvaltype_toadd=[['training_epoch', index_epoch, 'category']]
)
df_evaluation_result.to_pickle(
os.path.join(
path_ouptput_files,
'df_evaluation_result_epoch_{}.pkl'.format(index_epoch)
)
)
# save the checkpoint
mintflow.dump_checkpoint(
model=model,
data_mintflow=data_mintflow,
dict_all4_configs=dict_all4_configs,
path_dump=os.path.join(path_ouptput_files, "checkpoint_epoch_{}.pt".format(index_epoch)),
)
8. Select the best checkpoint
To perform the analysis you can either pick the last checkpoint, or alternatively, you can pick the best checkpoint by inspceting the dumped df_evaluation_result objects.
If the disentanglement is successful, the violinplots/boxplots that correspond to signalling genes should be skewed towards 1.0 (like the orange violin/box plots in the provided sample figures below) while for other genes they should be skewed towards 0.0 (like the blue violin/box plots in the provided sample figures below)
Sample violin/box plots:
a sample violinplot that shows succesful disentanglement: (link to the figure)
a sample boxplot that shows succesful disentanglement: (link to the figure)
To produce violin/box plots for a specific checkpoint, you can run the below cells.
To arrive at a metric, you can compute a discrepancy metric (e.g. Wasserstein distance) between the two groups specified in is_among_signalling_genes.
df_toinspect = pd.read_pickle(
os.path.join(path_ouptput_files, 'df_evaluation_result_epoch_{}.pkl'.format(10))
)
# TODO:MODIFY: change `10` to the checkpoint index that you want to inspect
sns.violinplot(
data=df_toinspect[
df_toinspect['read_count'] > 30.0
],
x='training_epoch',
y="fraction_assigned_to_Xmic",
hue="is_among_signalling_genes",
cut=0
)
sns.boxplot(
data=df_toinspect[
df_toinspect['read_count'] > 30.0
],
x='training_epoch',
y="fraction_assigned_to_Xmic",
hue="is_among_signalling_genes"
)
9. Use MintFlow predictions for analysis
Having selected the best (or a good) chcekpoint, set the variable index_selected_checkpoint below and run the following cells.
index_selected_checkpoint = 19
# TODO:MODIFY: the index of the best (or a good) checkpoint that you selected.
Load predictions for the selected checkpoint.
with open(os.path.join(path_ouptput_files, "predictions_epoch_{}.pkl".format(index_selected_checkpoint)), 'rb') as f:
predictions_selected_checkpoint = pickle.load(f)
MintFlow predictions are available in predictions_selected_checkpoint. In particular, the intrinsic- and microenvironment-induced components of expression are available in
predictions_selected_checkpoint['TissueSection 0 (zero-based)']['MintFlow_Xint']predictions_selected_checkpoint['TissueSection 0 (zero-based)']['MintFlow_Xmic']
For example we can compute MintFlow’s per-cell microenvironment signalling score as follows:
Xint = predictions_selected_checkpoint['TissueSection 0 (zero-based)']['MintFlow_Xint']
Xmic = predictions_selected_checkpoint['TissueSection 0 (zero-based)']['MintFlow_Xmic']
MintFlow_microenvironment_signalling_score = Xmic.sum(1) / (Xint+Xmic).sum(1)