diff --git a/COMEBin/cluster.py b/COMEBin/cluster.py index 2b5dc58..129d487 100644 --- a/COMEBin/cluster.py +++ b/COMEBin/cluster.py @@ -9,8 +9,11 @@ import logging from igraph import Graph +from sklearn.cluster import KMeans +from sklearn.metrics.pairwise import euclidean_distances from sklearn.preprocessing import normalize -from sklearn.cluster._kmeans import euclidean_distances, stable_cumsum, KMeans, check_random_state, row_norms, MiniBatchKMeans +from sklearn.utils import check_random_state +from sklearn.utils.extmath import row_norms from utils import get_length, calculateN50, save_result from scripts.gen_bins_from_tsv import gen_bins as gen_bins_from_tsv @@ -100,7 +103,7 @@ def seed_kmeans_full(logger, contig_file: str, namelist: List[str], out_path: st output_temp = out_path + '_k_' + str( bin_number) + '_result.tsv' if not (os.path.exists(output_temp)): - km = KMeans(n_clusters=bin_number, n_jobs=-1, random_state=7, algorithm="full", + km = KMeans(n_clusters=bin_number, random_state=7, algorithm="lloyd", n_init=10, init=functools.partial(partial_seed_init, seed_idx=seed_bacar_marker_idx)) km.fit(X_mat, sample_weight=length_weight) idx = km.labels_ @@ -192,8 +195,10 @@ def partial_seed_init(X, n_clusters: int, random_state, seed_idx, n_local_trials # Choose center candidates by sampling with probability proportional # to the squared distance to the closest existing center rand_vals = random_state.random_sample(n_local_trials) * current_pot - candidate_ids = np.searchsorted(stable_cumsum(closest_dist_sq), - rand_vals) + candidate_ids = np.searchsorted( + np.cumsum(closest_dist_sq, dtype=np.float64), + rand_vals, + ) # XXX: numerical imprecision can result in a candidate_id out of range np.clip(candidate_ids, None, closest_dist_sq.size - 1, out=candidate_ids) @@ -417,5 +422,3 @@ def cluster(logger, args, prefix=None): multiprocess.join() logger.info('multiprocess Done') - - diff --git a/COMEBin/get_augfeature.py b/COMEBin/get_augfeature.py index 906e320..4910046 100644 --- a/COMEBin/get_augfeature.py +++ b/COMEBin/get_augfeature.py @@ -38,7 +38,7 @@ def get_kmer_coverage(data_path: str, n_views: int = 2, kmer_model_path: str = ' shuffled_covMat = pd.read_csv(cov_file, sep='\t', usecols=range(1, covHeader.shape[1])).values shuffled_namelist = pd.read_csv(cov_file, sep='\t', usecols=range(1)).values[:, 0] - covIdxArr = np.empty(len(mapObj), dtype=np.int) + covIdxArr = np.empty(len(mapObj), dtype=np.int64) for contigIdx in range(len(shuffled_namelist)): if shuffled_namelist[contigIdx].split('_aug')[0] in mapObj: covIdxArr[mapObj[shuffled_namelist[contigIdx].split('_aug')[0]]] = contigIdx @@ -50,7 +50,7 @@ def get_kmer_coverage(data_path: str, n_views: int = 2, kmer_model_path: str = ' shuffled_compositMat = pd.read_csv(com_file, sep=',', usecols=range(1, compositHeader.shape[1])).values shuffled_namelist = pd.read_csv(com_file, sep=',', usecols=range(1)).values[:, 0] - covIdxArr = np.empty(len(mapObj), dtype=np.int) + covIdxArr = np.empty(len(mapObj), dtype=np.int64) for contigIdx in range(len(shuffled_namelist)): if shuffled_namelist[contigIdx].split('_aug')[0] in mapObj: covIdxArr[mapObj[shuffled_namelist[contigIdx].split('_aug')[0]]] = contigIdx @@ -63,7 +63,7 @@ def get_kmer_coverage(data_path: str, n_views: int = 2, kmer_model_path: str = ' shuffled_varsMat = pd.read_csv(vars_file, sep='\t', usecols=range(1, varsHeader.shape[1])).values shuffled_namelist = pd.read_csv(vars_file, sep='\t', usecols=range(1)).values[:, 0] - covIdxArr = np.empty(len(mapObj), dtype=np.int) + covIdxArr = np.empty(len(mapObj), dtype=np.int64) for contigIdx in range(len(shuffled_namelist)): if shuffled_namelist[contigIdx].split('_aug')[0] in mapObj: covIdxArr[mapObj[shuffled_namelist[contigIdx].split('_aug')[0]]] = contigIdx diff --git a/COMEBin/main.py b/COMEBin/main.py index d37fe6d..c3ea94c 100644 --- a/COMEBin/main.py +++ b/COMEBin/main.py @@ -2,10 +2,39 @@ import logging import os import pandas as pd +import sys +import warnings from comebin_version import __version__ as ver -from train_CLmodel import train_CLmodel -from cluster import cluster + + +def _resolve_checkm_data_path(): + """ + Resolve a CheckM data directory from the current runtime environment. + + Preference order: + 1) Existing CHECKM_DATA_PATH env var + 2) $CONDA_PREFIX/checkm_data + 3) sys.prefix/checkm_data + 4) Parent of interpreter prefix/checkm_data + """ + existing = os.environ.get("CHECKM_DATA_PATH") + if existing: + return existing + + candidates = [] + conda_prefix = os.environ.get("CONDA_PREFIX") + if conda_prefix: + candidates.append(os.path.join(conda_prefix, "checkm_data")) + + candidates.append(os.path.join(sys.prefix, "checkm_data")) + candidates.append(os.path.join(os.path.dirname(sys.prefix), "checkm_data")) + + for candidate in candidates: + if os.path.isfile(os.path.join(candidate, "taxon_marker_sets.tsv")): + return candidate + + return None def arguments(): @@ -278,6 +307,13 @@ def main(): args.output_path = args.out_augdata_path os.makedirs(args.output_path, exist_ok=True) + + # Avoid matplotlib trying to write under ~/.matplotlib in restricted runtimes. + if not os.environ.get("MPLCONFIGDIR"): + mpl_config_dir = os.path.join(args.output_path, ".mplconfig") + os.makedirs(mpl_config_dir, exist_ok=True) + os.environ["MPLCONFIGDIR"] = mpl_config_dir + handler = logging.FileHandler(args.output_path+'/comebin.log') handler.setLevel(logging.INFO) handler.setFormatter(formatter) @@ -285,11 +321,13 @@ def main(): ## training if args.subcmd == 'train': + from train_CLmodel import train_CLmodel logger.info('train') train_CLmodel(logger,args) ## clustering if args.subcmd == 'bin': + from cluster import cluster logger.info('bin') from utils import gen_seed @@ -301,6 +339,7 @@ def main(): ## clustering NoContrast if args.subcmd == 'nocontrast': + from cluster import cluster logger.info('NoContrast mode') from utils import get_kmer_coverage_aug0 @@ -356,6 +395,23 @@ def main(): ###Generate the final results from the Leiden clustering results if args.subcmd == 'get_result': logger.info('get_result') + # checkm currently imports pkg_resources, which emits a noisy deprecation warning. + warnings.filterwarnings( + "ignore", + message="pkg_resources is deprecated as an API.*", + category=UserWarning, + ) + # Keep CheckM from defaulting to ~/.checkm in restricted environments. + checkm_data_path = _resolve_checkm_data_path() + if checkm_data_path: + os.environ["CHECKM_DATA_PATH"] = checkm_data_path + logger.info("Using CHECKM_DATA_PATH:\t" + checkm_data_path) + else: + fallback_checkm_path = os.path.join(args.output_path, "checkm_data") + os.makedirs(fallback_checkm_path, exist_ok=True) + os.environ["CHECKM_DATA_PATH"] = fallback_checkm_path + logger.info("CHECKM_DATA_PATH was unset; using fallback path:\t" + fallback_checkm_path) + from utils import gen_seed from get_final_result import run_get_final_result @@ -367,4 +423,3 @@ def main(): if __name__ == '__main__': main() - diff --git a/COMEBin/simclr.py b/COMEBin/simclr.py index 4039b90..5d557d6 100644 --- a/COMEBin/simclr.py +++ b/COMEBin/simclr.py @@ -4,7 +4,7 @@ import torch import torch.nn.functional as F -from torch.cuda.amp import GradScaler, autocast +from torch.amp import GradScaler, autocast from torch.utils.tensorboard import SummaryWriter from tqdm import tqdm from utils import save_config_file, accuracy, save_checkpoint @@ -139,7 +139,8 @@ def train(self, train_loader, data, namelist): :param data: Input data. :param namelist: List of sequence names. """ - scaler = GradScaler(enabled=self.args.fp16_precision) + _device_type = 'cuda' if 'cuda' in str(self.args.device) else 'cpu' + scaler = GradScaler(_device_type, enabled=self.args.fp16_precision) # save config file save_config_file(self.args.output_path, self.args) @@ -155,7 +156,7 @@ def train(self, train_loader, data, namelist): contig_features = contig_features.to(self.args.device) - with autocast(enabled=self.args.fp16_precision): + with autocast(_device_type, enabled=self.args.fp16_precision): features = self.model(contig_features) logits, labels = self.info_nce_loss(features) loss = self.criterion(logits, labels) @@ -214,7 +215,8 @@ def train_addpretrain(self, train_loader, data, namelist): :param data: Input data. :param namelist: List of sequence names. """ - scaler = GradScaler(enabled=self.args.fp16_precision) + _device_type = 'cuda' if 'cuda' in str(self.args.device) else 'cpu' + scaler = GradScaler(_device_type, enabled=self.args.fp16_precision) # save config file save_config_file(self.args.output_path, self.args) @@ -241,7 +243,7 @@ def train_addpretrain(self, train_loader, data, namelist): contig_features = contig_features.to(self.args.device) # print(contig_features.shape) - with autocast(enabled=self.args.fp16_precision): + with autocast(_device_type, enabled=self.args.fp16_precision): if self.args.addcovloss and not self.args.addkmerloss: if self.args.pretrain_kmer_model_path !='no': features, covemb, kmeremb = self.model(contig_features[:, -kmer_len:], contig_features[:, :-kmer_len]) @@ -365,7 +367,8 @@ def covmodeltrain(self, train_loader): :param train_loader: Data loader for training. """ - scaler = GradScaler(enabled=self.args.fp16_precision) + _device_type = 'cuda' if 'cuda' in str(self.args.device) else 'cpu' + scaler = GradScaler(_device_type, enabled=self.args.fp16_precision) # save config file save_config_file(self.args.output_path, self.args) @@ -381,7 +384,7 @@ def covmodeltrain(self, train_loader): contig_features = contig_features.to(self.args.device) - with autocast(enabled=self.args.fp16_precision): + with autocast(_device_type, enabled=self.args.fp16_precision): features = self.model(contig_features[:, :-128]) logits, labels = self.info_nce_loss(features) loss = self.criterion(logits, labels) diff --git a/COMEBin/train_CLmodel.py b/COMEBin/train_CLmodel.py index da45c94..af5ce7d 100644 --- a/COMEBin/train_CLmodel.py +++ b/COMEBin/train_CLmodel.py @@ -88,11 +88,9 @@ def train_CLmodel(logger, args): config_file= os.path.dirname(args.pretrain_kmer_model_path)+'/kmerMetric_config.yaml' from ruamel.yaml import YAML - from pathlib import Path - yaml = YAML(typ='safe') - - cnf = yaml.load(Path(config_file)) + with open(config_file, "r", encoding="utf-8") as f: + cnf = yaml.load(f) ps = [cnf['dropout_value']]*(len(cnf['emb_szs'])-1) actn= nn.LeakyReLU() diff --git a/COMEBin/utils.py b/COMEBin/utils.py index ad9fe31..248d70b 100644 --- a/COMEBin/utils.py +++ b/COMEBin/utils.py @@ -76,7 +76,7 @@ def save_result(result, filepath, namelist): os.makedirs(filedir) f = open(filepath, 'w') for contigIdx in range(len(result)): - f.write(namelist[contigIdx] + "\t" + str(result[contigIdx].item(0)) + "\n") + f.write(namelist[contigIdx] + "\t" + str(result[contigIdx].item()) + "\n") f.close() @@ -106,7 +106,7 @@ def get_kmer_coverage_aug0(data_path): shuffled_covMat = pd.read_csv(cov_file, sep='\t', usecols=range(1, covHeader.shape[1])).values shuffled_namelist = pd.read_csv(cov_file, sep='\t', usecols=range(1)).values[:, 0] - covIdxArr = np.empty(len(mapObj), dtype=np.int) + covIdxArr = np.empty(len(mapObj), dtype=np.int64) for contigIdx in range(len(shuffled_namelist)): if shuffled_namelist[contigIdx].split('_aug')[0] in mapObj: covIdxArr[mapObj[shuffled_namelist[contigIdx].split('_aug')[0]]] = contigIdx @@ -116,7 +116,7 @@ def get_kmer_coverage_aug0(data_path): shuffled_compositMat = pd.read_csv(com_file, sep=',', usecols=range(1, compositHeader.shape[1])).values shuffled_namelist = pd.read_csv(com_file, sep=',', usecols=range(1)).values[:, 0] - covIdxArr = np.empty(len(mapObj), dtype=np.int) + covIdxArr = np.empty(len(mapObj), dtype=np.int64) for contigIdx in range(len(shuffled_namelist)): if shuffled_namelist[contigIdx].split('_aug')[0] in mapObj: covIdxArr[mapObj[shuffled_namelist[contigIdx].split('_aug')[0]]] = contigIdx @@ -153,7 +153,6 @@ def get_kmerMetric_emb(kmer_model_path,compositMats,device=torch.device('cpu'),k config_file = os.path.dirname(kmer_model_path) + '/kmerMetric_config.yaml' from ruamel.yaml import YAML - from pathlib import Path import torch.nn as nn from models.mlp import EmbeddingNet from sklearn.preprocessing import normalize @@ -161,7 +160,8 @@ def get_kmerMetric_emb(kmer_model_path,compositMats,device=torch.device('cpu'),k yaml = YAML(typ='safe') - cnf = yaml.load(Path(config_file)) + with open(config_file, "r", encoding="utf-8") as f: + cnf = yaml.load(f) ps = [cnf['dropout_value']] * (len(cnf['emb_szs']) - 1) actn = nn.LeakyReLU() @@ -225,13 +225,22 @@ def gen_seed(logger, contig_file: str, threads: int, contig_length_threshold: in os.system(fragCmd) if os.path.exists(fragResultURL): - if not (os.path.exists(hmmResultURL)): + if (not os.path.exists(hmmResultURL)) or os.path.getsize(hmmResultURL) == 0: hmmCmd = hmmExeURL + " --domtblout " + hmmResultURL + " --cut_tc --cpu " + str( threads) + " " + markerURL + " " + fragResultURL + " 1>" + hmmResultURL + ".out 2>" + hmmResultURL + ".err" logger.info("exec cmd: " + hmmCmd) - os.system(hmmCmd) - - if os.path.exists(hmmResultURL): + hmm_ret = os.system(hmmCmd) + if hmm_ret != 0 or (os.path.exists(hmmResultURL) and os.path.getsize(hmmResultURL) == 0): + # Newer HMMER builds can reject --cut_tc for marker files without TC thresholds. + logger.info("hmmsearch with --cut_tc failed; retrying with an E-value cutoff.") + if os.path.exists(hmmResultURL): + os.remove(hmmResultURL) + hmmCmd = hmmExeURL + " --domtblout " + hmmResultURL + " -E 1e-10 --cpu " + str( + threads) + " " + markerURL + " " + fragResultURL + " 1>" + hmmResultURL + ".fallback.out 2>" + hmmResultURL + ".fallback.err" + logger.info("exec cmd: " + hmmCmd) + os.system(hmmCmd) + + if os.path.exists(hmmResultURL) and os.path.getsize(hmmResultURL) > 0: if not (os.path.exists(seedURL)): markerCmd = markerExeURL + " " + hmmResultURL + " " + contig_file + " " + str( contig_length_threshold) + " " + seedURL @@ -244,9 +253,9 @@ def gen_seed(logger, contig_file: str, threads: int, contig_length_threshold: in logger.info("markerCmd failed! Not exist: " + markerCmd) candK = 0 else: - logger.info("Hmmsearch failed! Not exist: " + hmmResultURL) + logger.info("Hmmsearch failed! Not exist or empty: " + hmmResultURL) sys.exit() else: logger.info("FragGeneScan failed! Not exist: " + fragResultURL) sys.exit() - return candK \ No newline at end of file + return candK diff --git a/comebin_env.yaml b/comebin_env.yaml index 8c3bedc..05570ee 100644 --- a/comebin_env.yaml +++ b/comebin_env.yaml @@ -1,46 +1,42 @@ name: comebin_env channels: - - pytorch - conda-forge - bioconda + - pytorch - defaults dependencies: - - atomicwrites=1.4.0 - - bedtools=2.30.0 - - biolib=0.1.6 - - biopython=1.76 - - bwa=0.7.17 - - checkm-genome=1.1.3 - - cudatoolkit=11.1.1 - - fraggenescan=1.31 - - hmmer=3.1b2 - - matplotlib-base=3.5.1 - - matplotlib-inline=0.1.3 - - numpy=1.19.0 + - python>=3.10,<3.12 + - pip>=24 + - setuptools>=69,<81 + - numpy>=1.26,<2.0 + - pandas>=2.1 + - scipy>=1.11 + - scikit-learn>=1.4 + - pytorch>=2.2 + - tensorboard>=2.15 + - biopython>=1.83 + - click>=8.1 + - pyyaml>=6 + - ruamel.yaml>=0.17 + - atomicwrites>=1.4 + - tqdm>=4.66 + - wget + - bedtools>=2.31 + - biolib>=0.1.9 + - bwa>=0.7.17 + - checkm-genome>=1.2 + - fraggenescan>=1.31 + - hmmer>=3.3 - pplacer=1.1.alpha19 - - prodigal=2.6.3 - - python=3.7 - - python-dateutil=2.8.2 - - python-fastjsonschema=2.15.3 - - python_abi=3.7 - - pytorch=1.10.2=py3.7_cuda11.1_cudnn8.0.5_0 - - pytorch-mutex=1.0=cuda - - pip=22.0.4 - - samtools=1.15.1 - - scipy=1.7.3 - - setuptools=59.5.0 - - tensorboard=1.15.0 - - click=8.0.4 - - pip: - - hnswlib==0.6.2 - - igraph==0.9.9 - - joblib==1.1.0 - - leidenalg==0.8.10 - - networkx==2.6.3 - - numba==0.56.4 - - pandas==1.3.5 - - scanpy==1.9.1 - - scikit-learn==0.22.1 - - seaborn==0.12.1 - - statsmodels==0.13.5 - - pyyaml==6.0 + - prodigal>=2.6 + - samtools>=1.18 + - matplotlib-base>=3.8 + - networkx>=3.2 + - numba>=0.59 + - joblib>=1.3 + - seaborn>=0.13 + - statsmodels>=0.14 + - hnswlib>=0.8.0 + - python-igraph>=0.11 + - leidenalg>=0.10 + - scanpy>=1.10