aegis_sim.utilities.vcf

  1"""VCF export/import for AEGIS genomes (composite architecture).
  2
  3Each bit position is treated as a biallelic SNP (REF=A, ALT=T by convention).
  4Each individual is one sample column with phased genotype CHROMATID0|CHROMATID1.
  5Rows are ordered by logical (trait × age × bit-within-locus) position so that
  6adjacent rows are biologically adjacent loci — useful for downstream tools that
  7look at local linkage or sliding-window stats.
  8
  9The VCF header records the `locus_permutation` and architecture metadata so the
 10file is self-contained: `decode_vcf_to_genomes()` can reconstruct the exact
 11multi-dimensional genome array without re-initializing AEGIS.
 12
 13Only the composite architecture is supported here; modifying-architecture export
 14would need a different structural mapping.
 15"""
 16
 17import pathlib
 18from typing import Tuple, List, Dict
 19
 20import numpy as np
 21
 22REF_BASE = "A"
 23ALT_BASE = "T"
 24
 25
 26def encode_population_to_vcf(
 27    population,
 28    output_dir: pathlib.Path,
 29    name: str = "dump",
 30) -> pathlib.Path:
 31    """Write a VCF with one row per genome bit and one sample column per individual.
 32
 33    Pulls architecture metadata (locus_permutation, trait layout) from
 34    `aegis_sim.submodels.architect`, so the architect must be initialized.
 35
 36    Returns the path to the written VCF.
 37    """
 38    from aegis_sim import submodels, parameterization
 39
 40    output_dir = pathlib.Path(output_dir)
 41    output_dir.mkdir(parents=True, exist_ok=True)
 42
 43    architecture = submodels.architect.architecture
 44    if not hasattr(architecture, "locus_permutation"):
 45        raise RuntimeError("VCF export only supports the composite architecture")
 46
 47    genome_array = population.genomes.array
 48    n_individuals, ploidy, n_loci, bits_per_locus = genome_array.shape
 49    if ploidy != 2:
 50        raise ValueError(f"VCF export expects diploid genomes (ploidy=2), got ploidy={ploidy}")
 51
 52    locus_permutation = np.asarray(architecture.locus_permutation, dtype=int)
 53    traits = parameterization.traits
 54
 55    sample_names: List[str] = []
 56    birthdays = getattr(population, "birthdays", None)
 57    ages = getattr(population, "ages", None)
 58    for i in range(n_individuals):
 59        b = int(birthdays[i]) if birthdays is not None else -1
 60        a = int(ages[i]) if ages is not None else -1
 61        sample_names.append(f"ind_{i}_b{b}_a{a}")
 62
 63    vcf_path = output_dir / f"{name}.vcf"
 64    with open(vcf_path, "w") as fh:
 65        fh.write("##fileformat=VCFv4.2\n")
 66        fh.write("##source=AEGIS\n")
 67        fh.write(f"##contig=<ID=aegis_genome,length={n_loci * bits_per_locus}>\n")
 68        fh.write('##INFO=<ID=TRAIT,Number=1,Type=String,Description="Trait name">\n')
 69        fh.write('##INFO=<ID=AGE,Number=1,Type=Integer,Description="Age (0-indexed) for age-specific traits, -1 otherwise">\n')
 70        fh.write('##INFO=<ID=BIT,Number=1,Type=Integer,Description="Bit position within the BITS_PER_LOCUS locus (0-indexed)">\n')
 71        fh.write('##INFO=<ID=PHYS_LOCUS,Number=1,Type=Integer,Description="Physical locus index in storage">\n')
 72        fh.write('##FORMAT=<ID=GT,Number=1,Type=String,Description="Phased genotype: chromatid0|chromatid1">\n')
 73        fh.write(f"##AEGIS_BITS_PER_LOCUS={bits_per_locus}\n")
 74        fh.write(f"##AEGIS_N_LOCI={n_loci}\n")
 75        fh.write(f"##AEGIS_PLOIDY={ploidy}\n")
 76        fh.write(f"##AEGIS_N_INDIVIDUALS={n_individuals}\n")
 77        fh.write(f"##AEGIS_LOCUS_PERMUTATION={','.join(str(x) for x in locus_permutation.tolist())}\n")
 78        fh.write("#" + "\t".join(["CHROM", "POS", "ID", "REF", "ALT", "QUAL", "FILTER", "INFO", "FORMAT"] + sample_names) + "\n")
 79
 80        for logical_locus in range(n_loci):
 81            physical_locus = int(locus_permutation[logical_locus])
 82            trait_name, age = _logical_to_trait_age(logical_locus, traits)
 83            for bit_in_locus in range(bits_per_locus):
 84                pos = logical_locus * bits_per_locus + bit_in_locus + 1  # 1-indexed
 85                variant_id = f"{trait_name}_a{age}_b{bit_in_locus}" if age >= 0 else f"{trait_name}_b{bit_in_locus}"
 86                info = f"TRAIT={trait_name};AGE={age};BIT={bit_in_locus};PHYS_LOCUS={physical_locus}"
 87
 88                # Pull bits for all individuals at this physical position
 89                chr0 = genome_array[:, 0, physical_locus, bit_in_locus].astype(np.int8)
 90                chr1 = genome_array[:, 1, physical_locus, bit_in_locus].astype(np.int8)
 91                # Phased genotype strings
 92                gts = [f"{int(chr0[i])}|{int(chr1[i])}" for i in range(n_individuals)]
 93
 94                row = ["aegis_genome", str(pos), variant_id, REF_BASE, ALT_BASE, ".", "PASS", info, "GT"] + gts
 95                fh.write("\t".join(row) + "\n")
 96
 97    return vcf_path
 98
 99
100def _logical_to_trait_age(logical_locus: int, traits) -> Tuple[str, int]:
101    """Map a logical locus index to (trait_name, age). age=-1 for non-age-specific."""
102    for trait in traits.values():
103        if trait.length == 0:
104            continue
105        if trait.start <= logical_locus < trait.end:
106            age = logical_locus - trait.start if trait.agespecific is True else -1
107            return trait.name, age
108    raise IndexError(f"logical_locus {logical_locus} is outside any trait range")
109
110
111def decode_vcf_to_genomes(vcf_path: pathlib.Path) -> Tuple[np.ndarray, List[str]]:
112    """Parse an AEGIS-emitted VCF back to (genome_array, sample_names).
113
114    The VCF header contains the metadata needed to reconstruct the original
115    (n_individuals, ploidy, n_loci, bits_per_locus) array — no AEGIS init required.
116    """
117    vcf_path = pathlib.Path(vcf_path)
118
119    meta: Dict[str, str] = {}
120    sample_names: List[str] = []
121    data_rows: List[Tuple[int, int, int, List[Tuple[int, int]]]] = []  # phys_locus, bit, _, genotypes
122
123    with open(vcf_path) as fh:
124        for line in fh:
125            line = line.rstrip("\n")
126            if line.startswith("##"):
127                if "=" in line:
128                    key, _, value = line[2:].partition("=")
129                    meta[key] = value
130                continue
131            if line.startswith("#CHROM"):
132                fields = line.lstrip("#").split("\t")
133                sample_names = fields[9:]
134                continue
135            if not line:
136                continue
137            fields = line.split("\t")
138            info = dict(kv.split("=") for kv in fields[7].split(";"))
139            phys_locus = int(info["PHYS_LOCUS"])
140            bit_in_locus = int(info["BIT"])
141            gt_strings = fields[9:]
142            genotypes = []
143            for gt in gt_strings:
144                a, b = gt.split("|")
145                genotypes.append((int(a), int(b)))
146            data_rows.append((phys_locus, bit_in_locus, 0, genotypes))
147
148    n_individuals = int(meta["AEGIS_N_INDIVIDUALS"])
149    ploidy = int(meta["AEGIS_PLOIDY"])
150    n_loci = int(meta["AEGIS_N_LOCI"])
151    bits_per_locus = int(meta["AEGIS_BITS_PER_LOCUS"])
152    assert len(sample_names) == n_individuals
153    assert len(data_rows) == n_loci * bits_per_locus
154
155    genome_array = np.zeros((n_individuals, ploidy, n_loci, bits_per_locus), dtype=np.bool_)
156    for phys_locus, bit_in_locus, _, genotypes in data_rows:
157        for ind_idx, (a, b) in enumerate(genotypes):
158            genome_array[ind_idx, 0, phys_locus, bit_in_locus] = bool(a)
159            genome_array[ind_idx, 1, phys_locus, bit_in_locus] = bool(b)
160
161    return genome_array, sample_names