fix: don't count mutations at masked sites towards branch length

augur/refine.py

@@ -362,6 +362,21 @@ def run(args):
     if args.divergence_units=='mutations-per-site': #default
         pass
     elif args.divergence_units=='mutations':
+        # find column of the compressed alignment that maps to masked positions of all N
+        compressed_alignment = np.array([x for x in tt.data.compressed_alignment.values()]).T
+        ambiguous_positions = []
+        for ci, col in enumerate(compressed_alignment):
+            if np.unique(col).size==1 and col[0]==tt.gtr.ambiguous:
+                # the column with all Ns has only one unique value, and the first value is N
+                ambiguous_positions.extend(tt.data.compressed_to_full_sequence_map[ci])
+
+                # Since TreeTime represents all columns of all Ns with a single
+                # entry in its compressed alignment structure, we can stop
+                # searching the data structure.
+                break
+        # convert to set for faster lookup
+        ambiguous_positions = set(ambiguous_positions)
+
         if not args.timetree:
             # infer ancestral sequence for the purpose of counting mutations
             # sample mutations from the root profile, otherwise use most likely state.
@@ -385,7 +400,7 @@ def are_sequence_states_different(nuc1, nuc2):
             n_muts = len([
                 position
                 for ancestral, position, derived in node.mutations
-                if are_sequence_states_different(ancestral, derived)
+                if are_sequence_states_different(ancestral, derived) and position not in ambiguous_positions
             ])
 
             if args.timetree:

tests/functional/refine/cram/not-timetree-mutations-with-ambiguous-column.t

@@ -0,0 +1,34 @@
+Setup
+
+  $ source "$TESTDIR"/_setup.sh
+
+Run refine without a time tree, but request number of mutations per branch as the divergence unit.
+This approach only works when we provide an alignment FASTA.
+The given alignment has a column of all Ns which should get excluded from branch length calculations as an ambiguous column.
+
+  $ ${AUGUR} refine \
+  >  --tree "$TESTDIR/../data/tree_raw.nwk" \
+  >  --alignment "$TESTDIR/../data/aligned_with_ambiguous_column.fasta" \
+  >  --metadata "$TESTDIR/../data/metadata.tsv" \
+  >  --output-tree tree.nwk \
+  >  --output-node-data branch_lengths.json \
+  >  --coalescent opt \
+  >  --date-confidence \
+  >  --date-inference marginal \
+  >  --clock-filter-iqd 4 \
+  >  --seed 314159 \
+  >  --divergence-units mutations &> /dev/null
+
+Confirm that trees match expected topology and branch lengths, given that the output should not be a time tree.
+
+  $ python3 "$TESTDIR/../../../../scripts/diff_trees.py" \
+  >   "$TESTDIR/../data/not_time_tree.nwk" \
+  >   tree.nwk \
+  >   --significant-digits 2
+  {}
+  $ python3 "$TESTDIR/../../../../scripts/diff_jsons.py" \
+  >   "$TESTDIR/../data/integer_branch_lengths_with_ambiguous_column.json" \
+  >   branch_lengths.json \
+  >   --significant-digits 0 \
+  >   --exclude-paths "root['generated_by']['version']" "root['input_tree']" "root['alignment']"
+  {}