aegis_sim.utilities.fasta

  1"""FASTA export/import for AEGIS genomes.
  2
  3Encodes population genomes as DNA-like sequences for input to read simulators
  4(Badread, etc.) and provides a lossless round-trip back to the original bit
  5array via a sidecar mapping file.
  6
  7Encoding: 4-letter packing of 2 bits per base — 00→A, 01→C, 10→G, 11→T.
  8A global uniform-random XOR mask is applied to every individual's flattened
  9genome before packing, to debias base composition (otherwise fit individuals
 10with 1-heavy bitstrings would produce T-heavy sequences). XOR preserves
 11Hamming distance between every pair of individuals exactly, so all population
 12structure relevant to introgression detection is retained.
 13
 14The mapping JSON stores the XOR mask, the original genome shape, and the
 15encoding table — everything needed to invert the transform.
 16"""
 17
 18import json
 19import pathlib
 20from typing import Tuple, List
 21
 22import numpy as np
 23
 24BASES = np.array(["A", "C", "G", "T"])
 25BASE_TO_PAIR = {"A": (0, 0), "C": (0, 1), "G": (1, 0), "T": (1, 1)}
 26ENCODING_NAME = "4-letter-v1"
 27
 28
 29def encode_population_to_fasta(
 30    population,
 31    output_dir: pathlib.Path,
 32    name: str = "dump",
 33    mask_seed: int = 0,
 34    line_width: int = 80,
 35) -> Tuple[pathlib.Path, pathlib.Path]:
 36    """Write population genomes to <output_dir>/<name>.genome.fasta + <name>.mapping.json.
 37
 38    Returns (fasta_path, mapping_path).
 39    """
 40    output_dir = pathlib.Path(output_dir)
 41    output_dir.mkdir(parents=True, exist_ok=True)
 42
 43    genome_array = population.genomes.array
 44    n_individuals = len(genome_array)
 45    original_shape = tuple(int(x) for x in genome_array.shape)
 46
 47    flat = genome_array.reshape(n_individuals, -1).astype(np.uint8)
 48    bits_per_individual = int(flat.shape[1])
 49
 50    pad = bits_per_individual % 2
 51    if pad:
 52        flat = np.concatenate([flat, np.zeros((n_individuals, 1), dtype=np.uint8)], axis=1)
 53    padded_length = bits_per_individual + pad
 54
 55    mask_rng = np.random.default_rng(mask_seed)
 56    mask = mask_rng.integers(0, 2, size=padded_length, dtype=np.uint8)
 57    masked = flat ^ mask[np.newaxis, :]
 58
 59    pairs = masked.reshape(n_individuals, padded_length // 2, 2)
 60    base_indices = pairs[:, :, 0] * 2 + pairs[:, :, 1]
 61    sequences = BASES[base_indices]
 62
 63    birthdays = getattr(population, "birthdays", None)
 64    ages = getattr(population, "ages", None)
 65
 66    fasta_path = output_dir / f"{name}.genome.fasta"
 67    with open(fasta_path, "w") as fh:
 68        for i in range(n_individuals):
 69            birthday = int(birthdays[i]) if birthdays is not None else -1
 70            age = int(ages[i]) if ages is not None else -1
 71            header = f">ind_{i}|birthday={birthday}|age={age}"
 72            fh.write(header + "\n")
 73            seq = "".join(sequences[i].tolist())
 74            for j in range(0, len(seq), line_width):
 75                fh.write(seq[j : j + line_width] + "\n")
 76
 77    mapping = {
 78        "encoding": ENCODING_NAME,
 79        "bit_pair_to_base": {"00": "A", "01": "C", "10": "G", "11": "T"},
 80        "original_shape": list(original_shape),
 81        "bits_per_individual": bits_per_individual,
 82        "padding_bits": pad,
 83        "n_individuals": int(n_individuals),
 84        "xor_mask_hex": mask.tobytes().hex(),
 85        "mask_seed": int(mask_seed),
 86    }
 87    mapping_path = output_dir / f"{name}.mapping.json"
 88    with open(mapping_path, "w") as fh:
 89        json.dump(mapping, fh, indent=2)
 90
 91    return fasta_path, mapping_path
 92
 93
 94def decode_fasta_to_genomes(
 95    fasta_path: pathlib.Path,
 96    mapping_path: pathlib.Path,
 97) -> Tuple[np.ndarray, List[str]]:
 98    """Read FASTA + mapping JSON, return (genome_array, record_ids).
 99
100    genome_array has the same shape and dtype (bool) as the original
101    population.genomes.array.
102    """
103    fasta_path = pathlib.Path(fasta_path)
104    mapping_path = pathlib.Path(mapping_path)
105
106    with open(mapping_path) as fh:
107        mapping = json.load(fh)
108
109    if mapping["encoding"] != ENCODING_NAME:
110        raise ValueError(f"Unsupported encoding {mapping['encoding']!r}; expected {ENCODING_NAME!r}")
111
112    mask = np.frombuffer(bytes.fromhex(mapping["xor_mask_hex"]), dtype=np.uint8).copy()
113    original_shape = tuple(mapping["original_shape"])
114    bits_per_individual = int(mapping["bits_per_individual"])
115    pad = int(mapping["padding_bits"])
116    padded_length = bits_per_individual + pad
117    n_expected = int(mapping["n_individuals"])
118
119    assert mask.shape == (padded_length,), f"mask length {mask.shape} != padded length {padded_length}"
120
121    record_ids: List[str] = []
122    sequences: List[str] = []
123    current_id = None
124    current_chunks: List[str] = []
125    with open(fasta_path) as fh:
126        for line in fh:
127            line = line.strip()
128            if not line:
129                continue
130            if line.startswith(">"):
131                if current_id is not None:
132                    sequences.append("".join(current_chunks))
133                    record_ids.append(current_id)
134                current_id = line[1:]
135                current_chunks = []
136            else:
137                current_chunks.append(line)
138        if current_id is not None:
139            sequences.append("".join(current_chunks))
140            record_ids.append(current_id)
141
142    n_records = len(sequences)
143    assert n_records == n_expected, f"FASTA has {n_records} records but mapping says {n_expected}"
144
145    seq_len = padded_length // 2
146    flat = np.zeros((n_records, padded_length), dtype=np.uint8)
147    for i, seq in enumerate(sequences):
148        if len(seq) != seq_len:
149            raise ValueError(f"Record {record_ids[i]!r} has length {len(seq)}; expected {seq_len}")
150        for j, base in enumerate(seq):
151            b0, b1 = BASE_TO_PAIR[base]
152            flat[i, 2 * j] = b0
153            flat[i, 2 * j + 1] = b1
154
155    flat ^= mask[np.newaxis, :]
156
157    if pad:
158        flat = flat[:, :bits_per_individual]
159
160    decoded = flat.reshape((n_records,) + original_shape[1:]).astype(np.bool_)
161    return decoded, record_ids