File size: 8,334 Bytes
03677d0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | #!/usr/bin/env python3
"""Export lightweight, reproducible views of the paper's frozen million-point Atlas.
Requires numpy and Pillow in the research environment. No embeddings are refitted,
no remote dataset is downloaded, and full-corpus ID search is independent of these assets.
"""
from export_atlas_coverage import export_coverage
import csv
import gzip
import hashlib
import json
from pathlib import Path
import numpy as np
from PIL import Image, ImageChops
ROOT = Path(__file__).resolve().parents[2]
SOURCE = ROOT / 'Ouroboros_Atlas/visualization/20261107-atlas-million-r1'
LAYOUT = SOURCE / 'research/20261107-category-layout-r1'
PANEL = SOURCE / 'research/20261108-prot2text-panel-r5'
CASES = SOURCE / 'research/20261107-integrated-showcases-r1'
PUBLIC = ROOT / 'Ouroboros_Showcase/public'
SEED = 20261110
AA = dict(zip('ALA ARG ASN ASP CYS GLN GLU GLY HIS ILE LEU LYS MET PHE PRO SER THR TRP TYR VAL SEC PYL ASX GLX UNK'.split(), 'ARNDCQEGHILKMFPSTWYVUOBZX'))
def sha(data):
return hashlib.sha256(data).hexdigest()
def pdb_residues(data):
residues, seen = [], set()
for line in data.decode().splitlines():
if line.startswith('ENDMDL'):
break
if not line.startswith('ATOM ') or line[12:16].strip() != 'CA' or line[16] not in ' A':
continue
key = (line[21], line[22:27])
if key in seen:
continue
seen.add(key)
residues.append(dict(index=int(line[22:26]), chain=line[21].strip(),
code=AA[line[17:20].strip()], confidence=float(line[60:66])))
assert residues, 'No CA residues'
return residues
def sample(pool, labels, count, required):
"""Community-stratified sample; retain small communities and every paper case."""
rng = np.random.default_rng(SEED)
chosen = set(required)
for label in np.unique(labels[pool]):
group = pool[labels[pool] == label]
chosen.update(rng.choice(group, min(8, len(group)), replace=False).tolist())
remainder = np.array(sorted(set(pool.tolist()) - chosen))
chosen.update(rng.choice(remainder, count - len(chosen), replace=False).tolist())
# Random draw order avoids overpainting one community with another.
return rng.permutation(sorted(chosen))
def thumbnail(case):
data = Path(case['image_path']).read_bytes()
assert sha(data) == case['image_sha256']
im = Image.open(case['image_path']).convert('RGBA')
white = Image.new('RGBA', im.size, 'white')
white.alpha_composite(im)
im = white.convert('RGB')
mask = ImageChops.difference(im, Image.new('RGB', im.size, 'white')).convert('L').point(lambda x: 255 if x > 12 else 0)
bbox = mask.getbbox()
if bbox:
im = im.crop(bbox)
im.thumbnail((320, 220))
result = Image.new('RGB', (352, 244), 'white')
result.paste(im, ((352-im.width)//2, (244-im.height)//2))
name = 'atlas-' + case['accession'] + '.webp'
result.save(PUBLIC / 'images' / name, quality=86, method=6)
return '/images/' + name
def main():
export_coverage()
cluster = json.loads((PANEL / 'cluster_manifest.json').read_text())
coordinates = LAYOUT / 'coordinates_landmark.npy'
labels_path = PANEL / 'cluster_labels.npy'
assert sha(labels_path.read_bytes()) == cluster['labels_sha256']
assert sha(coordinates.read_bytes()) == json.loads((LAYOUT/'landmark_map_manifest.json').read_text())['coordinates_sha256']
assert sha((SOURCE/'map_points.csv').read_bytes()) == cluster['inputs']['map_points_sha256']
xy, labels = np.load(coordinates), np.load(labels_path)
assert xy.shape == (1000000, 2) and np.isfinite(xy).all()
assert np.bincount(labels).tolist() == cluster['cluster_sizes']
selected = json.loads((CASES/'main_cases.json').read_text()) + json.loads((CASES/'supplementary_cases.json').read_text())
assert len(selected) == 8
required = [c['index'] for c in selected]
main = sample(np.arange(len(xy)), labels, 12000, required)
preview = sample(main, labels, 4000, required)
keep = set(main.tolist())
rows = {}
with (SOURCE/'map_points.csv').open() as handle:
for i, row in enumerate(csv.DictReader(handle)):
if i in keep:
rows[i] = row
assert len(rows) == 12000
(PUBLIC/'images').mkdir(exist_ok=True)
output_cases = []
titles = ['C2H2 zinc finger', 'RACK1 scaffold', 'PPE family', 'Lysine transporter',
'C2H2 zinc finger', 'C2H2 zinc finger', 'PPE family', 'WXG100 family']
for number, case in enumerate(selected, 1):
row = rows[case['index']]
assert row['uniref_id'] == case['uniref_id']
assert np.array_equal(xy[case['index']], case['map_xy'])
pdb = Path(case['pdb_path']).read_bytes()
assert sha(pdb) == case['pdb_sha256']
residues = pdb_residues(pdb)
sequence = ''.join(r['code'] for r in residues)
assert sha(sequence.encode()) == case['sequence_sha256']
assert len(sequence) == case['sequence_length'] == int(row['sequence_length'])
assert abs(sum(r['confidence'] for r in residues)/len(residues) - case['mean_plddt']) < 1e-8
filename = 'atlas-' + case['accession'] + '.pdb'
(PUBLIC/'structures'/filename).write_bytes(pdb)
fields = ('uniref_id', 'accession', 'pred_text', 'sequence_length', 'map_xy', 'organism',
'source_url', 'afdb_entry', 'mean_plddt', 'sequence_sha256', 'pdb_sha256',
'original_name', 'domain_support_display')
output_cases.append({**{k: case[k] for k in fields}, 'sequence': sequence,
'residues': residues, 'pdbUrl': '/structures/'+filename, 'thumbnail': thumbnail(case),
'number': number, 'title': titles[number-1],
'prediction_status': row['prediction_status'], 'output_hit_limit': row['output_hit_limit'].lower() == 'true',
'model': 'Ouroboros FIX large',
'selection_note': 'Paper illustration; generated function has not been experimentally validated.'})
colors = json.loads((PANEL/'figures/cluster_colors.json').read_text())['cluster_colors']
sources = {str(p.relative_to(ROOT)): sha(p.read_bytes()) for p in [coordinates, labels_path,
PANEL/'figures/cluster_colors.json', CASES/'main_cases.json', CASES/'supplementary_cases.json']}
angle = -45 * np.pi / 180
display_xy = xy @ np.array([[np.cos(angle), np.sin(angle)], [-np.sin(angle), np.cos(angle)]])
def asset(indices):
return dict(ids=[rows[i]['uniref_id'] for i in indices],
xy=np.round(xy[indices].astype(float), 6).tolist(),
reference=[rows[i]['reference_status']=='swissprot2024_exact_match' for i in indices],
lengths=[int(rows[i]['sequence_length']) for i in indices],
clusters=labels[indices].tolist(), colors=colors,
bounds=[xy.min(axis=0).tolist(), xy.max(axis=0).tolist()],
display_bounds=[display_xy.min(axis=0).tolist(), display_xy.max(axis=0).tolist()],
rotation_degrees=-45, records=len(indices), population=51321839,
source_records=1000000, sample_seed=SEED,
sampling='Community-stratified display sample with all eight paper cases retained; display proportions are not corpus estimates.',
reference_label='Exact sequence match to Swiss-Prot 2024_01',
source_sample_sha256=cluster['inputs']['sample_sha256'], sources=sources)
files = {}
for name, value in [('atlas-map.json', asset(main)), ('atlas-preview.json', asset(preview)), ('atlas-cases.json', output_cases)]:
payload = json.dumps(value, ensure_ascii=False, separators=(',', ':')).encode()
(PUBLIC/'data'/name).write_bytes(payload)
if name != 'atlas-cases.json':
(PUBLIC/'data'/(name+'.gz')).write_bytes(gzip.compress(payload, compresslevel=9, mtime=0))
files[name] = dict(bytes=len(payload), sha256=sha(payload))
report = dict(status='passed', records=len(main), preview_records=len(preview), cases=len(output_cases),
communities=len(colors), full_lookup_records=51321839, files=files, sources=sources)
(PUBLIC/'data/atlas-provenance.json').write_text(json.dumps(report, indent=2)+'\n')
print(json.dumps(report))
if __name__ == '__main__':
main()
|