Protein trace: record what happens to every input protein

FastOMA.nf

@@ -146,6 +146,7 @@ process infer_roothogs{
     path "omamer_rhogs"
     path "gene_id_dic_xml.pickle"
     path "selected_isoforms" , optional: true
+    path "protein_events_*.events", optional: true, emit: protein_events
 
   script:
     """
@@ -154,6 +155,7 @@ process infer_roothogs{
                                --splice ${splice_folder} \
                                --out-rhog-folder "omamer_rhogs" \
                                --min-sequence-length ${params.min_sequence_length} \
+                               ${ params.protein_trace.toString().toBoolean() ? "--protein-trace" : "" } \
                                -vv
     """
 }
@@ -209,6 +211,7 @@ process hog_big{
     path "*.nwk" , optional: true  // gene trees  if write True
     path "*.tsv", optional: true
     path "*.tsv.gz", optional: true
+    path "protein_events_*.events", optional: true, emit: protein_events
 
   script:
     """
@@ -225,7 +228,8 @@ process hog_big{
                                ${ params.write_genetrees ? "--gene-trees-write" : ""} \
                                ${ params.disable_completeness_score.toString().toBoolean()   ? "--no-store-completeness-score" : "" } \
                                ${ params.disable_implied_losses_score.toString().toBoolean() ? "--no-store-implied-losses-score" : "" } \
-                               ${ params.disable_tcs_score.toString().toBoolean()            ? "--no-store-tcs-score" : "" }
+                               ${ params.disable_tcs_score.toString().toBoolean()            ? "--no-store-tcs-score" : "" } \
+                               ${ params.protein_trace.toString().toBoolean() ? "--protein-trace" : "" }
     """
 }
 
@@ -247,6 +251,7 @@ process hog_rest{
     path "*.nwk" , optional: true  // gene trees  if write True
     path "*.tsv", optional: true
     path "*.tsv.gz", optional: true
+    path "protein_events_*.events", optional: true, emit: protein_events
   script:
     """
         fastoma-infer-subhogs --input-rhog-folder ${rhogsrest}  \
@@ -261,7 +266,8 @@ process hog_rest{
                               ${ params.write_genetrees ? "--gene-trees-write" : ""} \
                               ${ params.disable_completeness_score.toString().toBoolean()   ? "--no-store-completeness-score" : "" } \
                               ${ params.disable_implied_losses_score.toString().toBoolean() ? "--no-store-implied-losses-score" : "" } \
-                              ${ params.disable_tcs_score.toString().toBoolean()            ? "--no-store-tcs-score" : "" }
+                              ${ params.disable_tcs_score.toString().toBoolean()            ? "--no-store-tcs-score" : "" } \
+                              ${ params.protein_trace.toString().toBoolean() ? "--protein-trace" : "" }
     """
 }
 
@@ -299,6 +305,30 @@ process collect_subhogs{
     """
 }
 
+process protein_trace {
+  label "process_single"
+  publishDir params.output_folder, mode: 'copy'
+
+  input:
+    path proteome_folder
+    path events, stageAs: "events/*"
+    path roothogs_tsv
+    val  id_transform
+
+  output:
+    path "protein_fates.tsv"
+    path "protein_events.tsv.gz"
+    path "protein_trace_summary.txt"
+
+  script:
+    """
+        python -m FastOMA.protein_trace --proteomes ${proteome_folder} \
+                                        --events events/ \
+                                        --roothog-tsv ${roothogs_tsv} \
+                                        --id-transform ${id_transform}
+    """
+}
+
 process extract_pairwise_ortholog_relations {
   label "process_medium"
   publishDir params.output_folder, mode: 'copy'
@@ -492,6 +522,14 @@ workflow {
     channel.empty().concat(pickle_big_rhog, pickle_rest_rhog).set{ all_rhog_pickle }
 
     (orthoxml_file, OrthologousGroupsFasta, OrthologousGroups_tsv, rootHOGs_tsv)  = collect_subhogs(all_rhog_pickle.collect(), gene_id_dic_xml, omamer_rhogs, species_tree_checked, params.fasta_header_id_transformer)
+    if (params.protein_trace.toString().toBoolean()) {
+        channel.empty()
+            .concat(infer_roothogs.out.protein_events, hog_big.out.protein_events, hog_rest.out.protein_events)
+            .flatten()
+            .collect()
+            .set { all_protein_events }
+        protein_trace(proteome_folder, all_protein_events, rootHOGs_tsv, params.fasta_header_id_transformer)
+    }
     c = hogmap.count().branch{ n->
         TRUE: (n<=25 || params.force_pairwise_ortholog_generation)
         FALSE: n>25

FastOMA/_infer_subhog.py

@@ -26,6 +26,7 @@
 from . import _wrappers, logger
 from . import _utils_subhog
 from . import _utils_frag_SO_detection
+from . import _protein_trace as ptrace
 from ._hog_class import HOG, Representative, split_hog, attach_scores, _member_species, _species_name_index, ScoreFlags
 from ._utils_subhog import MSAFilter, MSAFilterElbow, MSAFilterTrimAL
 
@@ -104,8 +105,14 @@ def read_infer_xml_rhog(rhogid, inferhog_concurrent_on, pickles_rhog_folder,  pi
 
     tot_genes, placed_genes = 0, 0
     hogs_rhogs_xml = []
+    members_final = set()
     for hog_i in hogs_a_rhog:
         tot_genes += len(hog_i)
+        if ptrace.is_enabled():
+            members_final |= {part for m in hog_i.get_members() for part in m.split("_|_")}  # merged fragments count per protein
+            final_event = "in_final_hog" if len(hog_i) >= inferhog_min_hog_size_xml else "final_singleton_not_written"
+            for member in hog_i.get_members():
+                ptrace.member_event(member, "infer_subhogs", final_event, rhogid, f"hog={hog_i.hogid} size={len(hog_i)}")
         if len(hog_i) >= inferhog_min_hog_size_xml:
             # could be improved   # hogs_a_rhog_xml = hog_i.to_orthoxml(**gene_id_name)
             hogs_a_rhog_xml_raw = hog_i.to_orthoxml(full_species_tree, score_flags)    # <generef  >      <paralg object >
@@ -135,6 +142,12 @@ def read_infer_xml_rhog(rhogid, inferhog_concurrent_on, pickles_rhog_folder,  pi
             if len(hog_i) > 1:
                 logger.warning("issue 166312309 this is not a singleton "+str(hog_i._members))
 
+    # any protein of the input rootHOG that is in no final HOG and was not reported as lost above
+    if ptrace.is_enabled():
+        for rec in rhog_i:
+            if rec.id not in members_final:
+                ptrace.member_event(rec.id, "infer_subhogs", "missing_at_end", rhogid)
+
     pickles_rhog_file = pickles_rhog_folder + '/file_' + rhogid + '.pickle'
     with open(pickles_rhog_file, 'wb') as handle:
         # dill_pickle.dump(hogs_rhogs_xml, handle, protocol=dill_pickle.HIGHEST_PROTOCOL)
@@ -330,6 +343,7 @@ def __init__(self, node_species_tree:TreeNode, subhogs:List[HOG], rhogid:str, co
         self._base_msa_tree_filename = "HOG_"+rhogid+"_"+str(node_species_tree.name)
         self._outname_step = 0
         self._removed_rep = set([])
+        self._lost_members = set([])   # members dropped by split_hog at this level (for the protein trace)
         self._rep_lookup = self._prepare_lookups()
         self._msa_filter = self._instantiate_msa_filter()
 
@@ -645,6 +659,13 @@ def merge_subhogs(self, reconciled_genetree:TreeNode, msa:MultipleSeqAlignment):
                     split_parts = [list(n.name for n in sub.iter_leaves() if n.hogid == hogid) for sub in subtrees]
                     split_hogs = split_hog(self.subhogs[hogid], self.node_species_tree.name, *split_parts)
                     if split_hogs and len(split_hogs) > 1:
+                        if ptrace.is_enabled():
+                            kept = set().union(*(h.get_members() for h in split_hogs))
+                            for member in self.subhogs[hogid].get_members() - kept:
+                                self._lost_members.add(member)
+                                ptrace.member_event(member, "infer_subhogs", "lost_in_split_hog", self.rhogid,
+                                                    f"level={self.node_species_tree.name} hog={hogid} "
+                                                    f"hog_size={len(self.subhogs[hogid])} parts={len(split_hogs)}")
                         # we could split the current hog.
                         self.subhogs.pop(hogid)
                         self.subhogs.update({h.hogid: h for h in split_hogs})
@@ -673,6 +694,17 @@ def merge_subhogs(self, reconciled_genetree:TreeNode, msa:MultipleSeqAlignment):
         return new_hogs
 
     def process(self) -> List[HOG]:
+        if not ptrace.is_enabled():
+            return self._process()
+        members_before = set().union(set(), *(h.get_members() for h in list(self.subhogs.values()) + list(self._frozen_hogs_before)))
+        new_hogs = self._process()
+        members_after = set().union(set(), *(h.get_members() for h in new_hogs))
+        for member in members_before - members_after - self._lost_members:
+            ptrace.member_event(member, "infer_subhogs", "lost_at_level", self.rhogid,
+                                f"level={self.node_species_tree.name}")
+        return new_hogs
+
+    def _process(self) -> List[HOG]:
         """Process the given roothog at the given taxonomic level.
 
         This method does all the processessing of the subhogs from the

FastOMA/_protein_trace.py

@@ -0,0 +1,101 @@
+"""Per-protein event log, so that the fate of every input protein can be traced.
+
+Each FastOMA step appends events to its own file in the task directory
+(protein_events_<step>_<uuid>.events, one per process so that parallel
+Nextflow tasks never collide). The columns are, tab separated and without header:
+
+    species  protein  step  event  roothog  detail
+
+`fastoma-protein-trace` (FastOMA/protein_trace.py) merges all event files at the
+end of the run into one row per input protein.
+
+Nothing is recorded unless enable() was called (the --protein-trace option of
+fastoma-infer-roothogs and fastoma-infer-subhogs). Expect a few events per protein,
+i.e. roughly 250-400 bytes per input protein in total.
+"""
+
+import atexit
+import threading
+import uuid
+
+EVENT_FILE_SUFFIX = ".events"
+COLUMNS = ["species", "protein", "step", "event", "roothog", "detail"]
+
+_lock = threading.Lock()
+_handles = {}
+_enabled = False
+
+
+def enable():
+    global _enabled
+    _enabled = True
+
+
+def is_enabled():
+    return _enabled
+
+
+def _handle(step):
+    fh = _handles.get(step)
+    if fh is None:
+        fh = open(f"protein_events_{step}_{uuid.uuid4().hex[:12]}{EVENT_FILE_SUFFIX}", "wt")
+        _handles[step] = fh
+    return fh
+
+
+def event(species, protein, step, event_name, roothog="", detail=""):
+    if not _enabled:
+        return
+    line = "\t".join(str(x).replace("\t", " ") for x in (species, protein, step, event_name, roothog, detail))
+    with _lock:
+        _handle(step).write(line + "\n")
+
+
+def split_member(member):
+    """Internal FastOMA member ids look like 'protId||species||numericId', or several of
+    those joined by '_|_' for merged fragments. Returns a list of (species, protId)."""
+    out = []
+    for part in str(member).split("_|_"):
+        fields = part.split("||")
+        if len(fields) >= 2:
+            out.append((fields[1], fields[0]))
+    return out
+
+
+def member_event(member, step, event_name, roothog="", detail=""):
+    if not _enabled:
+        return
+    for species, protein in split_member(member):
+        event(species, protein, step, event_name, roothog, detail)
+
+
+def rhog_assignment(rhogs_prots):
+    """(species, protId) -> rootHOG id, for diffing the rootHOG grouping steps (empty if tracing is off)."""
+    if not _enabled:
+        return {}
+    return {sp_prot: rhogid for rhogid, sp_prots in rhogs_prots.items() for sp_prot in sp_prots}
+
+
+def record_rhog_changes(before, after, step, event_prefix):
+    """Record proteins that changed rootHOG, or disappeared, between two assignments."""
+    if not _enabled:
+        return
+    for sp_prot, old in before.items():
+        new = after.get(sp_prot)
+        if new is None:
+            event(*sp_prot, step, f"{event_prefix}_removed", old)
+        elif new != old:
+            event(*sp_prot, step, f"{event_prefix}_moved", new, f"from={old}")
+    for sp_prot, new in after.items():
+        if sp_prot not in before:
+            event(*sp_prot, step, f"{event_prefix}_added", new)
+
+
+def close():
+    with _lock:
+        for fh in _handles.values():
+            fh.close()
+        _handles.clear()
+
+
+atexit.register(close)

FastOMA/_utils_roothog.py

@@ -14,6 +14,7 @@
 
 from .zoo.unionfind import UnionFind
 from . import logger
+from . import _protein_trace as ptrace
 import collections
 
 filter_nonchild_rootHOG= False
@@ -175,7 +176,12 @@ def parse_proteomes(folder=None, min_sequence_length=0):  # list_oma_species
     #smallprot_recs_lists ={}
     for species_name in species_names:
         prot_address = os.path.join(folder, species_name + "." + fasta_format_keep)
-        prots_record = list(rec for rec in SeqIO.parse(prot_address, "fasta") if len(rec) >= min_sequence_length)
+        prots_record = []
+        for rec in SeqIO.parse(prot_address, "fasta"):
+            if len(rec) >= min_sequence_length:
+                prots_record.append(rec)
+            else:
+                ptrace.event(species_name, rec.id, "infer_roothogs", "too_short", "", f"length={len(rec)} min={min_sequence_length}")
         #prots_record_small = list(rec for rec in SeqIO.parse(prot_address, "fasta") if len(rec) < min_sequence_length)
         # logger.debug(prots_record)
         logger.info(f"{species_name} contains {len(prots_record)} that are at least {min_sequence_length} long.")
@@ -193,12 +199,23 @@ def create_root_hogs(hogmaps, conf):
     # Initial grouping
     rhogs_prots = group_prots_by_roothogs(hogmaps)
     logger.info(f"Initial grouping: {len(rhogs_prots)} root HOGs")
+    assignment = ptrace.rhog_assignment(rhogs_prots)
+    for (species_name, prot_id), rhogid in assignment.items():
+        n_hits = len(hogmaps[species_name][prot_id])
+        ptrace.event(species_name, prot_id, "infer_roothogs", "omamer_best_hit", rhogid, f"n_hits={n_hits}")
 
     # Handle singletons
     rhogs_prots = resolve_singletons(rhogs_prots, hogmaps, conf)
+    new_assignment = ptrace.rhog_assignment(rhogs_prots)
+    ptrace.record_rhog_changes(assignment, new_assignment, "infer_roothogs", "resolve_singletons")
+    assignment = new_assignment
 
     rhogs_prots = merge_rhogs2(hogmaps, rhogs_prots, conf)
+    new_assignment = ptrace.rhog_assignment(rhogs_prots)
+    ptrace.record_rhog_changes(assignment, new_assignment, "infer_roothogs", "merge_rhogs")
+    assignment = new_assignment
     rhogs_prots = filter_big_roothogs(hogmaps, rhogs_prots, conf)
+    ptrace.record_rhog_changes(assignment, ptrace.rhog_assignment(rhogs_prots), "infer_roothogs", "filter_big_roothogs")
 
     # # Merge similar HOGs
     # if conf.enable_merging:
@@ -286,6 +303,7 @@ def parse_hogmap_omamer(proteomes, fasta_format_keep, folder=None):
             for row in reader:
                 if row['hogid'] == "N/A":
                     unmapped[species_name].append(row['qseqid'])
+                    ptrace.event(species_name, row['qseqid'], "infer_roothogs", "no_omamer_hit")
                 elif row['qseqid'] not in proteome:
                     logger.debug(f"Ignoring {row} [{species_name}] from hogmap as not in proteome")
                 else:
@@ -511,8 +529,12 @@ def write_rhog(rhogs_prot_records, prot_recs_all, address_rhogs_folder, min_rhog
             # todo add the release id   to file names  rhogids_list[:2] > ['HOG:C0884658', 'HOG:C0709155']
             SeqIO.write(rhog_recs, address_rhogs_folder + "/HOG_" + rhogid + ".fa", "fasta")
             rhogid_written_list.append(rhogid)
+            for rec in rhog_recs:
+                ptrace.member_event(rec.id, "infer_roothogs", "roothog_written", rhogid, f"size={len(rhog_recs)}")
         else:
             logger.debug("The roothog " +str(rhogid)+" was too small with size of "+str(len(rhog_recs))+" which is smaller than threshold "+str(min_rhog_size))
+            for rec in rhog_recs:
+                ptrace.member_event(rec.id, "infer_roothogs", "roothog_too_small", rhogid, f"size={len(rhog_recs)} min={min_rhog_size}")
             #logger.debug(" issue 12314405 to check" )
             #for prot1 in rhog_recs:
                 #logger.debug("we are removing due to omamer signleton hog |*|" + str(prot1.id))
@@ -833,6 +855,8 @@ def collect_unmapped_singleton(rhogs_prots, unmapped, prot_recs_all, unmapped_si
     logger.debug(f"collected sequence records of {len(singleton_recs)} singleton proteins.")
 
     singleton_unmapped_recs = unmapped_recs + singleton_recs
+    for rec in singleton_unmapped_recs:
+        ptrace.member_event(rec.id, "infer_roothogs", "sent_to_clustering")
     SeqIO.write(singleton_unmapped_recs, unmapped_singleton_fasta , "fasta")
     logger.debug(f"Wrote {len(singleton_unmapped_recs)} sequences of unmapped and singleton proteins to {unmapped_singleton_fasta}.")
     return len(singleton_unmapped_recs)
@@ -903,6 +927,7 @@ def write_clusters(fasta_clustered, address_rhogs_folder, min_rhog_size):
     for cluster in clusters:
         # cluster is like ['>sp|O67032|RF1_AQUAE', 'MLKEAYISRLDKLQEKYRKLQEELSKPEVIQD...', ...], i.e. an id and its sequence.
         # thus, check for double the min_rhog_size to have at least min_rhog_size proteins
+        member_ids = [cluster[pr_idi*2][1:].split(" ")[0] for pr_idi in range(int(len(cluster)/2))]
         if len(cluster) >= 2 * min_rhog_size:
             species_names = set([])     # it seems that genes from same species tend to cluster together, discard such clusters
             for pr_idi in range(int(len(cluster)/2)):
@@ -912,7 +937,15 @@ def write_clusters(fasta_clustered, address_rhogs_folder, min_rhog_size):
                 with open(address_rhogs_folder + "/HOG_clust" + str(cluster_iter) + ".fa", "w") as file_idx:
                     for line in cluster:
                         file_idx.write(line + "\n")
+                for member in member_ids:
+                    ptrace.member_event(member, "infer_roothogs", "roothog_written", f"clust{cluster_iter}", f"size={len(member_ids)} source=linclust")
                 cluster_iter += 1
+            else:
+                for member in member_ids:
+                    ptrace.member_event(member, "infer_roothogs", "cluster_single_species_discarded", "", f"size={len(member_ids)}")
+        else:
+            for member in member_ids:
+                ptrace.member_event(member, "infer_roothogs", "cluster_too_small", "", f"size={len(member_ids)}")
     return cluster_iter - 10000
 
 
@@ -977,6 +1010,9 @@ def select_best_isoforms(species_names, isoform_by_gene_all, hogmaps):
 
             isoform_selected[species_name].append(protname_best)
             isoform_not_selected[species_name].extend([i for i in isoform_list if i != protname_best])  # flatten for all
+            for i in isoform_list:
+                if i != protname_best:
+                    ptrace.event(species_name, i, "infer_roothogs", "isoform_not_selected", "", f"selected={protname_best}")
 
     return isoform_selected, isoform_not_selected