# Colab には torch と huggingface_hub が入っている. 手元で動かす場合は別途入れる
%pip install -q git+https://github.com/GenAIBio/evo2-sae-handson

import numpy as np
import pandas as pd
import torch
import matplotlib.pyplot as plt

import evo2_sae
evo2_sae.setup("https://huggingface.co/datasets/suzuki-2001/evo2-sae-handson/resolve/main")

d = torch.load(evo2_sae.download("region_emb.pt"), map_location="cpu")
x = d["emb"].float()
start, end = d["start"], d["end"]
chrom = d["accession"]

# まずは埋め込みのテンソルを見てみる
print(d.keys())
print(d["emb"])
print(f"{tuple(x.shape)}")
print(f"{chrom}:{start:,}-{end:,}")
print(f"Model: {d['model']}, Layer: {d['layer']}")

downloading region_emb.pt

dict_keys(['emb', 'accession', 'start', 'end', 'model', 'layer'])
tensor([[-0.2266, -0.3711,  0.1055,  ...,  0.2617,  0.1768, -0.2188],
        [ 0.6250, -0.0342, -0.3398,  ..., -1.1875,  0.9062, -0.4277],
        [ 0.1543, -0.5703, -0.1611,  ..., -0.3242,  0.1680, -0.1445],
        ...,
        [ 0.3516,  0.1260, -0.7734,  ..., -0.1494, -0.4551, -0.1846],
        [-0.1147,  0.4258, -0.5664,  ..., -0.6094, -0.4414, -0.4707],
        [-0.5078, -0.5117, -0.0923,  ..., -0.0117,  0.3301,  0.4023]],
       dtype=torch.bfloat16)
(30720, 4096)
NC_000913.3:4,162,560-4,193,280
Model: evo2_7b, Layer: blocks.26

value = x.flatten().numpy()

fig, ax = plt.subplots(figsize=(6.2, 3.4))
ax.set_facecolor("#f7f7f7")
ax.hist(value, bins=300, range=(-1.5, 1.5), color="#334155")
ax.set_title("Evo 2 blocks.26", fontsize=11)
ax.set_xlabel("Value")
ax.set_ylabel("Count")
for spine in ax.spines.values():
    spine.set(linewidth=1.4, color="black")
plt.show()

#@markdown 表示する次元. コンマ区切りで 5 本まで指定できる (0 から 4095)
dims = "0, 100, 2000, 4000"  #@param {type:"string"}

chosen = [int(i) for i in dims.split(",")][:5]

counts = {}
for r in evo2_sae.overlapping(0, 5_000_000):
    counts[r["kind"]] = counts.get(r["kind"], 0) + 1
print("ecoli.bed  " + ",".join(f"{k} {n}" for k, n in sorted(counts.items())))

embedding = evo2_sae.Region.from_matrix(chrom, start, x)
embedding.show(
    [(i, "#334155", f"embedding dim {i}") for i in chosen],
    floor=False,
)

downloading ecoli.bed

ecoli.bed  ORF 4340,prophage 6,rRNA 22,repeat 20,spacer 18,tRNA 86

from huggingface_hub import hf_hub_download

ckpt = hf_hub_download(
    "Goodfire/Evo-2-Layer-26-Mixed",
    "sae-layer26-mixed-expansion_8-k_64.pt"
)
sd = torch.load(ckpt, map_location="cpu")

W = sd["_orig_mod.W"].float()
b_enc = sd["_orig_mod.b_enc"].float()

print(f"W {tuple(W.shape)}   b_enc {tuple(b_enc.shape)}")

Warning: You are sending unauthenticated requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits and faster downloads.

sae-layer26-mixed-expansion_8-k_64.pt: reconstructing file:   0%|          |  0.00B /  537MB

sae-layer26-mixed-expansion_8-k_64.pt: downloading bytes:           |  0.00B

W (4096, 32768)   b_enc (32768,)

K = 64

def batch_topk(a, k):
    floor = torch.topk(a.flatten(), k * a.shape[0]).values.min()
    return torch.where(a >= floor, a, torch.zeros_like(a))

A = torch.relu(x @ W + b_enc)
Z = batch_topk(A, K)

region = evo2_sae.Region.from_genome(chrom, start, end)

# Evo 2 論文が挙げている特徴のインデックス
# Fig. 4c
F_ORF, F_INTERGENIC, F_TRNA, F_RRNA = 13606, 26069, 30262, 2812
F_HELIX, F_SHEET = 28741, 22326  # Fig. 4d
F_PROPHAGE = 19746  # Fig. 4b

# CRISPR のリピートに対応する特徴は論文に挙がっていないので, 第 2 部の探索で選んだ
F_REPEAT = 1261

region.show([(F_RRNA, "#059669", f"f/{F_RRNA}  rRNA")])

fig4c = [
    (F_RRNA, "#059669", f"f/{F_RRNA}  rRNA"),
    (F_TRNA, "#7c3aed", f"f/{F_TRNA}  tRNA"),
    (F_INTERGENIC, "#ea580c", f"f/{F_INTERGENIC}  intergenic"),
    (F_ORF, "#2563eb", f"f/{F_ORF}  ORF"),
]

region.show(fig4c)

region.show(fig4c, 4_166_000, 4_172_500)

rrnA = evo2_sae.Region.from_genome(chrom, 4_035_000, 4_041_200)
rrnA.show(fig4c)

fig4d = [
    (F_TRNA, "#1e3a8a", f"f/{F_TRNA}  tRNA"),
    (F_HELIX, "#2563eb", f"f/{F_HELIX}  alpha-helix"),
    (F_SHEET, "#f59e0b", f"f/{F_SHEET}  beta-sheet"),
]

region.show(fig4d, 4_175_250, 4_177_250)

region.show(fig4d, 4_181_000, 4_189_650)

span = {
    r["name"]: (r["start"] - start, r["end"] - start)
    for r in evo2_sae.overlapping(start, end)
}
loci = ["rrsB", "rrlB", "thrT", "tufB", "rpoB", "rpoC"]
paper = {
    F_RRNA: "rRNA",
    F_TRNA: "tRNA",
    F_HELIX: "alpha-helix",
    F_SHEET: "beta-sheet",
}

rows = []
for feature, label in paper.items():
    track = region.track(feature)
    row = {"特徴": f"f/{feature}", "対応": label}
    row.update({g: round(float(track[span[g][0]:span[g][1]].max()), 2) for g in loci})
    row["全域平均"] = round(float(track.mean()), 3)
    rows.append(row)
pd.DataFrame(rows)

PROTEINS = evo2_sae.proteins()  # ORF の座標と DSSP

# region.residue_acts は, 遺伝子が配布した 30 kb の外にあれば全ゲノム行列から引く.
# 4 つぶんをここでまとめて取り出しておく
for gene in PROTEINS:
    region.residue_acts(gene, F_HELIX)

downloading proteins.json

region.show_complex("tufB", F_HELIX, F_SHEET)

downloading eftu_trna.pdb

region.show_complex("rpoB-rpoC", F_HELIX, F_SHEET)

downloading rpoBC.pdb

protein = "crp"  #@param ["pspA","ompG","crp","groL"]

region.show_complex(protein, F_HELIX, F_SHEET)

downloading crp.pdb

cpz55 = evo2_sae.Region.from_genome(chrom, 2_552_000, 2_572_000)
cpz55.show([(F_PROPHAGE, "#dc2626", f"f/{F_PROPHAGE}  prophage")])

crispr = evo2_sae.Region.from_genome(chrom, 2_872_000, 2_890_000)
crispr.show(
    [
        (F_PROPHAGE, "#dc2626", f"f/{F_PROPHAGE}  spacer"),
        (F_REPEAT, "#334155", f"f/{F_REPEAT}  repeat"),
    ],
    2_877_800, 2_878_550,
)

def mean_on(feature, kind):
    track = crispr.track(feature)
    mask = np.zeros(len(crispr), bool)
    for r in evo2_sae.overlapping(crispr.start, crispr.end, kind):
        begin = max(r["start"], crispr.start) - crispr.start
        end = min(r["end"], crispr.end) - crispr.start
        mask[begin:end] = True
    return track[mask].mean()

rows = []
for feature in (F_PROPHAGE, F_REPEAT):
    row = {"特徴": f"f/{feature}"}
    row.update({
        kind: round(mean_on(feature, kind), 3)
        for kind in ("spacer", "repeat", "ORF")
    })
    row["全体"] = round(float(crispr.track(feature).mean()), 3)
    rows.append(row)
pd.DataFrame(rows)

LOCI = {
    "lac operon": (360_800, 367_900),
    "rrnH operon": (223_300, 229_400),
    "e14 prophage": (1_193_000, 1_216_000),
    "CPZ-55": (2_552_000, 2_572_000),
    "CRISPR array": (2_877_700, 2_878_650),
    "groL": (4_370_600, 4_373_100),
}

locus = "CPZ-55"  #@param ["lac operon","rrnH operon","e14 prophage","CPZ-55","CRISPR array","groL"]

columns = fig4c + [(F_PROPHAGE, "#dc2626", f"f/{F_PROPHAGE}  prophage")]

lo, hi = LOCI[locus]

field = evo2_sae.Region.from_genome(chrom, lo, hi)
field.show(columns)

zoom_from = 2_556_000  #@param {type:"integer"}
zoom_to = 2_560_000  #@param {type:"integer"}

lo = max(zoom_from, field.start)
hi = min(zoom_to, field.end)
if lo >= hi:  # 上のセルで選んだ領域と重ならないので, 領域全体を表示する
    lo, hi = field.start, field.end

field.show(columns, lo, hi)

#@markdown 任意の座標と特徴のインデックス
start_bp = 4_035_000  #@param {type:"integer"}
end_bp = 4_041_200  #@param {type:"integer"}
feature = 2812  #@param {type:"integer"}

evo2_sae.Region.from_genome(chrom, start_bp, end_bp).show(
    [(feature, "#111827", f"f/{feature}")],
)

節	題材	内容
04	Fig. 4c	ORF・遺伝子間・tRNA・rRNA に対応する特徴
05	Fig. 4d	tRNA アレイと tufB, rpoB-rpoC, および立体構造
06	Fig. 4b	プロファージと CRISPR スペーサー
07	—	任意の座位を自分で見る

ファイル	サイズ	内容
`region_emb.pt`	252 MB	Evo 2 の埋め込み. 大腸菌ゲノムの 30,720 塩基ぶん
`ecoli.bed`	0.2 MB	ゲノム全体の遺伝子座標
`proteins.json` と PDB 6 件	2.9 MB	7 つのタンパク質の座標, DSSP, 立体構造
`ecoli_values.f16` ほか 2 本	2.4 GB	SAE の活性行列

遺伝子	残基	H	E	産物
tufB	394	24%	37%	翻訳伸長因子 EF-Tu
rpoB	1342	33%	25%	RNA ポリメラーゼ $\beta$ 鎖
rpoC	1407	39%	19%	RNA ポリメラーゼ $\beta'$ 鎖
pspA	222	91%	0%	ファージショックタンパク質 A
ompG	301	1%	66%	外膜ポリン OmpG
crp	210	38%	30%	cAMP 受容体タンパク質
groL	548	51%	18%	シャペロニン GroEL

	特徴	spacer	repeat	ORF	全体
0	f/19746	1.111	0.000	0.0	0.026
1	f/1261	0.032	3.424	0.0	0.074

Evo 2 の内部表現をスパースオートエンコーダで読む¶

第 1 部 — ゲノムトラックと立体構造への読み出し¶

01. 準備¶

02. Evo 2 の埋め込みを見る¶

03. SAE を通して特徴を読む¶

04. ORF・遺伝子間・tRNA・rRNA に対応する特徴 (Fig. 4c)¶

05. tRNA とタンパク質の二次構造 (Fig. 4d)¶

ゲノムトラック¶

立体構造に載せる¶

06. プロファージと CRISPR スペーサー (Fig. 4b)¶

07. 自由に探索する¶

参考文献¶

	特徴	対応	rrsB	rrlB	thrT	tufB	rpoB	rpoC	全域平均
0	f/2812	rRNA	5.09	6.79	0.0	1.18	0.66	1.18	0.524
1	f/30262	tRNA	1.17	0.92	2.4	0.00	0.00	0.00	0.008
2	f/28741	alpha-helix	0.00	0.66	0.6	1.60	3.79	2.74	0.035
3	f/22326	beta-sheet	0.00	0.00	0.0	3.19	3.37	3.59	0.052