"""Reheader a SAM in a stream. When using the bam2fq -> minimap2 pattern for (re)aligning BAM data, we lose any existing RG and PG headers. This is particularly egregious when handling basecalled data as lines related to dorado basecalling settings as well as dorado RG headers are lost; orphaning RG tags in the reads. This is problematic for downstream anaylses that would like to read the XAM header to intelligently determine how to handle the reads based on the basecaller model and basecaller configuration. This script handles: - Inserting RG, PG and CO lines from an existing XAM header into the header of the SAM emitted from minimap2's alignment stream - Inserting a PG header to indicate that a call to bam2fq was made - Updating the first streamed PG.PP parent tag with the last PG.ID of the existing XAM header to maintain a chain of custody - Updating any streamed PG.ID (and PG.PP) tags to avoid collisions with inserted PG.ID Handling collisions may seem like overkill but it is anticipated that this script will be called immediately after minimap2, any previous attempt to use minimap2 will lead to ambiguity. This would be the expected case where users have used wf-basecalling or wf-alignment to align a set of reads, only to realign them to another reference (eg. via wf-human-variation). Arguably, we should remove older references to minimap2 as they will have been invalidated by the call to bam2fq but removing PG records and sticking the PG chain back together seems more fraught with annoying future bugs than simply resolving conflicts. This script will explode on a stream that contains: - PG lines in the original header where the last PG in the chain is ambiguous, or where the parent PP IDs are not injective - PG lines in the stream that do not appear in the order of their chain (that is if a PG.PP refers to a PG.ID that has not been encountered yet) SQ lines are retained after an HD line. That is to say, the most recent set of SQ lines observed after an HD will appear in the final output. SQ, RG, PG and CO lines are emitted as a group together, with elements written out in the order observed. PG lines are naively appended to the last PG element in the chain. No attempt is made to keep multiple program chains intact as this can lead to bloated headers. Broken PG metadata is a known problem (see samtools/hts-specs#275) but one that is preferable to headers that become unwieldly large to process: there IS an upper limit to a SAM header's size after all. This script takes advantage of minimap2's SAM output to immediately reheader the stream before any downstream calls to other programs pollute the PG header. This script is a little overkill but attempts to be robust with handling PG collisions and more obviously encapsulates reheadering behaviour, and leaves some room to do more clever things as necessary. """ from shutil import copyfileobj import sys from ..util import wf_parser # noqa: ABS101 class SamHeader: """An overkill container to manage merging PG lines in SAM headers. Collision handling is simple. If a PG.ID is duplicated by the stream then we add a suffix to its name and keep an eye out for the corresponding PG.PP later. We assume that headers emitted by the stream are chronological because this script should not be called as part of any complicated pipework other than immediately following minimap2. """ def __init__(self): """Initialise a collision aware PG container.""" self.remapped_pgids = {} self.collision_suffix = 0 # Default HD, in case the new stream does not provide one self.hd = "@HD\tVN:1.6\tSO:unknown" # We'll merge RG, CO and PG self.rg_records = [] self.co_records = [] self.pg_records = [] # We keep the most recently observed block of SQ records by # resetting SQ on the first SQ seen after non-SQ. We cannot # rely on HD being emitted (as minimap2 does not do this!) self.sq_records = [] self.reset_sq = False self.observed_rgids = set() self.observed_pgids = set() self.last_pgid = None @staticmethod def str_to_record(line): """Return an appropriate struct for a given string record.""" try: record_type, record_data = line.strip().split('\t', 1) except ValueError: raise Exception(f"Record type could not be determined: {line}") if len(record_type) > 3: raise Exception(f"Record type malformed: {record_type}") record = {} if record_type in ["@HD", "@CO", "@SQ"]: return record_type, record_data elif record_type in ["@RG", "@PG"]: for field in record_data.strip().split('\t'): k, v = field.split(':', 1) if len(k) == 2 and k[0].isalpha() and k[1].isalnum(): record[k] = v else: raise Exception(f"{record_type} with invalid tag: '{k}'") if "ID" not in record: raise Exception(f"{record_type} with no ID: {record_data}") return record_type, record else: raise Exception(f"Unknown record type: {line}") @staticmethod def record_to_str(record_type, record_data): """Form a string from a header record.""" if record_type in ["@PG", "@RG"]: tags = [f"{k}:{v}" for k, v in record_data.items()] return f"{record_type}\t" + '\t'.join(tags) elif record_type in ["@SQ", "@CO"]: return f"{record_type}\t{record_data}" @staticmethod def resolve_pg_chain(pg_dicts): """Check links between PG.ID and PP.ID, exploding if inconsistent.""" links = {} # Document links between all ID and their PP parent pgids_without_ppid = 0 for pgd in pg_dicts: pgid = pgd["ID"] pgpp = pgd.get("PP") links[pgid] = pgpp if pgpp is None: pgids_without_ppid += 1 if len(links) > 0: # If there are links, exactly one should have a None parent # to indicate the first PG in the chain. Explode if we see # no head or multiple heads. if pgids_without_ppid == 0: raise Exception("PG chain does not have a head.") elif pgids_without_ppid > 1: raise Exception("PG chain has multiple heads.") for source in links: head = source path = [head] while True: head = links[head] if head is None: break if head in path: path.append(head) raise Exception(f"PG chain appears to contain cycle: {path}") path.append(head) # This function is only really called to catch any explosions # but we'll return the links here as it is useful for testing return links def _bump_pg_collider(self): """Alter the collision suffix after determining a collision.""" self.collision_suffix += 1 def _uncollide_pgid(self, pgid): """Return an uncollided string for a given PG ID.""" new_pgid = f"{pgid}-{self.collision_suffix}" self.remapped_pgids[pgid] = new_pgid self._bump_pg_collider() return new_pgid def add_line(self, line): """Add a header line to the header.""" record_type, record = self.str_to_record(line) if record_type == "@HD": self.hd = f"@HD\t{record}" elif record_type == "@CO": self.co_records.append(record) elif record_type == "@SQ": if self.reset_sq: self.sq_records = [] self.reset_sq = False self.sq_records.append(record) elif record_type == "@RG": rgid = record["ID"] if rgid not in self.observed_rgids: self.observed_rgids.add(rgid) self.rg_records.append(record) elif record not in self.rg_records: # if rgid has been seen before, abort if this record is different raise Exception( f"Duplicate RG with ID '{rgid}' conflicts with previously seen RG with same ID." # noqa:E501 ) elif record_type == "@PG": pgid = record["ID"] if pgid in self.observed_pgids: # collision, rewrite the pgid pgid = self._uncollide_pgid(pgid) record["ID"] = pgid else: self.observed_pgids.add(pgid) # maintain chain ppid = record.get("PP") if not ppid: # record has no parent, this is either # - the first record (last_pgid is None) so is the tail # - an inserted record that needs its parent to be the current tail if not self.last_pgid: self.last_pgid = pgid else: record["PP"] = self.last_pgid self.last_pgid = pgid else: if ppid not in self.observed_pgids: raise Exception( f"Encountered PG.PP '{ppid}' before observing corresponding PG.ID" # noqa:E501 ) # remap parent id (if needed) record["PP"] = self.remapped_pgids.get(ppid, ppid) # set tail to this record self.last_pgid = pgid self.pg_records.append(record) if len(self.sq_records) > 0 and record_type != '@SQ': self.reset_sq = True return record def write_header(self, fh): """Write this header to a file handle.""" self.resolve_pg_chain(self.pg_records) # check PG header fh.write(f"{self.hd}\n") for sq in self.sq_records: fh.write(self.record_to_str("@SQ", sq) + '\n') for rg in self.rg_records: fh.write(self.record_to_str("@RG", rg) + '\n') for pg in self.pg_records: fh.write(self.record_to_str("@PG", pg) + '\n') for co in self.co_records: fh.write(self.record_to_str("@CO", co) + '\n') def reheader_samstream(header_in, stream_in, stream_out, args): """Run reheader_samstream.""" # read original header into container sh = SamHeader() for line in header_in: sh.add_line(line) # append user provided lines to container for line in args.insert: sh.add_line(line) # read the header portion of the minimap2 stream wrote_header = False for line in stream_in: if line[0] != '@': # write out header on first alignment sh.write_header(stream_out) wrote_header = True # and actually write the first alignment stream_out.write(line) break sh.add_line(line) # Pass through the rest of the alignments. # I toyed with a few ways of doing this: # - A trivial iter over the input file was slow. presumably as we incurred some # overhead calling read() and write() and decoding more than other methods. # - os.read/write avoids dealing with higher level python read/write but requires # file descriptors which rules out non-file-like objects. this made testing more # annoying as StringIO does not have a file descriptor. we could have mocked fds # but i was not happy with the discrepancy between real and test execution. # - copyfileobj with the stream_in.buffer would also avoid some of the higher # level text handling but would require all tests to provide inputs that have # an underlying binary buffer. it was also not possible to seek the buffer to # the position of the text stream as we've used next() to iterate over the # header lines, fixing this would have required rewriting of the header # handling or keeping track of the position in the stream ourselves which # just seemed unncessary overkill given how we expect this program to be used. # copyfileobj on the text streams is more efficient than merely iterating the file # and dumping the lines out and seems to do the job. this keeps the code and tests # simple with minimal additional cost to performance. i anticipate any overhead of # this program will be dwarfed by that of minimap2/samtools sort anyway. # increasing the buffer size gave worse performance in my limited testing so we # leave it as the default here. copyfileobj(stream_in, stream_out) # If there were no alignments, we won't have hit the != @ case in the first stdin, # and we won't have written the header out. Write a header if we haven't already. if not wrote_header: sh.write_header(stream_out) def argparser(): """Argument parser for entrypoint.""" parser = wf_parser("reheader_samstream") parser.add_argument("header_in") parser.add_argument("--insert", action="append", default=[]) return parser def main(args): """reheader_samstream default entry point.""" with open(args.header_in) as header_in: reheader_samstream(header_in, sys.stdin, sys.stdout, args)