Skip to content

Compatibility fixes for modern NumPy / sklearn / HMMER / CheckM #49

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rec3141 wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Compatibility fixes for modern NumPy / sklearn / HMMER / CheckM #49

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 9 additions & 6 deletions COMEBin/cluster.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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_
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -417,5 +422,3 @@ def cluster(logger, args, prefix=None):
multiprocess.join()
logger.info('multiprocess Done')



6 changes: 3 additions & 3 deletions COMEBin/get_augfeature.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
61 changes: 58 additions & 3 deletions COMEBin/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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():
Expand Down Expand Up @@ -278,18 +307,27 @@ 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)
logger.addHandler(handler)

## 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

Expand All @@ -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

Expand Down Expand Up @@ -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

Expand All @@ -367,4 +423,3 @@ def main():

if __name__ == '__main__':
main()

17 changes: 10 additions & 7 deletions COMEBin/simclr.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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])
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand Down
6 changes: 2 additions & 4 deletions COMEBin/train_CLmodel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
31 changes: 20 additions & 11 deletions COMEBin/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()


Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -153,15 +153,15 @@ 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


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()
Expand Down Expand Up @@ -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
Expand All @@ -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
return candK
Loading