main.py

import os
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["OPENBLAS_NUM_THREADS"] = "1"
os.environ["NUMEXPR_NUM_THREADS"] = "1"
os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
import torch
import pandas as pd
import scanpy as sc
import numpy as np
import argparse
import time
from pathlib import Path
from spamo.preprocess import fix_seed
from spamo.preprocess import clr_normalize_each_cell, pca
from spamo.preprocess import construct_neighbor_graph, lsi
from spamo.trainer import Train
from spamo.trainer_3m import Train_3M
from spamo.utils import clustering, spatial_smoothing
from cal_matrics import main as calculate_metrics
try:
    from spamo.preprocess_3m import construct_neighbor_graph as construct_neighbor_graph_3M
except ImportError:
    construct_neighbor_graph_3M = None


def resolve_output_paths(args):
    output_dir = Path(args.output_dir) if args.output_dir else None
    if output_dir is not None:
        output_dir.mkdir(parents=True, exist_ok=True)
        args.txt_out_path = args.labels_out or str(output_dir / 'labels.txt')
        args.vis_out_path = args.vis_out_path or str(output_dir / 'spamo.png')
        args.embedding_out = args.embedding_out or str(output_dir / 'embedding.csv')
        args.metrics_out = args.metrics_out or str(output_dir / 'metrics.txt')
    else:
        args.txt_out_path = args.labels_out or args.txt_out_path or 'results/labels.txt'
        args.vis_out_path = args.vis_out_path or 'results/spamo.png'
        if args.embedding_out is None:
            args.embedding_out = str(Path(args.txt_out_path).with_name('embedding.csv'))
        if args.gt_labels and args.metrics_out is None:
            args.metrics_out = str(Path(args.txt_out_path).with_name('metrics.txt'))
        Path(args.txt_out_path).parent.mkdir(parents=True, exist_ok=True)
        if args.vis_out_path:
            Path(args.vis_out_path).parent.mkdir(parents=True, exist_ok=True)
        if args.embedding_out:
            Path(args.embedding_out).parent.mkdir(parents=True, exist_ok=True)
        if args.metrics_out:
            Path(args.metrics_out).parent.mkdir(parents=True, exist_ok=True)


def infer_generic_data_type(args):
    if args.mod3_h5ad:
        return 'Simulation'
    if args.mod2_type == 'ATAC':
        return 'Spatial-epigenome-transcriptome'
    return '10x'


def load_inputs(args):
    if args.rna_h5ad or args.mod2_h5ad or args.mod3_h5ad:
        if not args.rna_h5ad or not args.mod2_h5ad:
            raise ValueError('Generic mode requires both --rna_h5ad and --mod2_h5ad.')
        data_type = args.data_type or infer_generic_data_type(args)
        adata_omics1 = sc.read_h5ad(args.rna_h5ad)
        adata_omics2 = sc.read_h5ad(args.mod2_h5ad)
        adata_omics3 = sc.read_h5ad(args.mod3_h5ad) if args.mod3_h5ad else None
        return adata_omics1, adata_omics2, adata_omics3, data_type

    if not args.file_fold or not args.data_type:
        raise ValueError('Provide either --rna_h5ad/--mod2_h5ad or the benchmark pair --file_fold/--data_type.')

    if args.data_type in ['10x', 'SPOTS', 'Stereo-CITE-seq']:
        adata_omics1 = sc.read_h5ad(args.file_fold + '/adata_RNA.h5ad')
        adata_omics2 = sc.read_h5ad(args.file_fold + '/adata_ADT.h5ad')
        adata_omics3 = None
    elif args.data_type == 'Spatial-epigenome-transcriptome':
        adata_omics1 = sc.read_h5ad(args.file_fold + '/adata_RNA.h5ad')
        adata_omics2 = sc.read_h5ad(args.file_fold + '/adata_peaks_normalized.h5ad')
        adata_omics3 = None
    elif args.data_type == 'Simulation':
        adata_omics1 = sc.read_h5ad(args.file_fold + '/adata_RNA.h5ad')
        adata_omics2 = sc.read_h5ad(args.file_fold + '/adata_ADT.h5ad')
        adata_omics3 = sc.read_h5ad(args.file_fold + '/adata_ATAC.h5ad')
    else:
        raise ValueError(f'Unsupported data_type: {args.data_type}')
    return adata_omics1, adata_omics2, adata_omics3, args.data_type


def align_modalities(*adatas):
    adatas = [adata for adata in adatas if adata is not None]
    common = adatas[0].obs_names
    for adata in adatas[1:]:
        common = common.intersection(adata.obs_names)
    if len(common) == 0:
        raise ValueError('Input modalities have no overlapping obs_names.')
    common = [name for name in adatas[0].obs_names if name in set(common)]
    aligned = [adata[common].copy() for adata in adatas]
    return aligned


def preprocess_inputs(args, adata_omics1, adata_omics2, adata_omics3, data_type):
    adata_omics1.var_names_make_unique()
    adata_omics2.var_names_make_unique()
    if adata_omics3 is not None:
        adata_omics3.var_names_make_unique()

    if data_type == '10x':
        sc.pp.filter_genes(adata_omics1, min_cells=10)
        sc.pp.highly_variable_genes(adata_omics1, flavor="seurat_v3", n_top_genes=min(3000, adata_omics1.n_vars))
        sc.pp.normalize_total(adata_omics1, target_sum=1e4)
        sc.pp.log1p(adata_omics1)
        sc.pp.scale(adata_omics1)
        adata_omics1_high = adata_omics1[:, adata_omics1.var['highly_variable']]
        adata_omics1.obsm['feat'] = pca(adata_omics1_high, n_comps=max(1, min(adata_omics2.n_vars - 1, adata_omics1.n_obs - 1)))
        adata_omics2 = clr_normalize_each_cell(adata_omics2)
        sc.pp.scale(adata_omics2)
        adata_omics2.obsm['feat'] = pca(adata_omics2, n_comps=max(1, min(adata_omics2.n_vars - 1, adata_omics2.n_obs - 1)))
        data = construct_neighbor_graph(adata_omics1, adata_omics2, datatype=data_type, Arg=args)
    elif data_type == 'Spatial-epigenome-transcriptome':
        sc.pp.filter_genes(adata_omics1, min_cells=10)
        sc.pp.filter_cells(adata_omics1, min_genes=200)
        adata_omics2 = adata_omics2[adata_omics1.obs_names].copy()
        sc.pp.highly_variable_genes(adata_omics1, flavor="seurat_v3", n_top_genes=min(3000, adata_omics1.n_vars))
        sc.pp.normalize_total(adata_omics1, target_sum=1e4)
        sc.pp.log1p(adata_omics1)
        sc.pp.scale(adata_omics1)
        adata_omics1_high = adata_omics1[:, adata_omics1.var['highly_variable']]
        adata_omics1.obsm['feat'] = pca(adata_omics1_high, n_comps=max(1, min(50, adata_omics1.n_obs - 1, adata_omics1_high.n_vars)))
        if 'X_lsi' not in adata_omics2.obsm.keys():
            sc.pp.highly_variable_genes(adata_omics2, flavor="seurat_v3", n_top_genes=min(3000, adata_omics2.n_vars))
            lsi(adata_omics2, use_highly_variable=False, n_components=max(2, min(51, adata_omics2.n_obs - 1, adata_omics2.n_vars - 1)))
        adata_omics2.obsm['feat'] = adata_omics2.obsm['X_lsi'].copy()
        data = construct_neighbor_graph(adata_omics1, adata_omics2, datatype=data_type, Arg=args)
    elif data_type == 'SPOTS':
        sc.pp.filter_genes(adata_omics1, min_cells=10)
        sc.pp.highly_variable_genes(adata_omics1, flavor="seurat_v3", n_top_genes=min(3000, adata_omics1.n_vars))
        sc.pp.normalize_total(adata_omics1, target_sum=1e4)
        sc.pp.log1p(adata_omics1)
        sc.pp.scale(adata_omics1)
        adata_omics1_high = adata_omics1[:, adata_omics1.var['highly_variable']]
        adata_omics1.obsm['feat'] = pca(adata_omics1_high, n_comps=max(1, min(adata_omics2.n_vars - 1, adata_omics1.n_obs - 1)))
        adata_omics2 = clr_normalize_each_cell(adata_omics2)
        sc.pp.scale(adata_omics2)
        adata_omics2.obsm['feat'] = pca(adata_omics2, n_comps=max(1, min(adata_omics2.n_vars - 1, adata_omics2.n_obs - 1)))
        data = construct_neighbor_graph(adata_omics1, adata_omics2, datatype=data_type, Arg=args)
    elif data_type == 'Stereo-CITE-seq':
        sc.pp.filter_genes(adata_omics1, min_cells=10)
        sc.pp.filter_cells(adata_omics1, min_genes=80)
        sc.pp.filter_genes(adata_omics2, min_cells=50)
        adata_omics2 = adata_omics2[adata_omics1.obs_names].copy()
        sc.pp.highly_variable_genes(adata_omics1, flavor="seurat_v3", n_top_genes=min(3000, adata_omics1.n_vars))
        sc.pp.normalize_total(adata_omics1, target_sum=1e4)
        sc.pp.log1p(adata_omics1)
        adata_omics1_high = adata_omics1[:, adata_omics1.var['highly_variable']]
        adata_omics1.obsm['feat'] = pca(adata_omics1_high, n_comps=max(1, min(adata_omics2.n_vars - 1, adata_omics1.n_obs - 1)))
        adata_omics2 = clr_normalize_each_cell(adata_omics2)
        adata_omics2.obsm['feat'] = pca(adata_omics2, n_comps=max(1, min(adata_omics2.n_vars - 1, adata_omics2.n_obs - 1)))
        data = construct_neighbor_graph(adata_omics1, adata_omics2, datatype=data_type, Arg=args)
    elif data_type == 'Simulation':
        if adata_omics3 is None:
            raise ValueError('Simulation/three-modality mode requires --mod3_h5ad or Data/Simulation/adata_ATAC.h5ad.')
        n_protein = adata_omics2.n_vars
        sc.pp.highly_variable_genes(adata_omics1, flavor="seurat_v3", n_top_genes=min(3000, adata_omics1.n_vars))
        sc.pp.normalize_total(adata_omics1, target_sum=1e4)
        sc.pp.log1p(adata_omics1)
        adata_omics1_high = adata_omics1[:, adata_omics1.var['highly_variable']]
        n_comps = max(1, min(n_protein, adata_omics1.n_obs - 1, adata_omics1_high.n_vars))
        adata_omics1.obsm['feat'] = pca(adata_omics1_high, n_comps=n_comps)
        adata_omics2 = clr_normalize_each_cell(adata_omics2)
        adata_omics2.obsm['feat'] = pca(adata_omics2, n_comps=max(1, min(n_protein, adata_omics2.n_obs - 1, adata_omics2.n_vars - 1)))
        sc.pp.highly_variable_genes(adata_omics3, flavor="seurat_v3", n_top_genes=min(3000, adata_omics3.n_vars))
        lsi(adata_omics3, use_highly_variable=False, n_components=max(2, min(n_protein + 1, adata_omics3.n_obs - 1, adata_omics3.n_vars - 1)))
        adata_omics3.obsm['feat'] = adata_omics3.obsm['X_lsi'].copy()
        data = construct_neighbor_graph_3M(adata_omics1, adata_omics2, adata_omics3)
    else:
        raise ValueError(f'Unsupported data_type: {data_type}')
    return data, adata_omics1


def train_model(args, data, data_type, device, random_seed):
    if data_type == 'Simulation':
        model = Train_3M(
            data, datatype=data_type, device=device, Arg=args,
            dgi_weight=args.dgi_weight,
            spatial_weight=args.spatial_weight,
            epochs_override=args.epochs_override,
            dropout=args.dropout,
            use_cross_attn=args.use_cross_attn,
            optimizer_type=args.optimizer_type,
            lr_scheduler_type=args.lr_scheduler_type,
            ordered_ablation_mode=args.ordered_ablation_mode,
        )
    else:
        model = Train(
            data,
            datatype=data_type,
            device=device,
            dim_output=args.dim_output,
            Arg=args,
            dgi_weight=args.dgi_weight,
            spatial_weight=args.spatial_weight,
            epochs_override=args.epochs_override,
            dropout=args.dropout,
            use_cross_attn=args.use_cross_attn,
            random_seed=random_seed,
            optimizer_type=args.optimizer_type,
            lr_scheduler_type=args.lr_scheduler_type,
        )

    start_time = time.time()
    output = model.train()
    end_time = time.time()
    print("Training time: ", end_time - start_time)
    return output


def write_labels(path, labels):
    with open(path, 'w') as f:
        for label in labels:
            f.write(f"{label}\n")


def write_embedding(path, obs_names, embedding):
    if not path:
        return
    columns = [f'emb_{i + 1}' for i in range(embedding.shape[1])]
    df = pd.DataFrame(embedding, index=obs_names, columns=columns)
    df.index.name = 'spot'
    df.to_csv(path)


def cluster_and_write_outputs(args, adata_omics1, output, data_type):
    adata = adata_omics1.copy()
    adata.obsm['emb_latent_omics1'] = output['emb_latent_omics1'].copy()
    adata.obsm['emb_latent_omics2'] = output['emb_latent_omics2'].copy()
    if 'emb_latent_omics3' in output:
        adata.obsm['emb_latent_omics3'] = output['emb_latent_omics3'].copy()
    adata.obsm['emb_combined'] = output['emb_combined'].copy()

    tool = args.clustering_method
    clustering(adata, key='emb_combined', add_key='emb_combined', n_clusters=args.n_clusters, method=tool, use_pca=True)

    label = adata.obs['emb_combined']

    if data_type == 'Simulation' and args.reorder_simulation_labels:
        ids = label.index.astype(str).str[:4]
        int_list = [int(num_str) for num_str in ids]
        labels = [-1 for _ in range(len(int_list))]
        for i in range(len(int_list)):
            labels[int_list[i]] = label.iloc[i]
        spot_size = 60
    else:
        labels = label.tolist()
        spot_size = 20

    write_labels(args.txt_out_path, labels)
    write_embedding(args.embedding_out, adata.obs_names, output['emb_combined'])

    if args.gt_labels:
        calculate_metrics(argparse.Namespace(
            GT_path=args.gt_labels,
            our_path=args.txt_out_path,
            save_path=args.metrics_out,
        ))

    if args.skip_plot:
        return

    if data_type == 'Stereo-CITE-seq':
        adata.obsm['spatial'][:, 1] = -1 * adata.obsm['spatial'][:, 1]
    elif data_type == 'SPOTS':
        adata.obsm['spatial'] = np.rot90(np.rot90(np.rot90(np.array(adata.obsm['spatial'])).T).T).T
        adata.obsm['spatial'][:, 1] = -1 * adata.obsm['spatial'][:, 1]

    import matplotlib.pyplot as plt
    fig, ax_list = plt.subplots(1, 2, figsize=(7, 3))
    sc.pp.neighbors(adata, use_rep='emb_combined', n_neighbors=min(500, max(2, adata.n_obs - 1)))
    sc.tl.umap(adata)

    sc.pl.umap(adata, color='emb_combined', ax=ax_list[0], title='SPAMO', s=spot_size, show=False)
    sc.pl.embedding(adata, basis='spatial', color='emb_combined', ax=ax_list[1], title='SPAMO', s=spot_size, show=False)

    plt.tight_layout(w_pad=0.3)
    plt.savefig(args.vis_out_path)


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

    resolve_output_paths(args)
    adata_omics1, adata_omics2, adata_omics3, data_type = load_inputs(args)
    aligned = align_modalities(adata_omics1, adata_omics2, adata_omics3)
    adata_omics1, adata_omics2 = aligned[0], aligned[1]
    adata_omics3 = aligned[2] if len(aligned) > 2 else None
    random_seed = args.random_seed
    fix_seed(random_seed)

    data, adata_omics1 = preprocess_inputs(args, adata_omics1, adata_omics2, adata_omics3, data_type)
    output = train_model(args, data, data_type, device, random_seed)
    cluster_and_write_outputs(args, adata_omics1, output, data_type)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='SpaMO: Spatial Multi-Omics Integration')
    parser.add_argument('--rna_h5ad', type=str, help='Path to RNA AnnData .h5ad file')
    parser.add_argument('--mod2_h5ad', type=str, help='Path to second-modality AnnData .h5ad file')
    parser.add_argument('--mod3_h5ad', type=str, help='Optional third-modality AnnData .h5ad file')
    parser.add_argument('--mod2_type', type=str, default='ADT', choices=['ADT', 'ATAC'],
                        help='Second modality type for generic input mode')
    parser.add_argument('--mod3_type', type=str, default='ATAC', choices=['ATAC'],
                        help='Third modality type for generic input mode')
    parser.add_argument('--output_dir', type=str, help='Directory for labels, embedding, metrics, and plots')
    parser.add_argument('--labels_out', type=str, help='Optional explicit labels output path')
    parser.add_argument('--embedding_out', type=str, help='Optional explicit embedding CSV output path')
    parser.add_argument('--metrics_out', type=str, help='Optional explicit metrics output path')
    parser.add_argument('--gt_labels', type=str, help='Optional ground-truth labels for metric calculation')
    parser.add_argument('--file_fold', type=str, help='Benchmark shorthand: path to a dataset folder')
    parser.add_argument('--data_type', type=str,
                        choices=['10x', 'Spatial-epigenome-transcriptome', 'SPOTS', 'Stereo-CITE-seq', 'Simulation'],
                        help='Dataset preprocessing preset')
    parser.add_argument('--n_clusters', type=int, help='n_clusters for clustering')
    parser.add_argument('--dim_output', type=int, default=64, help='Latent space dimension')
    parser.add_argument('--init_k', type=int, default=10, help='init k')
    parser.add_argument('--KNN_k', type=int, default=20, help='KNN_k')
    parser.add_argument('--alpha', type=float, default=0.9, help='EMA coefficient')
    parser.add_argument('--cl_weight', type=float, default=1, help='Clustering loss weight')
    parser.add_argument('--RNA_weight', type=float, default=5, help='RNA reconstruction weight')
    parser.add_argument('--ADT_weight', type=float, default=5, help='ADT/ATAC reconstruction weight')
    parser.add_argument('--tau', type=float, default=2, help='Temperature for prototype contrastive loss')
    parser.add_argument('--vis_out_path', type=str, help='Backward-compatible plot output path')
    parser.add_argument('--txt_out_path', type=str, help='Backward-compatible labels output path')
    parser.add_argument('--skip_plot', action='store_true',
                        help='Skip UMAP/spatial plotting after writing clustering labels')
    parser.add_argument('--clustering_method', type=str, default='mclust', choices=['mclust', 'leiden', 'louvain'],
                        help='Clustering method for fused embedding')
    parser.add_argument('--reorder_simulation_labels', action='store_true',
                        help='Reorder Simulation labels by numeric spot prefix for benchmark compatibility')
    parser.add_argument('--dgi_weight', type=float, default=0.1, help='DGI monitor scaling weight')
    parser.add_argument('--spatial_weight', type=float, default=0.01, help='Spatial smoothness monitor scaling weight')
    parser.add_argument('--epochs_override', type=int, default=0, help='Override training epochs (0=use dataset default)')
    parser.add_argument('--dropout', type=float, default=0.1, help='Dropout rate')
    parser.add_argument('--use_cross_attn', action='store_true', default=True,
                        help='Use cross-modal attention in fusion (default: True)')
    parser.add_argument('--no_cross_attn', dest='use_cross_attn', action='store_false',
                        help='Disable cross-modal attention in fusion')
    parser.add_argument('--random_seed', type=int, default=2025, help='Random seed')
    parser.add_argument('--optimizer_type', type=str, default='adamw', choices=['sgd', 'adam', 'adamw'],
                        help='Optimizer type')
    parser.add_argument('--lr_scheduler_type', type=str, default='none', choices=['none', 'cosine', 'plateau'],
                        help='LR scheduler type')
    parser.add_argument('--ordered_ablation_mode', type=str, default='full',
                        choices=[
                            'full',
                            'early_interaction',
                            'late_interaction',
                            'no_ordered_design',
                            'fusion_before_graph_calibration',
                            'regularization_before_fusion',
                        ],
                        help='Ordered-framework ablation variant')
    args = parser.parse_args()
    main(args)