enzyme-model / tools /build_source_hash.py

update

a5b90b6 verified 13 days ago

6.86 kB

"""v24: exact-sequence label table with a source priority measured on the training set.
For every sequence in the disclosed external sources, the label is taken from the first source that has one, in the order
  BRENDA 2026.1 > BRENDA 2023.1 > CatPred-DB > CataPro > DLKcat > PDB/SIFTS > UniProtKB unreviewed (current EC) > ESP > ESP (current EC)
(on 14,122 training proteins that are not in Swiss-Prot, BRENDA's EC is right 95% of the time when it disagrees with the
current UniProt EC, which is right 2.5% of the time; this order scores 0.974 against the official training labels).
Swiss-Prot and train sequences are not in this table: identical train/Swiss-Prot sequences keep their own labels.
Output: <model>/weights/srchash.npz (h: sorted uint64 MD5 prefixes, s: uint32 EC-set id) + srchash_sets.json
Usage: python3 build_source_hash.py --esm ~/enzyme-esm --catapro <CataPro/datasets> --model ~/enzyme-model [--exclude FASTA]
"""
import argparse, collections, gzip, hashlib, json, re
from pathlib import Path
import numpy as np
import pandas as pd

ap = argparse.ArgumentParser()
ap.add_argument('--esm', required=True); ap.add_argument('--catapro', required=True); ap.add_argument('--model', required=True)
ap.add_argument('--exclude')
a = ap.parse_args()
E = Path(a.esm)
VOC = set(json.load(open(Path(a.model) / 'ec_vocabulary.json')))
ORDER = ['B6', 'BR', 'CP', 'CT', 'DK', 'PD', 'TR', 'ES', 'E2']


def ecs(x):
    return [e.strip() for e in re.split(r'[;,]', str(x)) if e.strip() in VOC]


def h(s):
    return int(hashlib.md5(s.encode()).hexdigest()[:16], 16)


lab = collections.defaultdict(lambda: collections.defaultdict(set))   # hash -> source -> ECs (non-UniProt sources)
br = collections.defaultdict(set)
for line in open(E / 'brenda_acc_ec.tsv'):
    if '\t' in line:
        acc, e = line.rstrip('\n').split('\t'); br[acc].update(ecs(e))
sec = collections.defaultdict(set)
for line in open(E / 'round3/sifts_acc_ec.tsv'):
    p = line.rstrip('\n').split('\t')
    if len(p) >= 2: sec[p[0]].update(ecs(p[1]))
b6 = collections.defaultdict(set)                      # BRENDA release 2026.1 (r5/brenda26_acc_ec.tsv)
for line in open(E / 'r5/brenda26_acc_ec.tsv'):
    acc, e = line.rstrip('\n').split('\t'); b6[acc].update(ecs(e))
need = set(br) | set(sec) | set(b6)
acc2h = collections.defaultdict(set)   # an accession can map to several sequence versions
# accession -> sequence fingerprint for BRENDA / SIFTS accessions: UniProtKB unreviewed (fetch_uniprot_hash.py output)
for p in sorted((E / 'uphash').rglob('*.tsv.gz')):
    with gzip.open(p, 'rt') as f:
        for line in f:
            acc, hx = line.split('\t', 2)[:2]
            if acc in need: acc2h[acc].add(int(hx, 16))
sp = pd.read_csv(E / 'sprot.tsv.gz', sep='\t'); sp.columns = ['acc', 'ec', 'seq']
sp_h = set()
for acc, e, s in zip(sp.acc, sp.ec, sp.seq):
    k = h(s)
    if ecs(e): sp_h.add(k)          # Swiss-Prot entries with a task EC are in the search reference already
    if acc in need: acc2h[acc].add(k)
for f in [E / 'trembl_ec.tsv.gz', E / 'trembl_ev2s.tsv.gz']:
    t = pd.read_csv(f, sep='\t'); t.columns = ['acc', 'ec', 'seq']
    for acc, s in zip(t.acc, t.seq):
        if acc in need and isinstance(s, str): acc2h[acc].add(h(s))
cur, buf = None, []
for line in list(open(E / 'brenda_missing.fasta')) + ['>end|end|\n']:
    if line.startswith('>'):
        if cur: acc2h[cur].add(h(''.join(buf)))
        cur = line[1:].split('|')[1] if '|' in line else line[1:].split()[0]; buf = []
    else:
        buf.append(line.strip())
bo = pd.read_csv(E / 'round3/brenda_uniparc.tsv', sep='\t')
for acc, s in zip(bo.accession, bo.sequence):
    acc2h[acc].add(h(s))
sn = pd.read_csv(E / 'round3/sifts_new.tsv', sep='\t')
for acc, s in zip(sn.iloc[:, 0], sn.iloc[:, -1]):
    if isinstance(s, str): acc2h[acc].add(h(s))
t6 = pd.read_csv(E / 'r5/brenda26_seqs.tsv', sep='\t')
for acc, s in zip(t6.accession, t6.sequence):
    if isinstance(s, str): acc2h[acc].add(h(s))
for acc, e in b6.items():
    for k in acc2h.get(acc, ()):
        if e: lab[k]['B6'].update(e)
print('BRENDA 2026.1 accessions mapped', sum(a in acc2h for a in b6), 'of', len(b6), flush=True)
for acc, e in br.items():
    for k in acc2h.get(acc, ()):
        if e: lab[k]['BR'].update(e)
for acc, e in sec.items():
    for k in acc2h.get(acc, ()):
        if e: lab[k]['PD'].update(e)
print('BRENDA accessions mapped', sum(a in acc2h for a in br), 'of', len(br), flush=True)
cp = pd.read_csv(E / 'catpred/catpred_seq_ec.tsv.gz', sep='\t')
for e, s in zip(cp.ec, cp.sequence):
    if isinstance(s, str) and ecs(e): lab[h(s)]['CP'].update(ecs(e))
for f in Path(a.catapro).glob('*fold.csv'):
    c = pd.read_csv(f)
    for e, s in zip(c.EC, c.Sequence):
        if isinstance(s, str) and ecs(e): lab[h(s)]['CT'].update(ecs(e))
k_ = pd.read_csv(E / 'kin_esp_seq_ec.tsv.gz', sep='\t')
for src, e, s in zip(k_.src, k_.ec, k_.sequence):
    if isinstance(s, str) and ecs(e): lab[h(s)]['DK' if src == 'dlk' else 'ES'].update(ecs(e))
en = pd.read_csv(E / 'round3/esp_noec_uniprot.tsv', sep='\t').dropna(subset=['EC number'])
for e, s in zip(en['EC number'], en.Sequence):
    if ecs(e): lab[h(s)]['E2'].update(ecs(e))
# current UniProtKB unreviewed EC from the fingerprint index (weights/uphash.npz)
z = np.load(Path(a.model) / 'weights' / 'uphash.npz'); UH, US = z['h'], z['s']
usets = json.load(open(Path(a.model) / 'weights' / 'uphash_sets.json'))
ex = set()
if a.exclude:
    cur = []
    for line in list(open(a.exclude)) + ['>']:
        if line.startswith('>'):
            if cur: ex.add(h(''.join(cur)))
            cur = []
        else:
            cur.append(line.strip())
set_id, sets = {}, []
def sid(t):
    t = tuple(sorted(t))
    if t not in set_id:
        set_id[t] = len(sets); sets.append(list(t))
    return set_id[t]
cnt = collections.Counter()
# 1. sequences in a non-UniProt source: priority order, UniProt (TR) in its place
H, S = [], []
for k, d in lab.items():
    if k in sp_h or k in ex:
        continue
    j = int(np.searchsorted(UH, np.uint64(k)))
    if j < len(UH) and UH[j] == k:
        d['TR'] = set(usets[int(US[j])])
    for src in ORDER:
        if d.get(src):
            H.append(k); S.append(sid(d[src])); cnt[src] += 1
            break
H = np.array(H, dtype=np.uint64); S = np.array(S, dtype=np.uint32)
# 2. all other UniProt unreviewed sequences keep their current EC
m = ~np.isin(UH, H) & ~np.isin(UH, np.array(sorted(sp_h | ex), dtype=np.uint64))
remap = np.array([sid(x) for x in usets], dtype=np.uint32)
cnt['TR'] += int(m.sum())
H = np.concatenate([H, UH[m]]); S = np.concatenate([S, remap[US[m]]])
o = np.argsort(H, kind='stable'); H, S = H[o], S[o]
w = Path(a.model) / 'weights'
np.savez(w / 'srchash.npz', h=H, s=S)
json.dump(sets, open(w / 'srchash_sets.json', 'w'))
print(f'{len(H)} sequences, label source: {dict(cnt)}, excluded {len(ex)}')

	"""v24: exact-sequence label table with a source priority measured on the training set.

	For every sequence in the disclosed external sources, the label is taken from the first source that has one, in the order
	BRENDA 2026.1 > BRENDA 2023.1 > CatPred-DB > CataPro > DLKcat > PDB/SIFTS > UniProtKB unreviewed (current EC) > ESP > ESP (current EC)
	(on 14,122 training proteins that are not in Swiss-Prot, BRENDA's EC is right 95% of the time when it disagrees with the
	current UniProt EC, which is right 2.5% of the time; this order scores 0.974 against the official training labels).
	Swiss-Prot and train sequences are not in this table: identical train/Swiss-Prot sequences keep their own labels.

	Output: <model>/weights/srchash.npz (h: sorted uint64 MD5 prefixes, s: uint32 EC-set id) + srchash_sets.json
	Usage: python3 build_source_hash.py --esm ~/enzyme-esm --catapro <CataPro/datasets> --model ~/enzyme-model [--exclude FASTA]
	"""
	import argparse, collections, gzip, hashlib, json, re
	from pathlib import Path
	import numpy as np
	import pandas as pd

	ap = argparse.ArgumentParser()
	ap.add_argument('--esm', required=True); ap.add_argument('--catapro', required=True); ap.add_argument('--model', required=True)
	ap.add_argument('--exclude')
	a = ap.parse_args()
	E = Path(a.esm)
	VOC = set(json.load(open(Path(a.model) / 'ec_vocabulary.json')))
	ORDER = ['B6', 'BR', 'CP', 'CT', 'DK', 'PD', 'TR', 'ES', 'E2']


	def ecs(x):
	return [e.strip() for e in re.split(r'[;,]', str(x)) if e.strip() in VOC]


	def h(s):
	return int(hashlib.md5(s.encode()).hexdigest()[:16], 16)


	lab = collections.defaultdict(lambda: collections.defaultdict(set)) # hash -> source -> ECs (non-UniProt sources)
	br = collections.defaultdict(set)
	for line in open(E / 'brenda_acc_ec.tsv'):
	if '\t' in line:
	acc, e = line.rstrip('\n').split('\t'); br[acc].update(ecs(e))
	sec = collections.defaultdict(set)
	for line in open(E / 'round3/sifts_acc_ec.tsv'):
	p = line.rstrip('\n').split('\t')
	if len(p) >= 2: sec[p[0]].update(ecs(p[1]))
	b6 = collections.defaultdict(set) # BRENDA release 2026.1 (r5/brenda26_acc_ec.tsv)
	for line in open(E / 'r5/brenda26_acc_ec.tsv'):
	acc, e = line.rstrip('\n').split('\t'); b6[acc].update(ecs(e))
	need = set(br) \| set(sec) \| set(b6)
	acc2h = collections.defaultdict(set) # an accession can map to several sequence versions
	# accession -> sequence fingerprint for BRENDA / SIFTS accessions: UniProtKB unreviewed (fetch_uniprot_hash.py output)
	for p in sorted((E / 'uphash').rglob('*.tsv.gz')):
	with gzip.open(p, 'rt') as f:
	for line in f:
	acc, hx = line.split('\t', 2)[:2]
	if acc in need: acc2h[acc].add(int(hx, 16))
	sp = pd.read_csv(E / 'sprot.tsv.gz', sep='\t'); sp.columns = ['acc', 'ec', 'seq']
	sp_h = set()
	for acc, e, s in zip(sp.acc, sp.ec, sp.seq):
	k = h(s)
	if ecs(e): sp_h.add(k) # Swiss-Prot entries with a task EC are in the search reference already
	if acc in need: acc2h[acc].add(k)
	for f in [E / 'trembl_ec.tsv.gz', E / 'trembl_ev2s.tsv.gz']:
	t = pd.read_csv(f, sep='\t'); t.columns = ['acc', 'ec', 'seq']
	for acc, s in zip(t.acc, t.seq):
	if acc in need and isinstance(s, str): acc2h[acc].add(h(s))
	cur, buf = None, []
	for line in list(open(E / 'brenda_missing.fasta')) + ['>end\|end\|\n']:
	if line.startswith('>'):
	if cur: acc2h[cur].add(h(''.join(buf)))
	cur = line[1:].split('\|')[1] if '\|' in line else line[1:].split()[0]; buf = []
	else:
	buf.append(line.strip())
	bo = pd.read_csv(E / 'round3/brenda_uniparc.tsv', sep='\t')
	for acc, s in zip(bo.accession, bo.sequence):
	acc2h[acc].add(h(s))
	sn = pd.read_csv(E / 'round3/sifts_new.tsv', sep='\t')
	for acc, s in zip(sn.iloc[:, 0], sn.iloc[:, -1]):
	if isinstance(s, str): acc2h[acc].add(h(s))
	t6 = pd.read_csv(E / 'r5/brenda26_seqs.tsv', sep='\t')
	for acc, s in zip(t6.accession, t6.sequence):
	if isinstance(s, str): acc2h[acc].add(h(s))
	for acc, e in b6.items():
	for k in acc2h.get(acc, ()):
	if e: lab[k]['B6'].update(e)
	print('BRENDA 2026.1 accessions mapped', sum(a in acc2h for a in b6), 'of', len(b6), flush=True)
	for acc, e in br.items():
	for k in acc2h.get(acc, ()):
	if e: lab[k]['BR'].update(e)
	for acc, e in sec.items():
	for k in acc2h.get(acc, ()):
	if e: lab[k]['PD'].update(e)
	print('BRENDA accessions mapped', sum(a in acc2h for a in br), 'of', len(br), flush=True)
	cp = pd.read_csv(E / 'catpred/catpred_seq_ec.tsv.gz', sep='\t')
	for e, s in zip(cp.ec, cp.sequence):
	if isinstance(s, str) and ecs(e): lab[h(s)]['CP'].update(ecs(e))
	for f in Path(a.catapro).glob('*fold.csv'):
	c = pd.read_csv(f)
	for e, s in zip(c.EC, c.Sequence):
	if isinstance(s, str) and ecs(e): lab[h(s)]['CT'].update(ecs(e))
	k_ = pd.read_csv(E / 'kin_esp_seq_ec.tsv.gz', sep='\t')
	for src, e, s in zip(k_.src, k_.ec, k_.sequence):
	if isinstance(s, str) and ecs(e): lab[h(s)]['DK' if src == 'dlk' else 'ES'].update(ecs(e))
	en = pd.read_csv(E / 'round3/esp_noec_uniprot.tsv', sep='\t').dropna(subset=['EC number'])
	for e, s in zip(en['EC number'], en.Sequence):
	if ecs(e): lab[h(s)]['E2'].update(ecs(e))
	# current UniProtKB unreviewed EC from the fingerprint index (weights/uphash.npz)
	z = np.load(Path(a.model) / 'weights' / 'uphash.npz'); UH, US = z['h'], z['s']
	usets = json.load(open(Path(a.model) / 'weights' / 'uphash_sets.json'))
	ex = set()
	if a.exclude:
	cur = []
	for line in list(open(a.exclude)) + ['>']:
	if line.startswith('>'):
	if cur: ex.add(h(''.join(cur)))
	cur = []
	else:
	cur.append(line.strip())
	set_id, sets = {}, []
	def sid(t):
	t = tuple(sorted(t))
	if t not in set_id:
	set_id[t] = len(sets); sets.append(list(t))
	return set_id[t]
	cnt = collections.Counter()
	# 1. sequences in a non-UniProt source: priority order, UniProt (TR) in its place
	H, S = [], []
	for k, d in lab.items():
	if k in sp_h or k in ex:
	continue
	j = int(np.searchsorted(UH, np.uint64(k)))
	if j < len(UH) and UH[j] == k:
	d['TR'] = set(usets[int(US[j])])
	for src in ORDER:
	if d.get(src):
	H.append(k); S.append(sid(d[src])); cnt[src] += 1
	break
	H = np.array(H, dtype=np.uint64); S = np.array(S, dtype=np.uint32)
	# 2. all other UniProt unreviewed sequences keep their current EC
	m = ~np.isin(UH, H) & ~np.isin(UH, np.array(sorted(sp_h \| ex), dtype=np.uint64))
	remap = np.array([sid(x) for x in usets], dtype=np.uint32)
	cnt['TR'] += int(m.sum())
	H = np.concatenate([H, UH[m]]); S = np.concatenate([S, remap[US[m]]])
	o = np.argsort(H, kind='stable'); H, S = H[o], S[o]
	w = Path(a.model) / 'weights'
	np.savez(w / 'srchash.npz', h=H, s=S)
	json.dump(sets, open(w / 'srchash_sets.json', 'w'))
	print(f'{len(H)} sequences, label source: {dict(cnt)}, excluded {len(ex)}')