main.py

import argparse
import os
import random
import numpy as np
import torch

from clean_data import clean_up
from preprocessing.create_data import preprocessing, load_marker_genes
from train_hvae import check_no_incomplete_run, main as train_hvae_main

# Named parameter sets selected with `--preset`. A preset replaces the parser defaults of the
# listed arguments; any argument given explicitly on the command line still overrides it.
PRESETS = {
    # 1: contrastive setting (NT-Xent marker-gene head)
    "1": dict(
        epochs=50,
        latent_dimension=90,
        hidden_dimension=128,
        layers=1,
        learning_rate=0.001,
        batch_size=33,
        patience=15,
        encode_covariates=True,
        ln=False,
        kl_weight=0.001,
        mmd_weight=0.1,
        reconstruction_weight=1.0,
        kl_anneal_warmup_epochs=20,
        kl_anneal_max=0.05,
        library_kl="original",
        shared_base_prior=True,
        decoder_covariates=False,
        gene_factor=50,
        use_gene_contrastive=True,
        gene_contrastive_temperature=0.5,
        gene_contrastive_bandwidth=None,
        gene_contrastive_cross_batch_only=False,
        cluster_min=0.2,
        quantil=0.2,
        quantil_max=0.3,
    ),
    # 2: classifier setting (MSE head over binned marker-gene expression)
    "2": dict(
        epochs=50,
        latent_dimension=90,
        hidden_dimension=128,
        layers=1,
        learning_rate=0.001,
        batch_size=33,
        patience=15,
        encode_covariates=True,
        ln=False,
        kl_weight=0.001,
        mmd_weight=0.1,
        reconstruction_weight=1.0,
        kl_anneal_warmup_epochs=20,
        kl_anneal_max=0.05,
        library_kl="original",
        shared_base_prior=False,
        decoder_covariates=False,
        gene_factor=50,
        use_gene_contrastive=False,
        cluster_min=0.2,
        quantil=0.2,
        quantil_max=0.3,
    ),
}


def print_banner(message: str) -> None:
    """Print `message` centered in a row of `#` characters, to mark pipeline stages."""
    print(f"#################################### {message} #########################################")


def set_seed(seed: int = 42) -> None:
    """Set random seeds for reproducible runs."""
    random.seed(seed)
    np.random.seed(seed)

    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    os.environ["PYTHONHASHSEED"] = str(seed)

def get_obs_fields(obs_fields: list[str], dataset_symbol: str) -> list[str]:
    """
    Add the dataset/batch symbol to obs fields and remove duplicates.

    Kept list(set(...)) behavior from the original code.
    """
    obs_fields = obs_fields.copy()
    obs_fields.append(dataset_symbol)
    return list(set(obs_fields))


def build_parser() -> argparse.ArgumentParser:
    """
    Builds the CLI argument parser for the full pipeline, grouped into preprocessing,
    training, cleanup, and additional-load-fields argument groups.
    """
    parser = argparse.ArgumentParser(
        description="Handles data preprocessing, HVAE training and cleanup."
    )

    preprocessing_args = parser.add_argument_group("preprocessing")
    preprocessing_args.add_argument("--dataset_path", type=str, help="path to dataset")
    preprocessing_args.add_argument(
        "--marker_genes",
        type=str,
        default="marker_genes.txt",
        help="path to marker genes",
    )
    preprocessing_args.add_argument(
        "--output_h5ad",
        type=str,
        default="output",
        help="preprocessing output h5ad",
    )
    preprocessing_args.add_argument(
        "--dataset_symbol",
        type=str,
        default="batch",
        help="key to be used for batch correction",
    )
    preprocessing_args.add_argument(
        "--no_hvg",
        action=argparse.BooleanOptionalAction,
        default=False,
    )
    preprocessing_args.add_argument(
        "--n_top_genes",
        type=int,
        default=10000,
        help="number of genes if using hvg",
    )
    preprocessing_args.add_argument(
        "--split_data",
        action=argparse.BooleanOptionalAction,
        default=False,
    )
    preprocessing_args.add_argument(
        "--test_size",
        type=float,
        default=0.1,
        help="default test split if split_data is True",
    )

    training_args = parser.add_argument_group("training")
    training_args.add_argument(
        "--label_symbols",
        type=str,
        default="batch",
        help="labels used for plotting",
    )
    training_args.add_argument(
        "--data_symbol",
        type=str,
        default="data",
        help="data field",
    )
    training_args.add_argument(
        "--gene_symbol",
        type=str,
        default="gene_names",
        help="gene names",
    )
    training_args.add_argument(
        "--count_symbol",
        type=str,
        default=None,
        help="normalized data",
    )
    training_args.add_argument(
        "--output",
        type=str,
        default="embedding.npz",
        help="output embedding",
    )
    training_args.add_argument(
        "--model",
        type=str,
        default="model.ckpt",
        help="output model",
    )
    training_args.add_argument(
        "--resume",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="continue an aborted run from its checkpoint, with the same arguments (only --epochs may change)",
    )
    training_args.add_argument(
        "--checkpoint_every_n_epochs",
        type=int,
        default=1,
        help="how often to save the checkpoint used by --resume",
    )
    training_args.add_argument(
        "--epochs",
        type=int,
        default=50,
        help="number of epochs",
    )
    training_args.add_argument(
        "--latent_dimension",
        type=int,
        default=90,
        help="size of latent dim",
    )
    training_args.add_argument(
        "--hidden_dimension",
        type=int,
        default=128,
        help="size of hidden dim",
    )
    training_args.add_argument(
        "--layers",
        type=int,
        default=1,
        help="number of layers",
    )
    training_args.add_argument(
        "--learning_rate",
        type=float,
        default=0.001,
        help="learning rate",
    )
    training_args.add_argument(
        "--batch_size",
        type=int,
        default=33,
        help="minibatch size for training and embedding",
    )
    training_args.add_argument(
        "--patience",
        type=int,
        default=15,
        help="patience",
    )
    training_args.add_argument(
        "--gene_factor",
        type=int,
        default=10,
        help="weight for marker genes",
    )
    training_args.add_argument(
        "--use_gene_contrastive",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="use NT-Xent contrastive learning directly on log-normalized marker-gene expression, instead of the MSE-based classifier head over binned expression",
    )
    training_args.add_argument(
        "--gene_contrastive_temperature",
        type=float,
        default=0.5,
        help="NT-Xent temperature, only used with --use_gene_contrastive",
    )
    training_args.add_argument(
        "--gene_contrastive_bandwidth",
        type=float,
        default=None,
        help="kernel bandwidth for expression-similarity weighting, only used with --use_gene_contrastive (default: per-gene batch std)",
    )
    training_args.add_argument(
        "--gene_contrastive_cross_batch_only",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="only used with --use_gene_contrastive; if set (default), the contrastive loss only compares cross-batch pairs (batch-integration signal); pass --no-gene_contrastive_cross_batch_only to compare all pairs by marker-gene similarity regardless of batch, favoring cell-type structure over batch alignment",
    )
    training_args.add_argument(
        "--encode_covariates",
        type=bool,
        default=True,
        help="encoding covariates",
    )
    training_args.add_argument(
        "--ln",
        type=bool,
        default=False,
        help="layer normalization",
    )
    training_args.add_argument(
        "--val_data_path",
        type=str,
        default=None,
        help="alternative validation data to plot",
    )
    training_args.add_argument(
        "--random_assignment",
        type=bool,
        default=False,
        help="random assignment of clusters",
    )
    training_args.add_argument(
        "--kl_weight",
        type=float,
        default=0.001,
        help="initial KL weight",
    )
    training_args.add_argument(
        "--mmd_weight",
        type=float,
        default=0.1,
        help="initial MMD weight",
    )
    training_args.add_argument(
        "--reconstruction_weight",
        type=float,
        default=1.0,
        help="initial reconstruction weight",
    )
    training_args.add_argument(
        "--kl_anneal_warmup_epochs",
        type=int,
        default=20,
        help="initial warmup epochs",
    )
    training_args.add_argument(
        "--kl_anneal_max",
        type=float,
        default=0.05,
        help="maximum KL annealing weight",
    )
    training_args.add_argument(
        "--library_kl",
        type=str,
        choices=["original", "marginalized", "batch"],
        default="original",
        help="library-size KL variant: 'original' sums every component's KL (variance passed as std dev); 'marginalized' weights each component's KL by q_pi; 'batch' only uses the KL of each cell's batch-label component, matching how the library is sampled",
    )
    training_args.add_argument(
        "--shared_base_prior",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="parameterize the library and z GMM priors as a shared base plus per-component deviations (default); pass --no-shared_base_prior for independent per-component parameters",
    )
    training_args.add_argument(
        "--decoder_covariates",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="inject the batch covariate embedding into the decoder, so batch effects can be explained by the decoder instead of being encoded in z",
    )

    cleanup_args = parser.add_argument_group("cleanup")
    cleanup_args.add_argument(
        "--filtered_data_output",
        type=str,
        default="filtered_data.npz",
        help="cleaned up data",
    )
    cleanup_args.add_argument(
        "--conf_data",
        type=str,
        default="conf_data.npz",
        help="data with confidence scores",
    )
    cleanup_args.add_argument(
        "--plots",
        type=str,
        default="data",
        help="baseline plot name",
    )
    cleanup_args.add_argument(
        "--labels",
        type=str,
        default=None,
        help="label used for cleanup/leiden if None",
    )
    cleanup_args.add_argument(
        "--cluster_min",
        type=float,
        default=0.2,
        help="min confidence for per cluster cutoff",
    )
    cleanup_args.add_argument(
        "--quantil",
        type=float,
        default=0.2,
        help="confidence quantil",
    )
    cleanup_args.add_argument(
        "--quantil_max",
        type=float,
        default=0.3,
        help="max confidence for per cluster cutoff",
    )
    cleanup_args.add_argument(
        "--no_cleanup",
        action=argparse.BooleanOptionalAction,
        default=False,
    )
    cleanup_args.add_argument(
        "--no_plots",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="skip writing the TSNE scatter-plot HTML files during cleanup (the TSNE "
             "projection itself is still computed, since the confidence model is fit on "
             "it; this only skips the plot_tsne HTML writes, which dominate cleanup "
             "runtime for large datasets)",
    )

    additional_fields_args = parser.add_argument_group("additional load fields")
    additional_fields_args.add_argument(
        "--obs_fields",
        nargs="+",
        type=str,
        default=[],
        help="list of observation fields",
    )
    additional_fields_args.add_argument(
        "--layer_fields",
        nargs="+",
        type=str,
        default=[],
        help="list of layer fields",
    )

    parser.add_argument(
        "--seed",
        type=int,
        default=42,
        help="random seed for Python/NumPy/PyTorch (enables multi-seed runs)",
    )
    parser.add_argument(
        "--preset",
        type=str,
        choices=sorted(PRESETS),
        default=None,
        help="parameter set from PRESETS (1: contrastive, 2: classifier); explicitly given arguments override it",
    )

    return parser


def parse_args(parser: argparse.ArgumentParser, argv: list[str] | None = None) -> argparse.Namespace:
    """
    Parses the CLI args, with the defaults of `--preset` (if given) applied before the
    explicitly passed arguments, so those still take precedence over the preset.
    """
    preset = parser.parse_known_args(argv)[0].preset
    if preset is not None:
        parser.set_defaults(**PRESETS[preset])
        print(f"using preset {preset}: {PRESETS[preset]}")
    return parser.parse_args(argv)


def run_preprocessing(args: argparse.Namespace) -> None:
    """Runs marker-gene/HVG selection and optional train/test splitting via `preprocessing`."""
    print_banner("running pre-processing")

    preprocessing(
        marker_gene_list=args.marker_genes,
        dataset=args.dataset_path,
        full_name=args.output_h5ad,
        batch_key=args.dataset_symbol,
        no_hvg=args.no_hvg,
        n_top_genes=args.n_top_genes,
        split_data=args.split_data,
        test_size=args.test_size,
    )


def training_input_path(args: argparse.Namespace) -> str:
    """
    Returns the appropriate training input: the raw dataset when `no_hvg` and not
    `split_data`, the full HVG output when only `no_hvg` is False, or the reference split
    when `split_data` is True.
    """
    if args.no_hvg == True and args.split_data == False:
        return args.dataset_path
    elif args.split_data == False and args.no_hvg == False:
        return args.output_h5ad + ".h5ad"
    return args.output_h5ad + "_ref.h5ad"


def run_training(args: argparse.Namespace) -> None:
    """Trains the HVAE (`models/vae.py`) on the output of `training_input_path`."""
    print_banner("training (HVAE)")

    gene_names = load_marker_genes(args.marker_genes)
    train_data = training_input_path(args)

    train_hvae_main(
        train_data_name=train_data,
        gene_names=gene_names,
        dataset_symbol=args.dataset_symbol,
        label_symbols=args.label_symbols,
        data_symbol=args.data_symbol,
        gene_symbol=args.gene_symbol,
        val_data_path=args.val_data_path,
        num_epochs=args.epochs,
        latent_dim=args.latent_dimension,
        hidden_dim=args.hidden_dimension,
        layers=args.layers,
        learning_rate=args.learning_rate,
        batch_size=args.batch_size,
        patience=args.patience,
        encode_covariates=args.encode_covariates,
        ln=args.ln,
        gene_factor=args.gene_factor,
        kl_weight=args.kl_weight,
        mmd_weight=args.mmd_weight,
        reconstruction_weight=args.reconstruction_weight,
        save_embedding=args.output,
        kl_anneal_warmup_epochs=args.kl_anneal_warmup_epochs,
        kl_anneal_max=args.kl_anneal_max,
        model_checkpoint=args.model,
        resume=args.resume,
        checkpoint_every_n_epochs=args.checkpoint_every_n_epochs,
        count_symbol=args.count_symbol,
        random_assignment=args.random_assignment,
        use_gene_contrastive=args.use_gene_contrastive,
        gene_contrastive_temperature=args.gene_contrastive_temperature,
        gene_contrastive_bandwidth=args.gene_contrastive_bandwidth,
        gene_contrastive_cross_batch_only=args.gene_contrastive_cross_batch_only,
        library_kl=args.library_kl,
        shared_base_prior=args.shared_base_prior,
        decoder_covariates=args.decoder_covariates,
    )

def run_cleanup(args: argparse.Namespace) -> None:
    """Runs confidence-based embedding cleanup via `clean_up`, unless `--no_cleanup` is set."""
    print_banner("data cleanup")

    if not args.no_cleanup:
        clean_up(
            args.output,
            args.filtered_data_output,
            args.conf_data,
            plots=args.plots,
            labels=args.labels,
            cluster_min=args.cluster_min,
            quantil=args.quantil,
            quantil_max=args.quantil_max,
            make_plots=not args.no_plots,
        )


def main() -> None:
    """Entry point: parses CLI args and runs preprocessing, HVAE training, then cleanup in order."""
    parser = build_parser()
    args = parse_args(parser)
    check_no_incomplete_run(args.model, args.resume)
    set_seed(args.seed)
    if args.resume and os.path.exists(training_input_path(args)):
        print("Resuming: keeping the existing preprocessing output.")
    else:
        run_preprocessing(args)
    run_training(args)
    run_cleanup(args)

if __name__ == "__main__":
    main()