How to use the tsinfer.augment_ancestors function in tsinfer

def verify_example(self, subset, samples, ancestors, path_compression):
        ancestors_ts = tsinfer.match_ancestors(
            samples, ancestors, path_compression=path_compression
        )
        augmented_ancestors = tsinfer.augment_ancestors(
            samples, ancestors_ts, subset, path_compression=path_compression
        )
        self.verify_augmented_ancestors(
            subset, ancestors_ts, augmented_ancestors, path_compression
        )

        # Run the inference now
        final_ts = tsinfer.match_samples(samples, augmented_ancestors, simplify=False)
        t1 = ancestors_ts.dump_tables()
        tables = final_ts.tables
        for j, index in enumerate(subset):
            sample_id = final_ts.samples()[index]
            edges = [e for e in final_ts.edges() if e.child == sample_id]
            self.assertEqual(len(edges), 1)
            self.assertEqual(edges[0].left, 0)
            self.assertEqual(edges[0].right, final_ts.sequence_length)

def verify_example(self, full_subset, samples, ancestors, path_compression):
        ancestors_ts = tsinfer.match_ancestors(
            samples, ancestors, path_compression=path_compression
        )
        expected_sample_ancestors = 0
        for j in range(1, len(full_subset)):
            subset = full_subset[:j]
            expected_sample_ancestors += len(subset)
            augmented_ancestors = tsinfer.augment_ancestors(
                samples, ancestors_ts, subset, path_compression=path_compression
            )
            self.verify_augmented_ancestors(
                subset, ancestors_ts, augmented_ancestors, path_compression
            )
            # Run the inference now
            final_ts = tsinfer.match_samples(
                samples, augmented_ancestors, simplify=False
            )

            # Make sure metadata has been preserved in the final ts.
            num_sample_ancestors = 0
            for node in final_ts.nodes():
                if node.flags == tsinfer.NODE_IS_SAMPLE_ANCESTOR:
                    metadata = json.loads(node.metadata.decode())
                    self.assertIn(metadata["sample_data_id"], subset)

setup_logging(args)

    sample_data = tsinfer.SampleData.load(args.samples)
    ancestors_trees = get_ancestors_trees_path(args.ancestors_trees, args.samples)
    output_path = args.augmented_ancestors
    logger.info("Loading ancestral genealogies from {}".format(ancestors_trees))
    ancestors_trees = tskit.load(ancestors_trees)
    progress_monitor = ProgressMonitor(enabled=args.progress, augment_ancestors=True)
    # TODO Need some error checking on these values
    n = args.num_samples
    N = sample_data.num_samples
    if n is None:
        n = int(math.ceil(10 * N / 100))

    sample_indexes = np.linspace(0, N - 1, num=n).astype(int)
    ts = tsinfer.augment_ancestors(
        sample_data,
        ancestors_trees,
        sample_indexes,
        num_threads=args.num_threads,
        path_compression=not args.no_path_compression,
        progress_monitor=progress_monitor,
    )
    logger.info("Writing output tree sequence to {}".format(output_path))
    ts.dump(output_path)
    summarise_usage()

def run_augment(sample_data, ancestors_ts, subset, num_threads):
    progress_monitor = tsinfer.cli.ProgressMonitor(enabled=True, augment_ancestors=True)
    return tsinfer.augment_ancestors(
        sample_data, ancestors_ts, subset, num_threads=num_threads,
        progress_monitor=progress_monitor)