Remove daily-extend and add basic bash run script

Author: jeromekelleher

run.sh

@@ -0,0 +1,41 @@
+#!/bin/bash
+set -e
+set -u
+# set -x
+
+mismatches=3
+max_daily_samples=1000
+num_threads=8
+
+# Paths
+datadir=testrun
+run_id=tmp-dev
+# run_id=upgma-mds-$max_daily_samples-md-$max_submission_delay-mm-$mismatches
+resultsdir=results/$run_id
+results_prefix=$resultsdir/$run_id-
+logfile=logs/$run_id.log
+
+alignments=$datadir/alignments.db
+metadata=$datadir/metadata.db
+matches=$resultsdir/matces.db
+
+dates=`python3 -m sc2ts list-dates $metadata | grep -v 2021-12-31`
+echo $dates
+
+options="--num-threads $num_threads -vv -l $logfile "
+# options+="--max-submission-delay $max_submission_delay "
+# options+="--max-daily-samples $max_daily_samples "
+options+="--num-mismatches $mismatches"
+
+mkdir -p $resultsdir
+
+last_ts=$resultsdir/initial.ts
+
+python3 -m sc2ts initialise $last_ts $matches
+
+for date in $dates; do
+    out_ts="$results_prefix$date".ts
+    python3 -m sc2ts extend $last_ts $date $alignments $metadata \
+        $matches $out_ts $options
+    last_ts=$out_ts
+done

sc2ts/alignments.py

@@ -82,6 +82,7 @@ def __init__(self, path, mode="r"):
         self.env = lmdb.Environment(
             str(path), subdir=False, readonly=mode == "r", map_size=map_size
         )
+        logger.debug(f"Opened AlignmentStore at {path} mode={mode}")
 
     def __enter__(self):
         return self

sc2ts/cli.py

@@ -234,6 +234,132 @@ def list_dates(metadata, counts, verbose, log_file):
                 print(k)
 
 
+@click.command()
+@click.argument("base_ts", type=click.Path(exists=True, dir_okay=False))
+@click.argument("date")
+@click.argument("alignments", type=click.Path(exists=True, dir_okay=False))
+@click.argument("metadata", type=click.Path(exists=True, dir_okay=False))
+@click.argument("matches", type=click.Path(exists=True, dir_okay=False))
+@click.argument("output_ts", type=click.Path(dir_okay=False))
+@click.option(
+    "--num-mismatches",
+    default=3,
+    show_default=True,
+    type=float,
+    help="Number of mismatches to accept in favour of recombination",
+)
+@click.option(
+    "--hmm-cost-threshold",
+    default=5,
+    type=float,
+    show_default=True,
+    help="The maximum HMM cost for samples to be included unconditionally",
+)
+@click.option(
+    "--min-group-size",
+    default=10,
+    show_default=True,
+    type=int,
+    help="Minimum size of groups of reconsidered samples for inclusion",
+)
+@click.option(
+    "--retrospective-window",
+    default=30,
+    show_default=True,
+    type=int,
+    help="Number of days in the past to reconsider potential matches",
+)
+@click.option(
+    "--max-daily-samples",
+    default=None,
+    type=int,
+    help=(
+        "The maximum number of samples to match in a single day. If the total "
+        "is greater than this, randomly subsample."
+    ),
+)
+@click.option(
+    "--random-seed",
+    default=42,
+    type=int,
+    help="Random seed for subsampling",
+    show_default=True,
+)
+@click.option(
+    "--num-threads",
+    default=0,
+    type=int,
+    help="Number of match threads (default to one)",
+)
+@click.option("--progress/--no-progress", default=True)
+@click.option("-v", "--verbose", count=True)
+@click.option("-l", "--log-file", default=None, type=click.Path(dir_okay=False))
+@click.option(
+    "-f",
+    "--force",
+    is_flag=True,
+    flag_value=True,
+    help="Force clearing newer matches from DB",
+)
+def extend(
+    base_ts,
+    date,
+    alignments,
+    metadata,
+    matches,
+    output_ts,
+    num_mismatches,
+    hmm_cost_threshold,
+    min_group_size,
+    retrospective_window,
+    max_daily_samples,
+    num_threads,
+    random_seed,
+    progress,
+    verbose,
+    log_file,
+    force,
+):
+    """
+    Extend base_ts with sequences for the specified date, using specified
+    alignments and metadata databases, updating the specified matches
+    database, and outputting the result to the specified file.
+    """
+    setup_logging(verbose, log_file)
+    base = tskit.load(base_ts)
+    logger.debug(f"Loaded base ts from {base_ts}")
+    with contextlib.ExitStack() as exit_stack:
+        alignment_store = exit_stack.enter_context(sc2ts.AlignmentStore(alignments))
+        metadata_db = exit_stack.enter_context(sc2ts.MetadataDb(metadata))
+        match_db = exit_stack.enter_context(sc2ts.MatchDb(matches))
+
+        newer_matches = match_db.count_newer(date)
+        if newer_matches > 0:
+            if not force:
+                click.confirm(
+                    f"Do you want to remove {newer_matches} newer matches "
+                    f"from MatchDB > {date}?",
+                    abort=True,
+                )
+                match_db.delete_newer(date)
+        ts_out = inference.extend(
+            alignment_store=alignment_store,
+            metadata_db=metadata_db,
+            base_ts=base,
+            date=date,
+            match_db=match_db,
+            num_mismatches=num_mismatches,
+            hmm_cost_threshold=hmm_cost_threshold,
+            min_group_size=min_group_size,
+            retrospective_window=retrospective_window,
+            max_daily_samples=max_daily_samples,
+            random_seed=random_seed,
+            num_threads=num_threads,
+            show_progress=progress,
+        )
+        add_provenance(ts_out, output_ts)
+
+
 @click.command()
 @click.argument("alignments", type=click.Path(exists=True, dir_okay=False))
 @click.argument("metadata", type=click.Path(exists=True, dir_okay=False))
@@ -531,6 +657,7 @@ def cli():
 
 cli.add_command(initialise)
 cli.add_command(list_dates)
+cli.add_command(extend)
 cli.add_command(daily_extend)
 cli.add_command(validate)
 cli.add_command(annotate_recombinants)

sc2ts/inference.py

@@ -32,6 +32,7 @@ def __init__(self, path):
         self.uri = uri
         self.conn = sqlite3.connect(uri, uri=True)
         self.conn.row_factory = metadata.dict_factory
+        logger.debug(f"Opened MatchDb at {path} mode=rw")
 
     def __len__(self):
         sql = "SELECT COUNT(*) FROM samples"
@@ -52,6 +53,17 @@ def last_date(self):
             row = self.conn.execute(sql).fetchone()
             return row["MAX(match_date)"]
 
+    def count_newer(self, date):
+        with self.conn:
+            sql = "SELECT COUNT(*) FROM samples WHERE match_date >= ?"
+            row = self.conn.execute(sql, (date,)).fetchone()
+            return row["COUNT(*)"]
+
+    def delete_newer(self, date):
+        sql = "DELETE FROM samples WHERE match_date >= ?"
+        with self.conn:
+            self.conn.execute(sql, (date,))
+
     def __str__(self):
         return f"MatchDb at {self.uri} has {len(self)} samples"
 
@@ -376,49 +388,49 @@ def asdict(self):
         }
 
 
-def daily_extend(
-    *,
-    alignment_store,
-    metadata_db,
-    base_ts,
-    match_db,
-    num_mismatches=None,
-    max_hmm_cost=None,
-    min_group_size=None,
-    num_past_days=None,
-    show_progress=False,
-    max_submission_delay=None,
-    max_daily_samples=None,
-    num_threads=None,
-    precision=None,
-    rng=None,
-    excluded_sample_dir=None,
-):
-    assert num_past_days is None
-    assert max_submission_delay is None
-
-    start_day = last_date(base_ts)
-
-    last_ts = base_ts
-    for date in metadata_db.get_days(start_day):
-        ts = extend(
-            alignment_store=alignment_store,
-            metadata_db=metadata_db,
-            date=date,
-            base_ts=last_ts,
-            match_db=match_db,
-            num_mismatches=num_mismatches,
-            max_hmm_cost=max_hmm_cost,
-            min_group_size=min_group_size,
-            show_progress=show_progress,
-            max_submission_delay=max_submission_delay,
-            max_daily_samples=max_daily_samples,
-            num_threads=num_threads,
-            precision=precision,
-        )
-        yield ts, date
-
-        last_ts = ts
+# def daily_extend(
+#     *,
+#     alignment_store,
+#     metadata_db,
+#     base_ts,
+#     match_db,
+#     num_mismatches=None,
+#     max_hmm_cost=None,
+#     min_group_size=None,
+#     num_past_days=None,
+#     show_progress=False,
+#     max_submission_delay=None,
+#     max_daily_samples=None,
+#     num_threads=None,
+#     precision=None,
+#     rng=None,
+#     excluded_sample_dir=None,
+# ):
+#     assert num_past_days is None
+#     assert max_submission_delay is None
+
+#     start_day = last_date(base_ts)
+
+#     last_ts = base_ts
+#     for date in metadata_db.get_days(start_day):
+#         ts = extend(
+#             alignment_store=alignment_store,
+#             metadata_db=metadata_db,
+#             date=date,
+#             base_ts=last_ts,
+#             match_db=match_db,
+#             num_mismatches=num_mismatches,
+#             max_hmm_cost=max_hmm_cost,
+#             min_group_size=min_group_size,
+#             show_progress=show_progress,
+#             max_submission_delay=max_submission_delay,
+#             max_daily_samples=max_daily_samples,
+#             num_threads=num_threads,
+#             precision=precision,
+#         )
+#         yield ts, date
+
+#         last_ts = ts
 
 
 def preprocess(
@@ -559,30 +571,28 @@ def extend(
     base_ts,
     match_db,
     num_mismatches=None,
-    max_hmm_cost=None,
+    hmm_cost_threshold=None,
     min_group_size=None,
-    show_progress=False,
-    max_submission_delay=None,
     max_daily_samples=None,
+    show_progress=False,
+    retrospective_window=None,
+    random_seed=42,
     num_threads=0,
-    precision=None,
-    rng=None,
 ):
     if num_mismatches is None:
         num_mismatches = 3
-    if max_hmm_cost is None:
-        max_hmm_cost = 5
+    if hmm_cost_threshold is None:
+        hmm_cost_threshold = 5
     if min_group_size is None:
         min_group_size = 10
 
+    # TMP
+    precision = 6
     check_base_ts(base_ts)
     logger.info(
         f"Extend {date}; ts:nodes={base_ts.num_nodes};samples={base_ts.num_samples};"
         f"mutations={base_ts.num_mutations};date={base_ts.metadata['sc2ts']['date']}"
     )
-    # TODO not sure whether we'll keep these params. Making sure they're not
-    # used for now
-    assert max_submission_delay is None
 
     samples = preprocess(
         date,
@@ -615,7 +625,7 @@ def extend(
 
     logger.info(f"Update ARG with low-cost samples for {date}")
     ts = add_matching_results(
-        f"match_date=='{date}' and hmm_cost>0 and hmm_cost<={max_hmm_cost}",
+        f"match_date=='{date}' and hmm_cost>0 and hmm_cost<={hmm_cost_threshold}",
         ts=ts,
         match_db=match_db,
         date=date,

sc2ts/metadata.py

@@ -22,6 +22,7 @@ def __init__(self, path):
         self.path = path
         self.conn = sqlite3.connect(uri, uri=True)
         self.conn.row_factory = dict_factory
+        logger.debug(f"Opened MetadataDb at {path} mode=ro")
 
     @staticmethod
     def import_csv(csv_path, db_path):