Add is_reference option to AncestralState config

When is_reference is true, the REF allele (variant_allele[:, 0]) from the
VCZ store is used as the ancestral allele instead of reading a named field.
Useful for simulated data where REF equals the ancestral allele.

Exactly one of field or is_reference must be set; both default to None.

The genotype roundtrip tests now use is_reference=True throughout, since
all test inputs originate from tree sequences where REF is ancestral.

Author: jeromekelleher

docs/config.md

@@ -33,7 +33,8 @@ Specifies where to read the ancestral allele for each variant position.
 | Field | Type | Default | Description |
 |-------|------|---------|-------------|
 | `path` | string | (required) | Path to VCZ containing ancestral alleles |
-| `field` | string | (required) | Array name in the store (e.g. `"variant_AA"`) |
+| `field` | string | — | Array name in the store (e.g. `"variant_AA"`). Required unless `is_reference` is set. |
+| `is_reference` | bool | `false` | Use the REF allele (`variant_allele[:, 0]`) as the ancestral state. Useful for simulations. `field` must not be set when this is `true`. |
 
 
 ## `[[ancestors]]`

docs/quickstart.md

@@ -28,8 +28,9 @@ vcf2zarr convert mydata.vcf.gz mydata.vcz
 Each site used for inference requires a known **ancestral allele**. If your VCF
 has an `AA` INFO field, `vcf2zarr` stores it as `variant_AA` in the `.vcz`
 store and you can reference it directly in the config. Alternatively, ancestral
-alleles can come from a separate VCZ store. See the
-{ref}`config reference <sec_config_reference>` for details.
+alleles can come from a separate VCZ store. For simulated data where the REF
+allele is the ancestral allele, set `is_reference = true` instead of specifying
+a field. See the {ref}`config reference <sec_config_reference>` for details.
 
 
 ## Writing the config

example_config.toml

@@ -59,6 +59,11 @@ path = "data/1kgp_chr20.vcz"
 path = "data/homo_sapiens-chr20.vcz"
 field = "variant_AA"
 
+# For simulated data where REF is the ancestral allele, use:
+# [ancestral_state]
+# path = "data/simulated.vcz"
+# is_reference = true
+
 
 # ============================================================================
 # Ancestors

tests/test_genotype_roundtrip.py

@@ -39,12 +39,13 @@
 # ---------------------------------------------------------------------------
 
 
-def _anc_state(store):
-    return config.AncestralState(path=store, field="variant_ancestral_allele")
-
-
 def _run_pipeline(sample_store):
-    """Build config, run full pipeline, return output tree sequence."""
+    """Build config, run full pipeline, return output tree sequence.
+
+    Uses ``is_reference=True`` so the REF allele (variant_allele[:, 0])
+    is treated as the ancestral allele — the natural choice when input
+    data originates from a tree sequence.
+    """
     src = config.Source(path=sample_store, name="test")
     anc_src = config.Source(path=None, name="ancestors", sample_time="sample_time")
     cfg = config.Config(
@@ -62,7 +63,7 @@ def _run_pipeline(sample_store):
             output="output.trees",
         ),
         post_process=config.PostProcessConfig(),
-        ancestral_state=_anc_state(sample_store),
+        ancestral_state=config.AncestralState(path=sample_store, is_reference=True),
     )
     return pipeline.run(cfg)
 
@@ -103,7 +104,7 @@ def _run_pipeline_with_augment(sample_store, augment_store=None, ann_store=None)
         ),
         post_process=config.PostProcessConfig(),
         augment_sites=config.AugmentSitesConfig(sources=["augment"]),
-        ancestral_state=_anc_state(ann_store),
+        ancestral_state=config.AncestralState(path=ann_store, is_reference=True),
     )
     return pipeline.run(cfg)
 

tests/test_pipeline.py

@@ -355,6 +355,35 @@ def test_hand_constructed(self):
         assert out_ts.num_sites == 2
 
 
+# ---------------------------------------------------------------------------
+# TestIsReferenceConfig
+# ---------------------------------------------------------------------------
+
+
+class TestIsReferenceConfig:
+    def test_is_reference_with_field_raises(self):
+        """is_reference=True with field set raises ValueError."""
+        with pytest.raises(ValueError, match="field must not be set"):
+            config.AncestralState(path="x.vcz", field="variant_AA", is_reference=True)
+
+    def test_neither_field_nor_is_reference_raises(self):
+        """Neither field nor is_reference raises ValueError."""
+        with pytest.raises(ValueError, match="requires either"):
+            config.AncestralState(path="x.vcz")
+
+    def test_valid_field(self):
+        """field without is_reference is valid."""
+        a = config.AncestralState(path="x.vcz", field="variant_AA")
+        assert a.field == "variant_AA"
+        assert a.is_reference is None
+
+    def test_valid_is_reference(self):
+        """is_reference=True without field is valid."""
+        a = config.AncestralState(path="x.vcz", is_reference=True)
+        assert a.is_reference is True
+        assert a.field is None
+
+
 # ---------------------------------------------------------------------------
 # TestNodeMetadata
 # ---------------------------------------------------------------------------

tsinfer/config.py

@@ -113,10 +113,28 @@ def __post_init__(self):
 
 @dataclasses.dataclass
 class AncestralState:
-    """Specifies where to read the ancestral allele for each variant position."""
+    """Specifies where to read the ancestral allele for each variant position.
+
+    Exactly one of *field* or *is_reference* must be set.
+
+    If *is_reference* is ``True``, the REF allele (``variant_allele[:, 0]``)
+    from the store at *path* is used as the ancestral allele.  Otherwise
+    *field* names the array to read.
+    """
 
     path: str | pathlib.Path
-    field: str
+    field: str | None = None
+    is_reference: bool | None = None
+
+    def __post_init__(self):
+        if self.is_reference is None and self.field is None:
+            raise ValueError(
+                "[ancestral_state] requires either 'field' or 'is_reference = true'"
+            )
+        if self.is_reference is True and self.field is not None:
+            raise ValueError(
+                "[ancestral_state] field must not be set when is_reference is true"
+            )
 
 
 @dataclasses.dataclass
@@ -356,7 +374,7 @@ def from_toml(cls, path: str | pathlib.Path) -> Config:
     "sample_time",
 }
 
-_KNOWN_ANCESTRAL_STATE_KEYS = {"path", "field"}
+_KNOWN_ANCESTRAL_STATE_KEYS = {"path", "field", "is_reference"}
 
 _KNOWN_ANCESTORS_KEYS = {
     "name",
@@ -429,13 +447,13 @@ def _parse_ancestral_state(raw: dict) -> AncestralState:
     if entry is None:
         raise ValueError("Config must contain an [ancestral_state] section")
     _check_unknown_keys("ancestral_state", entry, _KNOWN_ANCESTRAL_STATE_KEYS)
-    try:
-        return AncestralState(
-            path=_resolve_path(entry["path"]),
-            field=entry["field"],
-        )
-    except KeyError as e:
-        raise ValueError(f"[ancestral_state] missing required key: {e}") from e
+    if "path" not in entry:
+        raise ValueError("[ancestral_state] missing required key: 'path'")
+    return AncestralState(
+        path=_resolve_path(entry["path"]),
+        field=entry.get("field"),
+        is_reference=entry.get("is_reference"),
+    )
 
 
 def _parse_one_ancestor(entry: dict) -> AncestorsConfig:

tsinfer/vcz.py

@@ -2368,7 +2368,10 @@ def __init__(
         # --- Ancestral state ---
         ann_store = open_store(ancestral_state.path)
         ann_positions = np.asarray(ann_store["variant_position"][:], dtype=np.int32)
-        ann_values = np.asarray(ann_store[ancestral_state.field][:])
+        if ancestral_state.is_reference:
+            ann_values = np.asarray(ann_store["variant_allele"][:, 0])
+        else:
+            ann_values = np.asarray(ann_store[ancestral_state.field][:])
 
         # --- Unified site set (all numpy) ---
         valid_per_source = []