aegis_sim.utilities.gvcf

  1"""FASTA-coordinate multi-sample gVCF export for AEGIS populations.
  2
  3Designed as a drop-in replacement for Clair3-emitted gVCFs in pipelines
  4that joint-genotype via GLnexus and then run population-genetics tests
  5like ABBA-BABA. Output is one gVCF per Population, each individual is a
  6sample column, each row is a FASTA base position with the same coordinate
  7system as the FASTA + Badread reads.
  8
  9Per position:
 10  REF = consensus (modal) base across all 2*N alleles
 11  ALT = sorted list of observed non-REF bases + <NON_REF>
 12  GT  = allele indices into [REF, *ALTs] for each diploid individual
 13  GQ  = 99 (synthetic; AEGIS has no read uncertainty)
 14  DP  = 30 (synthetic)
 15  AD  = per-allele depth, fixed at 15 for the called genotype's alleles
 16
 17Long runs of all-REF positions are compressed into reference blocks
 18(single record with END=<last_pos_of_block>, ALT=<NON_REF>) — standard
 19gVCF convention, expected by GLnexus.
 20
 21Only the composite architecture is supported. Modifying architecture
 22output would need a separate emitter (different locus layout).
 23"""
 24
 25import io
 26import pathlib
 27from typing import List, Tuple
 28
 29import numpy as np
 30
 31BASES = np.array(["A", "C", "G", "T"])
 32NON_REF = "<NON_REF>"
 33
 34
 35def encode_population_to_gvcf(
 36    population,
 37    output_dir: pathlib.Path,
 38    name: str = "dump",
 39) -> pathlib.Path:
 40    """Write a FASTA-coordinate multi-sample gVCF for the given Population.
 41
 42    Coordinates and bases match what `encode_population_to_fasta` produces:
 43    same XOR mask seed, same 4-letter packing, same physical/logical ordering.
 44    Returns the path to the written gVCF.
 45    """
 46    from aegis_sim import submodels
 47    from aegis_sim.parameterization import parametermanager
 48
 49    output_dir = pathlib.Path(output_dir)
 50    output_dir.mkdir(parents=True, exist_ok=True)
 51
 52    architecture = submodels.architect.architecture
 53    if not hasattr(architecture, "locus_permutation"):
 54        raise RuntimeError("gVCF export only supports the composite architecture")
 55
 56    base_indices_phys = _decode_to_base_indices(population, parametermanager)
 57    # base_indices_phys shape: (n_individuals, ploidy, n_base_positions) in PHYSICAL order
 58    # We emit rows in LOGICAL order so adjacent rows are biologically adjacent loci.
 59    base_indices = _physical_to_logical_bases(
 60        base_indices_phys,
 61        locus_permutation=np.asarray(architecture.locus_permutation, dtype=int),
 62        bits_per_locus=architecture.BITS_PER_LOCUS,
 63    )
 64
 65    n_individuals, ploidy, n_positions = base_indices.shape
 66    if ploidy != 2:
 67        raise ValueError(f"gVCF export expects diploid (ploidy=2), got ploidy={ploidy}")
 68
 69    # Per-position consensus REF + ALT list
 70    refs, alts_per_pos = _consensus_ref_and_alts(base_indices)
 71
 72    sample_names = _sample_names(population, n_individuals)
 73
 74    gvcf_path = output_dir / f"{name}.gvcf"
 75    with open(gvcf_path, "w") as fh:
 76        _write_header(fh, n_positions, sample_names)
 77        _write_body(fh, base_indices, refs, alts_per_pos, sample_names)
 78
 79    return gvcf_path
 80
 81
 82# ----- decode + per-position aggregation -------------------------------------
 83
 84def _decode_to_base_indices(population, parametermanager) -> np.ndarray:
 85    """Replicate the FASTA encoding pipeline up to the point where we know
 86    each individual's bases. Returns int8 array of shape (n_ind, 2, n_bases)
 87    in PHYSICAL locus order, with values in {0,1,2,3} for {A,C,G,T}."""
 88    genome_array = population.genomes.array  # (n, 2, n_loci, bits_per_locus) bool
 89    n_individuals, ploidy, n_loci, bits_per_locus = genome_array.shape
 90
 91    flat = genome_array.reshape(n_individuals, ploidy, -1).astype(np.uint8)
 92    total_bits = flat.shape[-1]
 93    pad = total_bits % 2
 94    if pad:
 95        flat = np.concatenate(
 96            [flat, np.zeros((n_individuals, ploidy, 1), dtype=np.uint8)], axis=-1
 97        )
 98    padded_length = total_bits + pad
 99
100    mask_seed = int(getattr(parametermanager.parameters, "FASTA_MASK_SEED", 0))
101    mask = np.random.default_rng(mask_seed).integers(0, 2, size=padded_length, dtype=np.uint8)
102    masked = flat ^ mask  # broadcasts over (n_ind, ploidy)
103
104    pairs = masked.reshape(n_individuals, ploidy, padded_length // 2, 2)
105    base_indices = (pairs[..., 0] * 2 + pairs[..., 1]).astype(np.int8)
106    return base_indices  # (n_ind, 2, n_bases)
107
108
109def _physical_to_logical_bases(
110    base_indices_phys: np.ndarray,
111    locus_permutation: np.ndarray,
112    bits_per_locus: int,
113) -> np.ndarray:
114    """Reorder bases from physical-locus order to logical-locus order.
115
116    Each logical locus has BITS_PER_LOCUS bits → bits_per_locus/2 bases.
117    """
118    n_individuals, ploidy, n_bases = base_indices_phys.shape
119    bases_per_locus = bits_per_locus // 2
120    if bases_per_locus * bits_per_locus // 2 != bases_per_locus:
121        # Handles odd bits_per_locus by leaving the last base in its physical slot;
122        # for the standard 8-bits-per-locus case, bases_per_locus=4 and this is exact.
123        pass
124    n_loci = n_bases // bases_per_locus
125    by_locus = base_indices_phys[:, :, : n_loci * bases_per_locus].reshape(
126        n_individuals, ploidy, n_loci, bases_per_locus
127    )
128    by_locus = by_locus[:, :, locus_permutation, :]
129    out = by_locus.reshape(n_individuals, ploidy, n_loci * bases_per_locus)
130    # Append any leftover trailing bases (rare; only if bits_per_locus is odd)
131    if n_bases > n_loci * bases_per_locus:
132        out = np.concatenate(
133            [out, base_indices_phys[:, :, n_loci * bases_per_locus :]], axis=-1
134        )
135    return out
136
137
138def _consensus_ref_and_alts(
139    base_indices: np.ndarray,
140) -> Tuple[np.ndarray, List[List[int]]]:
141    """Per-position consensus REF + sorted list of observed non-REF ALTs.
142
143    base_indices is (n_ind, 2, n_bases). Returns:
144      refs:          int8 (n_bases,) — modal base index per position
145      alts_per_pos:  list of lists; alts_per_pos[p] is the sorted ALT indices
146    """
147    n_individuals, ploidy, n_bases = base_indices.shape
148    flat = base_indices.transpose(2, 0, 1).reshape(n_bases, -1)  # (n_bases, n_ind*ploidy)
149
150    refs = np.zeros(n_bases, dtype=np.int8)
151    alts_per_pos: List[List[int]] = []
152    for p in range(n_bases):
153        counts = np.bincount(flat[p], minlength=4)
154        ref = int(np.argmax(counts))
155        refs[p] = ref
156        observed = sorted(int(b) for b in np.unique(flat[p]) if int(b) != ref)
157        alts_per_pos.append(observed)
158    return refs, alts_per_pos
159
160
161# ----- writers ---------------------------------------------------------------
162
163def _sample_names(population, n_individuals: int) -> List[str]:
164    birthdays = getattr(population, "birthdays", None)
165    ages = getattr(population, "ages", None)
166    names = []
167    for i in range(n_individuals):
168        b = int(birthdays[i]) if birthdays is not None else -1
169        a = int(ages[i]) if ages is not None else -1
170        names.append(f"ind_{i}_b{b}_a{a}")
171    return names
172
173
174def _write_header(fh, n_positions: int, sample_names: List[str]) -> None:
175    fh.write("##fileformat=VCFv4.2\n")
176    fh.write("##source=AEGIS\n")
177    fh.write(f"##contig=<ID=aegis_genome,length={n_positions}>\n")
178    fh.write('##ALT=<ID=NON_REF,Description="Represents any possible alternative allele not already in ALT">\n')
179    fh.write('##INFO=<ID=END,Number=1,Type=Integer,Description="End position of the reference block">\n')
180    fh.write('##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">\n')
181    fh.write('##FORMAT=<ID=GQ,Number=1,Type=Integer,Description="Genotype Quality">\n')
182    fh.write('##FORMAT=<ID=DP,Number=1,Type=Integer,Description="Read Depth">\n')
183    fh.write('##FORMAT=<ID=AD,Number=R,Type=Integer,Description="Allele Depth">\n')
184    fh.write("##GVCFBlock=AEGIS-synthetic\n")
185    columns = ["CHROM", "POS", "ID", "REF", "ALT", "QUAL", "FILTER", "INFO", "FORMAT"] + sample_names
186    fh.write("#" + "\t".join(columns) + "\n")
187
188
189def _write_body(
190    fh,
191    base_indices: np.ndarray,
192    refs: np.ndarray,
193    alts_per_pos: List[List[int]],
194    sample_names: List[str],
195) -> None:
196    """Walk positions emitting either a variant record or accumulating a
197    reference block. Variant: any position with a non-empty ALT list.
198    Reference block: a run of consecutive positions all with empty ALTs and
199    same REF — emitted as one record with END=<last>.
200    """
201    n_individuals, ploidy, n_positions = base_indices.shape
202
203    p = 0
204    while p < n_positions:
205        if not alts_per_pos[p]:
206            # Start of (or continuation of) a reference block.
207            block_ref = int(refs[p])
208            block_start = p
209            while p < n_positions and not alts_per_pos[p] and int(refs[p]) == block_ref:
210                p += 1
211            block_end = p  # exclusive
212            _emit_ref_block(
213                fh, start=block_start, end=block_end, ref=block_ref, n_individuals=n_individuals
214            )
215        else:
216            _emit_variant(
217                fh,
218                pos=p,
219                ref=int(refs[p]),
220                alts=alts_per_pos[p],
221                base_indices=base_indices,
222            )
223            p += 1
224
225
226def _emit_ref_block(fh, start: int, end: int, ref: int, n_individuals: int) -> None:
227    """One record covering positions [start, end), all individuals 0/0:99:30:30,0."""
228    pos = start + 1  # 1-indexed POS
229    end_pos = end  # END is 1-indexed inclusive == 0-indexed exclusive end
230    fields = [
231        "aegis_genome",
232        str(pos),
233        ".",
234        BASES[ref],
235        NON_REF,
236        ".",
237        ".",
238        f"END={end_pos}",
239        "GT:GQ:DP:AD",
240    ]
241    sample_strs = ["0/0:99:30:30,0"] * n_individuals
242    fh.write("\t".join(fields + sample_strs) + "\n")
243
244
245def _emit_variant(
246    fh,
247    pos: int,
248    ref: int,
249    alts: List[int],
250    base_indices: np.ndarray,
251) -> None:
252    """One variant record at position pos. Alleles = [REF, *alts, <NON_REF>]."""
253    n_individuals = base_indices.shape[0]
254    all_alleles = [ref, *alts]  # index 0..len(alts) maps to actual base indices
255    alt_strs = [BASES[a] for a in alts] + [NON_REF]
256    ref_str = BASES[ref]
257
258    base_to_allele_idx = {b: i for i, b in enumerate(all_alleles)}
259    n_distinct = len(all_alleles) + 1  # +1 for <NON_REF>
260
261    sample_strs: List[str] = []
262    for i in range(n_individuals):
263        a0 = int(base_indices[i, 0, pos])
264        a1 = int(base_indices[i, 1, pos])
265        # All observed alleles are in all_alleles by construction
266        gt0 = base_to_allele_idx[a0]
267        gt1 = base_to_allele_idx[a1]
268        # AD: one entry per allele in REF + ALTs + NON_REF (1+len(alts)+1 = n_distinct)
269        ad = [0] * n_distinct
270        # Synthetic 15-read AD for each chromatid's allele
271        ad[gt0] += 15
272        ad[gt1] += 15
273        ad_str = ",".join(str(x) for x in ad)
274        sample_strs.append(f"{gt0}/{gt1}:99:30:{ad_str}")
275
276    fields = [
277        "aegis_genome",
278        str(pos + 1),
279        ".",
280        ref_str,
281        ",".join(alt_strs),
282        ".",
283        "PASS",
284        ".",
285        "GT:GQ:DP:AD",
286    ]
287    fh.write("\t".join(fields + sample_strs) + "\n")