Add support for terminal site to generate_ancestors with PY engine

Author: duncanMR

tsinfer/algorithm.py

@@ -58,6 +58,7 @@ class Site:
     id = attr.ib()
     time = attr.ib()
     derived_count = attr.ib()
+    terminal = attr.ib()
 
 
 class AncestorBuilder:
@@ -137,21 +138,23 @@ def store_site_genotypes(self, site_id, genotypes):
         stop = start + self.encoded_genotypes_size
         self.genotype_store[start:stop] = genotypes
 
-    def add_site(self, time, genotypes):
+    def add_site(self, time, genotypes, terminal):
         """
         Adds a new site at the specified ID to the builder.
         """
         site_id = len(self.sites)
         derived_count = np.sum(genotypes == 1)
-        self.store_site_genotypes(site_id, genotypes)
-        self.sites.append(Site(site_id, time, derived_count))
-        sites_at_fixed_timepoint = self.time_map[time]
-        # Sites with an identical variant distribution (i.e. with the same
-        # genotypes.tobytes() value) and at the same time, are put into the same ancestor
-        # to which we allocate a unique ID (just use the genotypes value)
-        ancestor_uid = tuple(genotypes)
-        # Add each site to the list for this ancestor_uid at this timepoint
-        sites_at_fixed_timepoint[ancestor_uid].append(site_id)
+        self.sites.append(Site(site_id, time, derived_count, terminal))
+        if not terminal:
+            self.store_site_genotypes(site_id, genotypes)
+            sites_at_fixed_timepoint = self.time_map[time]
+            # Sites with an identical variant distribution (i.e. with the same
+            # genotypes.tobytes() value) and at the same time, are put into the
+            # same ancestor to which we allocate a unique ID (just use the genotypes
+            # value)
+            ancestor_uid = tuple(genotypes)
+            # Add each site to the list for this ancestor_uid at thigs timepoint
+            sites_at_fixed_timepoint[ancestor_uid].append(site_id)
 
     def print_state(self):
         print("Ancestor builder")
@@ -221,6 +224,8 @@ def compute_ancestral_states(self, a, focal_site, sites):
         disagree = np.zeros(self.num_samples, dtype=bool)
 
         for site_index in sites:
+            if self.sites[site_index].terminal:
+                break
             a[site_index] = 0
             last_site = site_index
             g_l = self.get_site_genotypes(site_index)

tsinfer/formats.py

@@ -3093,7 +3093,14 @@ class AncestorData(DataContainer):
     FORMAT_NAME = "tsinfer-ancestor-data"
     FORMAT_VERSION = (3, 0)
 
-    def __init__(self, position, sequence_length, chunk_size_sites=None, **kwargs):
+    def __init__(
+        self,
+        inference_position,
+        terminal_position,
+        sequence_length,
+        chunk_size_sites=None,
+        **kwargs,
+    ):
         super().__init__(**kwargs)
         self._last_time = 0
         self.inference_sites_set = False
@@ -3111,15 +3118,22 @@ def __init__(self, position, sequence_length, chunk_size_sites=None, **kwargs):
         self.create_dataset("sample_end", dtype=np.int32)
         self.create_dataset("sample_time", dtype=np.float64)
         self.create_dataset("sample_focal_sites", dtype="array:i4")
-
+        variant_position = np.concatenate([inference_position, terminal_position])
         self.create_dataset(
             "variant_position",
-            data=position,
-            shape=position.shape,
+            data=variant_position,
+            shape=variant_position.shape,
             chunks=self._chunk_size_sites,
             dtype=np.float64,
             dimensions=["variants"],
         )
+        self.create_dataset(
+            "terminal_position",
+            data=terminal_position,
+            shape=terminal_position.shape,
+            dtype=np.float64,
+            dimensions=["terminal_sites"],
+        )
 
         # We have to include a ploidy dimension sgkit compatibility
         a = self.create_dataset(
@@ -3277,10 +3291,17 @@ def num_sites(self):
     @property
     def sites_position(self):
         """
-        The positions of the inference sites used to generate the ancestors
+        The positions of the inference and terminal sites used to generate the ancestors
         """
         return self.data["variant_position"]
 
+    @property
+    def terminal_position(self):
+        """
+        The positions of the terminal sites used to generate the ancestors
+        """
+        return self.data["terminal_position"]
+
     @property
     def ancestors_start(self):
         return self.data["sample_start"]
@@ -3314,10 +3335,10 @@ def ancestors_length(self):
         """
         # Ancestor start and end are half-closed. The last site is assumed
         # to cover the region up to sequence length.
-        pos = np.hstack([self.sites_position[:], [self.sequence_length]])
+
         start = self.ancestors_start[:]
         end = self.ancestors_end[:]
-        return pos[end] - pos[start]
+        return self.sites_position[end] - self.sites_position[start]
 
     def insert_proxy_samples(
         self,
@@ -3683,6 +3704,7 @@ def add_ancestor(self, start, end, time, focal_sites, haplotype):
         if start < 0:
             raise ValueError("Start must be >= 0")
         if end > self.num_sites:
+            print(f"[INFO] {end}, {self.num_sites}")
             raise ValueError("end must be <= num_sites")
         if start >= end:
             raise ValueError("start must be < end")

tsinfer/inference.py

@@ -1807,6 +1807,8 @@ def __init__(
         self.num_samples = variant_data.num_samples
         self.num_threads = num_threads
         self.mmap_temp_file = None
+        self.sites_position = None
+        self.terminal_position = None
         mmap_fd = -1
 
         genotype_matrix_size = self.max_sites * self.num_samples
@@ -1865,6 +1867,8 @@ def add_sites(self, exclude_positions=None):
         logger.info(f"Starting addition of {self.max_sites} sites")
         progress = self.progress_monitor.get("ga_add_sites", self.max_sites)
         inference_site_id = []
+        last_position = 0
+
         for variant in self.variant_data.variants(recode_ancestral=True):
             # If there's missing data the last allele is None
             num_alleles = len(variant.alleles) - int(variant.alleles[-1] is None)
@@ -1879,6 +1883,7 @@ def add_sites(self, exclude_positions=None):
                 and site.ancestral_state is not None
             ):
                 use_site = True
+                last_position = site.position
                 time = site.time
                 if tskit.is_unknown_time(time):
                     # Non-variable sites have no obvious freq-as-time values
@@ -1888,12 +1893,18 @@ def add_sites(self, exclude_positions=None):
                 if np.isnan(time):
                     use_site = False  # Site with meaningless time value: skip inference
             if use_site:
-                self.ancestor_builder.add_site(time, variant.genotypes)
+                self.ancestor_builder.add_site(time, variant.genotypes, terminal=False)
                 inference_site_id.append(site.id)
                 self.num_sites += 1
             progress.update()
         progress.close()
         self.inference_site_ids = inference_site_id
+        # Add terminal site at end of sequence
+        zeros = np.zeros(self.num_samples, dtype=np.int8)
+        self.ancestor_builder.add_site(tskit.UNKNOWN_TIME, zeros, terminal=True)
+        self.num_sites += 1
+        self.terminal_position = np.array([last_position + 1], dtype=np.float64)
+
         logger.info("Finished adding sites")
 
     def _run_synchronous(self, progress):
@@ -2000,15 +2011,18 @@ def run(self):
             if t not in self.timepoint_to_epoch:
                 self.timepoint_to_epoch[t] = len(self.timepoint_to_epoch) + 1
         self.ancestor_data = formats.AncestorData(
-            self.variant_data.sites_position[:][self.inference_site_ids],
-            self.variant_data.sequence_length,
+            inference_position=self.variant_data.sites_position[:][
+                self.inference_site_ids
+            ],
+            terminal_position=self.terminal_position,
+            sequence_length=self.variant_data.sequence_length,
             path=self.ancestor_data_path,
             **self.ancestor_data_kwargs,
         )
         if self.num_ancestors > 0:
             logger.info(f"Starting build for {self.num_ancestors} ancestors")
             progress = self.progress_monitor.get("ga_generate", self.num_ancestors)
-            a = np.zeros(self.num_sites, dtype=np.int8)
+            a = np.zeros(self.num_sites - 1, dtype=np.int8)
             root_time = max(self.timepoint_to_epoch.keys())
             av_timestep = root_time / len(self.timepoint_to_epoch)
             root_time += av_timestep  # Add a root a bit older than the oldest ancestor
@@ -2017,15 +2031,15 @@ def run(self):
             # line up. It's normally removed when processing the final tree sequence.
             self.ancestor_data.add_ancestor(
                 start=0,
-                end=self.num_sites,
+                end=self.num_sites - 1,
                 time=root_time + av_timestep,
                 focal_sites=np.array([], dtype=np.int32),
                 haplotype=a,
             )
             # This is the the "ultimate ancestor" of all zeros
             self.ancestor_data.add_ancestor(
                 start=0,
-                end=self.num_sites,
+                end=self.num_sites - 1,
                 time=root_time,
                 focal_sites=np.array([], dtype=np.int32),
                 haplotype=a,