allelign 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
allelign/align.py ADDED
@@ -0,0 +1,1252 @@
1
+ """Align summary statistics or weights to a reference variant table.
2
+
3
+ The aligned ``beta`` counts the reference's ``a1`` allele, so it can be applied
4
+ directly to ``a1`` dosages. This is where a silent allele mix-up would flip
5
+ effect signs, so the rules are explicit:
6
+
7
+ 1. *Locate.* A row whose identifier the reference holds (``id`` or
8
+ ``alternate_id``) is matched by identifier, coordinate-checked: a unique
9
+ identifier at a disagreeing chrom/pos is unmatched (a silent build mix-up),
10
+ an identifier held at several loci is resolved by the row's coordinates or
11
+ rejected, and neither falls through to a different variant at the row's
12
+ coordinates. A row with no usable identifier matches by ``(chrom, pos)``.
13
+ Chromosome labels are normalised (:func:`~allelign.alleles.normalize_chrom`)
14
+ and ``pos <= 0`` means missing.
15
+ 2. *Ambiguity.* With ``drop_ambiguous`` (default) a strand-ambiguous pair
16
+ (A/T, C/G, and multi-base pairs such as ``AA``/``TT``) is dropped.
17
+ 3. *Orient.* Alleles are compared upper-cased: same order, swapped,
18
+ complemented, complemented and swapped; a swap negates ``beta``.
19
+ 4. *Claim.* Each reference variant is taken by at most one row, the first in
20
+ file order to reach it with compatible alleles and finite ``beta``/``se``.
21
+ Several rows at one locus (multi-allelic sites) are tried in order.
22
+
23
+ Output rows are in reference order; ``src_index`` points back to the input.
24
+ Two implementations of one rule set -- a vectorised join for plain string and
25
+ numeric NumPy columns, and the reference row loop for anything else -- give
26
+ identical results. Ported from ldpred3 (its harmonisation was checked against an
27
+ independent naive implementation on 5,500 cases).
28
+ """
29
+
30
+ from __future__ import annotations
31
+
32
+ import math
33
+ import threading
34
+ import warnings
35
+ from dataclasses import dataclass, field, replace
36
+ from itertools import repeat
37
+
38
+ import numpy as np
39
+
40
+ from . import status as _st
41
+ from .alleles import (allele_orientation as _allele_orientation,
42
+ is_ambiguous as _ambiguous_pair,
43
+ normalize_chrom as _norm_chrom,
44
+ reverse_complement as _complement)
45
+
46
+ __all__ = ["Harmonized", "harmonize", "diagnose_match", "reference_coverage",
47
+ "coverage_message", "reanchor_on_identifier", "HarmonizedWeights",
48
+ "harmonize_weights", "count_positional_only_matches",
49
+ "variant_indices", "REFERENCE_COVERAGE_WARN",
50
+ "SEVERE_REFERENCE_COVERAGE", "BUILD_MISMATCH_FRACTION"]
51
+
52
+ # Per-table matching indices, reused across calls on a frozen table: rebuilding
53
+ # the O(variants) dicts per file dominated runs that align many files to one
54
+ # reference. The cache lives on the table instance; only read-only columns are
55
+ # cached, and the key holds the column objects so replacement invalidates it.
56
+ _CACHE_ATTR = "_harmonize_index_cache"
57
+ _INDEX_CACHE_LOCK = threading.Lock()
58
+
59
+
60
+ def _readonly_column(column):
61
+ """A read-only view over a writable owner is still mutable metadata."""
62
+ if column is None:
63
+ return True
64
+ if not isinstance(column, np.ndarray):
65
+ return False
66
+ while isinstance(column, np.ndarray):
67
+ if column.flags.writeable:
68
+ return False
69
+ column = column.base
70
+ if column is None:
71
+ return True
72
+ try:
73
+ return memoryview(column).readonly
74
+ except TypeError:
75
+ return False
76
+
77
+
78
+ def _table_cache_key(variants):
79
+ columns = (variants.id, variants.chrom, variants.pos, variants.a1,
80
+ variants.a2, getattr(variants, "alternate_id", None))
81
+ if not all(_readonly_column(column) for column in columns):
82
+ return None
83
+ return (*columns, len(variants))
84
+
85
+
86
+ @dataclass
87
+ class Harmonized:
88
+ """Summary statistics aligned to a reference's ``a1`` alleles.
89
+
90
+ ``var_index`` indexes the reference table; ``beta``/``se``/``n_eff`` are
91
+ parallel to it and count ``a1``. ``src_index`` is the input row each match
92
+ came from, so extra source columns can be pulled positionally.
93
+ """
94
+
95
+ var_index: np.ndarray
96
+ beta: np.ndarray
97
+ se: np.ndarray
98
+ n_eff: np.ndarray
99
+ flipped: np.ndarray
100
+ src_index: np.ndarray = None
101
+ log: dict = field(default_factory=dict)
102
+
103
+ def __len__(self):
104
+ return len(self.var_index)
105
+
106
+
107
+ _TOLIST_KINDS = frozenset("biuUSO")
108
+
109
+
110
+ def _rows(column):
111
+ """``column`` as a Python list when that is value-identical, else as is."""
112
+ if isinstance(column, np.ndarray) and (
113
+ column.dtype.kind in _TOLIST_KINDS
114
+ or column.dtype == np.float64):
115
+ return column.tolist()
116
+ return column
117
+
118
+
119
+ def _build_id_index(variants):
120
+ v_id = variants.id
121
+ alternate_id = getattr(variants, "alternate_id", None)
122
+ by_id = {}
123
+ for i, vid in enumerate(v_id):
124
+ if vid and vid != ".":
125
+ by_id.setdefault(vid, []).append(i)
126
+ if alternate_id is not None:
127
+ aid = alternate_id[i]
128
+ if aid and aid != "." and aid != vid:
129
+ by_id.setdefault(aid, []).append(i)
130
+ return by_id
131
+
132
+
133
+ _STR_TYPES = frozenset((str, np.str_))
134
+
135
+
136
+ def _str_rows(column):
137
+ """``column`` as a list of ``str`` when it is a NumPy array of text only."""
138
+ if not isinstance(column, np.ndarray) or column.ndim != 1:
139
+ return None
140
+ if column.dtype.kind == "U":
141
+ return column.tolist()
142
+ if column.dtype.kind != "O":
143
+ return None
144
+ rows = column.tolist()
145
+ return rows if set(map(type, rows)) <= _STR_TYPES else None
146
+
147
+
148
+ def _query_locus(chrom, query_memo):
149
+ """``(usable, normalised)`` for a sumstats chromosome label, memoised."""
150
+ hit = query_memo.get(chrom)
151
+ if hit is None:
152
+ hit = query_memo[chrom] = (bool(str(chrom).strip()),
153
+ _norm_chrom(chrom))
154
+ return hit
155
+
156
+
157
+ def _coordinate_candidates(candidates, variants, chrom, pos, memo,
158
+ query_memo=None):
159
+ """Disambiguate an identifier, and reject a unique ID at the wrong locus."""
160
+ if query_memo is None:
161
+ query_memo = {}
162
+ loci = []
163
+ for gi in candidates:
164
+ raw = variants.chrom[gi]
165
+ candidate_norm = memo.get(raw)
166
+ if candidate_norm is None:
167
+ candidate_norm = memo[raw] = _norm_chrom(raw)
168
+ loci.append((candidate_norm, int(variants.pos[gi])))
169
+
170
+ unique_loci = len(set(loci)) == 1
171
+ try:
172
+ pos = int(pos)
173
+ except (TypeError, ValueError, OverflowError):
174
+ pos = 0
175
+ usable_query = False
176
+ if pos > 0:
177
+ usable_query, query_norm = _query_locus(chrom, query_memo)
178
+
179
+ if unique_loci:
180
+ if not usable_query:
181
+ return candidates
182
+ candidate_norm, candidate_pos = loci[0]
183
+ if candidate_pos <= 0:
184
+ return candidates
185
+ if candidate_pos == pos and candidate_norm == query_norm:
186
+ return candidates
187
+ return None
188
+
189
+ if not usable_query:
190
+ return None
191
+ hits = []
192
+ for gi, (candidate_norm, candidate_pos) in zip(candidates, loci):
193
+ if candidate_pos == pos and candidate_norm == query_norm:
194
+ hits.append(gi)
195
+ return hits or None
196
+
197
+
198
+ def variant_indices(variants):
199
+ """``[key, by_id, by_pos, upper, chrom_memo, vector]`` for one table.
200
+
201
+ The matching indices :func:`harmonize` builds once per frozen table and
202
+ reuses; ``by_id`` maps an identifier to its row(s). Callers that look up
203
+ many identifiers against a reference can borrow it instead of building
204
+ their own (a 1.4M-variant dict costs about 0.6 GiB).
205
+ """
206
+ return _variant_indices(variants)
207
+
208
+
209
+ def _variant_indices(variants):
210
+ """``[key, by_id, by_pos, upper, chrom_memo, vector]`` for one table."""
211
+ key = _table_cache_key(variants)
212
+ try:
213
+ namespace = variants.__dict__
214
+ except AttributeError: # slotted / C-implemented variant tables
215
+ return [key, _build_id_index(variants), None, {}, {}, {}]
216
+ if key is None:
217
+ with _INDEX_CACHE_LOCK:
218
+ namespace.pop(_CACHE_ATTR, None)
219
+ return [key, _build_id_index(variants), None, {}, {}, {}]
220
+ with _INDEX_CACHE_LOCK:
221
+ entry = namespace.get(_CACHE_ATTR)
222
+ if (entry is not None and entry[0][-1] == key[-1]
223
+ and all(old is new for old, new in zip(entry[0][:-1], key[:-1]))):
224
+ return entry
225
+ namespace.pop(_CACHE_ATTR, None)
226
+ del entry
227
+ entry = [key, _build_id_index(variants), None, {}, {}, {}]
228
+ try:
229
+ with _INDEX_CACHE_LOCK:
230
+ namespace[_CACHE_ATTR] = entry
231
+ except AttributeError: # pragma: no cover - frozen table
232
+ pass
233
+ return entry
234
+
235
+
236
+ def _build_pos_index(variants, memo):
237
+ v_chrom, v_pos = variants.chrom, variants.pos
238
+ by_pos = {}
239
+ for i in range(len(variants)):
240
+ pos = int(v_pos[i])
241
+ if pos <= 0: # the missing-coordinate sentinel, not a locus
242
+ continue
243
+ c = v_chrom[i]
244
+ nc = memo.get(c)
245
+ if nc is None:
246
+ nc = memo[c] = _norm_chrom(c)
247
+ by_pos.setdefault((nc, pos), []).append(i)
248
+ return by_pos
249
+
250
+
251
+ # ---------------------------------------------------------------------------
252
+ # Vectorised harmonisation: chromosome labels and alleles factorised to
253
+ # integer codes, the rsID join as one ``map`` over a dict, the position join
254
+ # as sorted ``chrom_code * span + pos`` keys. Only rows whose outcome depends on
255
+ # another row (multi-allelic sites, duplicate IDs, a reference row reached
256
+ # twice) go through the row loop.
257
+ # ---------------------------------------------------------------------------
258
+
259
+ _POS_LIMIT = 1 << 62
260
+
261
+
262
+ def _int_positions(column):
263
+ """``column`` as int64 when ``int(value)`` of every element is exact."""
264
+ if not isinstance(column, np.ndarray) or column.ndim != 1:
265
+ return None
266
+ kind = column.dtype.kind
267
+ if kind not in "iuf" or (kind == "f" and column.dtype.itemsize > 8):
268
+ return None
269
+ if column.size:
270
+ if kind == "f" and not np.isfinite(column).all():
271
+ return None
272
+ if kind != "i" and not (
273
+ -_POS_LIMIT < column.min() and column.max() < _POS_LIMIT):
274
+ return None
275
+ return column.astype(np.int64)
276
+
277
+
278
+ def _numeric(column):
279
+ return (isinstance(column, np.ndarray) and column.ndim == 1
280
+ and column.dtype.kind in "iuf" and column.dtype.itemsize <= 8)
281
+
282
+
283
+ def _factorize(rows):
284
+ """``(uniques, codes)``: first-appearance distinct values and int codes."""
285
+ index = dict.fromkeys(rows)
286
+ for code, value in enumerate(index):
287
+ index[value] = code
288
+ codes = np.fromiter(map(index.__getitem__, rows), dtype=np.int64,
289
+ count=len(rows))
290
+ return list(index), codes
291
+
292
+
293
+ def _genotype_columns(variants, entry):
294
+ """Per-table integer view of ``variants``, cached in ``entry[5]``."""
295
+ cache = entry[5]
296
+ cols = cache.get("columns")
297
+ if cols is not None:
298
+ return cols or None
299
+ cols = False
300
+ n = len(variants)
301
+ chrom, a1, a2 = (_str_rows(variants.chrom), _str_rows(variants.a1),
302
+ _str_rows(variants.a2))
303
+ pos = _int_positions(variants.pos)
304
+ if (chrom is not None and a1 is not None and a2 is not None
305
+ and pos is not None and len(variants.id) == n
306
+ and len(chrom) == len(a1) == len(a2) == len(pos) == n):
307
+ labels, label_codes = _factorize(chrom)
308
+ norm_code = {}
309
+ per_label = np.array(
310
+ [norm_code.setdefault(_norm_chrom(c), len(norm_code))
311
+ for c in labels], dtype=np.int64)
312
+ vocab = {}
313
+ allele_codes = []
314
+ for column in (a1, a2):
315
+ uniques, codes = _factorize(column)
316
+ per_value = np.array(
317
+ [vocab.setdefault(str(a).upper(), len(vocab))
318
+ for a in uniques], dtype=np.int64)
319
+ allele_codes.append(per_value[codes])
320
+ cols = {"pos": pos, "chrom": per_label[label_codes],
321
+ "norm_code": norm_code, "vocab": vocab,
322
+ "a1": allele_codes[0], "a2": allele_codes[1]}
323
+ cache["columns"] = cols
324
+ return cols or None
325
+
326
+
327
+ def _position_index(cols, cache):
328
+ """Sorted ``(chrom, pos)`` join keys, built on first use and cached."""
329
+ index = cache.get("pos_index")
330
+ if index is not None:
331
+ return index
332
+ pos, chrom = cols["pos"], cols["chrom"]
333
+ rows = np.flatnonzero(pos > 0)
334
+ span = int(pos[rows].max()) + 1 if rows.size else 1
335
+ if max(len(cols["norm_code"]), 1) * span >= _POS_LIMIT:
336
+ index = False
337
+ else:
338
+ keys = chrom[rows] * span + pos[rows]
339
+ order = np.argsort(keys, kind="stable")
340
+ index = (keys[order], rows[order], span)
341
+ cache["pos_index"] = index
342
+ return index
343
+
344
+
345
+ def _vector_columns(sumstats, variants, entry):
346
+ """Columns for the vectorised join, or ``None`` (defer to the row loop)."""
347
+ try:
348
+ n = len(sumstats)
349
+ s_id, s_chrom = _str_rows(sumstats.id), _str_rows(sumstats.chrom)
350
+ s_ea, s_oa = _str_rows(sumstats.ea), _str_rows(sumstats.oa)
351
+ if s_id is None or s_chrom is None or s_ea is None or s_oa is None:
352
+ return None
353
+ s_pos = _int_positions(sumstats.pos)
354
+ beta, se, n_eff = sumstats.beta, sumstats.se, sumstats.n_eff
355
+ if (s_pos is None or not (_numeric(beta) and _numeric(se)
356
+ and _numeric(n_eff))):
357
+ return None
358
+ if not (len(s_id) == len(s_chrom) == len(s_ea) == len(s_oa)
359
+ == len(s_pos) == len(beta) == len(se) == len(n_eff) == n):
360
+ return None
361
+ genotype = _genotype_columns(variants, entry)
362
+ except Exception: # noqa: BLE001 - defer every oddity to the row loop
363
+ return None
364
+ if genotype is None:
365
+ return None
366
+ return {"id": s_id, "chrom": s_chrom, "pos": s_pos, "raw_pos":
367
+ sumstats.pos, "ea": s_ea, "oa": s_oa, "beta": beta, "se": se,
368
+ "n_eff": n_eff, "genotype": genotype}
369
+
370
+
371
+ def _sumstats_alleles(rows, vocab, local):
372
+ """Upper-cased allele codes over the reference vocabulary (+ extras)."""
373
+ uniques, codes = _factorize(rows)
374
+ upper = [str(a).upper() for a in uniques]
375
+ base = len(vocab)
376
+ per_value = np.array(
377
+ [vocab[a] if a in vocab else local.setdefault(a, base + len(local))
378
+ for a in upper],
379
+ dtype=np.int64)
380
+ return upper, codes, per_value
381
+
382
+
383
+ def _complement_codes(upper, vocab, local):
384
+ out = []
385
+ for a in upper:
386
+ c = _complement(a)
387
+ if c is None:
388
+ out.append(-1)
389
+ else:
390
+ code = vocab.get(c)
391
+ out.append(code if code is not None else local.get(c, -2))
392
+ return np.array(out, dtype=np.int64)
393
+
394
+
395
+ def _harmonize_vector(sumstats, variants, entry, columns, drop_ambiguous):
396
+ """Vectorised :func:`harmonize`; identical output to the row loop."""
397
+ g = columns["genotype"]
398
+ cache = entry[5]
399
+ n = len(sumstats)
400
+ s_pos = columns["pos"]
401
+
402
+ labels, label_codes = _factorize(columns["chrom"])
403
+ norm_code = g["norm_code"]
404
+ usable = np.array([bool(str(c).strip()) for c in labels],
405
+ dtype=bool)[label_codes]
406
+ q_chrom = np.array([norm_code.get(_norm_chrom(c), -1) for c in labels],
407
+ dtype=np.int64)[label_codes]
408
+
409
+ by_id = entry[1]
410
+ none = (-1,)
411
+ id_hit = np.array(
412
+ [c[0] if len(c) == 1 else -2
413
+ for c in map(by_id.get, columns["id"], repeat(none))],
414
+ dtype=np.int64)
415
+
416
+ cand = np.full(n, -1, dtype=np.int64)
417
+ multi = {}
418
+ n_unmatched = n_ambiguous_id = 0
419
+
420
+ rows = np.flatnonzero(id_hit >= 0)
421
+ gi = id_hit[rows]
422
+ q = s_pos[rows]
423
+ g_pos = g["pos"][gi]
424
+ keep = (~((q > 0) & usable[rows]) | (g_pos <= 0)
425
+ | ((g_pos == q) & (g["chrom"][gi] == q_chrom[rows])))
426
+ cand[rows[keep]] = gi[keep]
427
+ n_unmatched += int(rows.size - np.count_nonzero(keep))
428
+
429
+ rows = np.flatnonzero(id_hit == -2)
430
+ if rows.size:
431
+ memo, query_memo = entry[4], {}
432
+ s_id, s_chrom, raw_pos = (columns["id"], columns["chrom"],
433
+ _rows(columns["raw_pos"]))
434
+ for k in rows.tolist():
435
+ hits = _coordinate_candidates(by_id[s_id[k]], variants,
436
+ s_chrom[k], raw_pos[k], memo,
437
+ query_memo)
438
+ if not hits:
439
+ n_ambiguous_id += 1
440
+ elif len(hits) == 1:
441
+ cand[k] = hits[0]
442
+ else:
443
+ multi[k] = hits
444
+
445
+ rows = np.flatnonzero(id_hit == -1)
446
+ if rows.size:
447
+ index = _position_index(g, cache)
448
+ if index is False:
449
+ return None
450
+ keys, key_rows, span = index
451
+ q = s_pos[rows]
452
+ c = q_chrom[rows]
453
+ ok = (q > 0) & (q < span) & (c >= 0)
454
+ rows_ok = rows[ok]
455
+ qk = c[ok] * span + q[ok]
456
+ lo = np.searchsorted(keys, qk, side="left")
457
+ hi = np.searchsorted(keys, qk, side="right")
458
+ width = hi - lo
459
+ one = width == 1
460
+ cand[rows_ok[one]] = key_rows[lo[one]]
461
+ for k, a, b in zip(rows_ok[width > 1].tolist(),
462
+ lo[width > 1].tolist(), hi[width > 1].tolist()):
463
+ multi[k] = key_rows[a:b].tolist()
464
+ n_unmatched += int(rows.size - np.count_nonzero(width))
465
+
466
+ vocab, local = g["vocab"], {}
467
+ ea_up, ea_idx, ea_val = _sumstats_alleles(columns["ea"], vocab, local)
468
+ oa_up, oa_idx, oa_val = _sumstats_alleles(columns["oa"], vocab, local)
469
+ ea_comp = _complement_codes(ea_up, vocab, local)
470
+ oa_comp = _complement_codes(oa_up, vocab, local)
471
+
472
+ has = cand >= 0
473
+ if multi:
474
+ has[np.fromiter(multi, dtype=np.int64, count=len(multi))] = True
475
+ stage = np.flatnonzero(has)
476
+ e, o = ea_val[ea_idx[stage]], oa_val[oa_idx[stage]]
477
+ ce, co = ea_comp[ea_idx[stage]], oa_comp[oa_idx[stage]]
478
+ n_ambiguous = 0
479
+ if drop_ambiguous:
480
+ amb = (ce >= 0) & (co >= 0) & (ce == o) & (co == e)
481
+ n_ambiguous = int(np.count_nonzero(amb))
482
+ live = ~amb
483
+ stage, e, o, ce, co = (stage[live], e[live], o[live], ce[live],
484
+ co[live])
485
+
486
+ reach = [cand[stage[cand[stage] >= 0]]]
487
+ reach += [np.asarray(multi[k], dtype=np.int64)
488
+ for k in stage.tolist() if k in multi] if multi else []
489
+ hits = np.bincount(np.concatenate(reach), minlength=len(variants))
490
+ single = cand[stage] >= 0
491
+ solo = single & (hits[np.where(single, cand[stage], 0)] == 1)
492
+
493
+ beta, se, n_eff = columns["beta"], columns["se"], columns["n_eff"]
494
+ finite = np.isfinite(beta) & np.isfinite(se)
495
+
496
+ k_solo = stage[solo]
497
+ gi = cand[k_solo]
498
+ g1, g2 = g["a1"][gi], g["a2"][gi]
499
+ e_, o_, ce_, co_ = e[solo], o[solo], ce[solo], co[solo]
500
+ same = (e_ == g1) & (o_ == g2)
501
+ swap = ~same & (e_ == g2) & (o_ == g1)
502
+ strand_same = ~same & ~swap & (ce_ == g1) & (co_ == g2)
503
+ strand_swap = ~same & ~swap & ~strand_same & (ce_ == g2) & (co_ == g1)
504
+ compatible = same | swap | strand_same | strand_swap
505
+ flip = swap | strand_swap
506
+ n_mismatch = int(k_solo.size - np.count_nonzero(compatible))
507
+ n_strand = int(np.count_nonzero(strand_same | strand_swap))
508
+ fin = finite[k_solo]
509
+ n_nonfinite = int(np.count_nonzero(compatible & ~fin))
510
+ kept = compatible & fin
511
+ n_flip = int(np.count_nonzero(kept & flip))
512
+ idx_parts, src_parts, flip_parts = [gi[kept]], [k_solo[kept]], [flip[kept]]
513
+
514
+ loop_rows = stage[~solo].tolist()
515
+ if loop_rows:
516
+ vocab_rows = list(vocab)
517
+ ga1, ga2 = g["a1"], g["a2"]
518
+ idx, srcs, flips = [], [], []
519
+ n_duplicate = 0
520
+ seen = set()
521
+ for k in loop_rows:
522
+ candidates = multi.get(k) or [int(cand[k])]
523
+ ea, oa = ea_up[ea_idx[k]], oa_up[oa_idx[k]]
524
+ gi_hit = None
525
+ n_seen_here = 0
526
+ for gi_k in candidates:
527
+ if gi_k in seen:
528
+ n_seen_here += 1
529
+ continue
530
+ f, strand = _allele_orientation(
531
+ ea, oa, vocab_rows[ga1[gi_k]], vocab_rows[ga2[gi_k]])
532
+ if f is None:
533
+ continue
534
+ if strand:
535
+ n_strand += 1
536
+ gi_hit = gi_k
537
+ break
538
+ if gi_hit is None:
539
+ if n_seen_here == len(candidates):
540
+ n_duplicate += 1
541
+ else:
542
+ n_mismatch += 1
543
+ continue
544
+ if not finite[k]:
545
+ n_nonfinite += 1
546
+ continue
547
+ if f:
548
+ n_flip += 1
549
+ seen.add(gi_hit)
550
+ idx.append(gi_hit); srcs.append(k); flips.append(f)
551
+ idx_parts.append(np.array(idx, dtype=np.int64))
552
+ src_parts.append(np.array(srcs, dtype=np.int64))
553
+ flip_parts.append(np.array(flips, dtype=bool))
554
+ else:
555
+ n_duplicate = 0
556
+
557
+ idx = np.concatenate(idx_parts)
558
+ src = np.concatenate(src_parts)
559
+ flipped = np.concatenate(flip_parts)
560
+ order = np.argsort(idx)
561
+ idx, src, flipped = idx[order], src[order], flipped[order]
562
+ # Negate after widening: exact, and no wrap-around for integer effects.
563
+ out_beta = beta[src].astype(float)
564
+ np.negative(out_beta, out=out_beta, where=flipped)
565
+ if beta.dtype.kind in "iu":
566
+ out_beta += 0.0
567
+ log = {
568
+ "n_sumstats": n,
569
+ "n_genotype_variants": len(variants),
570
+ "n_matched": len(idx),
571
+ "n_flipped": n_flip,
572
+ "n_strand_flipped": n_strand,
573
+ "n_dropped_ambiguous": n_ambiguous,
574
+ "n_dropped_mismatch": n_mismatch,
575
+ "n_dropped_ambiguous_id": n_ambiguous_id,
576
+ "n_dropped_nonfinite": n_nonfinite,
577
+ "n_dropped_duplicate": n_duplicate,
578
+ "n_unmatched": n_unmatched,
579
+ }
580
+ return Harmonized(var_index=idx, beta=out_beta, se=se[src].astype(float),
581
+ n_eff=n_eff[src].astype(float), flipped=flipped,
582
+ src_index=src, log=log)
583
+
584
+
585
+ def harmonize(sumstats, variants, *, drop_ambiguous=True, status=None):
586
+ """Align ``sumstats`` to ``variants``; see the module docstring for the rules.
587
+
588
+ ``sumstats`` is any object with ``id, chrom, pos, ea, oa, beta, se,
589
+ n_eff`` arrays (a :class:`~allelign.sumstats.Sumstats`, or ldpred3's).
590
+ ``variants`` has ``id, chrom, pos, a1, a2`` (and optionally
591
+ ``alternate_id``). Frozen tables (:meth:`Variants.freeze`) cache their
592
+ indices across calls.
593
+
594
+ ``status`` is an optional ``int8`` array of ``len(sumstats)`` codes
595
+ (:mod:`allelign.status`): every dropped row is stamped with its reason
596
+ and matched rows stay 0. Such calls take the row loop -- the same rule
597
+ set and the same result, at O(rows) Python.
598
+ """
599
+ entry = _variant_indices(variants)
600
+ if status is not None:
601
+ status = _st.check_status(status, len(sumstats))
602
+ return _harmonize_rows(sumstats, variants, entry, drop_ambiguous,
603
+ status)
604
+ columns = _vector_columns(sumstats, variants, entry)
605
+ if columns is not None:
606
+ out = _harmonize_vector(sumstats, variants, entry, columns,
607
+ drop_ambiguous)
608
+ if out is not None:
609
+ return out
610
+ return _harmonize_rows(sumstats, variants, entry, drop_ambiguous)
611
+
612
+
613
+ def _harmonize_rows(sumstats, variants, entry, drop_ambiguous, status=None):
614
+ """The reference row-by-row implementation of :func:`harmonize`."""
615
+
616
+ def _drop(k, code):
617
+ if status is not None:
618
+ status[k] = code
619
+
620
+ by_id = entry[1]
621
+ by_pos = entry[2]
622
+ _memo = entry[4]
623
+ _upper = entry[3]
624
+
625
+ s_id, s_chrom, s_pos = (_rows(sumstats.id), _rows(sumstats.chrom),
626
+ _rows(sumstats.pos))
627
+ s_ea, s_oa = _rows(sumstats.ea), _rows(sumstats.oa)
628
+ s_beta, s_se, s_n = (_rows(sumstats.beta), _rows(sumstats.se),
629
+ _rows(sumstats.n_eff))
630
+ v_a1, v_a2 = variants.a1, variants.a2
631
+ by_id_get = by_id.get
632
+ isfinite = math.isfinite
633
+
634
+ idx, betas, ses, ns, flips, srcs = [], [], [], [], [], []
635
+ n_unmatched = n_ambiguous = n_mismatch = n_flip = n_strand = 0
636
+ n_duplicate = n_ambiguous_id = n_nonfinite = 0
637
+ seen = set()
638
+ query_memo = {}
639
+ ambiguous_memo = {}
640
+
641
+ for k in range(len(sumstats)):
642
+ sid = s_id[k]
643
+ candidates = by_id_get(sid) if sid else None
644
+ if candidates:
645
+ n_id = len(candidates)
646
+ candidates = _coordinate_candidates(
647
+ candidates, variants, s_chrom[k], s_pos[k], _memo,
648
+ query_memo)
649
+ if not candidates:
650
+ if n_id > 1:
651
+ n_ambiguous_id += 1
652
+ _drop(k, _st.AMBIGUOUS_ID)
653
+ else:
654
+ n_unmatched += 1
655
+ _drop(k, _st.UNMATCHED)
656
+ continue
657
+ if not candidates:
658
+ if by_pos is None:
659
+ by_pos = _build_pos_index(variants, _memo)
660
+ entry[2] = by_pos
661
+ s_pos_k = int(s_pos[k])
662
+ if s_pos_k <= 0:
663
+ candidates = None
664
+ else:
665
+ c = s_chrom[k]
666
+ nc = _memo.get(c)
667
+ if nc is None:
668
+ nc = _memo[c] = _norm_chrom(c)
669
+ candidates = by_pos.get((nc, s_pos_k))
670
+ if not candidates:
671
+ n_unmatched += 1
672
+ _drop(k, _st.UNMATCHED)
673
+ continue
674
+
675
+ ea = _upper.get(s_ea[k])
676
+ if ea is None:
677
+ ea = _upper[s_ea[k]] = str(s_ea[k]).upper()
678
+ oa = _upper.get(s_oa[k])
679
+ if oa is None:
680
+ oa = _upper[s_oa[k]] = str(s_oa[k]).upper()
681
+ if drop_ambiguous:
682
+ pair = (ea, oa)
683
+ ambiguous = ambiguous_memo.get(pair)
684
+ if ambiguous is None:
685
+ ambiguous = ambiguous_memo[pair] = _ambiguous_pair(ea, oa)
686
+ if ambiguous:
687
+ n_ambiguous += 1
688
+ _drop(k, _st.AMBIGUOUS_ALLELES)
689
+ continue
690
+
691
+ gi_hit = None
692
+ flip = None
693
+ n_seen_here = 0
694
+ for gi in candidates:
695
+ if gi in seen:
696
+ n_seen_here += 1
697
+ continue
698
+ g1 = _upper.get(v_a1[gi])
699
+ if g1 is None:
700
+ g1 = _upper[v_a1[gi]] = str(v_a1[gi]).upper()
701
+ g2 = _upper.get(v_a2[gi])
702
+ if g2 is None:
703
+ g2 = _upper[v_a2[gi]] = str(v_a2[gi]).upper()
704
+ flip, strand = _allele_orientation(ea, oa, g1, g2)
705
+ if flip is None:
706
+ continue
707
+ if strand:
708
+ n_strand += 1
709
+ gi_hit = gi
710
+ break
711
+
712
+ if gi_hit is None:
713
+ if n_seen_here == len(candidates):
714
+ n_duplicate += 1
715
+ _drop(k, _st.DUPLICATE_MATCH)
716
+ else:
717
+ n_mismatch += 1
718
+ _drop(k, _st.ALLELE_MISMATCH)
719
+ continue
720
+
721
+ beta = s_beta[k]
722
+ if not isfinite(beta) or not isfinite(s_se[k]):
723
+ n_nonfinite += 1
724
+ _drop(k, _st.NONFINITE_BETA_SE)
725
+ continue
726
+ if flip:
727
+ beta = -beta
728
+ n_flip += 1
729
+
730
+ seen.add(gi_hit)
731
+ idx.append(gi_hit); betas.append(beta); ses.append(s_se[k])
732
+ ns.append(s_n[k]); flips.append(flip); srcs.append(k)
733
+
734
+ order = np.argsort(idx)
735
+ idx = np.array(idx, dtype=np.int64)[order]
736
+ log = {
737
+ "n_sumstats": len(sumstats),
738
+ "n_genotype_variants": len(variants),
739
+ "n_matched": len(idx),
740
+ "n_flipped": n_flip,
741
+ "n_strand_flipped": n_strand,
742
+ "n_dropped_ambiguous": n_ambiguous,
743
+ "n_dropped_mismatch": n_mismatch,
744
+ "n_dropped_ambiguous_id": n_ambiguous_id,
745
+ "n_dropped_nonfinite": n_nonfinite,
746
+ "n_dropped_duplicate": n_duplicate,
747
+ "n_unmatched": n_unmatched,
748
+ }
749
+ return Harmonized(
750
+ var_index=idx,
751
+ beta=np.array(betas, dtype=float)[order],
752
+ se=np.array(ses, dtype=float)[order],
753
+ n_eff=np.array(ns, dtype=float)[order],
754
+ flipped=np.array(flips, dtype=bool)[order],
755
+ src_index=np.array(srcs, dtype=np.int64)[order],
756
+ log=log,
757
+ )
758
+
759
+
760
+ # ---------------------------------------------------------------------------
761
+ # Coverage and match diagnostics
762
+ # ---------------------------------------------------------------------------
763
+
764
+ #: Flag an alignment covering less than this fraction of the reference: LDSC's
765
+ #: M, the h2 denominator and the LD screen's windows are defined on the full
766
+ #: reference, so sparse coverage changes what those numbers mean.
767
+ REFERENCE_COVERAGE_WARN = 0.5
768
+ #: Below this the analysis is no longer genome-wide on that reference.
769
+ SEVERE_REFERENCE_COVERAGE = 0.1
770
+ #: Call it a build/coordinate mismatch when at least this fraction of the
771
+ #: unmatched rows carry an identifier the reference holds elsewhere.
772
+ BUILD_MISMATCH_FRACTION = 0.2
773
+
774
+
775
+ class _ReferenceLoci:
776
+ """Loci carrying each reference identifier, over the cached index."""
777
+
778
+ __slots__ = ("_by_id", "_chrom", "_pos", "_memo")
779
+
780
+ def __init__(self, variants):
781
+ self._by_id = _variant_indices(variants)[1]
782
+ self._chrom = variants.chrom
783
+ self._pos = variants.pos
784
+ self._memo = {}
785
+
786
+ def __len__(self):
787
+ return len(self._by_id)
788
+
789
+ def _norm(self, raw):
790
+ norm = self._memo.get(raw)
791
+ if norm is None:
792
+ norm = self._memo[raw] = _norm_chrom(raw)
793
+ return norm
794
+
795
+ def loci(self, key):
796
+ """Distinct normalised ``(chrom, pos)`` pairs for ``key``, or ``None``."""
797
+ rows = self._by_id.get(key)
798
+ if not rows:
799
+ return None
800
+ return {(self._norm(self._chrom[gi]), int(self._pos[gi]))
801
+ for gi in rows}
802
+
803
+
804
+ def _diagnose_unmatched(sumstats, variants, harmonized):
805
+ """Split unmatched rows into build-mismatch and absent-variant classes."""
806
+ n_rows = len(sumstats.id)
807
+ matched_rows = np.zeros(n_rows, dtype=bool)
808
+ if harmonized.src_index is not None and len(harmonized):
809
+ matched_rows[np.asarray(harmonized.src_index, dtype=np.int64)] = True
810
+ unmatched = np.where(~matched_rows)[0]
811
+ report = {"n_unmatched_rows": int(unmatched.size),
812
+ "n_unmatched_id_elsewhere": 0,
813
+ "n_unmatched_id_absent": 0,
814
+ "n_unmatched_id_missing": 0,
815
+ "examples": []}
816
+ if unmatched.size == 0:
817
+ return report
818
+ index = _ReferenceLoci(variants)
819
+ ss_id, ss_chrom, ss_pos = sumstats.id, sumstats.chrom, sumstats.pos
820
+ memo = {}
821
+ for k in unmatched:
822
+ key = ss_id[k]
823
+ if not key:
824
+ report["n_unmatched_id_missing"] += 1
825
+ continue
826
+ loci = index.loci(key)
827
+ if loci is None:
828
+ report["n_unmatched_id_absent"] += 1
829
+ continue
830
+ raw = ss_chrom[k]
831
+ norm = memo.get(raw)
832
+ if norm is None:
833
+ norm = memo[raw] = _norm_chrom(raw)
834
+ try:
835
+ here = int(ss_pos[k])
836
+ except (TypeError, ValueError, OverflowError):
837
+ here = 0
838
+ if (norm, here) in loci:
839
+ continue
840
+ report["n_unmatched_id_elsewhere"] += 1
841
+ if len(report["examples"]) < 3:
842
+ ref_chrom, ref_pos = sorted(loci)[0]
843
+ report["examples"].append(
844
+ f"{key} GWAS {norm}:{here} vs reference "
845
+ f"{ref_chrom}:{ref_pos}")
846
+ return report
847
+
848
+
849
+ def reference_coverage(sumstats, variants, harmonized, *, n_reference=None):
850
+ """How much of the LD reference the aligned rows cover, and why not more.
851
+
852
+ Zero overlap is an error a caller catches elsewhere. What this reports is
853
+ the case that still produces numbers: a fit running on a small,
854
+ unrepresentative corner of the reference. Returns a dict with the two
855
+ fractions and a ``reason`` from a fixed vocabulary --
856
+ ``adequate_coverage``, ``no_rows_offered``, ``build_mismatch``,
857
+ ``absent_identifiers``, ``rows_rejected_at_matched_locus``,
858
+ ``sparse_sumstats`` -- plus the unmatched-row breakdown when it looked.
859
+
860
+ Preparation uses the same report and adds validity, QC and re-anchoring
861
+ counts, so CLI and web runs diagnose a thin alignment consistently.
862
+ """
863
+ return _reference_coverage(sumstats, variants, harmonized,
864
+ n_reference=n_reference)
865
+
866
+
867
+ def _reference_coverage(sumstats, variants, harmonized, *, n_reference=None,
868
+ n_matched=None, n_offered=None, n_invalid=0,
869
+ n_qc_dropped=0):
870
+ """Shared coverage classification, including preparation's row filters."""
871
+ n_rows = int(len(sumstats.id))
872
+ n_cache = int(len(variants.id) if n_reference is None else n_reference)
873
+ n_matched = int(len(harmonized) if n_matched is None else n_matched)
874
+ n_offered = int(n_rows if n_offered is None else n_offered)
875
+ out = {
876
+ "n_sumstats_offered": n_offered,
877
+ "n_cache": n_cache,
878
+ "n_matched": n_matched,
879
+ "frac_of_reference": (n_matched / n_cache) if n_cache else 0.0,
880
+ "frac_of_sumstats": (n_matched / n_offered) if n_offered else 0.0,
881
+ "diagnosed": False,
882
+ }
883
+ if out["frac_of_reference"] >= REFERENCE_COVERAGE_WARN:
884
+ out["reason"] = "adequate_coverage"
885
+ return out
886
+ if n_rows == 0:
887
+ # Nothing reached harmonization, so the reference says nothing about
888
+ # why. Naming the build here would be a guess.
889
+ out["reason"] = "no_rows_offered"
890
+ return out
891
+ diagnosis = _diagnose_unmatched(sumstats, variants, harmonized)
892
+ out.update(diagnosis, diagnosed=True)
893
+ n_unmatched = diagnosis["n_unmatched_rows"]
894
+ elsewhere = diagnosis["n_unmatched_id_elsewhere"]
895
+ absent = diagnosis["n_unmatched_id_absent"]
896
+ build_mismatch = bool(
897
+ n_unmatched and elsewhere / n_unmatched >= BUILD_MISMATCH_FRACTION)
898
+ out["build_mismatch_suspected"] = build_mismatch
899
+ out["severe"] = out["frac_of_reference"] < SEVERE_REFERENCE_COVERAGE
900
+ out["reason"] = (
901
+ "rows_failed_validity" if n_invalid and n_invalid >= max(n_matched, 1) else
902
+ "qc_removed_rows" if n_qc_dropped and n_qc_dropped >= max(n_matched, 1) else
903
+ "build_mismatch" if build_mismatch else
904
+ "absent_identifiers" if absent and absent >= elsewhere else
905
+ "rows_rejected_at_matched_locus" if n_unmatched else
906
+ "sparse_sumstats")
907
+ return out
908
+
909
+
910
+ def coverage_message(coverage, *, label="sumstats", n_invalid=0,
911
+ n_harmonized=0, n_qc_dropped=0):
912
+ """One paragraph naming the coverage, the cause, and the remedy.
913
+
914
+ Kept beside the thresholds it describes; the pipeline's alignment
915
+ warnings use it.
916
+ """
917
+ reason = coverage.get("reason")
918
+ if reason in (None, "adequate_coverage"):
919
+ return None
920
+ if reason == "reanchor_dropped_all":
921
+ return (
922
+ f"{label}: identifier re-anchoring dropped every one of "
923
+ f"{coverage['n_sumstats_offered']:,} rows against the LD reference "
924
+ f"({coverage['n_dropped_absent_identifier']:,} identifiers the reference "
925
+ f"does not hold, {coverage['n_dropped_ambiguous_locus']:,} held at "
926
+ "several loci). This is not a QC filter failure: the reference "
927
+ "was consulted and none of the offered identifiers could be "
928
+ "re-stamped onto it.")
929
+ if reason == "no_rows_offered":
930
+ return (f"{label}: QC removed every row before alignment to the LD "
931
+ "reference, so no variant could match. Check the QC filters "
932
+ "against the file's columns -- a missing or unparsed "
933
+ "sample-size, allele-frequency, or standard-error column "
934
+ "removes every row -- before suspecting the reference.")
935
+ lines = [
936
+ f"{label}: only {coverage['n_matched']:,} of {coverage['n_cache']:,} "
937
+ f"LD-reference variants are covered "
938
+ f"({coverage['frac_of_reference']:.1%}), from "
939
+ f"{coverage['n_sumstats_offered']:,} rows offered to harmonisation."]
940
+ if coverage.get("severe"):
941
+ lines.append(
942
+ "A fit on this few reference variants is not a genome-wide "
943
+ "analysis of that reference: the LD blocks are nearly empty, the "
944
+ "LD-consistency screen leaves most windows below its minimum "
945
+ "size, and LDSC keeps the full reference M, so h2 and "
946
+ "polygenicity are on a denominator the data do not populate.")
947
+ else:
948
+ lines.append(
949
+ "LDSC keeps the full reference M and the LD blocks stay sparse, "
950
+ "so h2, polygenicity, and screen coverage are all affected.")
951
+ if reason == "build_mismatch" or coverage.get("build_mismatch_suspected"):
952
+ lines.append(
953
+ f"{coverage['n_unmatched_id_elsewhere']:,} of the "
954
+ f"{coverage['n_unmatched_rows']:,} unmatched rows carry an identifier "
955
+ "the reference holds at coordinates that disagree. That is the "
956
+ "signature of a genome-build mismatch (GRCh38 summary statistics "
957
+ "against a GRCh37 reference, or the reverse) or of CHR/BP read "
958
+ "from the wrong columns. Check the column mapping; if the builds "
959
+ "differ, lift the summary statistics over to the reference's build "
960
+ "and keep the chain file and failed mappings with the results, "
961
+ "or use a reference on the GWAS build. Do not drop coordinates "
962
+ "to force an identifier-only match.")
963
+ if coverage.get("examples"):
964
+ lines.append("Examples: " + "; ".join(coverage["examples"]) + ".")
965
+ elif reason == "absent_identifiers":
966
+ lines.append(
967
+ f"{coverage['n_unmatched_id_absent']:,} unmatched rows carry an "
968
+ "identifier the reference does not hold at all, which points at a "
969
+ "different variant set or ancestry panel rather than at "
970
+ "coordinates.")
971
+ elif reason == "rows_rejected_at_matched_locus":
972
+ lines.append(
973
+ "The unmatched rows are present in the reference at the same "
974
+ "locus, so they were rejected on alleles, duplication, or a "
975
+ "non-finite effect; check the effect/other allele columns.")
976
+ elif reason == "rows_failed_validity":
977
+ lines.append(
978
+ f"{n_invalid:,} of the {n_harmonized:,} harmonized rows had a "
979
+ "non-finite or non-positive standard error or sample size and "
980
+ "were dropped after alignment. Check the SE and N columns -- a "
981
+ "mis-mapped or all-zero standard-error column is a common cause "
982
+ "-- before suspecting the deposition.")
983
+ elif reason == "qc_removed_rows":
984
+ lines.append(
985
+ f"Sumstats QC removed {n_qc_dropped:,} rows before alignment, so "
986
+ "only the survivors were offered to the reference. Relax the QC "
987
+ "filters (min_n_ratio, min_maf, min_info, max_chisq) if they are "
988
+ "tighter than intended, and check they point at the right columns.")
989
+ else:
990
+ lines.append(
991
+ "Every offered row matched, so the GWAS itself covers only this "
992
+ "part of the reference -- a significant-hits-only or "
993
+ "array-restricted deposition rather than a misalignment. "
994
+ "Genome-wide estimands cannot be read off such a subset.")
995
+ return " ".join(lines)
996
+
997
+
998
+ def diagnose_match(sumstats, variants, *, sample=2000):
999
+ """Explain *why* few or no variants matched (build, chr naming, alleles).
1000
+
1001
+ Returns a dict with overlap counts and a human-readable ``message``.
1002
+ """
1003
+ gid = {v for v in variants.id if v and v != "."}
1004
+ sid = {s for s in sumstats.id if s}
1005
+ rsid_overlap = len(sid & gid)
1006
+
1007
+ gpos = {(_norm_chrom(variants.chrom[i]), int(variants.pos[i]))
1008
+ for i in range(len(variants)) if int(variants.pos[i]) > 0}
1009
+ spos = {(_norm_chrom(sumstats.chrom[k]), int(sumstats.pos[k]))
1010
+ for k in range(len(sumstats)) if int(sumstats.pos[k]) > 0}
1011
+ pos_overlap = len(spos & gpos)
1012
+
1013
+ def _prefixed(labels):
1014
+ return any(str(c)[:3].lower() == "chr" for c in labels[:sample])
1015
+ s_pref, g_pref = _prefixed(sumstats.chrom), _prefixed(variants.chrom)
1016
+
1017
+ build_mismatch = False
1018
+ if rsid_overlap:
1019
+ g_at = {variants.id[i]: (_norm_chrom(variants.chrom[i]),
1020
+ int(variants.pos[i]))
1021
+ for i in range(len(variants)) if variants.id[i]}
1022
+ # sorted, not set order: set iteration varies with hash randomisation.
1023
+ shared = set(sorted(sid & gid)[:sample])
1024
+ disagree = tot = 0
1025
+ for k in range(len(sumstats)):
1026
+ r = sumstats.id[k]
1027
+ if r in g_at and r in shared:
1028
+ tot += 1
1029
+ if g_at[r] != (_norm_chrom(sumstats.chrom[k]),
1030
+ int(sumstats.pos[k])):
1031
+ disagree += 1
1032
+ build_mismatch = tot > 0 and disagree / tot > 0.5
1033
+
1034
+ msgs = []
1035
+ if build_mismatch:
1036
+ msgs.append("rsIDs overlap but their chrom:pos disagree -- likely a "
1037
+ "genome-build mismatch (e.g. hg19 vs hg38); lift over to a "
1038
+ "common build")
1039
+ if rsid_overlap == 0 and pos_overlap == 0:
1040
+ if s_pref != g_pref:
1041
+ msgs.append("chromosome labels differ in style (one side uses a "
1042
+ "'chr' prefix) -- normalisation handles this, so a "
1043
+ "remaining 0-overlap points to different builds or "
1044
+ "variant sets")
1045
+ else:
1046
+ msgs.append("no rsID and no chrom:pos overlap -- the sumstats and "
1047
+ "reference look like different variant sets or builds")
1048
+ if rsid_overlap and pos_overlap and not build_mismatch:
1049
+ msgs.append("variants overlap but harmonisation dropped them -- likely "
1050
+ "an allele/strand or effect-column problem (check ea/oa)")
1051
+ return {
1052
+ "n_sumstats": len(sumstats),
1053
+ "n_genotype_variants": len(variants),
1054
+ "rsid_overlap": rsid_overlap,
1055
+ "pos_overlap_normalized": pos_overlap,
1056
+ "sumstats_chr_prefixed": s_pref,
1057
+ "genotype_chr_prefixed": g_pref,
1058
+ "build_mismatch_suspected": build_mismatch,
1059
+ "message": "; ".join(msgs) if msgs else "no obvious cause detected",
1060
+ }
1061
+
1062
+
1063
+ def count_positional_only_matches(sumstats, variants, harmonized):
1064
+ """Matched rows whose identifier the reference lacks (position-only match).
1065
+
1066
+ A coordinate can land on a *different* variant that happens to carry a
1067
+ compatible allele pair, so these rows cannot be verified by identifier.
1068
+ """
1069
+ if not len(harmonized) or harmonized.src_index is None:
1070
+ return 0
1071
+ by_id = _variant_indices(variants)[1]
1072
+ ss_id = sumstats.id
1073
+ n = 0
1074
+ for k in np.asarray(harmonized.src_index, dtype=np.int64):
1075
+ key = ss_id[k]
1076
+ if key and key not in by_id:
1077
+ n += 1
1078
+ return n
1079
+
1080
+
1081
+ def reanchor_on_identifier(sumstats, variants, *, label="sumstats",
1082
+ warn=True, return_keep=False,
1083
+ warning_stacklevel=2):
1084
+ """Move each row onto the reference's own coordinates for its identifier.
1085
+
1086
+ The coordinate repair for a build mismatch when identifiers are shared and
1087
+ trustworthy (a GWAS Catalog harmonised file's ``hm_rsid`` against an
1088
+ rsID-keyed reference). It is *not* a chain-file liftover: no interval
1089
+ mapping is consulted and no coordinate is invented. A row whose identifier
1090
+ the reference holds at exactly one locus is re-stamped with that locus;
1091
+ alleles are still checked afterwards by :func:`harmonize`.
1092
+
1093
+ Dropped rather than re-anchored: identifiers the reference does not hold
1094
+ (their coordinates are on the other build and could match a *different*
1095
+ variant by position), identifiers held at several loci, and rows with no
1096
+ identifier. Returns ``(anchored, log)`` -- plus the boolean keep mask over
1097
+ the input rows with ``return_keep=True``. Prefer a properly lifted-over
1098
+ file where one exists; the log is the audit trail.
1099
+ """
1100
+ ids = sumstats.id
1101
+ n_rows = len(ids)
1102
+ index = _ReferenceLoci(variants)
1103
+ chrom = np.array(sumstats.chrom, dtype=object)
1104
+ pos = np.array(sumstats.pos, dtype=np.int64)
1105
+ keep = np.ones(n_rows, dtype=bool)
1106
+ memo = {}
1107
+ n_anchored = n_moved = n_ambiguous = n_absent = n_missing_id = 0
1108
+ for k in range(n_rows):
1109
+ key = ids[k]
1110
+ if not key:
1111
+ n_missing_id += 1
1112
+ keep[k] = False
1113
+ continue
1114
+ loci = index.loci(key)
1115
+ if loci is None:
1116
+ n_absent += 1
1117
+ keep[k] = False
1118
+ continue
1119
+ if len(loci) > 1:
1120
+ n_ambiguous += 1
1121
+ keep[k] = False
1122
+ continue
1123
+ ref_chrom, ref_pos = next(iter(loci))
1124
+ raw = chrom[k]
1125
+ norm = memo.get(raw)
1126
+ if norm is None:
1127
+ norm = memo[raw] = _norm_chrom(raw)
1128
+ if norm != ref_chrom or int(pos[k]) != ref_pos:
1129
+ n_moved += 1
1130
+ chrom[k] = ref_chrom
1131
+ pos[k] = ref_pos
1132
+ n_anchored += 1
1133
+ log = {"applied": True, "n_rows": int(n_rows),
1134
+ "n_anchored": int(n_anchored), "n_moved": int(n_moved),
1135
+ "n_dropped_ambiguous_locus": int(n_ambiguous),
1136
+ "n_dropped_absent_identifier": int(n_absent),
1137
+ "n_dropped_missing_identifier": int(n_missing_id)}
1138
+ anchored = replace(sumstats, chrom=chrom, pos=pos).subset(keep)
1139
+ n_dropped = n_absent + n_ambiguous + n_missing_id
1140
+ # Warn on drops as well as moves: identifiers retired between builds can
1141
+ # lose a large fraction with nothing moved.
1142
+ if warn and (n_moved or n_dropped):
1143
+ actions = []
1144
+ if n_moved:
1145
+ actions.append(
1146
+ f"re-anchored {n_moved:,} of {n_rows:,} rows onto the LD "
1147
+ "reference's coordinates for their identifiers, because the "
1148
+ "two sides disagreed on where those variants are")
1149
+ if n_dropped:
1150
+ actions.append(
1151
+ f"dropped {n_dropped:,} of {n_rows:,} rows it could not anchor "
1152
+ f"by identifier ({n_absent:,} identifiers the reference does "
1153
+ f"not hold, {n_ambiguous:,} held at several loci, "
1154
+ f"{n_missing_id:,} with no identifier at all)")
1155
+ warnings.warn(
1156
+ f"{label}: identifier-keyed re-anchoring: " + "; ".join(actions)
1157
+ + ". The fit is therefore on the reference's genome build, keyed "
1158
+ "on identifier rather than on position. This is not a chain-file "
1159
+ "liftover: record it, and prefer a properly lifted-over input "
1160
+ "where one is available.",
1161
+ RuntimeWarning, stacklevel=warning_stacklevel)
1162
+ return (anchored, log, keep) if return_keep else (anchored, log)
1163
+
1164
+
1165
+ # ---------------------------------------------------------------------------
1166
+ # Weights
1167
+ # ---------------------------------------------------------------------------
1168
+
1169
+ @dataclass
1170
+ class HarmonizedWeights:
1171
+ """Weights aligned to a reference's ``a1`` alleles.
1172
+
1173
+ ``af`` is the frequency of ``a1`` (``1 - af`` where the alleles swapped);
1174
+ ``sd`` is orientation-free. Both are ``None`` when the source had none.
1175
+ """
1176
+
1177
+ var_index: np.ndarray
1178
+ weight: np.ndarray
1179
+ flipped: np.ndarray
1180
+ src_index: np.ndarray
1181
+ af: np.ndarray = None
1182
+ sd: np.ndarray = None
1183
+ log: dict = field(default_factory=dict)
1184
+
1185
+ def __len__(self):
1186
+ return len(self.var_index)
1187
+
1188
+
1189
+ @dataclass
1190
+ class _EffectRows:
1191
+ id: np.ndarray
1192
+ chrom: np.ndarray
1193
+ pos: np.ndarray
1194
+ ea: np.ndarray
1195
+ oa: np.ndarray
1196
+ beta: np.ndarray
1197
+ se: np.ndarray
1198
+ n_eff: np.ndarray
1199
+
1200
+ def __len__(self):
1201
+ return len(self.beta)
1202
+
1203
+
1204
+ def harmonize_weights(weights, variants, *, drop_ambiguous=True, status=None):
1205
+ """Align a weights table or PGS Catalog score to ``variants``.
1206
+
1207
+ ``weights`` needs ``id, chrom, pos, weight`` and the effect/other alleles
1208
+ as ``a1``/``a2`` (:class:`~allelign.weights.WeightsTable`) or
1209
+ ``ea``/``oa`` (:class:`~allelign.weights.ScoringFile`); ``af_ref`` and
1210
+ ``sd_ref`` are carried along when present. The same rules as
1211
+ :func:`harmonize` apply. The log adds ``weight_mass_matched``, the share
1212
+ of the total squared weight that matched (over the whole file when the
1213
+ score was read with a ``keep=`` prefilter).
1214
+ """
1215
+ ea = getattr(weights, "a1", None)
1216
+ oa = getattr(weights, "a2", None)
1217
+ if ea is None or oa is None:
1218
+ ea, oa = weights.ea, weights.oa
1219
+ w = np.asarray(weights.weight, dtype=float)
1220
+ ones = np.ones(w.size)
1221
+ rows = _EffectRows(id=np.asarray(weights.id, dtype=object),
1222
+ chrom=np.asarray(weights.chrom, dtype=object),
1223
+ pos=np.asarray(weights.pos, dtype=np.int64),
1224
+ ea=np.asarray(ea, dtype=object),
1225
+ oa=np.asarray(oa, dtype=object),
1226
+ beta=w, se=ones, n_eff=ones)
1227
+ h = harmonize(rows, variants, drop_ambiguous=drop_ambiguous,
1228
+ status=status)
1229
+ af = getattr(weights, "af_ref", None)
1230
+ sd = getattr(weights, "sd_ref", None)
1231
+ af_aligned = sd_aligned = None
1232
+ if af is not None:
1233
+ af = np.asarray(af, dtype=float)[h.src_index]
1234
+ af_aligned = np.where(h.flipped, 1.0 - af, af)
1235
+ if sd is not None:
1236
+ sd_aligned = np.asarray(sd, dtype=float)[h.src_index]
1237
+ # A score read with read_scoring_file(keep=...) lacks the rows the
1238
+ # prefilter rejected: they count as unmatched, and the weight mass is
1239
+ # taken over the whole file.
1240
+ parse_log = getattr(weights, "log", None) or {}
1241
+ removed = int(parse_log.get("n_reference_prefiltered", 0))
1242
+ total = parse_log.get("weight_square_sum_all")
1243
+ total = float(w @ w) if total is None else float(total)
1244
+ log = dict(h.log)
1245
+ if removed:
1246
+ log["n_sumstats"] += removed
1247
+ log["n_unmatched"] += removed
1248
+ log["weight_mass_matched"] = (float(h.beta @ h.beta) / total
1249
+ if total > 0 else float("nan"))
1250
+ return HarmonizedWeights(var_index=h.var_index, weight=h.beta,
1251
+ flipped=h.flipped, src_index=h.src_index,
1252
+ af=af_aligned, sd=sd_aligned, log=log)