Refactor: HOGs store the name of their taxonomic level, not the ete3 node

FastOMA/_hog_class.py

@@ -89,14 +89,12 @@ def _count_untouched_clades(tree_node: TreeNode, reached_nodes: list) -> int:
 def _hog_implied_losses(hog: "HOG", tree_node: TreeNode) -> int:
     """Recursively counts implied gene losses for `hog` within `tree_node`'s subtree.
 
-    Mirrors _tax_overlap's _tax_now-based grouping of _subhogs, but sums losses per group member
-    instead of intersecting/folding lineages: a group of sub-hogs sharing the same _tax_now
+    Mirrors _tax_overlap's taxname-based grouping of _subhogs, but sums losses per group member
+    instead of intersecting/folding lineages: a group of sub-hogs sharing the same taxname
     represents a duplication, and each such paralogous copy's losses are counted independently
     so that one copy's presence in a species can't mask a sibling copy's loss there.
 
-    The sub-hogs' taxonomic levels are looked up by name in `tree_node`'s tree: their `_tax_now`
-    nodes belong to other tree objects (the pruned working tree, and a separate copy per pickled
-    level), and ete3 nodes only compare by identity."""
+    HOGs only store the name of their taxonomic level, which is looked up in `tree_node`'s tree."""
     subhogs = getattr(hog, "_subhogs", None)
     if not subhogs:
         return 0
@@ -105,7 +103,7 @@ def _hog_implied_losses(hog: "HOG", tree_node: TreeNode) -> int:
     losses = 0
     clade_groups = {}
     for sh in subhogs:
-        clade_groups.setdefault(sh._tax_now.name, []).append(sh)
+        clade_groups.setdefault(sh.taxname, []).append(sh)
     for clade_name, subhogs_of_clade in clade_groups.items():
         clade_node = index[clade_name]
         for sh in subhogs_of_clade:
@@ -159,8 +157,8 @@ def _combine_tax_overlap(parts):
 
 def _tax_overlap(hog: "HOG", species_lineage_index: dict):
     """Recursively computes (nset, leaf_size, leaf_acc, tax_score) for `hog`, mirroring the same
-    _tax_now-based grouping of _subhogs used by HOG.to_orthoxml() (a group of sub-hogs sharing a
-    _tax_now is a duplication -- its members are folded together as siblings, same as a group of
+    taxname-based grouping of _subhogs used by HOG.to_orthoxml() (a group of sub-hogs sharing a
+    taxname is a duplication -- its members are folded together as siblings, same as a group of
     literal gene-tree children would be). `leaf_size` is read off len(hog).
     See attach_scores for how the four returned values become TCSScore."""
     if not hog._subhogs:
@@ -169,7 +167,7 @@ def _tax_overlap(hog: "HOG", species_lineage_index: dict):
 
     groups = []
     for _, subhogs_of_clade in itertools.groupby(
-            sorted(hog._subhogs, key=lambda h: h._tax_now.name), key=lambda h: h._tax_now.name):
+            sorted(hog._subhogs, key=lambda h: h.taxname), key=lambda h: h.taxname):
         subhogs_of_clade = list(subhogs_of_clade)
         if len(subhogs_of_clade) == 1:
             groups.append(_tax_overlap(subhogs_of_clade[0], species_lineage_index))
@@ -282,13 +280,14 @@ def __repr__(self):
 class HOG:
     _hogid_iter = 10000
 
-    def __init__(self, input_instantiate, taxnomic_range:TreeNode, rhogid:str, msa:Optional[MultipleSeqAlignment] = None,
+    def __init__(self, input_instantiate, taxonomic_level: Union[str, TreeNode], rhogid:str, msa:Optional[MultipleSeqAlignment] = None,
                  representatives: Optional[List[Representative]] = None, conf_infer_subhhogs=None):
         """HOG class constructor.
 
         A HOG can be initialized with either a single SeqRecord or from a list of HOGs
-        (which will become its sub-hog elements). Furthermore, the taxonomic range
-        (a node from the species tree) needs to be provided.
+        (which will become its sub-hog elements). Furthermore, the taxonomic level
+        (the name of a species tree node, or the node itself) needs to be provided; only its
+        name is stored, so HOGs don't reference (and pickle) the species tree.
         Parameters:
         :param input_instantiate"""
         # fragment_list list of  sets , each set contains protein ID of fragments
@@ -301,7 +300,8 @@ def __init__(self, input_instantiate, taxnomic_range:TreeNode, rhogid:str, msa:O
         # 0070124
 
         self._hogid = "HOG_" + self._rhogid+ "_sub" + str(self.__class__._hogid_iter)
-        self._tax_now = taxnomic_range    # the taxnomic level that we are considering now, checking for duplication
+        # name of the taxonomic level that we are considering now, checking for duplication
+        self._taxname = taxonomic_level.name if isinstance(taxonomic_level, TreeNode) else taxonomic_level
         self.active = True
 
         if isinstance(input_instantiate, SeqRecord):  # if len(sub_hogs)==1:
@@ -349,19 +349,15 @@ def __init__(self, input_instantiate, taxnomic_range:TreeNode, rhogid:str, msa:O
             assert False
 
     def __repr__(self):
-        return f"<HOG:{self.hogid},size={len(self._members)},tax={self._tax_now.name}>"
+        return f"<HOG:{self.hogid},size={len(self._members)},tax={self._taxname}>"
 
     @property
     def hogid(self):
         return self._hogid
 
     @property
     def taxname(self):
-        return self._tax_now.name
-
-    @property
-    def taxlevel(self):
-        return self._tax_now
+        return self._taxname
 
     @property
     def rhogid(self):
@@ -480,11 +476,14 @@ def merge_prots_msa(self, merged_fragment_name, merged_msa_new):   # merged_frag
     #     self._msa = MultipleSeqAlignment(msa_new)
     #     return 1
 
-    def to_orthoxml(self, full_species_tree: Optional[TreeNode] = None, score_flags: ScoreFlags = ScoreFlags()):
+    def to_orthoxml(self, full_species_tree: TreeNode, score_flags: ScoreFlags = ScoreFlags()):
         """Public entry point: builds the orthoxml element tree for this HOG and, once the whole
         (recursive) build is done, logs how much of the total time went into computing scores
         (see _ScoreTimings and attach_scores). `score_flags` is forwarded to every attach_scores()
-        call made during the recursion."""
+        call made during the recursion.
+
+        `full_species_tree` is the species tree in which the TaxRange (MRCA of the members) and the
+        scores are computed, normally the unpruned tree from prepare_species_tree."""
         timings = _ScoreTimings()
         start = time.perf_counter()
         result = self._to_orthoxml(full_species_tree, timings, score_flags)
@@ -499,7 +498,7 @@ def to_orthoxml(self, full_species_tree: Optional[TreeNode] = None, score_flags:
             )
         return result
 
-    def _to_orthoxml(self, full_species_tree: Optional[TreeNode], timings: _ScoreTimings,
+    def _to_orthoxml(self, full_species_tree: TreeNode, timings: _ScoreTimings,
                       score_flags: ScoreFlags = ScoreFlags()):
         if len(self._subhogs) == 0:
             list_member = list(self._members)
@@ -538,7 +537,7 @@ def _to_orthoxml(self, full_species_tree: Optional[TreeNode], timings: _ScoreTim
         # We continue this function as an Implicit else :  len(self._subhogs) >=1
 
         def _sorter_key(sh):
-            return sh._tax_now.name
+            return sh.taxname
         self._subhogs.sort(key=_sorter_key)
 
         element_list = []
@@ -581,9 +580,9 @@ def _sorter_key(sh):
             #hog_elemnt = ET.Element('orthologGroup', attrib={"id": str(self._hogid)})
             hog_elemnt = ET.Element('orthologGroup', attrib={"id": str(self._hogid)}, )
             species_of_members = _member_species(self._members)  #  'tr|H2MU14|H2MU14_ORYLA||ORYLA||1056022282'
-            mrca = _species_mrca(full_species_tree if full_species_tree is not None else self.taxlevel, species_of_members)
-            if mrca.name != self.taxlevel.name:
-                logger.info(f"mrca ({mrca.name}) != self.taxlevel ({self.taxlevel.name})")
+            mrca = _species_mrca(full_species_tree, species_of_members)
+            if mrca.name != self.taxname:
+                logger.info(f"mrca ({mrca.name}) != self.taxname ({self.taxname})")
                 logger.info(f"<{hog_elemnt.tag} {hog_elemnt.attrib}>")
 
             attach_scores(hog_elemnt, self, mrca, species_of_members, timings=timings, score_flags=score_flags)
@@ -646,7 +645,7 @@ def split_hog(hog:HOG, level_name:str, *partitions):
         logger.info(f"Splitting {hog} into {len(first_level_subhogs)} subhogs: {first_level_subhogs}")
         hogs = []
         for p, subhogs in enumerate(first_level_subhogs):
-            h = HOG(subhogs, taxnomic_range=hog.taxlevel, rhogid=hog.rhogid, msa=hog.get_msa(), representatives=partitions[p])
+            h = HOG(subhogs, taxonomic_level=hog.taxname, rhogid=hog.rhogid, msa=hog.get_msa(), representatives=partitions[p])
             hogs.append(h)
         return hogs
     else:

FastOMA/_infer_subhog.py

@@ -120,9 +120,7 @@ def read_infer_xml_rhog(rhogid, inferhog_concurrent_on, pickles_rhog_folder,  pi
                 # in version v0.3 of orthoxml, there shouldn't be any paralogGroup at root level. Let's put them inside an orthogroup should be in
                 hog_elemnt = ET.Element('orthologGroup', attrib={"id": str(hog_i.hogid)})
                 species_of_members = _member_species(hog_i.get_members())
-                scoring_node = hog_i.taxlevel
-                if full_species_tree is not None:
-                    scoring_node = _species_name_index(full_species_tree)[hog_i.taxlevel.name]
+                scoring_node = _species_name_index(full_species_tree)[hog_i.taxname]
                 attach_scores(hog_elemnt, hog_i, scoring_node, species_of_members, score_flags=score_flags)
                 ET.SubElement(hog_elemnt, "property", attrib={"name": "TaxRange", "value": str(hog_i.taxname)})
                 hog_elemnt.append(hogs_a_rhog_xml_raw)
@@ -292,7 +290,7 @@ def singletone_hog_(node_species_tree, rhogid, pickles_subhog_folder_all, rhogs_
         species_seq_generator = (rec for rec in SeqIO.parse(fasta, format="fasta")
                                  if rec.id.split('||')[1] == node_species_name)
         for prot in species_seq_generator:
-            hog_leaf = HOG(prot, node_species_tree, rhogid)  # node_species_tree.name
+            hog_leaf = HOG(prot, node_species_tree.name, rhogid)
             hogs_this_level_list.append(hog_leaf)
     pickle_subhog_file = pickles_subhog_folder + str(this_level_node_name)+".pickle"
     with open(pickle_subhog_file, 'wb') as handle:
@@ -685,7 +683,7 @@ def merge_subhogs(self, reconciled_genetree:TreeNode, msa:MultipleSeqAlignment):
                 processed_nodes.add(subtree)
                 new_repr = self.find_most_divergent_representatives_from_genetree(subtree)
                 hog = HOG([self.subhogs[h] for h in subtree.hogids],
-                          self.node_species_tree,
+                          self.node_species_tree.name,
                           rhogid=self.rhogid,
                           msa=msa,
                           representatives=new_repr,

tests/test_hog_class.py

@@ -21,14 +21,22 @@ def test_leaf_hog(sp):
     h = leaf_hog(A1, sp)
     assert h.get_members() == {A1}
     assert len(h) == 1
-    assert h.taxname == "A" and h.taxlevel is sp & "A"
+    assert h.taxname == "A"
     assert h.hogid == "HOG_r_sub10001"
     assert [r.get_id() for r in h.get_representatives()] == [A1]
     assert [r.id for r in h.get_msa()] == [A1]
     assert h._subhogs == [] and h.get_dubious_members() == set()
     assert h.active
 
 
+def test_hog_stores_only_the_name_of_its_level(sp):
+    # HOGs are pickled after every taxonomic level; they must not drag a copy of the species tree along
+    import pickle
+    top = merged_hog([leaf_hog(A1, sp), leaf_hog(B1, sp)], sp & "AB")
+    assert top.taxname == "AB" and [h.taxname for h in top._subhogs] == ["A", "B"]
+    assert b"ete3" not in pickle.dumps(top)
+
+
 def test_hogids_are_sequential_per_process(sp):
     ids = [leaf_hog(x, sp).hogid for x in (A1, A2)]
     assert ids == ["HOG_r_sub10001", "HOG_r_sub10002"]
@@ -160,7 +168,7 @@ def test_split_hog_on_disjoint_first_level_subhogs(sp):
     assert len(parts) == 2
     assert [h.get_members() for h in parts] == [{A1, B1}, {C1}]
     assert [h._subhogs for h in parts] == [[ab], [c1]]
-    assert all(h.taxlevel is top.taxlevel for h in parts)
+    assert all(h.taxname == top.taxname for h in parts)
     assert [[r.get_id() for r in h.get_representatives()] for h in parts] == [[A1], [C1]]
 
 

tests/test_hog_scores.py

@@ -23,7 +23,7 @@ def leaf_hog(prot_id, species_node):
     hog = HOG.__new__(HOG)
     hog._rhogid = "test"
     hog._hogid = "HOG_test_" + prot_id
-    hog._tax_now = species_node
+    hog._taxname = species_node.name
     hog._members = {prot_id}
     hog._subhogs = []
     hog._dubious_members = set()
@@ -34,7 +34,7 @@ def merged_hog(subhogs, tax_node):
     hog = HOG.__new__(HOG)
     hog._rhogid = "test"
     hog._hogid = "HOG_test_" + tax_node.name
-    hog._tax_now = tax_node
+    hog._taxname = tax_node.name
     hog._subhogs = list(subhogs)
     hog._members = set().union(*(sh._members for sh in subhogs))
     hog._dubious_members = set()
@@ -452,10 +452,10 @@ class DuplicationAsymmetricRetentionTCSTests(TestCase):
                                                        |
                                                       SP4
 
-    HOG topology built in the test method below -- copy1 and copy2 both get _tax_now=7 (node
+    HOG topology built in the test method below -- copy1 and copy2 both get taxname "7" (node
     "7"), matching how the real pipeline's merge_subhogs() always bumps every surviving hog's
-    _tax_now to the level it was last processed at, whether or not it found a merge partner
-    there. That shared _tax_now is what marks them as a recognized duplication (paralogGroup):
+    taxonomic level to the one it was last processed at, whether or not it found a merge partner
+    there. That shared taxname is what marks them as a recognized duplication (paralogGroup):
 
                               1  (top HOG)
                              /            \

tests/test_level_processor.py

@@ -96,7 +96,7 @@ def test_merge_subhogs_two_orthogroups(make_lp, sp):
     t = _reconciled(lp, "((%s:1,%s:1):1,(%s:1,%s:1):1);" % q(A1, B1, A2, B2))
     new = lp.merge_subhogs(t, msa(rec(A1), rec(A2), rec(B1), rec(B2)))
     assert sorted(sorted(h.get_members()) for h in new) == [[A1, B1], [A2, B2]]
-    assert all(h.taxlevel is sp & "AB" for h in new)
+    assert all(h.taxname == "AB" for h in new)
     for h in new:
         assert {r.get_id() for r in h.get_representatives()} <= h.get_members()