deeptan-network 0.1.0__py3-none-any.whl

deeptan/__init__.py

@@ -0,0 +1,9 @@
+"""
+DeepTAN: A novel graph-based multi-task framework designed to infer large-scale multi-omics trait-associated networks (TANs) and reconstruct phenotype-specific omics states.
+"""
+
+__package_name__ = "deeptan"
+__distribution_name__ = "deeptan-network"
+__version__ = "0.1.0"
+__author__ = "Ying Wang, Zeming Meng"
+__credits__ = "Northwest A&F University"

deeptan/cli/__init__.py

@@ -0,0 +1,5 @@
+r"""
+CLI for DeepTAN.
+"""
+
+from deeptan.cli import deeptan_fit, deeptan_litdata, deeptan_perturb, deeptan_pkl2h5, deeptan_predict, hello, print_sc_h5

deeptan/cli/deeptan_fit.py

@@ -0,0 +1,77 @@
+import argparse
+import ast
+
+import deeptan.constants as const
+from deeptan.graph.model import DeepTANTune
+
+
+def deeptan_fit_tune():
+    parser = argparse.ArgumentParser(description="DeepTAN fitting and tuning pipeline.")
+    parser.add_argument("--auto_tune", "--atune", action="store_true", help="Whether to perform hyperparameter tuning")
+
+    parser.add_argument("--em", type=str, default="", help="Existing model checkpoint path for loading")
+    parser.add_argument("--focus", type=str, default="None", help="Focus on a specific task, choose from 'None', 'recon', 'label', 'recon_and_freeze', 'label_and_freeze'")
+    parser.add_argument("--no_guide_gat", "--nog", action="store_true", help="Whether to disable edge weights of guidance graphs on graph attentions")
+
+    parser.add_argument("--litdata", "--data", type=str, required=False, help="Path to litdata directory")
+    parser.add_argument("--bs", type=int, default=const.default.bs, help="Batch size for training")
+    parser.add_argument("--lr", type=float, default=const.default.lr, help="Learning rate")
+    parser.add_argument("--log_dir", "--logdir", type=str, default=".tmp_logs_tune", help="Directory for logging")
+    parser.add_argument("--input_node_emb_dim", "--indim", type=int, default=1, help="Input node embedding dimension")
+    parser.add_argument("--is_regression", "--ir", action="store_true", help="Whether the task is regression")
+    parser.add_argument("--acc_grad_batch", "--agb", type=int, default=const.default.accumulate_grad_batches, help="Accumulate gradients over multiple batches")
+    parser.add_argument("--chunk_size", "--ck", type=int, default=const.default.chunk_size, help="A proper chunk size can balance memory usage and speed")
+    parser.add_argument("--accelerator", "--ac", type=str, default=const.default.accelerator, help="cpu, gpu, tpu, hpu, mps, auto")
+    parser.add_argument("--devices", "--dev", type=str, default=const.default.devices, help="Devices to use")
+    parser.add_argument("--ntrials", "--nt", type=int, default=const.default.n_trials, help="Number of trials for hyperparameter tuning")
+    parser.add_argument("--njobs", "--nj", type=int, default=const.default.n_jobs, help="The number of parallel jobs for Optuna. If this argument is set to -1, the number is set to CPU count.")
+    args = parser.parse_args()
+
+    if args.no_guide_gat:
+        guide_gat = False
+    else:
+        guide_gat = True
+
+    _dev = args.devices
+    if args.devices.startswith("["):
+        if not args.devices.endswith("]"):
+            raise ValueError("Devices argument must be a valid list or a single device string.")
+        # Convert string representation of list to actual list
+        _dev = ast.literal_eval(args.devices)
+
+    _config = const.default.model_config.copy()
+    _config.update(
+        {
+            "es": const.default.es,
+            "max_ep": const.default.max_epoch,
+            "min_ep": const.default.min_epoch,
+            "log_dir": args.log_dir,
+            "litdata": args.litdata,
+            "bs": args.bs,
+            "lr": args.lr,
+            "chunk_size": args.chunk_size,
+            "is_regression": args.is_regression,
+            "accelerator": args.accelerator,
+            "devices": _dev,
+            "input_node_emb_dim": args.input_node_emb_dim,
+            "acc_grad_batch": args.acc_grad_batch,
+            "guide_gat": guide_gat,
+        }
+    )
+    print(f"\n🔧Configuration⚙️\n{_config}\n")
+
+    if len(args.em) < 3:
+        ckpt = None
+    else:
+        ckpt = args.em
+
+    if args.focus == "None":
+        focus = None
+    else:
+        focus = args.focus
+
+    trainer = DeepTANTune(_config, ckpt, focus)
+    if args.auto_tune:
+        trainer.optimize(n_trials=args.ntrials, n_jobs=args.njobs)
+    else:
+        trainer._train_on_args()

deeptan/cli/deeptan_litdata.py

@@ -0,0 +1,113 @@
+import argparse
+import json
+import os
+import pickle
+import shutil
+
+import litdata
+import numpy as np
+import polars as pl
+
+import deeptan.constants as const
+from deeptan.utils.data import DeepTANDataModule, read_nmic_npz
+
+
+def deeptan_litdata():
+    parser = argparse.ArgumentParser(description="Optimize data for Deeptan model")
+    parser.add_argument("--labels", type=str, default="", help="Path to label data in .parquet format")
+    parser.add_argument("--bs", type=int, default=const.default.bs, help="Batch size for training")
+    parser.add_argument("--trn_npz", type=str, required=True, help="Path to training data in .npz format (generated by mi2graph)")
+    parser.add_argument("--val_parquet", type=str, required=True, help="Path to validation data in .parquet format")
+    parser.add_argument("--tst_parquet", type=str, required=True, help="Path to test data in .parquet format")
+    parser.add_argument("--output_dir", type=str, default=".tmp_data_optimized", help="Directory for logging")
+    parser.add_argument("--thre_mi", type=float, default=const.default.threshold_nmic, help="Threshold for edge attribute")
+    parser.add_argument("--in_feat", type=str, default="", help="Path to a .csv with header, containing a list of features to specify. If None, all features are used")
+    parser.add_argument("--in_obs", type=str, default="", help="")
+    parser.add_argument("--onlytest", action="store_true", help="Only run the test phase")
+    parser.add_argument("--n_workers", type=int, default=const.default.n_threads, help="Number of workers for data loading")
+    args = parser.parse_args()
+
+    if len(args.labels) < 2:
+        labels = None
+    else:
+        labels = args.labels
+
+    if len(args.in_feat) < 2:
+        specify_features = None
+    else:
+        specify_features = args.in_feat
+
+    if len(args.in_obs) < 2:
+        specify_trn_obs = None
+    else:
+        specify_trn_obs = args.in_obs
+
+    files_fit = {
+        const.dkey.abbr_train: args.trn_npz,
+        const.dkey.abbr_val: args.val_parquet,
+        const.dkey.abbr_test: args.tst_parquet,
+    }
+    datamodule = DeepTANDataModule(
+        files_fit,
+        labels,
+        batch_size=args.bs,
+        edge_attr_threshold=args.thre_mi,
+        specify_features=specify_features,
+    )
+    datamodule.setup()
+
+    # Copy original training data to output directory for saving node_names and g_label_dim
+    # shutil.copy(args.trn_npz, os.path.join(args.output_dir, "trn.npz"))
+    others2save = {
+        "dict_node_names": datamodule.dict_node_names,
+        "output_g_label_dim": datamodule.label_dim,
+    }
+    os.makedirs(args.output_dir, exist_ok=True)
+    # Save as json
+    with open(os.path.join(args.output_dir, const.fname.litdata_others2save_json), "w") as f:
+        json.dump(others2save, f)
+    # Save as pickle
+    with open(os.path.join(args.output_dir, const.fname.litdata_others2save_pkl), "wb") as f:
+        pickle.dump(others2save, f)
+
+    if labels is not None:
+        shutil.copy(labels, os.path.join(args.output_dir, const.fname.label_class_onehot))
+
+    # Check obs_names filter
+    if specify_trn_obs is not None:
+        _edge_attr, _edge_index, _mat, _mat_feat_indices, _obs_names, _node_names = read_nmic_npz(args.trn_npz)
+        _obs_names_goal = pl.read_parquet(specify_trn_obs)["obs_names"].to_list()
+        # Get intersection of _obs_names and _obs_names_goal
+        _obs_names_filtered = list(set(_obs_names) & set(_obs_names_goal))
+        # Get available indices in _obs_names for the following litdata getting indices
+        _obs_names_indices = np.where(np.isin(_obs_names, _obs_names_filtered))[0]
+        _trn_indices = _obs_names_indices.tolist()
+    else:
+        _trn_indices = list(range(datamodule.train.len()))
+
+    # Optimize
+    if not args.onlytest:
+        litdata.optimize(
+            fn=datamodule.train.get,
+            inputs=_trn_indices,
+            output_dir=os.path.join(args.output_dir, const.dkey.abbr_train),
+            chunk_bytes=const.default.lit_chunk_bytes,
+            compression=const.default.lit_compression,
+            num_workers=min(args.n_workers, const.default.n_threads),
+        )
+        litdata.optimize(
+            fn=datamodule.val.get,
+            inputs=list(range(datamodule.val.len())),
+            output_dir=os.path.join(args.output_dir, const.dkey.abbr_val),
+            chunk_bytes=const.default.lit_chunk_bytes,
+            compression=const.default.lit_compression,
+            num_workers=min(args.n_workers, const.default.n_threads),
+        )
+    litdata.optimize(
+        fn=datamodule.test.get,
+        inputs=list(range(datamodule.test.len())),
+        output_dir=os.path.join(args.output_dir, const.dkey.abbr_test),
+        chunk_bytes=const.default.lit_chunk_bytes,
+        compression=const.default.lit_compression,
+        num_workers=min(args.n_workers, const.default.n_threads),
+    )

deeptan/cli/deeptan_perturb.py

@@ -0,0 +1,31 @@
+import argparse
+import os
+
+from deeptan.graph.recon import predict_perturbation
+
+
+def deeptan_perturb():
+    parser = argparse.ArgumentParser(description="DeepTAN perturbation script.")
+    parser.add_argument("--em", type=str, required=True, help="Existing model checkpoint path.")
+    parser.add_argument("--litdata", "--data", type=str, required=True, help="Path to litdata directory")
+    parser.add_argument("--output", "--out", type=str, required=True, help="Path to output file")
+    parser.add_argument("--maplocation", "--maploc", type=str, default=None, help="Map location for model loading")
+    parser.add_argument("--overwrite", action="store_true", help="Overwrite existing output directory")
+    args = parser.parse_args()
+
+    model_path = args.em
+    litdata_dir = args.litdata
+
+    # Create output directory
+    output_dir = os.path.dirname(args.output)
+    os.makedirs(output_dir, exist_ok=True)
+
+    print(f"Predicting with model {model_path} on data {litdata_dir}")
+    predict_perturbation(
+        model_ckpt_path=model_path,
+        litdata_dir=litdata_dir,
+        output_path=args.output,
+        n_perturbations=5,
+        map_location=args.maplocation,
+        overwrite_files=args.overwrite,
+    )

deeptan/cli/deeptan_pkl2h5.py

@@ -0,0 +1,41 @@
+import argparse
+import os
+
+from deeptan.utils.uni import convert_pickle_to_h5
+
+
+def convert_pkl_to_h5():
+    parser = argparse.ArgumentParser(description="Convert a pickle file to HDF5 format.")
+    parser.add_argument("-i", "--input", required=True, help="Path to a pickle file or a directory containing pickle files.")
+    parser.add_argument("-o", "--output", required=False, help="Path to the output HDF5 file or directory. If not provided, output will be in the same location as the input file with .h5 extension.")
+    parser.add_argument("-f", "--force", action="store_true", help="Overwrite existing HDF5 files without prompting.")
+    args = parser.parse_args()
+
+    if os.path.isdir(args.input):
+        for root, _, files in os.walk(args.input):
+            for file in files:
+                if file.endswith(".pkl"):
+                    input_pkl = os.path.join(root, file)
+                    if args.output is None:
+                        output_h5 = input_pkl.replace(".pkl", ".h5")
+                    else:
+                        output_h5 = os.path.join(args.output, file.replace(".pkl", ".h5"))
+                    if os.path.exists(output_h5) and not args.force:
+                        print(f"File already exists: {output_h5}. Skipping.")
+                    else:
+                        convert_pickle_to_h5(input_pkl, output_h5)
+                        print(f"Converted {input_pkl} to {output_h5}.")
+    else:
+        if args.output is None:
+            output_h5 = args.input.replace(".pkl", ".h5")
+        else:
+            if os.path.isdir(args.output):
+                output_h5 = os.path.join(args.output, os.path.basename(args.input).replace(".pkl", ".h5"))
+            else:
+                output_h5 = args.output.replace(".pkl", ".h5")
+
+        if os.path.exists(output_h5) and not args.force:
+            print(f"File already exists: {output_h5}. Skipping.")
+        else:
+            convert_pickle_to_h5(args.input, output_h5)
+            print(f"Conversion complete: {output_h5}.")

deeptan/cli/deeptan_predict.py

@@ -0,0 +1,36 @@
+import argparse
+import os
+
+from deeptan.graph.recon import predict
+
+
+def deeptan_predict():
+    parser = argparse.ArgumentParser(description="DeepTAN prediction script.")
+    parser.add_argument("--em", type=str, required=True, help="Existing model checkpoint path.")
+    parser.add_argument("--litdata", "--data", type=str, required=True, help="Path to litdata directory")
+    parser.add_argument("--output", "--out", type=str, required=True, help="Path to output file")
+    parser.add_argument("--bs", "--batch_size", type=int, default=8, help="Batch size for prediction")
+    parser.add_argument("--maplocation", "--maploc", type=str, default=None, help="Map location for model loading")
+    parser.add_argument("--overwrite", action="store_true", help="Overwrite existing output directory")
+    args = parser.parse_args()
+
+    model_path = args.em
+    litdata_dir = args.litdata
+
+    # Create output directory
+    output_dir = os.path.dirname(args.output)
+    os.makedirs(output_dir, exist_ok=True)
+    output_path = args.output if args.output.endswith(".h5") else f"{args.output}.h5"
+
+    if not os.path.exists(output_path) or args.overwrite:
+        print(f"Predicting with model {model_path} on data {litdata_dir}")
+        predict(
+            model_ckpt_path=model_path,
+            litdata_dir=litdata_dir,
+            output_path=output_path,
+            map_location=args.map_location,
+            batch_size=args.bs,
+            save_h5=True,
+        )
+    else:
+        print(f"Results already exist at {output_path}")

deeptan/cli/hello.py

@@ -0,0 +1,6 @@
+import deeptan.constants as const
+
+
+def hello():
+    print(const.art.ascii_art)
+    print("\nHello, DeepTAN!\n")

deeptan/cli/print_sc_h5.py

@@ -0,0 +1,32 @@
+import argparse
+import sys
+
+import scanpy as sc
+
+
+def print_sc_h5():
+    parser = argparse.ArgumentParser(description="Read an H5AD or H5 file and print its contents.")
+    parser.add_argument("-p", "--path", type=str, help="Path to the H5AD or H5 file.")
+    args = parser.parse_args()
+    if args.path.endswith(".h5ad"):
+        adata = sc.read_h5ad(args.path)
+    elif args.path.endswith(".h5"):
+        adata = sc.read_10x_h5(args.path)
+    else:
+        print("Unsupported file format. Please provide a .h5ad or .h5 file.")
+        sys.exit(1)
+    print(adata)
+    print("\nobs keys: ")
+    print(adata.obs.keys())
+    for _key in adata.obs.keys():
+        print(f"\n{_key}:")
+        print(adata.obs[_key])
+    print("\nvar keys: ")
+    for _key in adata.var_keys():
+        print(f"\n{_key}:")
+        print(adata.var[_key])
+    if "Celltype" in adata.obs.keys():
+        print("\nCelltype:")
+        print(type(adata.obs["Celltype"]))
+    if "Orig.ident" in adata.obs.keys():
+        sc.pl.umap(adata, color="Orig.ident", title="UMAP by Experiment")

deeptan/constants/__init__.py

@@ -0,0 +1,9 @@
+r"""
+Constants.
+"""
+
+import deeptan.constants.art as art
+import deeptan.constants.default as default
+import deeptan.constants.dict_key as dkey
+import deeptan.constants.fname as fname
+import deeptan.constants.hparam_candidates as hparam_candidates

deeptan/constants/art.py

@@ -0,0 +1,9 @@
+ascii_art = r"""
+    ____                _________    _   __
+   / __ \___  ___  ____/_  __/   |  / | / /
+  / / / / _ \/ _ \/ __ \/ / / /| | /  |/ / 
+ / /_/ /  __/  __/ /_/ / / / ___ |/ /|  /  
+/_____/\___/\___/ .___/_/ /_/  |_/_/ |_/   
+               /_/                         
+
+"""

deeptan/constants/default.py

@@ -0,0 +1,75 @@
+r"""
+Default values.
+"""
+
+from multiprocessing import cpu_count
+from os import getenv
+
+from numpy import ceil
+
+bs = 8
+accumulate_grad_batches = 4
+lr = 0.0002
+es = 5
+min_epoch = 3
+max_epoch = 50
+dropout = 0.2
+negative_slope = 0.2
+node_emb_dim = 128
+g_emb_dim = 192
+label_pred_hidden_dims = [1024, 256]
+fusion_dims_node_emb = [256, 128]
+n_heads_pooling = 4
+n_heads_node_emb = 4
+n_heads_ge_decoder = 4
+n_heads_label_pred = 4
+n_hop = 1
+
+chunk_size = 256
+mem_safety_factor = 0.6
+operation_overhead = 2.5
+
+threshold_nmic = 0.01
+threshold_subg_overlap = 0.85
+threshold_edge_exist = 0.03
+
+matmul_precision = "high"
+accelerator = "auto"
+devices = "auto"
+precision = "32-true"
+gradient_clip_val = 1.0
+
+n_threads = int(getenv("NUM_THREADS", ceil(cpu_count() * 0.9)))
+
+time_format = "%Y%m%d%H%M%S"
+time_delay = 11.7
+ckpt_fname_format = "best_model"
+optuna_db = "sqlite:///optuna.db"
+n_jobs = 1
+n_trials = 30
+n_workers = 1
+
+lit_chunk_bytes = "256MB"
+lit_compression = "zstd"
+lit_max_cache_size = "26GB"
+
+model_config = {
+    "guide_gat": True,
+    "class_weights": None,
+    "use_focal_loss": True,
+    "focal_alpha": None,
+    "node_emb_dim": node_emb_dim,
+    "fusion_dims_node_emb": fusion_dims_node_emb,
+    "output_dim_g_emb": g_emb_dim,
+    "n_hop": n_hop,
+    "threshold_edge_exist": threshold_edge_exist,
+    "threshold_subgraph_overlap": threshold_subg_overlap,
+    "n_heads_node_emb": n_heads_node_emb,
+    "n_heads_pooling": n_heads_pooling,
+    "n_heads_ge_decoder": n_heads_ge_decoder,
+    "n_heads_label_pred": n_heads_label_pred,
+    "dropout": dropout,
+    "lr": lr,
+    "chunk_size": chunk_size,
+    "n_workers": n_workers,
+}

deeptan/constants/dict_key.py

@@ -0,0 +1,98 @@
+r"""
+Dictionary keys and column names.
+"""
+
+title_train = "train"
+title_val = "val"
+title_test = "test"
+title_predict = "pred"
+
+abbr_train = "trn"
+abbr_val = "val"
+abbr_test = "tst"
+
+splits = [abbr_train, abbr_val, abbr_test]
+
+title_trn_loss = "trn/loss"
+title_val_loss = "val/loss"
+title_tst_loss = "tst/loss"
+
+tsb_keys2pick = [
+    "test/recon_MSE",
+    "test/recon_RMSE",
+    "test/recon_MAE",
+    "test/recon_PCC",
+    "test/label_MSE",
+    "test/label_RMSE",
+    "test/label_MAE",
+    "test/label_PCC",
+    "test/label_F1_weighted",
+    "test/label_F1_macro",
+    "test/label_F1_micro",
+    "test/label_AUROC",
+    "test/label_Accuracy",
+    "test/label_Precision",
+    "test/label_Recall",
+    "test/loss",
+    "test/loss_unweighted",
+    "val/loss",
+    "val/loss_unweighted",
+    "val/recon_MSE",
+    "val/recon_RMSE",
+    "val/recon_MAE",
+    "val/recon_PCC",
+    "val/label_MSE",
+    "val/label_RMSE",
+    "val/label_MAE",
+    "val/label_PCC",
+    "val/label_F1_weighted",
+    "val/label_F1_macro",
+    "val/label_F1_micro",
+    "val/label_AUROC",
+    "val/label_Accuracy",
+    "val/label_Precision",
+    "val/label_Recall",
+]
+
+title_metric_mapping = {
+    "jsd": "1 - JSD",
+    "mae": "1 - MAE",
+    "mse": "1 - MSE",
+    "pcc": "PCC",
+    "spearman": "Spearman",
+    "weighted_recall": "Recall (weighted)",
+    "weighted_precision": "Precision (weighted)",
+    "weighted_f1": "F1 Score (weighted)",
+    "macro_f1": "F1 Score (macro)",
+    "micro_f1": "F1 Score (micro)",
+    "auprc": "AUPRC",
+    "auroc": "AUROC",
+    "accuracy": "ACC",
+    "kbet_true_label": "kBET (true labels)",
+    "kbet_pred_label": "kBET (predicted labels)",
+    "asw_true_label": "ASW (true labels)",
+    "asw_pred_label": "ASW (predicted labels)",
+    "kbet": "kBET",
+    "asw": "ASW",
+    "ari": "ARI",
+    "nmi": "NMI",
+    "ami": "AMI",
+    "ari_leiden": "ARI (Leiden)",
+    "nmi_leiden": "NMI (Leiden)",
+    "ami_leiden": "AMI (Leiden)",
+    "method": "Method",
+    "train_size": "Train size",
+    "n_feat": "Number of features for test",
+}
+title_task_mapping = {
+    "multitask": "Multitask",
+    "multitask_noguide": "Multitask (no SGG)",
+    "focus_recon": "Focus on reconstruction",
+    "focus_label": "Focus on labelling",
+    "focus_label_on_focus_recon": "Fine-tuned on reconstruction",
+}
+title_colnameC2_mapping = {
+    "task": "Task",
+    "metric": "Metric",
+    "value": "Value",
+}

deeptan/constants/fname.py

@@ -0,0 +1,7 @@
+r"""
+File names.
+"""
+
+litdata_others2save_pkl = "others2save.pkl"
+litdata_others2save_json = "others2save.json"
+label_class_onehot = "label_onehot.parquet"

deeptan/constants/hparam_candidates.py

@@ -0,0 +1,8 @@
+r"""
+Hyperparameter candidates.
+"""
+
+batch_size = [16, 32, 64, 128]
+lr = [5e-3, 1e-3, 5e-4, 1e-4, 5e-5]
+dropout_high = 0.8
+dropout_step = 0.2

deeptan/graph/__init__.py

@@ -0,0 +1,5 @@
+r"""
+The backbone graph module.
+"""
+
+from deeptan.graph import model, modules, recon