Spaces:
Sleeping
Sleeping
"""In-silico knockout of the "Coupled regions" links with Bacformer (causal): the library.
pg_region_sets.py found 533 links between 290 accessory elements from presence/absence within lineages,
without the model. Here each link is put to the model: in every complete chromosome where the link can be
tested, the upstream element (dnaA reading order, decided per genome) is deleted and the model's belief in
the downstream element is read, in ONE causal fp32 pass over the whole chromosome (<= 5,584 proteins).
units (link, genome) where the link is testable; carriers are defined AT the element's spot (the
validated construction of the scouts' ko_units3.py, ported here unchanged: panRGP spot block,
consensus backbone flanks for loose regions, per-genome orientation, entry gene, insertion
point of an absent element)
co-occurrence genome carries both expected: deleting U LOWERS logit P(D entry)
avoidance carries U, not D (cleanly: expected: deleting U RAISES
no D family anywhere, none log P(D entry family) - log P(resident)
at D's spot) at D's insertion point
skipped (reason recorded): U or D overlapping, < 100 genes apart (MIN_APART; > MAX_APART when set),
wrapping the origin, resident = D's entry family, no entry family, no position.
deletion U's member genes (families of the region, never a backbone family) inside its element span
readout decoded family probability P(f | ctx) = sum_c P(c | ctx) P(f | c) over the FULL softmax at the
readout row, the decoder P(f | c) rebuilt from fp32 full-chromosome baseline passes (top 64 per
gene), cross-validated by lineage: a genome is always read with the decoder of the other half
controls per genome a pool of <= POOL_MAX (72) WHOLE accessory elements, each deleted by its member genes
(control_pool): coupled elements present at their spot (per unit, those linked to U or D excluded)
and, a small second arm, whole panRGP runs at spots holding no coupled region; kept to match the U
sizes and spread along the chromosome. Matched null per unit: controls of size x0.5-2 of U's at a
distance to the readout within +-35 % (WIDEN: x1/3-3, +-60 % when fewer than 3 match; MATCH_LOG a
third tier for close range)
numerics every pass of a genome has the same padded length and the head is applied on fixed-size row
chunks, so identical prefixes give bit-identical readouts: deleting nothing gives delta == 0
exactly and deleting genes after a readout row leaves it exactly unchanged (asserted in every
pass, not only in the dedicated check passes)
link sets Data(sets=PATH) reads any file in the format of pgb/region_sets.json (default); the units cache
is keyed by the md5 of the whole file (for the default file, < 1 MiB, the key is unchanged)
close range (pgb/region_sets_close.json: pairs 50-500 genes apart in the genomes carrying both, two separate
insertions, replicated in lineage-disjoint halves) MAX_APART bounds a unit's U-D distance (600 for these
sets; a unit is one genome's pair in dnaA reading order and can be thousands of genes apart, across the
origin); near_controls() adds whole-element controls at about a unit's own distance upstream of D, and
MATCH_LOG = F a third matching tier (distance within x1/F..xF): at 50-100 genes a genome rarely holds 3
whole elements in the +-35 % / +-60 % windows. Defaults leave the long-range behaviour unchanged.
influence what the model itself knows, independent of any link: after a deletion, the decoded logit
P(true family) (same lineage-CV decoder, fp32 pass, float64 decoding, read through a sparse
copy of the decoder rows: Decoder.srow, Runner.decode_true) at every gene downstream of the
deleted element. Data.elements(g) lists the accessory elements of a genome (panRGP runs at
every spot, by spot id; coupled regions present at their spot, by region id) and their entry
genes. influence_readout() turns one pass into 25-gene bins of the mean delta (and mean
|delta|, and the mean over panRGP genes) out to 1,500 genes, plus the delta at every element
entry out to 2,400 genes (1,500 x 1.6, so that the matched null of an entry near 1,500 genes,
distance +-35 % widened to +-60 %, is on file too). The 25 genes before the first deleted gene
are read as well and must be exactly unchanged (causality, check (g)).
load_influence(g*.npz) reads them back as tables; influence_matrix(g*.npz) is a reference
aggregation of the element-to-element matrix against the units' matched whole-element null (same
entry, size x0.5-2, distance +-35 %). influence(g, element) does it all on demand for one genome
(baseline + deletion + 8 size-matched whole-element controls), served by serve_live.py (/influence,
/elements: the page's own server, locally and in the Space) and by serve_influence.py (a process of its
own: its own port and GPU cap).
compact decoder influence() reads only the sparse decoder rows (Decoder.srow) of the families of the genome, with
the decoder of the other lineage half. export_decoder() writes exactly those rows, both halves, to
pgb/knock/decoder_halves.npz (cluster ids uint16, P(f | c) float64: the same rows bit for bit, 1/12 of
base_top64_{f,p}.npy), and Decoder uses that file when base_top64_{f,p}.npy are absent (the Space).
python3 -c "import pg_knock as K; K.export_decoder()"
Reads the project's data read-only; writes only under pgb/knock/ (through pg_knockout.py, the units cache and
export_decoder()).
"""
import os, json, time, zlib, pickle, hashlib, collections
import numpy as np
PROJ = os.path.dirname(os.path.abspath(__file__)) # the project (also the Space's app directory)
KDIR = f"{PROJ}/pgb/knock"
MODEL = "macwiatrak/bacformer-causal-complete-genomes"
NCL = 50001 # softmax width of gm_head (50,000 clusters + END)
TOPK = 64
CLS, PROT = 2, 4
SLACK, NEAR, GAP, CARRY = 5, 200, 30, 0.5 # as ko_units3.py
MIN_APART = 100 # genes between U's last gene and D's readout gene
MAX_APART = None # at most this many (None: no bound, as before); 600 for the close-range sets
FLOOR = 1e-12 # probability floor of every log
HEAD_CHUNK = 64 # rows per head call (fixed shape -> bit-identical rows across passes)
TRUE_CHUNK = HEAD_CHUNK # rows per head call of decode_true (256 was no faster and moved the logits by 2.5e-5)
PANEL_MAX = 40
ANCHOR_WIN, ANCHOR_SHARE = 40, 0.6 # where an absent element's entry would sit: carriers' flanking families
POOL_MAX = 72 # control deletions (whole accessory elements) per genome
MATCH = (0.5, 2.0, 0.35) # matched null: size x0.5-2 of U's, distance to the readout within +-35 %
WIDEN = (1 / 3.0, 3.0, 0.6) # adaptive widening when fewer than 3 controls match
MATCH_LOG = None # a third tier when still < 3: distance within x1/F..xF (log scale), size x1/3-3 (None: off,
# as before). At 50-100 genes a genome rarely holds 3 whole elements within +-60 % of the
# unit's distance (10 % of the close units); x2 reaches 21 %, x3 31 % (pg_knockout --match_log)
OFF_TOL = 0.35 # off-target readouts: within +-35 % of D's distance from U
SETS_DEFAULT = f"{PROJ}/pgb/region_sets.json"
DECODER_NPZ = f"{KDIR}/decoder_halves.npz" # the compact decoder (export_decoder), used when base_top64 is absent
INF_WIN, INF_BIN = 1500, 25 # influence profile: genes downstream of the deleted element, bin width
INF_ENT = int(INF_WIN * (1 + WIDEN[2])) # element entries read out to 2,400 genes: the widened matched null
INF_UP = 25 # genes read before the first deleted gene: must be exactly unchanged (check g)
MEM_FRAC = 0.35 # per-process GPU memory cap (two runs share the 8 GB card)
def log(*a):
print(time.strftime("%H:%M:%S"), *a, flush=True)
# ============================================================================ data and units
def units_cache_file(sets=None):
"""the units cache that Data(sets) reads or writes (the same key as Data.__init__): space/assemble.py ships it so that
the Space's first /elements does not rebuild it (~5-10 s)"""
raw = open(os.path.abspath(sets) if sets else SETS_DEFAULT, "rb").read()
NP = int(np.load(f"{PROJ}/pgb/chrom.npz")["offsets"][-1])
return f"{KDIR}/units_cache_{hashlib.md5(raw + str((NP, SLACK, NEAR, GAP, CARRY)).encode()).hexdigest()[:12]}.pkl"
class Data:
"""chromosomes, regions, links, and the per (region, genome) element / insertion point of ko_units3.py"""
def __init__(self, cache=True, sets=None):
t0 = time.time()
self.SETS = os.path.abspath(sets) if sets else SETS_DEFAULT
raw = open(self.SETS, "rb").read()
self.SETS_MD5 = hashlib.md5(raw).hexdigest()
self.RS = RS = json.loads(raw)
Z = np.load(f"{PROJ}/pgb/chrom.npz")
self.OFF = OFF = Z["offsets"]; self.FAM = FAM = Z["fam"]; self.SPOT = Z["spot"]; self.RGP = Z["rgp"]; self.PID = Z["pid"]
self.NG = NG = len(OFF) - 1; self.GL = np.diff(OFF); self.NP = int(OFF[-1])
CJ = json.load(open(f"{PROJ}/pgb/chrom_genomes.json"))
self.FN = FN = CJ["families"]; self.NF = len(FN); self.FIDX = {n: i for i, n in enumerate(FN)}
self.ACC = [g["acc"] for g in CJ["genomes"]]
self.gene_names = CJ.get("gene_names")
self.REG = RS["regions"]; self.LINKS = RS["links"]; self.NR = len(self.REG)
self.FAMS = [[self.FIDX[n] for n in r["families"]] for r in self.REG]
rsg = {g["acc"]: g for g in RS["genomes"]}
self.CL = np.array([rsg[a]["cl"] for a in self.ACC])
self.STRAIN = [rsg[a].get("strain", "") for a in self.ACC]
self.first = np.zeros(self.NP, bool); self.first[OFF[:-1]] = True
self.GID = np.repeat(np.arange(NG), self.GL)
self.LOC = np.arange(self.NP) - OFF[self.GID]
NPROT = os.path.getsize(f"{PROJ}/pgb/chrom_emb.f16") // (480 * 2)
self.EMB = np.memmap(f"{PROJ}/pgb/chrom_emb.f16", dtype=np.float16, mode="r", shape=(NPROT, 480))
self._lineages()
# keyed by the WHOLE sets file (the key used to hash its first MiB only: identical for files < 1 MiB,
# as the default one is, so its existing cache stays valid)
key = hashlib.md5(raw + str((self.NP, SLACK, NEAR, GAP, CARRY)).encode()).hexdigest()[:12]
cf = f"{KDIR}/units_cache_{key}.pkl"
loaded = False
if cache and os.path.exists(cf):
try: self.__dict__.update(pickle.load(open(cf, "rb"))); loaded = True
except Exception as e: log(f"units cache {os.path.basename(cf)} unreadable ({type(e).__name__}: {e}): rebuilt")
if not loaded:
self._locate()
if cache:
keep = ("BB", "PRES", "HAS", "TEST", "EL", "INS", "ENTRY", "FLANK", "how")
try:
os.makedirs(KDIR, exist_ok=True)
with open(cf + ".tmp", "wb") as fh: pickle.dump({k: self.__dict__[k] for k in keep}, fh)
os.replace(cf + ".tmp", cf)
except OSError as e: log(f"units cache {os.path.basename(cf)} not written ({e}): kept in memory")
self.COUPLED_SPOTS = {r["spot"] for r in self.REG if r["spot"] is not None}
self.cache_file = cf
log(f"data ready in {time.time()-t0:.0f}s: {NG} chromosomes, {self.NR} regions, {len(self.LINKS)} links"
+ ("" if self.SETS == SETS_DEFAULT else f" (sets {os.path.relpath(self.SETS, PROJ)}, md5 {self.SETS_MD5[:12]})"))
# ---------------------------------------------------------------- lineage clusters and their halves
def _lineages(self):
"""lineage = region_sets.json genomes[i]['cl'] (cut 0.4); cl <= 0 counts as its own lineage (singletons
already are). Halves: a fixed hash (crc32) of the lineage id, so a lineage is always on one side."""
lin = []
for g in range(self.NG):
c = int(self.CL[g]); lin.append(f"cl{c}" if c > 0 else f"acc:{self.ACC[g]}")
self.LIN = np.array(lin, dtype=object)
self.HALF = np.array([zlib.crc32(l.encode()) % 2 for l in lin])
self.LIN_HALF = {l: int(h) for l, h in zip(lin, self.HALF)}
def release_dates(self):
"""ncbi_seq_rel_date from the PanGBank metadata in the pangenome file (pgb/genome_meta.json lacks it)"""
import tables, warnings
warnings.simplefilter("ignore")
with tables.open_file(f"{PROJ}/pgb/ecoli_11587.h5") as h:
t = h.root.metadata.genomes.pangbank_wf.read()
d = {i.decode(): r.decode() for i, r in zip(t["ID"], t["ncbi_seq_rel_date"])}
return [d.get(a, "") for a in self.ACC]
# ---------------------------------------------------------------- ko_units3.py, ported
def _circ_blocks(self, ps, L, gap=GAP):
ps = sorted(ps)
if len(ps) == 1: return [(ps[0], ps[0], 1)]
d = [(ps[(i + 1) % len(ps)] - ps[i]) % L for i in range(len(ps))]
cuts = [i for i, x in enumerate(d) if x > gap]
if not cuts: return [(ps[0], ps[-1], len(ps))]
out = []
for k in range(len(cuts)):
i = (cuts[k] + 1) % len(ps); j = cuts[(k + 1) % len(cuts)]; idx = []
while True:
idx.append(ps[i])
if i == j: break
i = (i + 1) % len(ps)
out.append((idx[0], idx[-1], len(idx)))
return out
def _locate(self):
OFF, FAM, SPOT, NG, GL, NF, REG, NR, FAMS = self.OFF, self.FAM, self.SPOT, self.NG, self.GL, self.NF, self.REG, self.NR, self.FAMS
fampos, spotpos = [], []
for g in range(NG):
a = collections.defaultdict(list); b = collections.defaultdict(list)
for p, (f, s) in enumerate(zip(FAM[OFF[g]:OFF[g + 1]].tolist(), SPOT[OFF[g]:OFF[g + 1]].tolist())):
a[f].append(p)
if s >= 0: b[s].append(p)
fampos.append(a); spotpos.append(b)
uk, cnt = np.unique(self.GID.astype(np.int64) * NF + FAM, return_counts=True)
self.BB = BB = np.bincount((uk % NF)[cnt == 1], minlength=NF) >= 0.95 * NG
share = np.zeros((NR, NG))
for ri, fs in enumerate(FAMS):
for g in range(NG): share[ri, g] = sum(1 for f in fs if fampos[g].get(f)) / len(fs)
self.PRES = PRES = share >= CARRY
leaf = self._leaf()
bad = sum(str(min(9, int(share[ri, g] * 9.999))) != REG[ri]["presence"][leaf[g]]
for ri in range(NR) for g in range(NG))
assert bad == 0, f"presence strings of region_sets.json not reproduced ({bad} digits differ)"
def bb_flanks(g, s, e):
L = int(GL[g]); sl = FAM[OFF[g]:OFF[g + 1]]; out = []
for step, p0 in ((-1, s), (+1, e)):
p = p0
for _ in range(500):
p = (p + step) % L; f = int(sl[p])
if BB[f] and len(fampos[g].get(f, ())) == 1: out.append(f); break
else: out.append(None)
return out
FLANK = []
for ri, r in enumerate(REG):
lc = collections.Counter(); rc = collections.Counter(); n = 0
for g in range(NG):
if not PRES[ri, g]: continue
ps = [p for f in FAMS[ri] for p in fampos[g].get(f, ())]
if not ps: continue
ref = (min(spotpos[g][r["spot"]]) if (r["spot"] is not None and spotpos[g].get(r["spot"])) else r["center"])
bl = sorted(self._circ_blocks(ps, int(GL[g])), key=lambda b: (abs((b[0] + b[1]) / 2 - ref), -b[2]))[0]
lf, rf = bb_flanks(g, bl[0], bl[1]); n += 1
if lf is not None: lc[lf] += 1
if rf is not None: rc[rf] += 1
FLANK.append(dict(left=int(lc.most_common(1)[0][0]) if lc else None, right=int(rc.most_common(1)[0][0]) if rc else None,
lfrac=(lc.most_common(1)[0][1] / n) if lc and n else 0., n=n))
EL = {}; INS = np.full((NR, NG), -1); how = collections.Counter()
for ri, r in enumerate(REG):
s = r["spot"]
for g in range(NG):
blk = spotpos[g].get(s) if s is not None else None
member = [p for f in FAMS[ri] for p in fampos[g].get(f, ())]
if blk:
lo, hi = min(blk), max(blk); INS[ri, g] = lo; how["spot block"] += 1
inside = sorted(p for p in member if lo - SLACK <= p <= hi + SLACK)
if inside:
cl = sorted(self._circ_blocks(inside, int(GL[g])), key=lambda b: (-b[2], abs(b[0] - lo)))[0]
EL[(ri, g)] = dict(start=cl[0], end=cl[1], members=cl[2], span=(cl[1] - cl[0]) % int(GL[g]) + 1,
where="spot", spot_lo=lo, spot_hi=hi)
continue
lf = FLANK[ri]["left"]; lp = fampos[g].get(lf, ()) if lf is not None else ()
rf = FLANK[ri]["right"]; rp = fampos[g].get(rf, ()) if rf is not None else ()
pt = lp[0] + 1 if len(lp) == 1 else (rp[0] if len(rp) == 1 else None)
if pt is None: how["nowhere"] += 1; continue
INS[ri, g] = pt; how["flanks"] += 1
if member:
bl = sorted(self._circ_blocks(member, int(GL[g])), key=lambda b: (abs((b[0] + b[1]) / 2 - pt), -b[2]))[0]
if abs(bl[0] - pt) <= NEAR:
EL[(ri, g)] = dict(start=bl[0], end=bl[1], members=bl[2], span=(bl[1] - bl[0]) % int(GL[g]) + 1,
where="flanks", spot_lo=None, spot_hi=None)
HAS = np.zeros((NR, NG), bool)
for (ri, g) in EL: HAS[ri, g] = True
TEST = PRES & HAS
ENTRY = []
for ri in range(NR):
c = collections.Counter(); n = 0
for g in range(NG):
if TEST[ri, g]: c[int(FAM[OFF[g] + EL[(ri, g)]["start"]])] += 1; n += 1
ENTRY.append(dict(family=int(c.most_common(1)[0][0]) if c else None, frac=(c.most_common(1)[0][1] / n) if n else 0., n=n))
self.EL, self.INS, self.HAS, self.TEST, self.ENTRY, self.FLANK, self.how = EL, INS, HAS, TEST, ENTRY, FLANK, dict(how)
def _leaf(self):
m = {g["acc"]: i for i, g in enumerate(self.RS["genomes"])}
return [m[a] for a in self.ACC]
# ---------------------------------------------------------------- where an ABSENT element's entry would sit
def anchors(self, ri):
"""The families that flank region ri in its carriers: (pred family, share, succ family, share).
pred = the family of the gene just before the element's entry, succ = just after its last gene."""
if not hasattr(self, "_anch"): self._anch = {}
if ri not in self._anch:
pc, sc, n = collections.Counter(), collections.Counter(), 0
for g in range(self.NG):
if not self.TEST[ri, g]: continue
e = self.EL[(ri, g)]; a = self.OFF[g]; L = int(self.GL[g]); n += 1
pc[int(self.FAM[a + (e["start"] - 1) % L])] += 1
sc[int(self.FAM[a + (e["end"] + 1) % L])] += 1
self._anch[ri] = ((int(pc.most_common(1)[0][0]), pc.most_common(1)[0][1] / n) if n else (None, 0.0),
(int(sc.most_common(1)[0][0]), sc.most_common(1)[0][1] / n) if n else (None, 0.0))
return self._anch[ri]
def ins_row(self, ri, g):
"""Where region ri's ENTRY gene would sit in genome g if the element were inserted at its spot, read from the
CARRIERS' context. INS (the spot-block start, or the gene after the left backbone flank) is NOT where the
element starts in carriers -- for nanS it is 12 genes off, for the T6SS the preceding gene is itself absent --
so reading P(D entry) there lands in the far tail (1e-6) and measures nothing (both reviews' blocker).
1. row = 1 + the copy nearest INS (within ANCHOR_WIN) of the family preceding the entry in >= 60 % of carriers
2. else row = the copy nearest INS of the family following the last gene in >= 60 % of carriers
(the element would be inserted before that gene, so its entry would take that gene's index)
3. else no insertion anchor -> the unit is skipped
Returns (row, how) or (None, why)."""
p0 = int(self.INS[ri, g])
if p0 < 0: return None, "no insertion point"
rules = self.anchor_rules(ri)
if not rules: return None, "no insertion anchor (no flanking family of D that is consensus in its carriers " \
"and reproduces D's entry when the rule is applied to a carrier)"
a, L = self.OFF[g], int(self.GL[g])
sl = self.FAM[a:a + L]
for fam, off, how in rules:
c = [p for p in range(max(0, p0 - ANCHOR_WIN), min(L, p0 + ANCHOR_WIN + 1)) if int(sl[p]) == fam]
if not c: continue
r = min(c, key=lambda p: (abs(p - p0), p)) + off
if 0 <= r < L: return int(r), how
return None, "no insertion anchor (D's flanking family is not within 40 genes of the spot here)"
def anchor_rules(self, ri):
"""The flanking families that can place region ri's entry, in preference order: the family preceding the entry
in >= 60 % of carriers, then the family following the last gene, keeping only those whose rule reproduces D's
real entry in >= 50 % of the carriers. Frequency alone is not enough -- for 13 % of the regions on the absent
side of an avoidance link the most frequent neighbour does not mark where the element starts, and reading there
would repeat the error this fix is about. Each entry is (family, offset, name)."""
if not hasattr(self, "_arule"): self._arule = {}
if ri not in self._arule:
(fp, sp), (fs, ss) = self.anchors(ri)
out = []
for fam, share, off, how in ((fp, sp, 1, "pred"), (fs, ss, 0, "succ")):
if fam is None or share < ANCHOR_SHARE: continue
if self._anchor_validate(ri, fam, off) >= 0.5: out.append((int(fam), off, how))
self._arule[ri] = out
return self._arule[ri]
def _anchor_validate(self, ri, fam, off):
"""share of ri's carriers in which the rule returns D's real entry gene (off = 1) / the gene after its last
gene (off = 0), which is where the entry would land if the element were inserted before that gene"""
ok = n = 0
for g in range(self.NG):
if not self.TEST[ri, g]: continue
e = self.EL[(ri, g)]; a = self.OFF[g]; L = int(self.GL[g]); sl = self.FAM[a:a + L]
p0 = int(self.INS[ri, g])
if p0 < 0: continue
want = e["start"] if off == 1 else (e["end"] + 1) % L
c = [p for p in range(max(0, p0 - ANCHOR_WIN), min(L, p0 + ANCHOR_WIN + 1)) if int(sl[p]) == fam]
n += 1
if c and min(c, key=lambda p: (abs(p - p0), p)) + off == want: ok += 1
return (ok / n) if n else 0.0
def anchor_check(self, ri):
"""Validation of the rule actually used for region ri: (n carriers, share of them where it returns D's real
entry, rule name). Both candidate rules are also reported, to show what was rejected."""
rules = self.anchor_rules(ri)
n = int(self.TEST[ri].sum())
(fp, sp), (fs, ss) = self.anchors(ri)
both = {}
for fam, share, off, how in ((fp, sp, 1, "pred"), (fs, ss, 0, "succ")):
both[how] = (self._anchor_validate(ri, fam, off) if (fam is not None and share >= ANCHOR_SHARE) else np.nan)
if not rules: return n, np.nan, None, both
return n, self._anchor_validate(ri, rules[0][0], rules[0][1]), "+".join(r[2] for r in rules), both
def pos_of(self, ri, g):
if self.TEST[ri, g]: return int(self.EL[(ri, g)]["start"])
return int(self.INS[ri, g]) if self.INS[ri, g] >= 0 else None
def members(self, ri, g):
"""U's member genes inside its element span (0-based positions in genome g), never a backbone family"""
e = self.EL[(ri, g)]; a = self.OFF[g]
fs = set(self.FAMS[ri])
return [p for p in range(e["start"], e["end"] + 1) if int(self.FAM[a + p]) in fs and not self.BB[self.FAM[a + p]]]
# ---------------------------------------------------------------- testable units
def units(self, links):
"""one unit = one link in one genome: delete U (upstream in this genome), read D.
Returns (units, skip counter, skipped list)."""
OFF, FAM = self.OFF, self.FAM
units, skipped = [], []
skip = collections.Counter()
for li in links:
l = self.LINKS[li]; a, b, sg = l["a"], l["b"], l["sign"]
for g in range(self.NG):
ca, cb = bool(self.TEST[a, g]), bool(self.TEST[b, g])
if sg > 0 and not (ca and cb): continue # not a co-occurrence genome: not a unit
if sg < 0 and ca == cb: continue # both or neither: not an avoidance genome
pa, pb = self.pos_of(a, g), self.pos_of(b, g)
why = None
if pa is None or pb is None: why = "no position for an element"
else:
up, dn = (a, b) if pa < pb else (b, a)
if sg < 0 and not self.TEST[up, g]: why = "avoidance: the absent element is upstream (would need an insertion)"
elif sg < 0 and (self.PRES[dn, g] or self.HAS[dn, g]):
why = "avoidance: D not cleanly absent (its families elsewhere or a minority at its spot)"
if why is None:
eu = self.EL[(up, g)]
if eu["end"] < eu["start"]: why = "U wraps the origin"
if why is None and sg > 0:
ed = self.EL[(dn, g)]
if ed["end"] < ed["start"]: why = "D wraps the origin"
how = ""
if why is None:
ue = eu["end"]
if sg > 0:
rd = int(self.EL[(dn, g)]["start"]); f1 = int(FAM[OFF[g] + rd]); f2 = -1
else:
f1 = self.ENTRY[dn]["family"]
rd, how = self.ins_row(dn, g)
if f1 is None: why = "avoidance: D has no entry family (no carrier)"
elif rd is None: why = "avoidance: " + how
else:
f2 = int(FAM[OFF[g] + rd])
if f1 == f2: why = "avoidance: the resident gene is D's entry family"
if why is None:
if rd <= ue: why = "U and D overlap (D's readout gene not after U's last gene)"
elif rd - ue - 1 < MIN_APART: why = f"U and D < {MIN_APART} genes apart"
elif MAX_APART is not None and rd - ue - 1 > MAX_APART: why = f"U and D > {MAX_APART} genes apart"
if why is None:
dels = self.members(up, g)
if not dels: why = "U has no member gene outside the backbone"
if why is not None:
skip[why] += 1; skipped.append(dict(li=li, g=g, sign=sg, why=why)); continue
sec = []
if sg > 0: # D's member genes (secondary readout)
sec = self.members(dn, g)
# how complete is the knockout? U's families left elsewhere in the chromosome (multi-copy elements:
# for some links the "knockout" leaves every family in place -- mechanics review finding 2)
sl = FAM[OFF[g]:OFF[g + 1]]; left = np.delete(sl, np.asarray(dels, np.int64))
fs = set(self.FAMS[up]); pres_f = {f for f in fs if (sl == f).any()}
u_left = sum(1 for f in pres_f if (left == f).any())
units.append(dict(li=li, g=g, sign=sg, up=up, dn=dn, u_start=int(eu["start"]), u_end=int(ue),
u_dels=dels, u_size=len(dels), u_span=int(eu["span"]), row=rd, f1=int(f1), f2=int(f2),
dist=int(rd - ue), d_sec=sec, row_how=how,
u_fam=len(pres_f), u_left=int(u_left),
f1_up=int((sl[:rd] == f1).sum()), # is the readout family specific to D?
d_start=int(self.EL[(dn, g)]["start"]) if sg > 0 else rd,
d_end=int(self.EL[(dn, g)]["end"]) if sg > 0 else rd))
return units, skip, skipped
# ---------------------------------------------------------------- the exchangeable null and the off-targets
def links_of(self, ri):
"""regions linked to ri in region_sets.json (either sign)"""
if not hasattr(self, "_lof"):
self._lof = collections.defaultdict(set)
for l in self.LINKS: self._lof[l["a"]].add(l["b"]); self._lof[l["b"]].add(l["a"])
return self._lof[ri]
def control_pool(self, g, want=None, sizes=(), uncapped=False):
"""Whole accessory elements of genome g, each deleted by its member genes: the null for a U deletion.
The pilot's null was a contiguous SUB-RUN of a natural element, which is not the same kind of object as a
whole element (whole-vs-part shifts the readout as much as the effects being tested, inference review
finding 3), and it matched only 4 controls per unit (31 % of units unusable, finding 6). Whole COUPLED
elements are the right null: same construction as U, same size distribution (77 % are 1-2 genes), ~93 per
genome. Whole natural elements at spots holding no coupled region are kept as a second arm.
arm 'cpl': a coupled element present at its spot here; its region id is kept so that, per unit, controls
linked to U or D (or U/D themselves) are excluded
arm 'nat': a maximal run of panRGP genes of one spot at a spot with no coupled region, whole
`want`: rows that must not be deleted (readout genes). Capped at POOL_MAX, keeping a spread of sizes and
positions (uncapped=True: every candidate, for near_controls).
The 'nat' arm is small in practice: the cap keeps the controls that match most U sizes, and the natural runs
are longer than the 1-2-gene coupled elements, so ~4 % of the kept controls are natural (pilot: 1,714 of
38,654 control passes; 0 in some genomes). The null is in effect "delete another coupled element"."""
a, b = int(self.OFF[g]), int(self.OFF[g + 1]); n = b - a
fam = self.FAM[a:b]; rgp = self.RGP[a:b]; sp = self.SPOT[a:b]
want = set() if want is None else set(int(x) for x in want)
pool = []
for ri in range(self.NR):
if not (self.HAS[ri, g] and self.PRES[ri, g]): continue
e = self.EL[(ri, g)]
if e["end"] < e["start"]: continue # wraps the origin
dels = self.members(ri, g)
if dels and not (set(dels) & want):
pool.append(dict(arm="cpl", ri=ri, start=int(e["start"]), end=int(e["end"]), dels=dels,
size=len(dels), spot=(e["spot_lo"] if e["where"] == "spot" else -1)))
i = 0
while i < n:
if rgp[i] and sp[i] >= 0:
j = i
while j + 1 < n and rgp[j + 1] and sp[j + 1] == sp[i]: j += 1
if int(sp[i]) not in self.COUPLED_SPOTS:
dels = [p for p in range(i, j + 1) if not self.BB[fam[p]]]
if dels and not (set(dels) & want):
pool.append(dict(arm="nat", ri=-1, start=i, end=j, dels=dels, size=len(dels), spot=int(sp[i])))
i = j + 1
else: i += 1
if not uncapped and len(pool) > POOL_MAX:
# keep the controls that can actually serve as a matched null here: score each by the number of this
# genome's U sizes it matches (x0.5-2), then by a spread of sizes and positions
sz = np.asarray(sizes, float)
def score(c):
n_ = int(np.sum((c["size"] >= MATCH[0] * sz) & (c["size"] <= MATCH[1] * sz))) if len(sz) else 0
return (-n_, c["size"], c["start"])
pool.sort(key=score)
head = [c for c in pool if score(c)[0] < 0][:POOL_MAX]
if len(head) < POOL_MAX:
rest = [c for c in pool if score(c)[0] == 0]
idx = np.unique(np.linspace(0, len(rest) - 1, POOL_MAX - len(head)).round().astype(int)) if rest else []
head += [rest[k] for k in idx]
pool = head
return sorted(pool, key=lambda c: c["start"])
def near_controls(self, g, units, pool, want=5):
"""Extra whole-element controls for units whose matched null is thin (close range: with <= 72 controls spread
over ~5,000 genes, 0 % of the units 50-100 genes apart and 37 % of those 100-250 apart had >= 3 matched
controls in the close smoke run). For each unit of genome g, every candidate whole element (control_pool,
uncapped: coupled elements present at their spot, natural panRGP runs at spots without a coupled region) that
is usable in its null (no overlap with U or D, not U, D or a region linked to either, ending before D's
readout gene) and matches it under the widened rule (size x1/3-3, distance to the readout within +-60 %, or
within x1/F..xF with MATCH_LOG = F) is added, nearest in log distance and log size first, while it raises the unit's matched count as
match_controls() will compute it in the aggregation, until `want`. Returns the extra controls (not already in
`pool`), sorted by start."""
allc = self.control_pool(g, uncapped=True)
key = lambda c: (c["start"], c["end"], c["arm"], c["ri"])
have = {key(c) for c in pool}; extra = {}
for u in units:
banned = {u["up"], u["dn"]} | self.links_of(u["up"]) | self.links_of(u["dn"])
def usable(c):
return (c["end"] < u["row"] and not (c["start"] <= u["u_end"] and c["end"] >= u["u_start"])
and not (c["start"] <= u["d_end"] and c["end"] >= u["d_start"])
and not (c["ri"] >= 0 and c["ri"] in banned))
def nmatch(cs):
if not cs: return 0
idx, _w = match_controls(np.array([c["end"] for c in cs]), np.array([c["size"] for c in cs]),
np.ones(len(cs), bool), u["u_size"], u["dist"], u["row"])
return len(idx)
cur = [c for c in pool if usable(c)] + [c for c in extra.values() if usable(c)]
n0 = nmatch(cur)
if n0 >= want: continue
cand = [c for c in allc if key(c) not in have and key(c) not in extra and usable(c)
and WIDEN[0] * u["u_size"] <= c["size"] <= WIDEN[1] * u["u_size"]
and (abs((u["row"] - c["end"]) - u["dist"]) <= WIDEN[2] * u["dist"] or
(MATCH_LOG is not None and abs(np.log(max(u["row"] - c["end"], 1) / u["dist"])) <= np.log(MATCH_LOG)))]
cand.sort(key=lambda c: (abs(np.log((u["row"] - c["end"]) / u["dist"])), abs(np.log(c["size"] / u["u_size"])), c["start"]))
for c in cand:
n1 = nmatch(cur + [c])
if n1 > n0: cur.append(c); extra[key(c)] = c; n0 = n1
if n0 >= want: break
return sorted(extra.values(), key=lambda c: c["start"])
def offtargets(self, g, rows_needed):
"""Readouts of OTHER coupled elements in genome g, to test whether a U deletion moves D specifically or moves
every element downstream (inference review blocker): a present element is read at its entry gene (co kind),
an absent element with an insertion anchor at that row (avoidance kind, resident = the gene there).
`rows_needed`: only positions in this set are built (the union of the units' distance windows)."""
a = int(self.OFF[g]); out = []
for ri in range(self.NR):
if self.TEST[ri, g]:
e = self.EL[(ri, g)]
if e["end"] < e["start"]: continue
r = int(e["start"])
if r in rows_needed: out.append(dict(ri=ri, kind=1, row=r, f1=int(self.FAM[a + r]), f2=-1))
elif not self.PRES[ri, g] and not self.HAS[ri, g]:
f1 = self.ENTRY[ri]["family"]
if f1 is None: continue
r, _how = self.ins_row(ri, g)
if r is None or r not in rows_needed: continue
f2 = int(self.FAM[a + r])
if f2 != f1: out.append(dict(ri=ri, kind=-1, row=int(r), f1=int(f1), f2=f2))
return out
def phi(self, ri, rj):
"""across-genome (largely phylogenetic) correlation of two regions' presence: a region that merely marks a
strain background will move together with U without any within-lineage coupling"""
x = self.PRES[ri].astype(float); y = self.PRES[rj].astype(float)
if x.std() == 0 or y.std() == 0: return np.nan
return float(np.corrcoef(x, y)[0, 1])
# ---------------------------------------------------------------- the accessory elements of a genome (influence)
def elements(self, g):
"""Every accessory element present in genome g, with its entry gene (the first gene in dnaA reading order):
kind 0 a maximal run of panRGP genes sharing one spot id (every spot, coupled or not), id = spot id;
size = its non-backbone genes (what a deletion of the whole run removes, as control_pool's 'nat')
kind 1 a coupled region present at its spot (PRES and HAS, the carriers of the units), id = region id,
element span EL, size = its member genes; elements wrapping the origin are left out
panRGP genes outside every spot are not listed (no id to follow them across genomes). A coupled element
usually lies inside a spot run, so both can share an entry gene. Sorted by entry; cached per genome."""
if not hasattr(self, "_elts"): self._elts = {}
if g in self._elts: return self._elts[g]
a, b = int(self.OFF[g]), int(self.OFF[g + 1]); n = b - a
fam = self.FAM[a:b]; rgp = self.RGP[a:b]; sp = self.SPOT[a:b]
out = []
i = 0
while i < n:
if rgp[i] and sp[i] >= 0:
j = i
while j + 1 < n and rgp[j + 1] and sp[j + 1] == sp[i]: j += 1
out.append((0, int(sp[i]), i, j, i, sum(1 for p in range(i, j + 1) if not self.BB[fam[p]])))
i = j + 1
else: i += 1
for ri in range(self.NR):
if not (self.HAS[ri, g] and self.PRES[ri, g]): continue
e = self.EL[(ri, g)]
if e["end"] < e["start"]: continue
out.append((1, ri, int(e["start"]), int(e["end"]), int(e["start"]), len(self.members(ri, g))))
out.sort(key=lambda t: (t[4], t[0], t[1]))
A = np.array(out, np.int64).reshape(-1, 6)
E = dict(kind=A[:, 0].astype(np.int8), id=A[:, 1].astype(np.int32), start=A[:, 2].astype(np.int32),
end=A[:, 3].astype(np.int32), entry=A[:, 4].astype(np.int32), size=A[:, 5].astype(np.int32))
E["fam"] = fam[E["entry"]].astype(np.int32) if len(A) else np.zeros(0, np.int32)
self._elts[g] = E
return E
def element_label(self, kind, eid):
return (self.REG[int(eid)]["label"] if int(kind) == 1 else f"spot {int(eid)}")
# ---------------------------------------------------------------- control panel
def panel(self, g, target_sizes, avoid, rng):
"""<= PANEL_MAX control deletions in genome g, one per target size.
Natural accessory elements = maximal runs of panRGP genes sharing one spot id, at spots holding none of
the 290 coupled regions, not touching any coupled element here (`avoid` = positions). They are ~24 per
genome with >= 3-4 member genes each, while half the U deletions are 1-2 genes, so a control deletion is
a contiguous sub-run of t member genes (t = target size) of such an element -- itself a contiguous run of
panRGP genes of one spot, as a U element is a sub-run of its own spot's RGP. An element hosts several
non-overlapping sub-runs only when every element has been used (fewest uses first); among those,
farthest-point choice spreads the panel along the chromosome."""
a, b = self.OFF[g], self.OFF[g + 1]
rgp = self.RGP[a:b]; sp = self.SPOT[a:b]; fam = self.FAM[a:b]
cand = []
i, n = 0, b - a
while i < n:
if rgp[i] and sp[i] >= 0:
j = i
while j + 1 < n and rgp[j + 1] and sp[j + 1] == sp[i]: j += 1
if int(sp[i]) not in self.COUPLED_SPOTS:
pos = [p for p in range(i, j + 1) if not self.BB[fam[p]]]
if pos and not (set(range(i, j + 1)) & avoid):
cand.append(dict(start=i, end=j, pos=pos, spot=int(sp[i]), free=[True] * len(pos), uses=0))
i = j + 1
else: i += 1
chosen = []
order = sorted(range(len(target_sizes)), key=lambda k: -target_sizes[k]) # large first: rarer
for k in order:
t = int(target_sizes[k]); opts = []
for ci, c in enumerate(cand): # free windows of t member genes
w = [s0 for s0 in range(len(c["pos"]) - t + 1) if all(c["free"][s0:s0 + t])]
if w: opts.append((ci, w))
if not opts: continue
mu = min(cand[ci]["uses"] for ci, _ in opts); opts = [(ci, w) for ci, w in opts if cand[ci]["uses"] == mu]
if chosen:
mids = np.array([(e["start"] + e["end"]) / 2 for e in chosen])
dmin = [np.abs(mids - (cand[ci]["start"] + cand[ci]["end"]) / 2).min() for ci, _ in opts]
best = max(dmin); opts = [o for o, d in zip(opts, dmin) if d >= best - 1e-9]
ci, w = opts[int(rng.integers(len(opts)))]
s0 = w[int(rng.integers(len(w)))]; c = cand[ci]
for q in range(s0, s0 + t): c["free"][q] = False
c["uses"] += 1
dels = c["pos"][s0:s0 + t]
chosen.append(dict(start=dels[0], end=dels[-1], dels=dels, size=t, spot=c["spot"],
elem_start=c["start"], elem_end=c["end"], elem_size=len(c["pos"])))
return sorted(chosen, key=lambda e: e["start"]), len(cand)
# ============================================================================ decoder P(family | cluster)
class Decoder:
"""A[f, c] = sum over genes of family f (not the first gene of a genome) of P(c | context), top 64 of the
fp32 full-chromosome baseline calls, one per lineage half; P(f | c) = A[f, c] / sum_f A[f, c].
Genome g (half h) is read with the decoder of half 1 - h, built from genomes of the other lineages only.
Two sources: the fp32 baseline calls base_top64_{f,p}.npy (~0.9 GB; every method), or, when they are absent (or
compact=True), the export of export_decoder() (DECODER_NPZ, or `path`): the sparse rows srow() of every family of
the genomes each half reads, bit for bit (float64; a float32 export moves the logit of a gene predicted at
P > 0.99 by up to 2e-4, the deltas by < 1e-6). row() is then rebuilt from srow(); TF, TP and colsum do not exist,
and a family outside those genomes raises KeyError."""
def __init__(self, data, compact=None, path=None):
self.d = data
NP = data.NP
ff, fp = f"{KDIR}/base_top64_f.npy", f"{KDIR}/base_top64_p.npy"
if compact is None: compact = not (os.path.exists(ff) and os.path.exists(fp))
self.compact = bool(compact)
half_gene = data.HALF[data.GID]
self.mask = {h: (half_gene == h) & ~data.first for h in (0, 1)}
self.lineages = {h: set(data.LIN[data.HALF == h].tolist()) for h in (0, 1)}
self._rows = {}
if self.compact:
self._load_compact(path or DECODER_NPZ); return
self.TF = np.load(ff, mmap_mode="r")
self.TP = np.load(fp, mmap_mode="r")
assert self.TF.shape == (NP, TOPK)
self.colsum = {}
for h in (0, 1):
m = np.flatnonzero(self.mask[h]); cs = np.zeros(NCL)
for s in range(0, len(m), 200000):
mm = m[s:s + 200000]
cs += np.bincount(np.asarray(self.TF[mm]).ravel(), weights=np.asarray(self.TP[mm], np.float64).ravel(), minlength=NCL)
self.colsum[h] = cs
# genes of each family, per half (sorted global indices)
order = np.argsort(data.FAM, kind="stable"); fs = data.FAM[order]
self._ord = order; self._fs = fs
def _load_compact(self, path):
"""the arrays of export_decoder(), checked against this Data (same genes, same lineage halves)"""
if not os.path.exists(path):
raise FileNotFoundError(f"no decoder: neither pgb/knock/base_top64_{{f,p}}.npy nor {os.path.relpath(path, PROJ)}")
with np.load(path) as Z:
meta = json.loads(str(Z["meta"]))
assert meta["np"] == self.d.NP and meta["ncl"] == NCL, f"{path}: built for other data ({meta})"
assert np.array_equal(Z["half"], self.d.HALF), f"{path}: its lineage halves are not this Data's"
self.sp = {h: tuple(np.asarray(Z[f"{k}{h}"]) for k in ("fam", "ptr", "col", "val", "cover")) for h in (0, 1)}
self.meta = meta
def half_for(self, g):
"""the decoder that reads genome g: the other half. Self-check (d): it never saw g's lineage."""
h = 1 - int(self.d.HALF[g])
assert self.d.LIN[g] not in self.lineages[h], f"genome {g}: decoder {h} saw its lineage {self.d.LIN[g]}"
assert not self.mask[h][self.d.OFF[g]:self.d.OFF[g + 1]].any(), f"genome {g}: its own genes are in decoder {h}"
return h
def row(self, f, h):
"""P(f | c) for every cluster c (float64, length NCL), decoder half h"""
k = (int(f), h)
if k not in self._rows and self.compact:
c, w = self.srow(f, h); a = np.zeros(NCL); a[c] = w; self._rows[k] = a
if k not in self._rows:
lo, hi = np.searchsorted(self._fs, [f, f + 1]); idx = np.sort(self._ord[lo:hi])
idx = idx[self.mask[h][idx]]
a = np.zeros(NCL)
if len(idx):
a = np.bincount(np.asarray(self.TF[idx]).ravel(), weights=np.asarray(self.TP[idx], np.float64).ravel(), minlength=NCL)
with np.errstate(invalid="ignore", divide="ignore"):
w = np.where(self.colsum[h] > 0, a / np.maximum(self.colsum[h], 1e-300), 0.0)
self._rows[k] = np.clip(w, 0.0, 1.0)
return self._rows[k]
def srow(self, f, h):
"""row(f, h) as (clusters, values): the same float64 values, only the non-zero ones (median ~380 of 50,001
clusters per family, at most ~12,700), so that the true family of EVERY gene of a genome can be decoded
(~4,500 families per genome: dense rows would take 1.8 GB). Cached (~5 kB per family and half); the dense
row() cache is not touched. Compact decoder: the exported row."""
if self.compact:
fam, ptr, col, val, cover = self.sp[h]; f = int(f)
i = int(np.searchsorted(fam, f))
if i < len(fam) and fam[i] == f:
return col[ptr[i]:ptr[i + 1]].astype(np.int64), val[ptr[i]:ptr[i + 1]].astype(np.float64)
j = int(np.searchsorted(cover, f))
if j < len(cover) and cover[j] == f: return np.zeros(0, np.int64), np.zeros(0)
raise KeyError(f"family {f}: not in the compact decoder's half {h} (only the families of the genomes it reads)")
if not hasattr(self, "_srows"): self._srows = {}
k = (int(f), h)
if k not in self._srows:
lo, hi = np.searchsorted(self._fs, [f, f + 1]); idx = np.sort(self._ord[lo:hi])
idx = idx[self.mask[h][idx]]
if len(idx):
a = np.bincount(np.asarray(self.TF[idx]).ravel(), weights=np.asarray(self.TP[idx], np.float64).ravel(), minlength=NCL)
c = np.flatnonzero(a > 0)
with np.errstate(invalid="ignore", divide="ignore"):
w = np.where(self.colsum[h][c] > 0, a[c] / np.maximum(self.colsum[h][c], 1e-300), 0.0)
w = np.clip(w, 0.0, 1.0); keep = w > 0
self._srows[k] = (c[keep].astype(np.int64), w[keep])
else:
self._srows[k] = (np.zeros(0, np.int64), np.zeros(0))
return self._srows[k]
def sparse_table(self, fams, h):
"""CSR of the decoder rows of `fams` (half h): cols, vals, ptr (len(fams)+1), nnz"""
rs = [self.srow(f, h) for f in fams]
nnz = np.array([len(c) for c, _ in rs], np.int64)
ptr = np.concatenate([[0], np.cumsum(nnz)])
cols = np.concatenate([c for c, _ in rs]) if len(rs) else np.zeros(0, np.int64)
vals = np.concatenate([v for _, v in rs]) if len(rs) else np.zeros(0)
return cols, vals, ptr, nnz
def export_decoder(path=None, dtype=np.float64, data=None, compress=True):
"""The compact decoder for influence() where base_top64_{f,p}.npy (~0.9 GB) cannot go (the Space): for each half h,
the sparse rows srow(f, h) of every family f of the genomes read with h (half_for: the genomes of half 1 - h), the
only rows influence() reads. Per half: cover{h} (those families, int32), fam{h} (the ones with a non-empty row),
ptr{h} (int64 CSR offsets), col{h} (cluster ids, uint16: NCL = 50,001 < 65,536), val{h} (P(f | c), `dtype`:
float64 keeps srow() bit for bit, ~75 MB; float32 ~45 MB); `half` (the lineage half of each genome, checked on
load) and `meta`. ~90 s. Returns (path, bytes)."""
assert NCL <= 65536
path = path or DECODER_NPZ
D = data if data is not None else Data(); dec = Decoder(D, compact=False)
out = dict(half=D.HALF.astype(np.int8)); info = {}
for h in (0, 1):
gs = np.flatnonzero(D.HALF == 1 - h)
cover = np.unique(np.concatenate([D.FAM[D.OFF[g]:D.OFF[g + 1]] for g in gs]))
fams, ptr, cols, vals = [], [0], [], []
for f in cover.tolist():
c, w = dec.srow(f, h)
if not len(c): continue
fams.append(f); cols.append(c.astype(np.uint16)); vals.append(w.astype(dtype)); ptr.append(ptr[-1] + len(c))
out[f"cover{h}"] = cover.astype(np.int32); out[f"fam{h}"] = np.array(fams, np.int32)
out[f"ptr{h}"] = np.array(ptr, np.int64); out[f"col{h}"] = np.concatenate(cols); out[f"val{h}"] = np.concatenate(vals)
info[h] = dict(genomes_read=int(len(gs)), families=int(len(cover)), rows=len(fams), nnz=int(ptr[-1]))
dec._srows = {}
out["meta"] = np.array(json.dumps(dict(np=int(D.NP), ncl=NCL, topk=TOPK, dtype=np.dtype(dtype).name, halves=info,
sets_md5=D.SETS_MD5, built=time.strftime("%Y-%m-%d %H:%M"),
source="pgb/knock/base_top64_{f,p}.npy (fp32 baseline calls, top 64)",
source_files={os.path.basename(f): dict(size=os.path.getsize(f), mtime=os.path.getmtime(f))
for f in (f"{KDIR}/base_top64_f.npy", f"{KDIR}/base_top64_p.npy")})))
tmp = path + ".tmp.npz"
(np.savez_compressed if compress else np.savez)(tmp, **out)
os.replace(tmp, path)
log(f"decoder -> {path}: {os.path.getsize(path) / 1e6:.1f} MB, {info}")
return path, os.path.getsize(path)
# ============================================================================ the model
class Runner:
"""Bacformer causal, fp32, one chromosome per pass (batch 1), padded to a fixed length per genome.
Readouts go through the decoder in float64 on the GPU:
S[r, f] = sum_c p[r, c] W[f, c] = P(f | ctx)
C[r, f] = sum_c p[r, c] (1 - W[f, c]) = 1 - P(f | ctx), computed directly (no cancellation)"""
def __init__(self, dev="cuda:0", frac=MEM_FRAC, threads=None):
"""frac: per-process GPU memory cap (0.35 so that two runs share the 8 GB card; None leaves the process's
cap alone, e.g. inside serve_live). dev='cpu' runs on the CPU with `threads` threads."""
import torch
from transformers import AutoModelForCausalLM
self.torch = torch
if str(dev).startswith("cuda") and frac is not None:
torch.cuda.set_per_process_memory_fraction(frac, torch.device(dev).index or 0)
if not str(dev).startswith("cuda") and threads: torch.set_num_threads(int(threads))
torch.backends.cuda.matmul.allow_tf32 = False; torch.backends.cudnn.allow_tf32 = False
self.dev = dev
self.cm = AutoModelForCausalLM.from_pretrained(MODEL, trust_remote_code=True).float().eval().to(dev)
self.E = None
def sync(self):
if str(self.dev).startswith("cuda"): self.torch.cuda.synchronize()
def set_genome(self, emb_rows, L):
"""emb_rows: (n, 480) float array of the genome's proteins; L: padded sequence length for all passes"""
t = self.torch
self.E = t.tensor(np.asarray(emb_rows, np.float32), device=self.dev)
self.L = int(L)
self.stm = t.full((1, self.L), PROT, device=self.dev, dtype=t.long); self.stm[0, 0] = CLS
self.tt = t.zeros((1, self.L), device=self.dev, dtype=t.long)
def set_families(self, W):
"""W: (F, NCL) float64 decoder rows used for this genome (fixed for all its passes)"""
t = self.torch
self.W = t.tensor(W, device=self.dev, dtype=t.float64)
self.Wc = 1.0 - self.W
def hidden(self, src, extra=None):
"""src: int array, new position -> row of self.E (>= 0) or -(1 + k) for extra[k] (inserted proteins)"""
t = self.torch
m = len(src); assert m + 1 <= self.L
x = t.zeros((1, self.L, 480), device=self.dev)
s = t.as_tensor(np.asarray(src), device=self.dev)
if extra is None:
x[0, 1:1 + m] = self.E[s]
else:
X = t.cat([self.E, t.tensor(np.asarray(extra, np.float32), device=self.dev)])
s = t.where(s >= 0, s, self.E.shape[0] + (-s - 1))
x[0, 1:1 + m] = X[s]
with t.no_grad():
return self.cm.bacformer(protein_embeddings=x, special_tokens_mask=self.stm, token_type_ids=self.tt,
attention_mask=None, return_attn_weights=False, return_dict=True,
is_causal=True).last_hidden_state[0]
def decode(self, h, rows):
"""rows (new positions) -> (S, C) float64 numpy arrays (len(rows), F); fixed-size head chunks"""
t = self.torch
S, C = [], []
rows = np.asarray(rows, np.int64)
with t.no_grad():
for s in range(0, len(rows), HEAD_CHUNK):
r = rows[s:s + HEAD_CHUNK]; k = len(r)
rr = np.zeros(HEAD_CHUNK, np.int64); rr[:k] = r
p = t.softmax(self.cm.gm_head(h[t.as_tensor(rr, device=self.dev)]), -1).double()
S.append((p @ self.W.T)[:k].cpu().numpy()); C.append((p @ self.Wc.T)[:k].cpu().numpy())
if not S:
F = self.W.shape[0]; return np.zeros((0, F)), np.zeros((0, F))
return np.concatenate(S), np.concatenate(C)
def set_sparse(self, cols, vals, ptr, nnz):
"""the decoder rows of every family of the genome (Decoder.sparse_table), for decode_true. One fixed width per
genome (its largest row), so that every head chunk has the same shape in every pass."""
t = self.torch
self.sp_cols = t.as_tensor(np.concatenate([cols, [0]]), dtype=t.long, device=self.dev) # + a zero slot
self.sp_vals = t.as_tensor(np.concatenate([vals, [0.0]]), dtype=t.float64, device=self.dev)
self.sp_ptr = t.as_tensor(np.asarray(ptr[:-1], np.int64), device=self.dev)
self.sp_nnz = t.as_tensor(np.asarray(nnz, np.int64), device=self.dev)
self.sp_zero = len(cols)
w = max(1, int(np.max(nnz)) if len(nnz) else 1)
self.sp_ar = t.arange(-(-w // 64) * 64, device=self.dev)[None, :] # a multiple of 64: aligned rows
def decode_true(self, h, rows, fidx):
"""rows (new positions), fidx (index into the set_sparse table: each row's own family) -> (S, C) float64 numpy
arrays (len(rows),): S = P(f | ctx) = sum_c p[r, c] W[f, c] over the full softmax (the same float64 sum as
decode, over the non-zero W only; agrees with decode to ~2e-7 in logit, the fp32 rounding of the softmax),
C = sum_c p[r, c] - S = sum_c p[r, c] (1 - W[f, c]). Fixed-size head chunks (TRUE_CHUNK rows)
and a fixed gather width: a row's value does not depend on the other rows of its chunk nor on its slot in it
(asserted per genome in check (g) by reading every row of the baseline again in other slots). CUDA's row
reductions (softmax, sum) take another summation path when a row does not start on an aligned address, and a
50,001-wide row starts misaligned in half the slots, so the softmax runs on rows padded to 50,048 with -inf
(probability exactly 0) and the gather width is a multiple of 64: every row starts 256-byte aligned."""
t = self.torch
rows = np.asarray(rows, np.int64); fidx = np.asarray(fidx, np.int64)
S, C = [], []
NPAD = -(-NCL // 64) * 64; CH = TRUE_CHUNK
with t.no_grad():
zp = t.full((CH, NPAD), float("-inf"), device=self.dev)
for s in range(0, len(rows), CH):
r = rows[s:s + CH]; k = len(r)
rr = np.zeros(CH, np.int64); rr[:k] = r
ff = np.zeros(CH, np.int64); ff[:k] = fidx[s:s + CH]
zp[:, :NCL] = self.cm.gm_head(h[t.as_tensor(rr, device=self.dev)])
p = t.softmax(zp, -1).double()
ft = t.as_tensor(ff, device=self.dev)
ok = self.sp_ar < self.sp_nnz[ft][:, None]
ix = t.where(ok, self.sp_ptr[ft][:, None] + self.sp_ar, self.sp_zero)
sv = (p.gather(1, self.sp_cols[ix]) * self.sp_vals[ix]).sum(1)
S.append(sv[:k].cpu().numpy()); C.append((p.sum(1) - sv)[:k].cpu().numpy())
if not S: return np.zeros(0), np.zeros(0)
return np.concatenate(S), np.concatenate(C)
def full_calls(self, h, n, chunk=512):
"""top-64 clusters, probabilities and entropy (bits) at rows 0..n-1 (row k = the call for gene k)"""
t = self.torch
TF, TPp, EN = [], [], []
with t.no_grad():
for s in range(0, n, chunk):
r = t.arange(s, min(n, s + chunk), device=self.dev)
p = t.softmax(self.cm.gm_head(h[r]), -1)
tp, tf = p.topk(TOPK, 1)
TF.append(tf.int().cpu().numpy()); TPp.append(tp.cpu().numpy())
EN.append((-(p * t.log2(p.clamp(min=1e-30))).sum(1)).cpu().numpy())
return np.concatenate(TF), np.concatenate(TPp), np.concatenate(EN)
# ============================================================================ readouts
def readout_values(kind, S, C, i1, i2):
"""primary value per readout, from decoded S = P(f|ctx), C = 1 - P(f|ctx) (arrays over readouts)
co: logit P(f1) = log P(f1) - log(1 - P(f1))
av: log P(f1) - log P(f2) (f1 = D's entry family, f2 = the resident gene's family)
each probability floored at FLOOR"""
lp1 = np.log(np.maximum(S[np.arange(len(i1)), i1], FLOOR))
lc1 = np.log(np.maximum(C[np.arange(len(i1)), i1], FLOOR))
lp2 = np.log(np.maximum(S[np.arange(len(i2)), np.maximum(i2, 0)], FLOOR))
v = np.where(kind == 1, lp1 - lc1, lp1 - lp2)
return v, lp1, np.where(kind == 1, lc1, lp2)
# ============================================================================ influence readouts
def logit_true(S, C):
"""logit P(true family) from decode_true, each probability floored at FLOOR (as readout_values' co kind)"""
return np.log(np.maximum(S, FLOOR)) - np.log(np.maximum(C, FLOOR))
def inf_bins(dd, dl, ok, rgp, win=INF_WIN, width=INF_BIN):
"""dd: distance of each read gene from the deleted element's last gene (1 = the next gene), dl: its delta, ok: usable
(the decoder saw its family), rgp: panRGP gene. -> per bin of `width` genes out to `win`: mean delta, mean |delta|,
mean delta over panRGP genes, genes (all, panRGP). NaN where a bin is empty."""
nb = win // width
m = ok & (dd >= 1) & (dd <= win)
b = (dd[m] - 1) // width; x = dl[m]; rg = rgp[m]
n = np.bincount(b, minlength=nb); nr = np.bincount(b[rg], minlength=nb)
with np.errstate(invalid="ignore", divide="ignore"):
mean = np.bincount(b, x, nb) / n
mabs = np.bincount(b, np.abs(x), nb) / n
mr = np.bincount(b[rg], x[rg], nb) / nr
return mean, mabs, mr, n, nr
def influence_readout(R, h, keep, n, dels, pend, base_lg, tfi, seen, rgp, entry, lim=None,
up=INF_UP, win=INF_WIN, ent_max=INF_ENT):
"""The influence readouts of one deletion pass, from its hidden states h.
keep original genes in the pass, in order (new position = index in keep); n genes in the chromosome
dels the deleted genes; pend = the deleted element's last gene: distances count from it, as a unit's
distance counts from U's end
base_lg baseline logit P(true family) per original gene; tfi each gene's row in the sparse decoder table;
seen: the decoder saw its family; rgp: panRGP gene; entry: entry genes of Data.elements(g)
lim genes >= lim are not in the pass (a truncated sequence, influence()); default n
Reads every gene from `up` genes before the first deleted gene to `win` genes after pend, and the entries of
elements out to `ent_max` genes. Check (g), asserted here: every gene before the first deleted gene is
bit-identical to the baseline. Returns the bins (inf_bins), the entries (index into the element table, delta)
and the check counters."""
lim = n if lim is None else lim
orig_pos = np.full(n, -1, np.int64); orig_pos[keep] = np.arange(len(keep))
d0 = int(np.min(dels)); pend = int(pend)
lo = max(0, d0 - up); hi = min(lim - 1, pend + win)
rows = np.arange(lo, hi + 1); rows = rows[orig_pos[rows] >= 0]
ent_ok = (entry > pend) & (entry - pend <= ent_max) & (entry < lim)
ent_ok[ent_ok] &= orig_pos[entry[ent_ok]] >= 0 # an entry deleted by this pass is not read
far = np.unique(entry[ent_ok & (entry > hi)])
rows = np.concatenate([rows, far])
S, C = R.decode_true(h, orig_pos[rows], tfi[rows])
dl = logit_true(S, C) - base_lg[rows]
u = rows < d0
if u.any():
assert np.all(dl[u] == 0.0), f"check (g): a gene before the deletion changed (max |delta| {np.max(np.abs(dl[u]))})"
mean, mabs, mr, nn, nr = inf_bins(rows - pend, dl, seen[rows], rgp[rows], win)
val = np.full(n, np.nan); val[rows] = dl
ei = np.flatnonzero(ent_ok)
return dict(mean=mean, abs=mabs, rgp=mr, n=nn, n_rgp=nr, ent=ei, ent_delta=val[entry[ei]],
n_rows=len(rows), n_up=int(u.sum()), up_max=float(np.max(np.abs(dl[u]))) if u.any() else 0.0)
def load_influence(path, checks=False):
"""The influence readouts of one genome file written by pg_knockout.py --influence (g*.npz) as two tables:
passes one row per deletion pass: pass index, kind (U / ctrl_cpl / ctrl_nat / check_b), deleted region id
(-1 for a natural run), spot id, start, end, size (deleted genes), and the 60-bin profiles as arrays
(mean, abs, rgp, n, n_rgp; bin k = genes 25k+1 .. 25k+25 after `end`)
entries one row per (pass, element entry read): the affected element (kind 0 = panRGP run by spot id,
1 = coupled region by region id, its entry gene), distance = entry - pass end, delta (float16 ->
float64), the baseline logit at that entry and whether the decoder saw its family
The element-to-element matrix is `entries` with the deleted element (pass columns) joined in; the matched null of
one (pass, entry) is the control passes of the same genome with size x0.5-2 of the pass's and distance to the
same entry within +-35 % (match_controls, as for the units).
checks=False leaves out the check_b passes (10 arbitrary genes deleted after the last readout row, possibly
backbone genes: a self-check, not an element); checks=True keeps them."""
import pandas as pd
Z = np.load(path, allow_pickle=False)
assert "inf_e_off" in Z, f"{path}: no influence readouts (run with --influence)"
pk = Z["pass_kind"]; off = Z["inf_e_off"]
dp = np.flatnonzero((pk != "base") & (pk != "check_a") & (checks | (pk != "check_b")))
g = int(os.path.basename(path)[1:].split(".")[0])
P = pd.DataFrame(dict(g=g, pass_i=dp, kind=pk[dp], del_id=Z["pass_id"][dp], del_spot=Z["pass_spot"][dp],
start=Z["pass_start"][dp], end=Z["pass_end"][dp], size=Z["pass_size"][dp]))
for k in ("mean", "abs", "rgp"): P[k] = list(Z[f"inf_{k}"][dp].astype(np.float64))
for k in ("n", "n_rgp"): P[k] = list(Z[f"inf_{k}"][dp].astype(np.int64))
cnt = np.diff(off); pi = np.repeat(np.arange(len(cnt)), cnt); el = Z["inf_e_el"].astype(np.int64)
if not checks:
k_ = pk[pi] != "check_b"; pi = pi[k_]; el = el[k_]; de = Z["inf_e_delta"][k_]
else:
de = Z["inf_e_delta"]
ent = Z["el_entry"][el].astype(np.int64)
Ent = pd.DataFrame(dict(g=g, pass_i=pi, kind=pk[pi], del_id=Z["pass_id"][pi], del_spot=Z["pass_spot"][pi],
del_size=Z["pass_size"][pi], del_end=Z["pass_end"][pi],
el=el, el_kind=Z["el_kind"][el], el_id=Z["el_id"][el], el_entry=ent,
dist=ent - Z["pass_end"][pi], delta=de.astype(np.float64),
base_logit=Z["inf_base"][ent].astype(np.float64), seen=Z["inf_seen"][ent]))
return P, Ent
def influence_matrix(path, kinds=("U",), win=INF_WIN, floor=1e-5):
"""Reference aggregation of one genome file: the element-to-element influence matrix of the passes of `kinds`
against the whole-element null already used for the units. One row per (deleted element, affected element entry
within `win` genes): delta, and the matched null = the control passes (ctrl_cpl / ctrl_nat) of the same genome that
read the SAME entry, do not overlap the deleted element, and match it in size (x0.5-2) and distance to the entry
(+-35 %), widened to x1/3-3 / +-60 % when fewer than 3 match (match_controls); robust z floored at `floor` (the
fp32 floor, ~1e-5). A control pass put in `kinds` is scored against the OTHER controls (a pseudo-U)."""
import pandas as pd
P, E = load_influence(path)
E = E[E.seen.values]
ctl = E[E.kind.str.startswith("ctrl_")]
pst = dict(zip(P.pass_i, P.start)); pen = dict(zip(P.pass_i, P.end)); psz = dict(zip(P.pass_i, P["size"]))
by_el = {e: (d.pass_i.values, d.del_end.values, d.del_size.values, d.delta.values,
np.array([pst[i] for i in d.pass_i.values])) for e, d in ctl.groupby("el")}
rows = []
for r in E[E.kind.isin(kinds) & (E.dist >= 1) & (E.dist <= win)].itertuples(index=False):
n_c, z = 0, np.nan; med = mad = np.nan
if r.el in by_el:
ci, ce, cs, cd, cst = by_el[r.el]
ok = (ci != r.pass_i) & ~((cst <= pen[r.pass_i]) & (ce >= pst[r.pass_i]))
idx, _w = match_controls(ce, cs, ok, psz[r.pass_i], r.dist, r.el_entry)
n_c = len(idx)
if n_c >= 3: z, med, mad = robust_z(r.delta, cd[idx], floor)
rows.append(dict(g=r.g, pass_i=r.pass_i, kind=r.kind, del_id=r.del_id, del_spot=r.del_spot, del_size=r.del_size,
el_kind=r.el_kind, el_id=r.el_id, el_entry=r.el_entry, dist=r.dist, delta=r.delta,
n_ctrl=n_c, ctrl_med=med, ctrl_mad=mad, z=z))
return pd.DataFrame(rows)
# ---------------------------------------------------------------------------- on demand, for serve_live
_CTX = {}
def context(dev=None, threads=2, frac=MEM_FRAC, sets=None, data=None, compact=None):
"""(Data, Decoder, Runner) for influence(), built once per process. dev None: cuda:0 when available, else the CPU
with `threads` threads. Loads an fp32 copy of the model (the assay's precision; serve_live's copy is bf16). `frac`
caps the GPU memory of the WHOLE calling process (torch.cuda.set_per_process_memory_fraction): serve_influence.py
sets it (a process of its own), serve_live.py passes None (its process is not capped here). `data`: a Data already
loaded (serve_live's /elements), else Data(sets=sets). `compact`: Decoder's (None: the compact decoder only when
base_top64 is absent)."""
if "ctx" not in _CTX:
import torch
dev = dev or ("cuda:0" if torch.cuda.is_available() else "cpu")
D = data if data is not None else Data(sets=sets); dec = Decoder(D, compact=compact)
R = Runner(dev=dev, frac=frac, threads=threads)
_CTX["ctx"] = (D, dec, R); _CTX["base"] = collections.OrderedDict()
return _CTX["ctx"]
def resolve_element(D, g, element):
"""element: a region id (int) or ('region', ri): the coupled element of region ri in genome g (must be present at
its spot); ('spot', s): the panRGP run at spot s (the one with most non-backbone genes if the spot holds several);
('genes', [positions]): those genes. The deleted genes never include a backbone family."""
if isinstance(element, (int, np.integer)): element = ("region", int(element))
kind, x = element[0], element[1]
a = int(D.OFF[g]); fam = D.FAM[a:int(D.OFF[g + 1])]
if kind == "region":
ri = int(x)
assert D.HAS[ri, g] and D.PRES[ri, g], f"region {ri} is not present at its spot in genome {g}"
e = D.EL[(ri, g)]; assert e["end"] >= e["start"], f"region {ri} wraps the origin in genome {g}"
dels = D.members(ri, g)
return dict(kind=1, id=ri, start=int(e["start"]), end=int(e["end"]), dels=dels, size=len(dels),
label=D.REG[ri]["label"], spot=D.REG[ri]["spot"] if D.REG[ri]["spot"] is not None else -1)
if kind == "spot":
E = D.elements(g); m = np.flatnonzero((E["kind"] == 0) & (E["id"] == int(x)) & (E["size"] > 0))
assert len(m), f"no panRGP run at spot {x} in genome {g}"
k = m[np.argmax(E["size"][m])]; s, t = int(E["start"][k]), int(E["end"][k])
dels = [p for p in range(s, t + 1) if not D.BB[fam[p]]]
return dict(kind=0, id=int(x), start=s, end=t, dels=dels, size=len(dels), label=f"spot {int(x)}", spot=int(x))
if kind == "genes":
dels = sorted(int(p) for p in x if not D.BB[fam[int(p)]])
assert dels, "('genes', ...): every gene given is of a backbone family (never deleted)"
return dict(kind=-1, id=-1, start=dels[0], end=dels[-1], dels=dels, size=len(dels), label="genes", spot=-1)
raise ValueError(f"element {element!r}")
def influence_controls(D, g, el, n_ctrl=8):
"""n_ctrl whole accessory elements of genome g (control_pool: the null of the knockouts) not overlapping the
element, size x0.5-2 of its (widened to x1/3-3, then any size, when fewer match), nearest to it first, so that
their distances to the element's neighbours are close to the element's own"""
pool = D.control_pool(g, sizes=[el["size"]])
ok = [c for c in pool if not (c["start"] <= el["end"] and c["end"] >= el["start"])]
for lo_, hi_ in ((MATCH[0], MATCH[1]), (WIDEN[0], WIDEN[1]), (0.0, np.inf)):
cand = [c for c in ok if lo_ * el["size"] <= c["size"] <= hi_ * el["size"]]
if len(cand) >= n_ctrl: break
cand.sort(key=lambda c: (abs(c["end"] - el["end"]), c["start"]))
return cand[:n_ctrl], (lo_, hi_)
def _nan2none(x):
if isinstance(x, dict): return {k: _nan2none(v) for k, v in x.items()}
if isinstance(x, (list, tuple)): return [_nan2none(v) for v in x]
if isinstance(x, np.ndarray): return _nan2none(x.tolist())
if isinstance(x, (float, np.floating)): return None if not np.isfinite(x) else float(x)
if isinstance(x, np.integer): return int(x)
return x
def _spec_scores(r, entry, end, win=INF_WIN, floor=1e-5):
"""Specificity of one pass at each element entry it read (1..win genes after the deleted element's end `end`):
the entry's delta against the SAME pass's deltas at the other entries at a comparable distance (+-35 %, widened to
+-60 % when < 3; entries sharing the gene excluded) -- the off-target logic of the verdicts. A deletion that moves
everything downstream moves every entry alike and scores ~0 everywhere. Returns {element index: (dist, excess,
z, n_ref)} (excess = delta - median of the others; z = excess / (1.4826 MAD), floored)."""
ei = np.asarray(r["ent"], np.int64); dv = np.asarray(r["ent_delta"], float)
gene = entry[ei]; d = gene - end
m = np.isfinite(dv) & (d >= 1) & (d <= win)
ug, first = np.unique(gene[m], return_index=True) # one value per entry gene
gi = ei[m][first]; gd = d[m][first]; gv = dv[m][first]
out = {}
for k in range(len(ug)):
oth = np.array([], float)
for tol in (OFF_TOL, WIDEN[2]):
sel = (np.arange(len(ug)) != k) & (np.abs(gd - gd[k]) <= tol * gd[k])
oth = gv[sel]
if len(oth) >= 3: break
if len(oth) < 3: continue
med = float(np.median(oth)); mad = 1.4826 * float(np.median(np.abs(oth - med)))
out_v = (int(gd[k]), float(gv[k] - med), float((gv[k] - med) / max(mad, floor)), int(len(oth)))
for i in ei[m][gene[m] == ug[k]]: out[int(i)] = out_v
return out
def influence(g, element, n_ctrl=8, truncate=None, ctx=None, z_floor=1e-5, progress=None):
"""The radius of influence of one accessory element of genome g, as the model sees it: baseline, deletion of the
element (its member genes, never a backbone family), and n_ctrl size-matched whole-element control deletions of
the same genome (influence_controls), one fp32 pass each at one padded length, read as in --all
(influence_readout; check (g) asserted in every pass).
truncate run the passes only as far as the last gene read (the element's or a control's end + 2,400 genes);
exact for a causal model up to the fp32 floor (~1e-5, pilot check f). Default: on the CPU only.
Returns a JSON-ready dict:
element, controls what was deleted
bins [lo, hi] genes after the element's last gene, 25 per bin, out to 1,500
profile the element's mean delta, mean |delta|, mean delta over panRGP genes, genes per bin
ctrl_mean, ctrl_abs the same per control (n_ctrl x bins)
z_abs, z_mean robust z of the element's bin against its controls' same bin (median, 1.4826 MAD)
radius upper edge (genes) of the last bin with z_abs >= 3; radius_contig: of the first run of
such bins from the element on; n_bins_z3: how many of the 60 bins (with 8 controls and a
robust z, a few bins reach 3 by chance: read the profile against ctrl_abs, not the radius)
ctrl_abs_median the controls' median |delta| profile (the reference curve: every deletion moves the
model's expectations downstream, decaying with distance)
ctrl_abs_q per bin, min / median / max of the controls' mean |delta|: the band to draw the profile in
n_bins how many of the 60 bins are read for the element (fewer when the chromosome ends first)
null the same summaries with each control in turn as the element, against the other controls,
over the element's n_bins bins only (n_bins_z3, radius, radius_contig, broad): what 8
controls give by chance, like for like. radius_p /
n_bins_z3_p = share of these pseudo-elements with at least the element's value (+1 smoothed)
broad median over bins of the element's mean |delta| / the controls' median: how much more (or
less) than a typical deletion this one moves the model's expectations downstream
neighbours every element (Data.elements) whose entry lies 1-1,500 genes downstream: its delta
(logit P(its entry family)), and
spec_excess, z_spec SPECIFICITY: the neighbour's delta minus the median delta of the
element's own deletion at the other entries at a comparable distance
(+-35 %, widened to +-60 %; _spec_scores): does the deletion move THIS
neighbour more than it moves everything else at that distance?
p_spec, q_spec calibrated: z_spec against the same score computed for every entry of
every control deletion at a comparable distance (the controls scored
as if deleted elements); two-sided, +1 smoothed; q = BH over the
neighbours' distinct entry genes (an element and the panRGP run of
its spot often share one: one test, listed twice)
z_same, z_dist NOT specificity: the neighbour's delta against the controls' deltas at
the same entry (z_same) or at any entry at that distance (z_dist). An
element whose deletion moves everything downstream more than a typical
deletion flags every neighbour with z_same (yidK/yidJ, genome 1: 35 of
39 neighbours |z_same| > 3, the linked one included, while z_dist = -0.3)
and whether it is linked to the element in the sets file. Floors: 1e-5 (the fp32 floor).
Close to the element (< 100 genes) entries are few, so many scores are missing there.
seconds per pass and in total
progress: called as progress(done, total, what) before each pass (what: the pass about to run) and at the end,
for a server's job status (serve_live.py); total = the passes of this call (baseline included when not cached)."""
D, dec, R = ctx or context()
tick = progress or (lambda *a: None)
t0 = time.time()
a, b = int(D.OFF[g]), int(D.OFF[g + 1]); n = b - a
FAMg = D.FAM[a:b]; rgp = D.RGP[a:b]
E = D.elements(g); entry = E["entry"].astype(np.int64)
el = resolve_element(D, g, element)
ctrls, widen = influence_controls(D, g, el, n_ctrl)
dev_cuda = str(R.dev).startswith("cuda")
if truncate is None: truncate = not dev_cuda
lim = n if not truncate else min(n, max([el["end"]] + [c["end"] for c in ctrls]) + INF_ENT + 1)
hh = dec.half_for(g)
key = (g, lim, hh)
if key in _CTX.get("base", {}):
fams, tfi, seen, table, base_lg, t_base = _CTX["base"][key]
else:
fams = np.unique(FAMg); tfi = np.searchsorted(fams, FAMg)
table = dec.sparse_table(fams, hh); seen = table[3][tfi] > 0
base_lg = None
total = 1 + len(ctrls) + (base_lg is None)
R.set_sparse(*table)
R.set_genome(D.EMB[D.PID[a:a + lim]], lim + 1)
if base_lg is None:
tick(0, total, "baseline pass (the genome as it is)")
R.sync(); t = time.time()
h = R.hidden(np.arange(lim)); S, C = R.decode_true(h, np.arange(lim), tfi[:lim])
base_lg = np.full(n, np.nan); base_lg[:lim] = logit_true(S, C)
R.sync(); t_base = time.time() - t
_CTX.setdefault("base", collections.OrderedDict())[key] = (fams, tfi, seen, table, base_lg, t_base)
while len(_CTX["base"]) > 4: _CTX["base"].popitem(last=False)
res = []; secs = []
for k_, p in enumerate([el] + ctrls):
tick(total - 1 - len(ctrls) + k_, total, "deletion of the element" if k_ == 0 else f"control deletion {k_} of {len(ctrls)}")
R.sync(); t = time.time()
keep = np.setdiff1d(np.arange(lim), np.asarray(p["dels"], np.int64))
h = R.hidden(keep)
r = influence_readout(R, h, keep, n, p["dels"], p["end"], base_lg, tfi, seen, rgp, entry, lim=lim)
R.sync(); secs.append(time.time() - t); res.append(r)
nb = INF_WIN // INF_BIN
r0 = res[0]; cr = res[1:]
fin = np.isfinite(np.asarray(r0["abs"], float)) # the bins read for the element (fewer near a chromosome's end)
cm = np.array([r["mean"] for r in cr]).reshape(-1, nb); ca = np.array([r["abs"] for r in cr]).reshape(-1, nb)
z_abs = np.full(nb, np.nan); z_mean = np.full(nb, np.nan)
for k in range(nb):
x = ca[:, k][np.isfinite(ca[:, k])]
if np.isfinite(r0["abs"][k]) and len(x) >= 3: z_abs[k] = robust_z(r0["abs"][k], x, z_floor)[0]
x = cm[:, k][np.isfinite(cm[:, k])]
if np.isfinite(r0["mean"][k]) and len(x) >= 3: z_mean[k] = robust_z(r0["mean"][k], x, z_floor)[0]
def radii(z):
sig_ = np.flatnonzero(z >= 3); k_ = 0
while k_ < nb and np.isfinite(z[k_]) and z[k_] >= 3: k_ += 1
return (int((sig_[-1] + 1) * INF_BIN) if len(sig_) else 0), k_ * INF_BIN, int(len(sig_))
radius, radius_contig, n_z3 = radii(z_abs)
sig = np.arange(n_z3)
# the same with each control as the element, against the others, over the element's own bins (fin): what the
# summaries give by chance, compared like for like when the chromosome ends before the element's 60th bin
cmed = np.nanmedian(ca, 0) if len(ca) else np.full(nb, np.nan)
with np.errstate(invalid="ignore", divide="ignore"):
broad = float(np.nanmedian(np.asarray(r0["abs"], float) / cmed))
null = dict(radius=[], radius_contig=[], n_bins_z3=[], broad=[])
for c in range(len(ca)):
oth = np.delete(np.arange(len(ca)), c); zc = np.full(nb, np.nan)
for k_ in range(nb):
x = ca[oth, k_][np.isfinite(ca[oth, k_])]
if fin[k_] and np.isfinite(ca[c, k_]) and len(x) >= 3: zc[k_] = robust_z(ca[c, k_], x, z_floor)[0]
rr = radii(zc)
null["radius"].append(rr[0]); null["radius_contig"].append(rr[1]); null["n_bins_z3"].append(rr[2])
with np.errstate(invalid="ignore", divide="ignore"):
rc = (ca[c] / np.nanmedian(ca[oth], 0))[fin]
null["broad"].append(float(np.nanmedian(rc)) if np.isfinite(rc).any() else np.nan)
pshare = lambda v, xs: float((1 + sum(x >= v for x in xs)) / (1 + len(xs))) if xs else np.nan
# specificity at each neighbour, and its null: every control deletion scored the same way
spec0 = _spec_scores(r0, entry, el["end"])
spec_null = [v for c, r in zip(ctrls, cr) for v in _spec_scores(r, entry, c["end"]).values()]
sn_d = np.array([v[0] for v in spec_null], float); sn_z = np.array([v[2] for v in spec_null], float)
# the neighbours: the element's entries and each control's, by element index
ctl_ent = [dict(zip(r["ent"].tolist(), r["ent_delta"].tolist())) for r in cr]
linked = D.links_of(el["id"]) if el["kind"] == 1 else set()
nbrs = []
for i, dv in zip(r0["ent"].tolist(), r0["ent_delta"].tolist()):
d = int(entry[i] - el["end"])
if d > INF_WIN: continue
same = []
for tol in (OFF_TOL, WIDEN[2]):
same = [ce[i] for c, ce in zip(ctrls, ctl_ent)
if i in ce and np.isfinite(ce[i]) and abs((entry[i] - c["end"]) - d) <= tol * d]
if len(same) >= 3: break
pooled = [v for c, ce in zip(ctrls, ctl_ent) for j, v in ce.items()
if np.isfinite(v) and abs((entry[j] - c["end"]) - d) <= OFF_TOL * d]
zs = robust_z(dv, same, z_floor)[0] if len(same) >= 3 else np.nan
zd = robust_z(dv, pooled, z_floor)[0] if len(pooled) >= 3 else np.nan
sp_ = spec0.get(int(i)); ex_ = zsp = psp = np.nan; nref = nnull = 0
if sp_ is not None:
ex_, zsp, nref = sp_[1], sp_[2], sp_[3]
for tol in (OFF_TOL, WIDEN[2], np.inf):
mm = np.abs(sn_d - d) <= tol * d
if mm.sum() >= 20 or tol == np.inf: break
nnull = int(mm.sum())
if nnull: psp = float((1 + np.sum(np.abs(sn_z[mm]) >= abs(zsp))) / (1 + nnull))
kd, idd = int(E["kind"][i]), int(E["id"][i])
nbrs.append(dict(kind=("region" if kd == 1 else "spot"), id=idd, label=D.element_label(kd, idd)[:60],
entry=int(entry[i]), start=int(E["start"][i]), end=int(E["end"][i]), dist=d,
family=D.FN[int(E["fam"][i])], seen=bool(seen[int(entry[i])]),
base_p=float(1.0 / (1.0 + np.exp(-base_lg[int(entry[i])]))), delta=dv,
spec_excess=ex_, z_spec=zsp, n_spec_ref=nref, p_spec=psp, n_spec_null=nnull,
z_same=zs, n_same=len(same), z_dist=zd, n_dist=len(pooled),
linked=bool(kd == 1 and idd in linked)))
# BH over the distinct entry genes: an element and its spot's panRGP run often share one (the same readout, the same
# p), and counting it twice would inflate the number of tests
pe = {}
for x in nbrs: pe.setdefault(x["entry"], x["p_spec"])
qe = dict(zip(pe, bh(list(pe.values()))))
for x in nbrs: x["q_spec"] = qe[x["entry"]]
tick(total, total, "done")
out = dict(g=int(g), acc=D.ACC[g], n_genes=n, decoder_half=hh, device=str(R.dev), truncated_at=int(lim),
element={k_: el[k_] for k_ in ("kind", "id", "label", "start", "end", "size", "spot")},
controls=[dict(arm=c["arm"], ri=int(c["ri"]), spot=int(c["spot"]), start=int(c["start"]), end=int(c["end"]),
size=int(c["size"])) for c in ctrls], control_size_window=list(widen),
bins=[[k_ * INF_BIN + 1, (k_ + 1) * INF_BIN] for k_ in range(nb)],
profile=dict(mean=r0["mean"], abs=r0["abs"], rgp=r0["rgp"], n=r0["n"], n_rgp=r0["n_rgp"]),
ctrl_mean=cm, ctrl_abs=ca, z_abs=z_abs, z_mean=z_mean, radius=radius, radius_contig=radius_contig,
n_bins_z3=int(len(sig)), n_bins=int(fin.sum()), ctrl_abs_median=cmed,
ctrl_abs_q=(np.stack([np.nanmin(ca, 0), cmed, np.nanmax(ca, 0)]) if len(ca) else np.full((3, nb), np.nan)),
broad=broad, null=null, radius_p=pshare(radius, null["radius"]), n_bins_z3_p=pshare(int(len(sig)), null["n_bins_z3"]),
broad_p=pshare(broad, null["broad"]),
spec_null=dict(n=len(spec_null), abs_z_q=[float(x) for x in np.nanpercentile(np.abs(sn_z), [50, 90, 99])]
if len(sn_z) else None),
neighbours=nbrs,
checks=dict(upstream_rows_exact=int(sum(r["n_up"] for r in res)), rows_read=int(sum(r["n_rows"] for r in res))),
seconds=dict(baseline=t_base, passes=secs, total=time.time() - t0))
return _nan2none(out)
# ============================================================================ statistics
def bh(p):
p = np.asarray(p, float); q = np.full(len(p), np.nan); ok = np.isfinite(p)
if ok.sum():
pp = p[ok]; o = np.argsort(pp); r = pp[o] * ok.sum() / (np.arange(ok.sum()) + 1)
r = np.minimum.accumulate(r[::-1])[::-1]; qq = np.empty_like(pp); qq[o] = np.minimum(r, 1); q[ok] = qq
return q
def rank_share_null(n_ctrl, thr=0.95):
"""expected share of units with rank percentile >= thr among n controls, under exchangeability"""
n = np.asarray(n_ctrl, float)
return float(np.mean([(k - np.ceil(thr * k) + 1) / (k + 1) for k in n])) if len(n) else np.nan
def match_controls(pool_end, pool_size, ok, size, dist, row, widen=True):
"""indices of the controls matched to a deletion of `size` genes lying `dist` genes before `row`.
`ok` already excludes overlaps with U and D and controls linked to either. Returns (idx, widened).
With MATCH_LOG = F, a third tier when the widened one still has < 3: distance within x1/F..xF of the unit's."""
d = row - pool_end
for (lo, hi, tol), w in ((MATCH, False), (WIDEN, True)):
m = ok & (pool_end < row) & (pool_size >= lo * size) & (pool_size <= hi * size) & (np.abs(d - dist) <= tol * dist)
idx = np.flatnonzero(m)
if len(idx) >= 3 or not widen: return idx, w
if MATCH_LOG is not None:
with np.errstate(divide="ignore", invalid="ignore"):
m = ok & (pool_end < row) & (pool_size >= WIDEN[0] * size) & (pool_size <= WIDEN[1] * size) & \
(np.abs(np.log(np.maximum(d, 1) / dist)) <= np.log(MATCH_LOG))
idx = np.flatnonzero(m)
return idx, True
def robust_z(delta, ctrl, floor):
"""(delta - median) / (1.4826 x MAD) of the control deltas, floored. The pilot used mean/sd of 3-5 skewed
control deltas, which is what made the tails miscalibrated (inference review finding 2)."""
if len(ctrl) < 3: return np.nan, np.nan, np.nan
c = np.asarray(ctrl, float); med = float(np.median(c))
mad = 1.4826 * float(np.median(np.abs(c - med)))
return (delta - med) / max(mad, floor), med, mad
def rank_pct(delta, ctrl, e_sign):
"""percentile of delta among its controls in the expected direction (0.5 = a tie)"""
c = np.asarray(ctrl, float) * e_sign; x = delta * e_sign
return float((np.sum(c < x) + 0.5 * np.sum(c == x)) / len(c)) if len(c) else np.nan
def mde_wilcoxon(x, alpha=0.05, power=0.8):
"""smallest location shift (same units as x) a one-sided signed-rank test on the values x (one per lineage: the
lineage means) detects with the given power: normal approximation with the Wilcoxon ARE 3/pi, on a ROBUST scale,
1.4826 x MAD of x. A signed-rank test works on ranks, so a few extreme lineage means that inflate an SD do not blunt
it; an SD-based MDE overstated what the test cannot see. `alpha` is the per-link alpha: pg_knockout.py passes
0.05 / (number of TESTED links), Bonferroni. NaN with < 2 values or a zero MAD."""
from scipy.stats import norm
x = np.asarray(x, float); x = x[np.isfinite(x)]; n = len(x)
if n < 2: return np.nan
s = 1.4826 * float(np.median(np.abs(x - np.median(x))))
if not s > 0: return np.nan
return float((norm.ppf(1 - alpha) + norm.ppf(power)) * s / np.sqrt(n) * np.sqrt(np.pi / 3))
"""In-silico knockout of the "Coupled regions" links with Bacformer (causal): the library.
pg_region_sets.py found 533 links between 290 accessory elements from presence/absence within lineages,
without the model. Here each link is put to the model: in every complete chromosome where the link can be
tested, the upstream element (dnaA reading order, decided per genome) is deleted and the model's belief in
the downstream element is read, in ONE causal fp32 pass over the whole chromosome (<= 5,584 proteins).
units (link, genome) where the link is testable; carriers are defined AT the element's spot (the
validated construction of the scouts' ko_units3.py, ported here unchanged: panRGP spot block,
consensus backbone flanks for loose regions, per-genome orientation, entry gene, insertion
point of an absent element)
co-occurrence genome carries both expected: deleting U LOWERS logit P(D entry)
avoidance carries U, not D (cleanly: expected: deleting U RAISES
no D family anywhere, none log P(D entry family) - log P(resident)
at D's spot) at D's insertion point
skipped (reason recorded): U or D overlapping, < 100 genes apart (MIN_APART; > MAX_APART when set),
wrapping the origin, resident = D's entry family, no entry family, no position.
deletion U's member genes (families of the region, never a backbone family) inside its element span
readout decoded family probability P(f | ctx) = sum_c P(c | ctx) P(f | c) over the FULL softmax at the
readout row, the decoder P(f | c) rebuilt from fp32 full-chromosome baseline passes (top 64 per
gene), cross-validated by lineage: a genome is always read with the decoder of the other half
controls per genome a pool of <= POOL_MAX (72) WHOLE accessory elements, each deleted by its member genes
(control_pool): coupled elements present at their spot (per unit, those linked to U or D excluded)
and, a small second arm, whole panRGP runs at spots holding no coupled region; kept to match the U
sizes and spread along the chromosome. Matched null per unit: controls of size x0.5-2 of U's at a
distance to the readout within +-35 % (WIDEN: x1/3-3, +-60 % when fewer than 3 match; MATCH_LOG a
third tier for close range)
numerics every pass of a genome has the same padded length and the head is applied on fixed-size row
chunks, so identical prefixes give bit-identical readouts: deleting nothing gives delta == 0
exactly and deleting genes after a readout row leaves it exactly unchanged (asserted in every
pass, not only in the dedicated check passes)
link sets Data(sets=PATH) reads any file in the format of pgb/region_sets.json (default); the units cache
is keyed by the md5 of the whole file (for the default file, < 1 MiB, the key is unchanged)
close range (pgb/region_sets_close.json: pairs 50-500 genes apart in the genomes carrying both, two separate
insertions, replicated in lineage-disjoint halves) MAX_APART bounds a unit's U-D distance (600 for these
sets; a unit is one genome's pair in dnaA reading order and can be thousands of genes apart, across the
origin); near_controls() adds whole-element controls at about a unit's own distance upstream of D, and
MATCH_LOG = F a third matching tier (distance within x1/F..xF): at 50-100 genes a genome rarely holds 3
whole elements in the +-35 % / +-60 % windows. Defaults leave the long-range behaviour unchanged.
influence what the model itself knows, independent of any link: after a deletion, the decoded logit
P(true family) (same lineage-CV decoder, fp32 pass, float64 decoding, read through a sparse
copy of the decoder rows: Decoder.srow, Runner.decode_true) at every gene downstream of the
deleted element. Data.elements(g) lists the accessory elements of a genome (panRGP runs at
every spot, by spot id; coupled regions present at their spot, by region id) and their entry
genes. influence_readout() turns one pass into 25-gene bins of the mean delta (and mean
|delta|, and the mean over panRGP genes) out to 1,500 genes, plus the delta at every element
entry out to 2,400 genes (1,500 x 1.6, so that the matched null of an entry near 1,500 genes,
distance +-35 % widened to +-60 %, is on file too). The 25 genes before the first deleted gene
are read as well and must be exactly unchanged (causality, check (g)).
load_influence(g*.npz) reads them back as tables; influence_matrix(g*.npz) is a reference
aggregation of the element-to-element matrix against the units' matched whole-element null (same
entry, size x0.5-2, distance +-35 %). influence(g, element) does it all on demand for one genome
(baseline + deletion + 8 size-matched whole-element controls), served by serve_live.py (/influence,
/elements: the page's own server, locally and in the Space) and by serve_influence.py (a process of its
own: its own port and GPU cap).
compact decoder influence() reads only the sparse decoder rows (Decoder.srow) of the families of the genome, with
the decoder of the other lineage half. export_decoder() writes exactly those rows, both halves, to
pgb/knock/decoder_halves.npz (cluster ids uint16, P(f | c) float64: the same rows bit for bit, 1/12 of
base_top64_{f,p}.npy), and Decoder uses that file when base_top64_{f,p}.npy are absent (the Space).
python3 -c "import pg_knock as K; K.export_decoder()"
Reads the project's data read-only; writes only under pgb/knock/ (through pg_knockout.py, the units cache and
export_decoder()).
"""
import os, json, time, zlib, pickle, hashlib, collections
import numpy as np
PROJ = os.path.dirname(os.path.abspath(__file__)) # the project (also the Space's app directory)
KDIR = f"{PROJ}/pgb/knock"
MODEL = "macwiatrak/bacformer-causal-complete-genomes"
NCL = 50001 # softmax width of gm_head (50,000 clusters + END)
TOPK = 64
CLS, PROT = 2, 4
SLACK, NEAR, GAP, CARRY = 5, 200, 30, 0.5 # as ko_units3.py
MIN_APART = 100 # genes between U's last gene and D's readout gene
MAX_APART = None # at most this many (None: no bound, as before); 600 for the close-range sets
FLOOR = 1e-12 # probability floor of every log
HEAD_CHUNK = 64 # rows per head call (fixed shape -> bit-identical rows across passes)
TRUE_CHUNK = HEAD_CHUNK # rows per head call of decode_true (256 was no faster and moved the logits by 2.5e-5)
PANEL_MAX = 40
ANCHOR_WIN, ANCHOR_SHARE = 40, 0.6 # where an absent element's entry would sit: carriers' flanking families
POOL_MAX = 72 # control deletions (whole accessory elements) per genome
MATCH = (0.5, 2.0, 0.35) # matched null: size x0.5-2 of U's, distance to the readout within +-35 %
WIDEN = (1 / 3.0, 3.0, 0.6) # adaptive widening when fewer than 3 controls match
MATCH_LOG = None # a third tier when still < 3: distance within x1/F..xF (log scale), size x1/3-3 (None: off,
# as before). At 50-100 genes a genome rarely holds 3 whole elements within +-60 % of the
# unit's distance (10 % of the close units); x2 reaches 21 %, x3 31 % (pg_knockout --match_log)
OFF_TOL = 0.35 # off-target readouts: within +-35 % of D's distance from U
SETS_DEFAULT = f"{PROJ}/pgb/region_sets.json"
DECODER_NPZ = f"{KDIR}/decoder_halves.npz" # the compact decoder (export_decoder), used when base_top64 is absent
INF_WIN, INF_BIN = 1500, 25 # influence profile: genes downstream of the deleted element, bin width
INF_ENT = int(INF_WIN * (1 + WIDEN[2])) # element entries read out to 2,400 genes: the widened matched null
INF_UP = 25 # genes read before the first deleted gene: must be exactly unchanged (check g)
MEM_FRAC = 0.35 # per-process GPU memory cap (two runs share the 8 GB card)
def log(*a):
print(time.strftime("%H:%M:%S"), *a, flush=True)