Repository: cogent3/cogent3
Pull request #2096 (2024-11-05 00:18:10)

More control over annotation display with dotplots
## Summary by Sourcery

Enhance the dotplot functionality by adding support for biotype filtering, allowing users to control annotation display based on biotype. Refactor the dotplot function to accommodate this feature and add corresponding tests to verify its correct operation.

New Features:
- Introduce the ability to filter dotplots by biotype, allowing users to display only specific biotypes when sequences are annotated.

Enhancements:
- Refactor the dotplot function to support biotype filtering, improving the flexibility of annotation display in dotplots.

Tests:
- Add tests for the new biotype filtering feature in dotplots, ensuring correct functionality when sequences are annotated with specific biotypes.


Patch:
diff --git a/src/cogent3/core/new_alignment.py b/src/cogent3/core/new_alignment.py
index 41b68d8c6..d115b9c97 100644
--- a/src/cogent3/core/new_alignment.py
+++ b/src/cogent3/core/new_alignment.py
@@ -1342,6 +1342,7 @@ class SequenceCollection:
         width: int = 500,
         title: OptStr = None,
         rc: bool = False,
+        biotype: typing.Union[str, tuple[str]] = "gene",
         show_progress: bool = False,
     ):
         """make a dotplot between specified sequences. Random sequences
@@ -1371,6 +1372,8 @@ class SequenceCollection:
         rc
             include dotplot of reverse compliment also. Only applies to Nucleic
             acids moltypes
+        biotype
+            if selected sequences are annotated, display only these biotypes
 
         Returns
         -------
@@ -1398,16 +1401,9 @@ class SequenceCollection:
 
         seq1 = self.get_seq(seqname=name1, copy_annotations=False)
         seq2 = self.get_seq(seqname=name2, copy_annotations=False)
-
-        if seq1.is_annotated() or seq2.is_annotated():
-            annotated = True
-            data = getattr(seq1, "data", seq1)
-            bottom = data.get_drawable()
-            data = getattr(seq2, "data", seq2)
-            left = data.get_drawable(vertical=True)
-        else:
-            annotated = False
-
+        annotated = seq1.is_annotated(biotype=biotype) or seq2.is_annotated(
+            biotype=biotype
+        )
         dotplot = Dotplot(
             seq1,
             seq2,
@@ -1426,6 +1422,10 @@ class SequenceCollection:
         )
 
         if annotated:
+            data = getattr(seq1, "data", seq1)
+            bottom = data.get_drawable(biotype=biotype)
+            data = getattr(seq2, "data", seq2)
+            left = data.get_drawable(biotype=biotype, vertical=True)
             dotplot = AnnotatedDrawable(
                 dotplot,
                 left_track=left,
@@ -1436,6 +1436,7 @@ class SequenceCollection:
                 xrange=[0, len(seq1)],
                 yrange=[0, len(seq2)],
             )
+
         return dotplot
 
     @UI.display_wrap
diff --git a/src/cogent3/core/new_sequence.py b/src/cogent3/core/new_sequence.py
index 9eeacd0da..6efa1be68 100644
--- a/src/cogent3/core/new_sequence.py
+++ b/src/cogent3/core/new_sequence.py
@@ -1529,10 +1529,22 @@ class Sequence:
         seq.annotation_db = self.annotation_db
         return indel_map, seq
 
-    def is_annotated(self):
-        """returns True if sequence parent name has any annotations"""
+    def is_annotated(
+        self, biotype: typing.Optional[typing.Union[str, tuple[str]]] = None
+    ) -> bool:
+        """returns True if sequence parent name has any annotations
+
+        Parameters
+        ----------
+        biotype
+            amend condition to return True only if the sequence is
+            annotated with one of provided biotypes.
+        """
         with contextlib.suppress(AttributeError):
-            return self.annotation_db.num_matches(seqid=self._seq.seqid) != 0
+            return (
+                self.annotation_db.num_matches(seqid=self._seq.seqid, biotype=biotype)
+                != 0
+            )
         return False
 
     def annotate_matches_to(


Test patch:
diff --git a/tests/test_core/test_new_alignment.py b/tests/test_core/test_new_alignment.py
index 9e98df901..108fb2824 100644
--- a/tests/test_core/test_new_alignment.py
+++ b/tests/test_core/test_new_alignment.py
@@ -7,7 +7,7 @@ from warnings import catch_warnings, filterwarnings
 import numpy
 import pytest
 
-from cogent3 import get_app, load_seq, load_unaligned_seqs, open_
+from cogent3 import get_app, load_unaligned_seqs, open_
 from cogent3._version import __version__
 from cogent3.core import new_alignment, new_alphabet, new_moltype, new_sequence
 from cogent3.core.annotation import Feature
diff --git a/tests/test_core/test_new_sequence.py b/tests/test_core/test_new_sequence.py
index 66df2bef3..c8ec41f62 100644
--- a/tests/test_core/test_new_sequence.py
+++ b/tests/test_core/test_new_sequence.py
@@ -872,6 +872,15 @@ def test_is_annotated():
     assert s.is_annotated()
 
 
+@pytest.mark.parametrize("biotype", ("gene", "exon", ("gene", "exon")))
+def test_is_annotated_biotype(biotype):
+    """is_annotated operates correctly"""
+    s = new_moltype.DNA.make_seq(seq="ACGGCTGAAGCGCTCCGGGTTTAAAACG", name="s1")
+    _ = s.add_feature(biotype="gene", name="blah", spans=[(0, 10)])
+    _ = s.add_feature(biotype="exon", name="blah", spans=[(0, 10)])
+    assert s.is_annotated(biotype=biotype)
+
+
 def test_not_is_annotated():
     """is_annotated operates correctly"""
     s = new_moltype.DNA.make_seq(seq="ACGGCTGAAGCGCTCCGGGTTTAAAACG", name="s1")
@@ -881,6 +890,11 @@ def test_not_is_annotated():
         seqid="s2", biotype="gene", name="blah", spans=[(0, 10)]
     )
     assert not s.is_annotated()
+    # annotation wrong biotype
+    s.annotation_db.add_feature(
+        seqid="s1", biotype="exon", name="blah", spans=[(0, 10)]
+    )
+    assert not s.is_annotated(biotype="gene")
     s.annotation_db = None
     assert not s.is_annotated()
 
diff --git a/tests/test_draw/test_new_draw_integration.py b/tests/test_draw/test_new_draw_integration.py
index 8b5980ed3..b5e4ac90e 100644
--- a/tests/test_draw/test_new_draw_integration.py
+++ b/tests/test_draw/test_new_draw_integration.py
@@ -224,6 +224,39 @@ def test_annotated_dotplot_remove_tracks(load_alignment):
     assert isinstance(dp.figure, UnionDict)
 
 
+def test_sequence_collection_dotplot_annotated_no_matcing_seqids():
+    """has annotation_db but not for a selected sequence, returns plain DotPlot"""
+    db = GffAnnotationDb()
+    # dog not present
+    db.add_feature(seqid="Dog", biotype="exon", name="fred", spans=[(10, 15)])
+    data = {"Human": "CAGATTTGGCAGTT-", "Mouse": "CAGATTCAGCAGGTG"}
+    seqs = new_alignment.make_unaligned_seqs(data, moltype="dna")
+    seqs.annotation_db = db
+    seqs = seqs.take_seqs(["Human", "Mouse"], copy_annotations=True)
+    got = seqs.dotplot(show_progress=False, biotype="exon")
+    assert not isinstance(got, AnnotatedDrawable)
+    assert isinstance(got, Dotplot)
+
+
+def test_sequence_collection_dotplot_annotated_select_biotype():
+    """has annotation_db but not for a selected sequence, returns plain DotPlot"""
+    db = GffAnnotationDb()
+    # dog not present
+    db.add_feature(seqid="Human", biotype="exon", name="fred", spans=[(10, 15)])
+    data = {"Human": "CAGATTTGGCAGTT-", "Mouse": "CAGATTCAGCAGGTG"}
+    seqs = new_alignment.make_unaligned_seqs(data, moltype="dna")
+    seqs.annotation_db = db
+    seqs = seqs.take_seqs(["Human", "Mouse"])
+    got = seqs.dotplot(show_progress=False, biotype="gene")
+    assert not isinstance(got, AnnotatedDrawable)
+    assert isinstance(got, Dotplot)
+    # but if we select an available biotype
+    got = seqs.dotplot(show_progress=False, biotype="exon")
+    assert isinstance(got, AnnotatedDrawable)
+    # we get one trace with that biotype as the name
+    assert sum(tr.name == "exon" for tr in got.figure.data)
+
+
 def test_count_gaps_per_seq(load_alignment):
     """creation of drawables works"""
     styles = "bar", "box", "violin"