from anndata import AnnData
import pandas as pd # Please install pandas and matplotlib before you run this example
import matplotlib.pyplot as plt
import matplotlib
# Set the Matplotlib backend to 'Agg'
matplotlib.use('Agg')
import numpy as np
import scipy.sparse as sp
import sklearn.metrics as mpd
import gc
import os
import glob
from scipy.spatial.distance import pdist, squareform
from scMEDAL.utils.utils import read_adata, min_max_scaling,save_adata,calculate_zscores
# I run it with Aixa_genomap
"""Adapted genomap functions from https://github.com/xinglab-ai/genomap/blob/main/genomap/genomap.py and created utils for plotting"""
def createInteractionMatrix(data, metric='correlation'):
"""
Function from genomap github: https://github.com/xinglab-ai/genomap/blob/main/genomap/genomap.py
I added here because sometimes it throws nan values
Returns the interaction matrix among the genes
Parameters
----------
data : ndarray, shape (cellNum, geneNum)
gene expression data in cell X gene format. Each row corresponds
to one cell, whereas each column represents one gene
metric : 'string'
Metric for computing the genetic interaction
Returns
-------
interactMat : ndarray, shape (geneNum, geneNum)
pairwise interaction matrix among genes
"""
interactMat=mpd.pairwise_distances(data.T,metric=metric)
return interactMat
def create_gene_coordinates_mapping(projMat, gene_names, num_genes=2916, rowNum=54, colNum=54):
"""
Applies a projection matrix to a diagnostic array of gene indices, reshapes the
result into a 54×54 grid, and maps each gene to an (x,y) coordinate. Note that
if 'projMat' contains fractional assignments (e.g., from Gromov-Wasserstein),
rounding to int can cause collisions and potentially overwrite or miss certain
gene indices. Any missing genes are reported in the console.
"""
# Create and transform the diagnostic matrix
diagnostic_matrix = np.arange(num_genes).reshape(1, -1)
transformed_indices = np.matmul(diagnostic_matrix, projMat).flatten()
px = np.round(transformed_indices, 2)
# Reshape into a rowNum, colNum matrix
# Order ='F' returns transposed indexes
# genomaps_diagnostic = np.reshape(px, (rowNum, colNum), order='F')
# Reshape into a rowNum, colNum matrix
genomaps_diagnostic = np.reshape(px, (rowNum, colNum), order='C') # Default is 'C' order (Not transposed)
# Map genes to coordinates
gene_to_coordinates = {}
count = 0
for x in range(rowNum):
for y in range(colNum):
gene_index = int(genomaps_diagnostic[x, y])
if 0 <= gene_index < len(gene_names):
gene_name = gene_names[gene_index]
gene_to_coordinates[gene_name] = (x, y)
count += 1
else:
print(f"Index {gene_index} not assigned to a gene name")
expected_indices = set(range(len(gene_names)))
found_indices = set(int(genomaps_diagnostic[i, j]) for i in range(rowNum) for j in range(colNum) if 0 <= int(genomaps_diagnostic[i, j]) < len(gene_names))
missing_indices = expected_indices - found_indices
print("Missing indices:", missing_indices)
return gene_to_coordinates
# Example usage:
# projMat = np.random.rand(2916, 2916) # Example initialization; replace with actual matrix
# gene_names = ['Gene1', 'Gene2', ..., 'Gene2916'] # Define a list of gene names
# mapping = create_gene_name_coordinates_mapping(projMat, gene_names)
def construct_genomap(data,rowNum,colNum,epsilon=0,num_iter=1000):
from genomap.genomapOPT import create_space_distributions, gromov_wasserstein_adjusted_norm
from genomap.genomap import createMeshDistance
"""
Adapted function from genomap github: https://github.com/xinglab-ai/genomap/blob/main/genomap/genomap.py
Constructs 2D "genomaps" by coupling a gene-gene interaction matrix with
a grid distance matrix using Gromov-Wasserstein. Note that GW transport
yields fractional assignments, so forcibly rounding positions to integer
grid coordinates can cause collisions (duplicate integer indices) and
result in missing or overwritten gene indices.
Returns the constructed genomaps
I added code to avoid nan in interactions matrix
Parameters
----------
data : ndarray, shape (cellNum, geneNum)
gene expression data in cell X gene format. Each row corresponds
to one cell, whereas each column represents one gene
rowNum : int,
number of rows in a genomap
colNum : int,
number of columns in a genomap
Returns
-------
genomaps : ndarray, shape (rowNum, colNum, zAxisDimension, cell number)
genomaps are formed for each cell. zAxisDimension is more than
1 when 3D genomaps are created.
"""
sizeData=data.shape
numCell=sizeData[0]
numGene=sizeData[1]
# distance matrix of 2D genomap grid
distMat = createMeshDistance(rowNum,colNum)
# gene-gene interaction matrix
interactMat = createInteractionMatrix(data, metric='correlation')
# I added the following line to avoid nan in interactions matrix
interactMat = np.nan_to_num(interactMat)
totalGridPoint=rowNum*colNum
if (numGene<totalGridPoint):
totalGridPointEff=numGene
else:
totalGridPointEff=totalGridPoint
M = np.zeros((totalGridPointEff, totalGridPointEff))
p, q = create_space_distributions(totalGridPointEff, totalGridPointEff)
# Coupling matrix
T = gromov_wasserstein_adjusted_norm(
M, interactMat, distMat[:totalGridPointEff,:totalGridPointEff], p, q, loss_fun='kl_loss', epsilon=epsilon,max_iter=num_iter)
projMat = T*totalGridPoint
# Data projected onto the couping matrix
projM = np.matmul(data, projMat)
genomaps = np.zeros((numCell,rowNum, colNum, 1))
px = np.asmatrix(projM)
# Formation of genomaps from the projected data
for i in range(0, numCell):
dx = px[i, :]
fullVec = np.zeros((1,rowNum*colNum))
fullVec[:dx.shape[0],:dx.shape[1]] = dx
#ex = np.reshape(fullVec, (rowNum, colNum), order='F').copy()
ex = np.reshape(fullVec, (rowNum, colNum), order='C').copy()
genomaps[i, :, :, 0] = ex
geno_dict = {"genomaps":genomaps,"T":T,"totalGridPoint":totalGridPoint}
return geno_dict
#return genomaps
def process_data_genomap(inputs_path, recon_path=None, ncells=50000, ngenes=2916, return_input_zscores=False,issparse=True):
"""
Process input and recon data, calculate z-scores, and return a subset of the dataset.
Args:
inputs_path (str): Path to the input data file.
recon_path (str,optional): Path to the reconstruction data file.
ncells (int): Number of cells to subset. Default is 50000.
ngenes (int): Number of genes to subset. Default is 2916.
return_input_zscores (bool): Flag to return z-scores from input data. Default is False.
issparse (bool): Flag to read adata as sparse matrix. Default is True.
Returns:
dict: Nested dictionary containing subsets of z-scores and AnnData objects for input and recon data.
"""
# Load input data
X_input, var, obs = read_adata(inputs_path, issparse=issparse)
adata_input_train = AnnData(X_input, obs=obs, var=var)
# Process input data: Get z-scores
X_input_dense = adata_input_train.X.toarray() if sp.issparse(adata_input_train.X) else adata_input_train.X
input_z_scores = calculate_zscores(X_input_dense) if return_input_zscores else None
# Subset dataset to ncells, ngenes
subset_input_z_scores = input_z_scores[:ncells, :ngenes] if input_z_scores is not None else None
if subset_input_z_scores is not None:
print("subset_input_z_scores shape:", subset_input_z_scores.shape)
out_dict = {"input": {"z_scores": subset_input_z_scores,
"adata": adata_input_train}}
if recon_path is not None:
# Load recon data
loaded_data = np.load(recon_path, allow_pickle=False)
X_recon = loaded_data.item() if isinstance(loaded_data, np.ndarray) and loaded_data.dtype == object else loaded_data
# Load cell metadata
meta_path = glob.glob(inputs_path+'/*meta*.csv')[0]
obs = pd.read_csv(meta_path)
adata_recon_train = AnnData(X_recon, obs=obs, var=var)
# Process recon data: Get z-scores
X_recon_dense = adata_recon_train.X.toarray() if sp.issparse(adata_recon_train.X) else adata_recon_train.X
recon_z_scores = calculate_zscores(X_recon_dense)
# Subset dataset to ncells, ngenes
subset_recon_z_scores = recon_z_scores[:ncells, :ngenes]
print("subset_recon_z_scores shape:", subset_recon_z_scores.shape)
out_dict["recon"] = {"z_scores": subset_recon_z_scores,"adata": adata_recon_train}
return out_dict
def plot_genomaps(genoMaps, labels, filename):
"""
Plots genoMaps with a global color scale and includes a color bar.
Args:
genoMaps (numpy.ndarray): Array of genoMaps to plot, expected shape (num_cells, height, width, channels).
labels (list): List of labels for each cell to be used in the plot titles.
filename (str): The file path to save the figure.
Returns:
None
"""
# Determine the global min and max for the color scale
global_min = np.percentile(genoMaps, 1)
global_max = np.percentile(genoMaps, 98)
# Create a figure with 10x5 subplots
fig, axes = plt.subplots(10, 5, figsize=(20, 40))
# Iterate over the subplots and plot each cell
for i, ax in enumerate(axes.flat):
if i < len(genoMaps):
findI = genoMaps[i, :, :, :]
#for now i wont clip data
#clipped_data = np.clip(findI, global_min, global_max)
im = ax.imshow(findI, origin='lower', extent=[0, 10, 0, 10], aspect=1, vmin=global_min, vmax=global_max, cmap='viridis')
ax.set_title(f'Genomap of cell {i+1}\n' + labels[i])
else:
ax.axis('off')
# Add a color bar below the figure
cbar_ax = fig.add_axes([0.15, 0.05, 0.7, 0.02]) # [left, bottom, width, height]
fig.colorbar(im, cax=cbar_ax, orientation='horizontal', shrink=0.8)
plt.tight_layout(rect=[0, 0.1, 1, 1]) # Adjust layout to make room for colorbar
# Save the figure
plt.savefig(filename, bbox_inches='tight')
plt.close(fig)
gc.collect()
def process_and_plot_genomaps_singlepath(inputs_path, ncells, ngenes, rowNum, colNum, epsilon, num_iter, output_folder,genomap_name,gene_names=None):
"""
Process the input data to generate z-scores and genoMaps, then plot and save the results.
Args:
inputs_path (str): Path to the input data file.
ncells (int): Number of cells to subset.
ngenes (int): Number of genes to subset.
rowNum (int): Number of rows for the genoMap.
colNum (int): Number of columns for the genoMap.
epsilon (float): Epsilon value for the genoMap construction.
num_iter (int): Number of iterations for the genoMap construction.
output_folder (str): Path to the folder where outputs will be saved.
genomap_name (str): Name identifier for the genoMap.
gene_names (list,optional): List with the gene names of the genomap. The number of gene_names has to be equal to ngenes
"""
# Create output folder if it doesn't exist
if not os.path.exists(output_folder):
os.makedirs(output_folder)
# Process the data, returning input z-scores only for the first iteration
result = process_data_genomap(inputs_path, recon_path=None, ncells=ncells, ngenes=ngenes, return_input_zscores=True)
# Get input z-scores and construct genoMaps for the input data
input_z_scores = result['input']['z_scores']
genoMaps_input = construct_genomap(input_z_scores, rowNum, colNum, epsilon=epsilon, num_iter=num_iter)
projMat = genoMaps_input["T"]*genoMaps_input["totalGridPoint"]
# Save genoMaps and transition matrix (T) for input data
np.save(os.path.join(output_folder, f'genomap_{genomap_name}.npy'), genoMaps_input["genomaps"])
np.save(os.path.join(output_folder, f'T_input_{genomap_name}.npy'), genoMaps_input["T"])
# Plot and save the genoMaps for input data
input_ct_labels = result['input']['adata'].obs['celltype'].tolist()
input_batch_labels = result['input']['adata'].obs['batch'].tolist()
input_labels = [f'{ct}_{b}' for ct, b in zip(input_batch_labels, input_ct_labels)]
plot_genomaps(genoMaps_input["genomaps"], input_labels, os.path.join(output_folder, f'{genomap_name}.png'))
# Map genes to genomap
if gene_names is not None:
gene_to_coordinates = create_gene_coordinates_mapping(projMat, gene_names, ngenes, rowNum, colNum)
gene_to_coordinates_df = pd.DataFrame(gene_to_coordinates).T
gene_to_coordinates_df.columns=["pixel_i","pixel_j"]
gene_to_coordinates_df.to_csv(os.path.join(output_folder, f'gene_coordinates_{genomap_name}.csv'))
def process_and_plot_genomaps(df, models, types, splits, rowNum, colNum, output_folder, return_input_zscores=True, ncells=50000, ngenes=2916, epsilon=0.0, num_iter=200):
"""
Processes data, constructs genoMaps, and plots/saves the results for specified models, types, and splits.
Args:
df (DataFrame): DataFrame containing the paths to input and reconstruction data.
models (list): List of model names.
types (list): List of data types (e.g., 'train', 'test', validate').
splits (list): List of split indices.
rowNum (int): Number of rows for genoMap construction.
colNum (int): Number of columns for genoMap construction.
output_folder (str): Path to the folder where the results will be saved.
return_input_zscores (bool): Flag to determine if input z-scores should be returned.
Returns:
dict: Dictionary containing the processed results.
"""
# Create output folder if it doesn't exist
if not os.path.exists(output_folder):
os.makedirs(output_folder)
# Initialize the results dictionary
results = {}
# Loop through each model
for model in models:
results[model] = {}
# Loop through each data type (train, test, validate)
for Type in types:
results[model][Type] = {}
# Loop through each split
for Split in splits:
print("\nProcessing:",model,Type,Split)
# All the batches have the same input for the same type and split
k = 0
results[model][Type][Split] = {
'process_data_result': {},
'genoMaps_input_train': "",
'genoMaps_recon_train': {}
}
# Get the input and reconstruction paths from the DataFrame
inputs_path = df.loc[
(df["Key"] == model) &
(df["Type"] == Type) &
(df["Split"] == Split),
"InputsPath"
].values[0]
recon_paths = df.loc[
(df["Key"] == model) &
(df["Type"] == Type) &
(df["Split"] == Split),
"ReconPath"
].values
# Process each reconstruction path
print("\nrecon paths: ",recon_paths)
for recon_path in recon_paths:
recon_prefix = recon_path.split("/")[-1].split(".npy")[0]
print("\nProcessing ",recon_prefix)
# Process the data, returning input z-scores only for the first iteration
result = process_data_genomap(
inputs_path, recon_path,
ncells=ncells, ngenes=ngenes,
return_input_zscores=(k == 0)
)
# Create genomaps if return_input_zscores=True and k==0, ie first recon for the same type and split
if (k == 0) & return_input_zscores:
# Get input z-scores and construct genoMaps for the input data
input_z_scores = result['input']['z_scores']
print("Computing genomaps for input..")
genoMaps_input = construct_genomap(
input_z_scores, rowNum, colNum,
epsilon=epsilon, num_iter=num_iter
) if input_z_scores is not None else None
results[model][Type][Split]['genoMaps_input_train'] = genoMaps_input["genomaps"]
if genoMaps_input["genomaps"] is not None:
# Save genoMaps and transition matrix (T) for input data
save_genomaps(genoMaps_input, output_folder, model, Type, Split, data_type="input")
# Plot and save the genoMaps for input data
input_ct_labels = result['input']['adata'].obs['celltype'].tolist()
input_batch_labels = result['input']['adata'].obs['batch'].tolist()
input_labels = [f'{ct}_{b}' for ct, b in zip(input_batch_labels, input_ct_labels)]
plot_genomaps(genoMaps_input["genomaps"], input_labels, os.path.join(output_folder, f'{model}_{Type}_{Split}_input.png'))
# Get reconstruction z-scores and construct genoMaps for the reconstruction data
recon_z_scores = result['recon']['z_scores']
print("Computing genomaps for recon..")
genoMaps_recon = construct_genomap(
recon_z_scores, rowNum, colNum,
epsilon=epsilon, num_iter=num_iter
)
results[model][Type][Split]['process_data_result'][recon_prefix] = result
results[model][Type][Split]['genoMaps_recon_train'][recon_prefix] = genoMaps_recon["genomaps"]
if genoMaps_recon["genomaps"] is not None:
# Save genoMaps and transition matrix (T) for reconstruction data
save_genomaps(genoMaps_recon, output_folder, model,Type, Split=Split,data_type="recon",recon_prefix = recon_prefix)
# Plot and save the genoMaps for reconstruction data
recon_ct_labels = result['recon']['adata'].obs['celltype'].tolist()
n_labels = len(recon_ct_labels)
# recon original
if recon_prefix in ["recon_train", "recon_val"]:
recon_batch_labels = result['recon']['adata'].obs['batch'].tolist()
# cf recon
else:
recon_batch_labels = [recon_prefix ]*n_labels
recon_labels = [f'{ct}_{b}' for ct, b in zip(recon_batch_labels, recon_ct_labels)]
plot_genomaps(genoMaps_recon["genomaps"], recon_labels, os.path.join(output_folder, f'genomap_{recon_prefix}_{model}_{Split}.png'))
k += 1
gc.collect()
return results
def save_genomaps(genoMaps, output_folder, model, Type, Split, data_type, recon_prefix=None):
"""
Saves the genoMaps and transition matrices as .npy files.
Args:
genoMaps (dict): Dictionary containing the genoMaps and transition matrices.
output_folder (str): Path to the folder where the results will be saved.
model (str): Model name.
Type (str): Data type (e.g., 'train', 'test', validate').
Split (int): Split index.
data_type (str): Type of data ('input' or 'recon').
recon_prefix (str): Prefix for reconstruction files (optional).
"""
if data_type == "input":
np.save(os.path.join(output_folder, f'genomap_input_{model}_{Type}_{Split}.npy'), genoMaps["genomaps"])
np.save(os.path.join(output_folder, f'T_input_{model}_{Type}_{Split}.npy'), genoMaps["T"])
else:
np.save(os.path.join(output_folder, f'genomap_{recon_prefix}_{model}_{Split}.npy'), genoMaps["genomaps"])
np.save(os.path.join(output_folder, f'T_{recon_prefix}_{model}_{Split}.npy'), genoMaps["T"])
import numpy as np
def get_genomapfromT(t_files_path, inputs_path, recon_path,ncells=50000, ngenes=2916,colNum = 54 ,rowNum = 54,gene_names=None):
"""
Generates genomaps from given input data and projection matrix.
Parameters:
- t_files_path (str): Path to the T matrix file.
- inputs_path (str): Path to the input data.
- recon_path (str): Path to the reconstructed data.
- ncells (int, optional): Number of cells in the data. Default is 50000.
- ngenes (int, optional): Number of genes in the data. Default is 2916.
- colNum (int, optional): Number of columns in the genomap. Default is 54.
- rowNum (int, optional): Number of rows in the genomap. Default is 54.
- gene_names (list, optional): list of gene names for coordinates mapping
Returns:
dict: A dictionary containing:
- 'genomaps': A 4D numpy array representing the genomaps.
- 'genes_coor_map': A dictionary mapping gene names to their coordinates in the genomap.
"""
# Load the T matrix
T = np.load(t_files_path)
# Process data and get z scores
out_dict = process_data_genomap(inputs_path, recon_path, ncells, ngenes, return_input_zscores=False)
# Get genomap without perturbation
data = out_dict['recon']["z_scores"].copy()
sizeData = data.shape
print("sizeData", sizeData)
# Define parameters for genomap construction
numCell = sizeData[0]
numGene = sizeData[1]
totalGridPoint = ngenes
# Project data onto the coupling matrix
projMat = T * totalGridPoint
projM = np.matmul(data, projMat)
# Initialize genomaps
genomaps = np.zeros((numCell, rowNum, colNum, 1))
px = np.asmatrix(projM)
# Formation of genomaps from the projected data
for i in range(numCell):
dx = px[i, :]
fullVec = np.zeros((1, rowNum * colNum))
fullVec[:dx.shape[0], :dx.shape[1]] = dx
#ex = np.reshape(fullVec, (rowNum, colNum), order='F').copy()
ex = np.reshape(fullVec, (rowNum, colNum), order='C').copy()
genomaps[i, :, :, 0] = ex
# Map genes to genomap
if gene_names is not None:
gene_to_coordinates = create_gene_coordinates_mapping(projMat, gene_names, ngenes, rowNum, colNum)
return {"genomaps": genomaps, "genes_coor_map": gene_to_coordinates}
else:
return {"genomaps": genomaps}
import matplotlib.pyplot as plt
def plot_genomap_with_genes(cell_geno, top10hvg_genes, gene_to_coordinates):
"""
Plots a genomap with identified genes highlighted in red.
Parameters:
- cell_geno: numpy array representing the genomap data matrix.
- top10hvg_genes: list of top 10 highly variable genes.
- gene_to_coordinates: dictionary mapping gene names to their coordinates.
"""
# Set the size of the figure
fig, ax = plt.subplots(figsize=(10, 10)) # Width and height in inches
# Plot the genomap using an appropriate colormap
cax = ax.imshow(cell_geno, cmap='viridis')
# Coordinates for the first 20 genes, assuming you have a list of gene names in the correct order
first_n_genes = top10hvg_genes
coordinates = [gene_to_coordinates[gene] for gene in first_n_genes]
# Highlight each gene coordinate
for coord, gene in zip(coordinates, first_n_genes):
x, y = coord
ax.scatter(y, x, color='red', s=100, alpha=0.3) # Adjust size and color as needed
gene_str = gene.decode('utf-8') if isinstance(gene, bytes) else gene
ax.text(y, x, gene_str, color='white', ha='right', va='top')
# Add a colorbar to the plot
fig.colorbar(cax, orientation='vertical', fraction=0.02, pad=0.04)
# Adjust layout and show the plot
plt.tight_layout()
plt.show()
def sample_cells(obs, n_cells, celltype=None, batch_col="batch", force_batches=None, seed=None):
"""
Sample cells from a dataframe while ensuring representation from specified batches and cell types.
Parameters:
obs (pd.DataFrame): The dataframe containing cell data with columns including 'celltype' and 'batch'.
n_cells (int): The total number of cells to sample.
celltype (str or list, optional): The cell type(s) to be included in the sampling. Can be a single cell type (str) or a list of cell types (list).
batch_col (str): The name of the column representing batches in the dataframe.
force_batches (list, optional): A list of batch names to ensure at least one cell from each batch is included in the sample.
seed (int, optional): Random seed for reproducibility.
Returns:
pd.DataFrame: A dataframe containing the sampled cells.
Raises:
ValueError: If a specified batch or cell type does not exist in the dataframe.
"""
# Initialize selected_indices list to keep track of selected cells
selected_indices = []
# Ensure at least one cell from each forced batch and specified celltype (if provided) is included
if force_batches:
for batch in force_batches:
# Filter cells belonging to the current batch
batch_cells = obs[obs[batch_col] == batch]
if celltype is not None:
if isinstance(celltype, str):
# Further filter cells of the specified celltype within the batch
batch_cells = batch_cells[batch_cells["celltype"] == celltype]
if batch_cells.empty:
print(f"No cells available for celltype {celltype} in batch {batch}. Skipping batch.")
continue
# Randomly select one cell from the batch (and celltype if specified)
selected_index = batch_cells.sample(n=1, random_state=seed).index[0]
selected_indices.append(selected_index)
print(f"Selected 1 cell from batch {batch} with celltype {celltype}.")
elif isinstance(celltype, list):
for ct in celltype:
ct_batch_cells = batch_cells[batch_cells["celltype"] == ct]
if ct_batch_cells.empty:
print(f"No cells available for celltype {celltype} in batch {batch}. Skipping batch.")
continue
# Randomly select one cell from the batch for each celltype
selected_index = ct_batch_cells.sample(n=1, random_state=seed).index[0]
selected_indices.append(selected_index)
print(f"Selected 1 cell from batch {batch} with celltype {ct}.")
else:
if batch_cells.empty:
print(f"No cells available in batch {batch}. Skipping batch.")
continue
# Randomly select one cell from the batch without celltype constraint
selected_index = batch_cells.sample(n=1, random_state=seed).index[0]
selected_indices.append(selected_index)
print(f"Selected 1 cell from batch {batch} without specific celltype constraint.")
# Remove selected cells to avoid re-sampling them
obs_remaining = obs.drop(selected_indices)
# Sample remaining cells based on celltype
if celltype is not None:
if isinstance(celltype, str):
# Filter cells of the specified celltype
celltype_obs = obs_remaining[obs_remaining["celltype"] == celltype]
# Randomly select the remaining required cells from the filtered cells
additional_cells = celltype_obs.sample(n=n_cells - len(selected_indices), random_state=seed)
print(f"Selected {n_cells - len(selected_indices)} additional cells of celltype {celltype}.")
elif isinstance(celltype, list):
additional_cells_list = []
cells_needed_per_type = (n_cells - len(selected_indices)) // len(celltype)
for ct in celltype:
ct_cells = obs_remaining[obs_remaining["celltype"] == ct]
additional_cells_list.append(ct_cells.sample(n=cells_needed_per_type, random_state=seed))
additional_cells = pd.concat(additional_cells_list)
print(f"Selected {n_cells - len(selected_indices)} additional cells from specified celltypes.")
else:
# Randomly select the remaining required cells from the entire dataset
additional_cells = obs_remaining.sample(n=n_cells - len(selected_indices), random_state=seed)
print(f"Selected {n_cells - len(selected_indices)} additional cells without specific celltype constraint.")
# Combine the selected indices from forced batches and the additional sampled cells
final_selected_indices = pd.Index(selected_indices).append(additional_cells.index)
obs_final = obs.loc[final_selected_indices]
return obs_final
def create_count_matrix_multibatch(recon_prefix, recon_paths, obs, var, n_genes, n_cells, n_batches, out_path, add_inputs_fe=True, n_inputs_fe=2, celltype=None, save_data=False, scaling="min_max", issparse=True,seed=42,force_batches=None):
"""
Create a concatenated count matrix from cell reconstructions from multiple batches, optionally filtered by cell type.
This concatenated matrix will be used to generate a genomap.
Parameters:
recon_prefix (list of str): List of prefixes indicating the type of data in recon_paths.
recon_paths (list of str): List of paths to the reconstructed data files.
obs (DataFrame): DataFrame containing observation (cell) metadata.
var (DataFrame): DataFrame containing variable (gene) metadata.
n_genes (int): Number of genes to include in the count matrix.
n_cells (int): Number of cells to include from each batch.
n_batches (int): Number of batches to process.
out_path (str): Output directory path to save the concatenated AnnData object.
add_inputs_fe (bool, optional): Whether to add cells from the original input data and from the fixed effects recon. Default is True.
n_inputs_fe (int, optional): If add_inputs_fe = True. Specify how many reconstructions you want to add to the count matrix. Default is 2.
One for inputs and another one for fixed effects (fe). You could also add a fe classifier recon and a base autoencoder recon.
celltype (str or list, optional): Specific cell type(s) to filter for. If None, include all cell types. Default is None.
save_data (bool, optional): Whether to save the resulting AnnData object to disk. Default is False.
issparse (bool, optional): Flag to read sparse count matrix. Default is True.
seed (int,optional): seed for reproducibility. Default = 42.
force_batches (list, optional): A list of batch names to ensure at least one cell from each batch is included in the sample.
Returns:
None
"""
import numpy as np
import pandas as pd
from anndata import AnnData
import os
# Initialize empty arrays and DataFrame
x = np.empty((0, n_genes))
obs_combined = pd.DataFrame()
# Use sample_cells function to select the appropriate cells
if celltype is not None:
obs = sample_cells(obs, n_cells, celltype=celltype, batch_col="batch", force_batches=force_batches, seed=seed)
selected_indices = obs.index
if add_inputs_fe:
n_recon_prefix = n_batches + n_inputs_fe
else:
n_recon_prefix = n_batches
for recon_pref, recon_path in zip(recon_prefix[0:n_recon_prefix], recon_paths[0:n_recon_prefix]):
print(recon_pref)
if "input" in recon_pref:
X, _, _ = read_adata(recon_path, issparse=issparse)
if issparse:
X = X.toarray()
# Input needs to be scaled because reconstructions are in the scaled space
# Apply scaling to X based on the scaling parameter.
if scaling == "min_max":
X = min_max_scaling(X)
elif scaling == "z_scores":
X = calculate_zscores(X)
else:
X = np.load(recon_path)
if celltype is None:
x_i = X[:n_cells, :n_genes]
obs_i = obs.iloc[:n_cells].copy()
else:
x_i = X[selected_indices, :n_genes]
obs_i = obs.copy()
# Add a new column that indicates the recon prefix
obs_i["recon_prefix"] = recon_pref
x = np.concatenate((x, x_i), axis=0)
obs_combined = pd.concat([obs_combined, obs_i], axis=0)
obs_combined.reset_index(drop=True, inplace=True)
print("Shape of concatenated x:", x.shape)
print("Shape of concatenated obs:", obs_combined.shape)
if isinstance(celltype, str):
celltype_name = celltype.replace("/", "")
elif isinstance(celltype, list):
celltype_name = "_".join([ct.replace("/", "") for ct in celltype])
else:
celltype_name = None
folder_name = f"CMmultibatch_{n_cells}_cells_per_batch_{n_batches}batches" if not celltype_name else f"CMmultibatch_{n_cells}_cells_per_batch_{n_batches}batches_{celltype_name}"
if add_inputs_fe:
folder_name += f"_with_{n_inputs_fe}fe_input"
folder_path = os.path.join(out_path, folder_name)
if not os.path.exists(folder_path):
print("created folder path:",folder_path)
os.makedirs(folder_path)
adata_multibatch = AnnData(x, obs=obs_combined, var=var.iloc[:n_genes, :])
if save_data:
save_adata(adata_multibatch, folder_path)
return adata_multibatch
def find_intersection_batches(obs_combined, celltypes):
"""
Find the intersection of batches that are present in at least one cell from each of the selected celltypes.
Parameters:
obs_combined (DataFrame): DataFrame containing observation (cell) metadata with a 'batch' column.
celltypes (list): List of cell types to check for overlapping batch values.
Returns:
intersection_batches (set): Set of batches that are common across all specified cell types.
"""
# Initialize the set with batches from the first cell type
intersection_batches = set(obs_combined[obs_combined["celltype"] == celltypes[0]]["batch"].unique())
# Iterate through the remaining cell types and find the intersection
for celltype in celltypes[1:]:
batches = set(obs_combined[obs_combined["celltype"] == celltype]["batch"].unique())
intersection_batches.intersection_update(batches)
return intersection_batches
def select_cells_from_batches(obs_combined, celltypes, batches_to_select_from,seed = 42,cell_id_col="_index"):
"""
Select one cell from each cell type within the specified batches.
Parameters:
obs_combined (DataFrame): DataFrame containing observation (cell) metadata.
celltypes (str or list): A single cell type (str) or list of cell types to select cells from.
batches_to_select_from (list): List of batches to select cells from.
cell_id_col(str): column with the cell ids
Returns:
cell_ids_2plot (list): List of selected cell IDs.
"""
import random
random.seed(seed)
cell_ids_2plot = []
if isinstance(celltypes, str):
celltypes = [celltypes] # Convert to list for uniform processing
for batch in batches_to_select_from:
for celltype in celltypes:
cells_in_batch = obs_combined[(obs_combined['batch'] == batch) & (obs_combined['celltype'] == celltype)][cell_id_col].values
if len(cells_in_batch) > 0:
selected_cell = random.choice(cells_in_batch)
cell_ids_2plot.append(selected_cell)
return cell_ids_2plot
def compute_cell_stats_acrossbatchrecon(genomap, cell_indexes_batch_cf, genomap_coordinates, statistic='std', n_top_genes=10, order='C',path_2_genomap='',file_name='cell_id'):
"""
Calculate standard deviation or variance for genomaps of a single cell across batch reconstructions.
Update genomap_coordinates DataFrame.
Args:
genomap: 4D numpy array with genomap data. Axis = 0 indicates individual cells.
cell_indexes_batch_cf: List or array of cell indexes for batch reconstructions of single cell. The same cell was reconstructed for multiple batches.
genomap_coordinates: DataFrame containing gene names and pixel coordinates.
statistic: 'std' for standard deviation, 'var' for variance.
n_top_genes: Number of top genes to identify.
order: 'C' for default coordinate order, 'F' for transposed coordinates (i.e., pixel_i and pixel_j are swapped).
path_2_genomap (str): path to genomap directory. To save stats df.
file_name (str): Name to save the file. Default:"cell_id".
Returns:
Updated genomap_coordinates DataFrame with standard deviation/variance and rank.
Explanation:
- Statistics (standard deviation or variance) are calculated based on the genomap data indexed by cell_indexes_batch_cf.
- The cell_indexes parameter includes the cell_indexes_batch_cf (they are a supersubset) but they can also be equal.
- If order is 'C', the function uses pixel_i and pixel_j as they are to index the standard deviation/variance array.
- If order is 'F', the function transposes pixel_i and pixel_j when indexing the standard deviation/variance array, effectively swapping the coordinates.
"""
# Calculate the standard deviation or variance along axis 0 (batches)
if statistic == 'std':
stat_across_batches = np.std(genomap[cell_indexes_batch_cf, :, :, 0], axis=0, ddof=1)
elif statistic == 'var':
stat_across_batches = np.var(genomap[cell_indexes_batch_cf, :, :, 0], axis=0, ddof=1)
else:
raise ValueError("Statistic must be 'std' or 'var'")
print(f"Shape of {statistic}_across_batches:", stat_across_batches.shape)
# Add the standard deviation/variance to the genomap_coordinates DataFrame
genomap_coordinates[statistic] = np.nan
for idx, row in genomap_coordinates.iterrows():
pixel_i, pixel_j = int(row['pixel_i']), int(row['pixel_j'])
if order == 'C':
genomap_coordinates.at[idx, statistic] = stat_across_batches[pixel_i, pixel_j]
elif order == 'F':
genomap_coordinates.at[idx, statistic] = stat_across_batches[pixel_j, pixel_i]
# Rank the pixels based on the absolute value of standard deviation/variance
genomap_coordinates['Rank'] = genomap_coordinates[statistic].abs().rank(ascending=False)
# Convert gene_names from bytes to strings
genomap_coordinates['gene_names'] = genomap_coordinates['gene_names'].apply(lambda x: x.strip('b').strip("'"))
# Add a "Top N" column with True/False
genomap_coordinates['Top_N'] = False
top_n_indices = genomap_coordinates.nsmallest(n_top_genes, 'Rank').index
genomap_coordinates.loc[top_n_indices, 'Top_N'] = True
# Save
genomap_coordinates.to_csv(os.path.join(path_2_genomap, f"genomap_{n_top_genes}topvariablegenesacrossbatches_{file_name}_{statistic}.csv"))
return genomap_coordinates
def adjust_text_positions(x, y, threshold=0.5, offset=0.2):
"""
Adjust text positions to avoid overlap.
Args:
x: List or array of x coordinates.
y: List or array of y coordinates.
threshold: Minimum distance to maintain between points.
offset: Distance to shift overlapping points.
Returns:
List of adjusted (x, y) coordinates.
"""
# Calculate pairwise distances
points = np.array(list(zip(x, y)))
dists = squareform(pdist(points))
# Adjust positions to avoid overlap
adjusted_positions = []
for i, (x_i, y_i) in enumerate(points):
shift_x, shift_y = 0, 0
for j, (x_j, y_j) in enumerate(points):
if i != j and dists[i, j] < threshold:
shift_x += offset if x_i <= x_j else -offset
shift_y += offset if y_i <= y_j else -offset
adjusted_positions.append((x_i + shift_x, y_i + shift_y))
return adjusted_positions
def plot_cell_recon_genomap(genomap, cell_indexes, genomap_coordinates, obs, original_batch=None, n_top_genes=10, min_val=-5, max_val=10, n_cols = 3,order='C',path_2_genomap='',file_name="cell_id",remove_ticks=False):
"""
Plot the genomap with the top variable genes highlighted./
Args:
genomap: 4D numpy array with genomap data.
cell_indexes: List or array of cell indexes.
genomap_coordinates: DataFrame containing gene names and pixel coordinates.
obs: DataFrame containing cell metadata including the column 'recon_prefix' that indicates the name of the reconstruction for a cell id.
original_batch: Identifier for the original batch, if any. Default is None.
n_top_genes: Number of top genes to highlight.
min_val: Minimum value for color scale.
max_val: Maximum value for color scale.
n_cols (int): Number of columns for the subplot. Default is 3.
order: 'C' for default coordinate order, 'F' for transposed coordinates (i.e., pixel_i and pixel_j are swapped).
path_2_genomap (str): path to genomap directory. To save plot.
file_name (str): Name to save the file. Default:"cell_id"
remove_ticks (bool,optional): Flag to remove ticks from the plot. Default=False.
"""
geno_slices_cell_id = genomap[cell_indexes, :, :, 0]
n_images = len(cell_indexes)
n_rows = (n_images + n_cols - 1) // n_cols
fig, axes = plt.subplots(n_rows, n_cols, figsize=(5*n_cols+1, 5 * n_rows))
# If axes is 1D, convert to 2D
if n_rows == 1:
axes = np.expand_dims(axes, 0)
if n_cols == 1:
axes = np.expand_dims(axes, 1)
if genomap_coordinates is not None:
top_n_coordinates = genomap_coordinates[genomap_coordinates['Top_N']]
for i, (cell_index, cell_geno) in enumerate(zip(cell_indexes, geno_slices_cell_id)):
ax = axes[i // n_cols, i % n_cols]
im = ax.imshow(cell_geno, cmap='viridis', vmin=min_val, vmax=max_val)
# Remove ticks
if remove_ticks:
ax.set_xticks([])
ax.set_yticks([])
# # Add gene labels if genomap_coordinates is provided
if genomap_coordinates is not None:
# Get coordinates
if order == 'C':
x, y = top_n_coordinates['pixel_i'], top_n_coordinates['pixel_j']
elif order == 'F':
x, y = top_n_coordinates['pixel_j'], top_n_coordinates['pixel_i']
# Adjust text positions
adjusted_positions = adjust_text_positions(x, y, threshold=3, offset=6)
for (adj_x, adj_y), (pixel_i, pixel_j), gene in zip(adjusted_positions, zip(x, y), top_n_coordinates['gene_names']):
ax.plot(pixel_i, pixel_j, 'o', markerfacecolor='none', markeredgecolor='red', markersize=6, markeredgewidth=2)
ax.text(adj_x, adj_y, gene, color='black', ha='left', va='center', fontweight='bold', fontsize=12, clip_on=False)
# print("cell_index",cell_index)
# print("obs_index",obs.index)
recon_prefix = obs.loc[cell_index, "recon_prefix"]
ax.set_title(recon_prefix)
if 'input' in recon_prefix:
ax.set_title(f'{recon_prefix}\noriginal batch: {original_batch}')
elif original_batch in recon_prefix:
ax.set_title(f'{recon_prefix}\n(original batch)', color='red')
else:
ax.set_title(f'{recon_prefix}')
# Hide any unused subplots
for j in range(i + 1, n_rows * n_cols):
fig.delaxes(axes.flatten()[j])
cbar_ax = fig.add_axes([0.94, 0.15, 0.005, 0.7]) # Adjusted for vertical color bar
fig.colorbar(im, cax=cbar_ax, orientation='vertical')
fig.suptitle(file_name, fontsize=16)
plt.tight_layout(rect=[0, 0.03, 0.9, 0.97])
fig.savefig(os.path.join(path_2_genomap, f"genomap_{n_top_genes}topvariablegenesacrossbatches_{file_name}"))