gaftools 1.1.2__py3-none-any.whl

gaftools/cli/index.py

@@ -0,0 +1,181 @@
+"""
+Index a GAF file for the view functionality.
+
+This script creates an inverse look-up table where:
+    - key: the node information
+    - value: the offsets in the GAF file where the node is present
+"""
+
+import re
+import pickle
+from pysam import libcbgzf
+from collections import defaultdict
+import logging
+
+import gaftools.utils as utils
+from gaftools.gaf import GAF
+from gaftools.gfa import GFA
+from gaftools.timer import StageTimer
+from gaftools.cli import log_memory_usage
+
+logger = logging.getLogger(__name__)
+
+
+def run(gaf_path, gfa_path, output=None):
+    timers = StageTimer()
+    if output is None:
+        output = gaf_path + ".gvi"
+
+    # Detecting if GAF has stable or unstable coordinate
+    gaf = GAF(gaf_path)
+    stable = None
+    # checking format in the first 10 lines.
+    for i, gaf_line in enumerate(gaf.read_file()):
+        if i == 10:
+            break
+        if i == 0:
+            stable = gaf_line.detect_path_format()
+        assert stable == gaf_line.detect_path_format()
+    gaf.close()
+
+    nodes = {}
+    reference = defaultdict(lambda: [])
+    ref_contig = []
+
+    gfa_file = GFA(graph_file=gfa_path, low_memory=True)
+    contigs = list(gfa_file.contigs.keys())
+    ref_contig = [contig for contig in gfa_file.contigs if gfa_file.contigs[contig] == 0]
+    nodes = gfa_file.nodes
+    if stable:
+        for contig in contigs:
+            path = gfa_file.get_path(contig)
+            for node in path:
+                reference[contig].append(gfa_file[node])
+        nodes = gfa_file.nodes
+    del gfa_file
+
+    if utils.is_file_gzipped(gaf_path):
+        gaf_file = libcbgzf.BGZFile(gaf_path, "rb")
+    else:
+        gaf_file = open(gaf_path, "rt")
+
+    out_dict = {}
+    offset = 0
+    while True:
+        offset = gaf_file.tell()
+        mapping = gaf_file.readline()
+        if not mapping:
+            break
+        try:
+            val = mapping.rstrip().split("\t")
+        except TypeError:
+            val = mapping.decode("utf-8").rstrip().split("\t")
+
+        if stable:
+            with timers("convert_coord"):
+                alignment = convert_coord(val, reference)
+        else:
+            alignment = list(re.split(">|<", val[5]))[1:]
+        for a in alignment:
+            try:
+                out_dict[
+                    (
+                        nodes[a].id,
+                        nodes[a].tags["SN"][1],
+                        int(nodes[a].tags["SO"][1]),
+                        int(nodes[a].tags["SO"][1]) + int(nodes[a].tags["LN"][1]),
+                    )
+                ].append(offset)
+            except KeyError:
+                out_dict[
+                    (
+                        nodes[a].id,
+                        nodes[a].tags["SN"][1],
+                        int(nodes[a].tags["SO"][1]),
+                        int(nodes[a].tags["SO"][1]) + int(nodes[a].tags["LN"][1]),
+                    )
+                ] = [offset]
+    out_dict["ref_contig"] = ref_contig
+
+    gaf_file.close()
+
+    with open(output, "wb") as handle:
+        with timers("write_file"):
+            pickle.dump(out_dict, handle, protocol=pickle.HIGHEST_PROTOCOL)
+
+    logger.info("\n== SUMMARY ==")
+    total_time = timers.total()
+    log_memory_usage()
+    logger.info(
+        "Time to convert coordinates:                 %9.2f s", timers.elapsed("convert_coord")
+    )
+    logger.info(
+        "Time to write the file:                      %9.2f s", timers.elapsed("write_file")
+    )
+    logger.info("Time spent on rest:                          %9.2f s", total_time - timers.sum())
+    logger.info("Total time:                                  %9.2f s", total_time)
+
+
+def convert_coord(line, ref):
+    unstable_coord = []
+    gaf_contigs = list(filter(None, re.split("(>)|(<)", line[5])))
+    for nd in gaf_contigs:
+        if nd == ">" or nd == "<":
+            continue
+        if ":" in nd and "-" in nd:
+            tmp = nd.rstrip().split(":")
+            query_contig_name = tmp[0]
+            (query_start, query_end) = tmp[1].rstrip().split("-")
+        else:
+            query_start = line[7]
+            query_end = line[8]
+            query_contig_name = nd
+
+        """Find the matching nodes from the reference genome here"""
+        start, end = utils.search_intervals(
+            ref[query_contig_name], int(query_start), int(query_end), 0, len(ref[query_contig_name])
+        )
+
+        for node in ref[query_contig_name][start : end + 1]:
+            cases = -1
+            if (
+                int(node.tags["SO"][1])
+                <= int(query_start)
+                < int(node.tags["SO"][1]) + int(node.tags["LN"][1])
+            ):
+                cases = 1
+            elif (
+                int(node.tags["SO"][1])
+                < int(query_end)
+                <= int(node.tags["SO"][1]) + int(node.tags["LN"][1])
+            ):
+                cases = 2
+            elif (
+                int(query_start)
+                < int(node.tags["SO"][1])
+                < int(node.tags["SO"][1]) + int(node.tags["LN"][1])
+                < int(query_end)
+            ):
+                cases = 3
+
+            if cases != -1:
+                unstable_coord.append(node.id)
+
+    return unstable_coord
+
+
+# fmt: off
+def add_arguments(parser):
+    arg = parser.add_argument
+    # Positional arguments
+    arg('gaf_path', metavar='GAF',
+        help='GAF file (can be bgzip-compressed)')
+    arg('gfa_path', metavar='rGFA',
+        help='rGFA file (can be bgzip-compressed)')
+    arg('-o', '--output', default=None,
+        help='Output GAF View Index (GVI) file. Default: <GAF file>.gvi')
+# fmt: on
+
+
+def main(args):
+    run(**vars(args))

gaftools/cli/order_gfa.py

@@ -0,0 +1,335 @@
+"""
+Order bubbles in the GFA by adding BO and NO tags.
+
+The BO (Bubble Order) tags order bubbles in the GFA.
+The NO (Node Order) tags order the nodes in a bubble (in a lexicographic order).
+"""
+
+import sys
+import os
+import logging
+import time
+from collections import defaultdict
+from gaftools.gfa import GFA
+from gaftools.utils import DEFAULT_CHROMOSOME
+
+logger = logging.getLogger(__name__)
+
+
+def run_order_gfa(
+    gfa_filename,
+    outdir,
+    by_chrom,
+    chromosome_order=None,
+    with_sequence=False,
+):
+    if chromosome_order is not None:
+        chromosome_order = chromosome_order.split(sep=",")
+
+    if not os.path.isdir(outdir):
+        logging.warning(f"The directory {outdir} does not exist, creating one")
+        try:
+            os.makedirs(outdir)
+        except PermissionError:
+            logging.error(f"were not able to create directory {outdir}. Permission Denied")
+            sys.exit()
+        except FileNotFoundError:
+            logging.error(f"were not able to create directory {outdir}, File not found")
+            sys.exit()
+        except OSError:
+            logging.error(f"were not able to create directory {outdir}, OSError")
+            sys.exit()
+
+    logger.info(f"Reading {gfa_filename}")
+    if with_sequence:
+        graph = GFA(gfa_filename, low_memory=False)
+    else:
+        graph = GFA(gfa_filename, low_memory=True)
+    # the __str__ functions print the number of nodes and edges
+    logger.info("The graph has:")
+    logger.info(graph.__str__())
+
+    components = graph.all_components()
+    logger.info(f"Connected components: {len(components)}")
+
+    # name_comps checks the most frequent SN tag for the node in the component
+    # and returns a dict of chromosome_name: {nodes...}
+    components = name_comps(graph, components)
+    # if not chromosome_order is None:  # user gave a list
+    if chromosome_order != [""]:  # user gave a list
+        for c in chromosome_order:
+            if c not in set(components.keys()):
+                logger.error(
+                    f"The chromosome name provided {c} did not match with a component in the graph"
+                )
+                logger.error(f" What was Found: {','.join(sorted(components.keys()))}")
+                sys.exit(1)
+
+    else:  # user did not give a
+        try:
+            assert set(components.keys()) == set(DEFAULT_CHROMOSOME)
+        except AssertionError:
+            logger.error(
+                f"chromosome order was not provided, so the default was taken, but the default did not match"
+                f" what was found in the graph, which is {','.join(sorted(components.keys()))}"
+            )
+            sys.exit(1)
+        chromosome_order = DEFAULT_CHROMOSOME
+    # running index for the bubble index (BO) already used
+    bo = 0
+    total_bubbles = 0
+    # todo output final GFA with all the chromosomes ordered
+    out_gfa = []
+    out_csv = []
+    for chromosome in chromosome_order:
+        logger.info("Processing %s", chromosome)
+        component_nodes = components[chromosome]
+        # Initialize files
+        # f_gfa = open(outdir+'/'+gfa_filename.split("/")[-1][:-4]+'-'+chromosome+'.gfa', 'w')
+
+        scaffold_nodes, inside_nodes, node_order, bo, bubble_count = decompose_and_order(
+            graph, component_nodes, chromosome, bo
+        )
+
+        # skip a chromosome if something went wrong
+        if scaffold_nodes:
+            f_gfa = (
+                outdir
+                + os.sep
+                + gfa_filename.split(os.sep)[-1].split(".")[0]
+                + "-"
+                + chromosome
+                + ".gfa"
+            )
+            out_gfa.append(f_gfa)
+            csv_file = (
+                outdir + os.sep + gfa_filename.split(os.sep)[-1][:-4] + "-" + chromosome + ".csv"
+            )
+            out_csv.append(csv_file)
+            f_colors = open(csv_file, "w")
+            f_colors.write("Name,Color,SN,SO,BO,NO\n")
+            total_bubbles += bubble_count
+            for node_name in sorted(component_nodes):
+                node = graph.nodes[node_name]
+                bo_tag, no_tag = node_order[node_name]
+
+                node.tags["BO"] = ("i", bo_tag)
+                node.tags["NO"] = ("i", no_tag)
+
+                if node_name in scaffold_nodes:
+                    color = "orange"
+                elif node_name in inside_nodes:
+                    color = "blue"
+                else:
+                    color = "gray"
+
+                if "SN" in node.tags:
+                    sn_tag = node.tags["SN"][1]
+                else:
+                    sn_tag = "NA"
+                if "SO" in node.tags:
+                    so_tag = node.tags["SO"][1]
+                else:
+                    so_tag = "NA"
+                f_colors.write(
+                    "{},{},{},{},{},{}\n".format(node_name, color, sn_tag, so_tag, bo_tag, no_tag)
+                )
+
+            graph.write_gfa(
+                set_of_nodes=component_nodes,
+                output_file=f_gfa,
+                append=False,
+                order_bo=True,
+            )
+
+            f_colors.close()
+
+        else:
+            logger.warning(f"Chromosome {chromosome} was skipped")
+    if not by_chrom:
+        final_gfa = (
+            outdir + os.sep + gfa_filename.split(os.sep)[-1].split(".")[0] + "-complete" + ".gfa"
+        )
+        final_csv = (
+            outdir + os.sep + gfa_filename.split(os.sep)[-1].split(".")[0] + "-complete" + ".csv"
+        )
+        with open(final_gfa, "w") as outfile:
+            # outputting all the S lines first
+            for f in out_gfa:
+                with open(f, "r") as infile:
+                    for l in infile:
+                        if l.startswith("S"):
+                            outfile.write(l)
+            # outputting all the S lines
+            for f in out_gfa:
+                with open(f, "r") as infile:
+                    for l in infile:
+                        if l.startswith("L"):
+                            outfile.write(l)
+            for f in out_gfa:
+                os.remove(f)
+
+        with open(final_csv, "w") as outfile:
+            for f in out_csv:
+                with open(f, "r") as infile:
+                    for l in infile:
+                        outfile.write(l)
+            for f in out_csv:
+                os.remove(f)
+
+    logger.info("Total bubbles: %d", total_bubbles)
+
+
+def decompose_and_order(graph, component, component_name, bo_start=0):
+    """
+    This function takes the graph and a component
+    detects all biconnected components, order the scaffold nodes and starts ordering
+    """
+    logger.info(f" Input component: {len(component)} nodes")
+    logger.info(f" Finding Biconnected Components of the component {component_name}")
+    if len(component) == 1:
+        node = list(component)[0]
+        return component, set(), {node: (bo_start, 0)}, bo_start + 1, 0
+
+    new_graph = graph.graph_from_comp(component)
+    start = time.perf_counter()
+    all_biccs, artic_points = new_graph.biccs()
+    logger.info(
+        f" It took {time.perf_counter() - start} seconds to find the Biconnected Components"
+    )
+    bubbles = []
+    scaffold_graph = GFA()
+    scaffold_node_types = dict()
+    for n in artic_points:
+        scaffold_graph.add_node(n)
+        scaffold_node_types[n] = "s"
+    inside_nodes = set()
+
+    for bc in all_biccs:
+        # the components still have the articulation nodes in, need to be removed
+        bc_inside_nodes = bc.difference(artic_points)
+        bc_end_nodes = bc.intersection(artic_points)
+        inside_nodes.update(bc_inside_nodes)
+
+        if len(bc_inside_nodes) == 0:
+            assert len(bc_end_nodes) == 2
+            node1, node2 = tuple(bc_end_nodes)
+            scaffold_graph.add_edge(node1, "+", node2, "+", 0)
+
+        else:
+            bubble_index = len(bubbles)
+            bubbles.append(bc_inside_nodes)
+            scaffold_graph.add_node(str(bubble_index))
+            scaffold_node_types[str(bubble_index)] = "b"
+            for end_node in bc_end_nodes:
+                scaffold_graph.add_edge(str(bubble_index), "+", end_node, "+", 0)
+
+    logger.info(f"  Bubbles: {len(bubbles)}")
+    logger.info(f"  Scaffold graph: {len(scaffold_graph)} nodes")
+
+    # Find start/end points of the line by looking for nodes with degree 1
+    degree_one = [x.id for x in scaffold_graph.nodes.values() if len(x.neighbors()) == 1]
+    degree_two = [x.id for x in scaffold_graph.nodes.values() if len(x.neighbors()) == 2]
+
+    try:
+        assert len(degree_one) == 2
+    except AssertionError:
+        logger.warning(
+            f"Error: In Chromosome {component_name}, we found more or less than two nodes with degree 1. Skipping this chromosome"
+        )
+        # hacky but for now maybe ok
+        return None, None, None, None, None
+
+    try:
+        assert len(degree_two) == len(scaffold_graph) - 2
+    except AssertionError:
+        logger.warning(
+            f"Error: In Chromosome {component_name}, the number of nodes with degree 2 did not mach the expected number"
+        )
+        return None, None, None, None, None
+
+    # the scaffold graph should be a line graph here
+    traversal = scaffold_graph.dfs(degree_one[0])
+    traversal_scaffold_only = [
+        node_name for node_name in traversal if scaffold_node_types[node_name] == "s"
+    ]
+    # check that all scaffold nodes carry the same sequence name (SN), i.e. all came for the linear reference
+    assert len(set(new_graph[n].tags["SN"] for n in traversal_scaffold_only)) == 1
+    # I save tags as key:(type, value), so "SO":(i, '123')
+    coordinates = list(int(new_graph[n].tags["SO"][1]) for n in traversal_scaffold_only)
+
+    # make sure that the traversal is in ascending order
+    if coordinates[0] > coordinates[-1]:
+        traversal.reverse()
+        traversal_scaffold_only.reverse()
+        coordinates.reverse()
+    for i in range(len(coordinates) - 1):
+        assert coordinates[i] < coordinates[i + 1]
+    # compute dictionary mapping each node name to the corresponding "bubble order" and "node order" (BO,NO)
+    node_order = dict()
+    bo = bo_start
+    for node in traversal:
+        node_type = scaffold_node_types[node]
+        if node_type == "s":
+            node_order[node] = (bo, 0)
+        elif node_type == "b":
+            for i, n in enumerate(sorted(bubbles[int(node)])):
+                node_order[n] = (bo, i + 1)
+        else:
+            assert False
+        bo += 1
+    return artic_points, inside_nodes, node_order, bo, len(bubbles)
+
+
+def count_sn(graph, comp):
+    """
+    counts which SN tag is the majority in that chromosome to name the chromosome accordingly
+    """
+    counts = defaultdict(int)
+    for n in comp:
+        if "SN" not in graph[n].tags:
+            continue
+        counts[graph[n].tags["SN"][1]] += 1
+    return counts
+
+
+def name_comps(graph, components):
+    """
+    the graph is GFA object and components is a list of sets of the node ids of each component
+    for each component we take a majority vote of the SN tage and name the component accordingly
+    """
+    named_comps = dict()
+    current_tag = ""
+    for comp in components:
+        counts = count_sn(graph, comp)
+        most_freq = 0
+        for tag, count in counts.items():
+            if most_freq <= count:
+                current_tag, most_freq = tag, count
+        if current_tag == "":
+            raise ValueError(
+                "Were not able to assign a chromosome to component, SN tags could be missing"
+            )
+        named_comps[current_tag] = comp
+    return named_comps
+
+
+# fmt: off
+def add_arguments(parser):
+    arg = parser.add_argument
+    arg("--chromosome_order", default="",
+        help="Order in which to arrange chromosomes in terms of BO sorting. "
+        "Expecting comma-separated list. Default: chr1,...,chr22,chrX,chrY,chrM")
+    arg("--with-sequence", default=False, action="store_true",
+        help="Retain sequences in output (default is to strip sequences)")
+    arg("gfa_filename", metavar="GRAPH",
+        help="Input rGFA file")
+    arg("--outdir", default="./out",
+        help='Output Directory to store all the GFA and CSV files. Default location is a "out" folder from the directory of execution.')
+    arg("--by-chrom", default=False, action="store_true",
+        help="Outputs each chromosome as a separate GFA, otherwise, all chromosomes in one GFA file")
+# fmt: on
+
+
+def main(args):
+    run_order_gfa(**vars(args))

gaftools/cli/phase.py

@@ -0,0 +1,135 @@
+"""
+Add phasing information to the GAF file from a haplotag TSV file.
+
+The script uses the TSV file containing the haplotag information generated from WhatsHap's haplotag command.
+The H1 and H2 labels for each read are then added to the reads in the GAF file.
+"""
+
+import logging
+
+from gaftools.utils import FileWriter
+from gaftools.cli import log_memory_usage
+from gaftools.timer import StageTimer
+from gaftools.gaf import GAF
+
+
+class Node:
+    def __init__(self, chr_name, haplotype, phase_set):
+        self.chr_name = chr_name
+        self.phase_set = phase_set
+        self.haplotype = haplotype
+
+
+logger = logging.getLogger(__name__)
+timers = StageTimer()
+
+
+def run(gaf_file, tsv_file, output=None):
+    add_phase_info(gaf_file, tsv_file, output)
+    logger.info("\n== SUMMARY ==")
+    total_time = timers.total()
+    log_memory_usage()
+    logger.info("Time to read TSV                             %9.2f s", timers.elapsed("read_tsv"))
+    logger.info(
+        "Time to read and write GAF                   %9.2f s", timers.elapsed("read_write_gaf")
+    )
+    logger.info("Time spent on rest:                          %9.2f s", total_time - timers.sum())
+    logger.info("Total time:                                  %9.2f s", total_time)
+
+
+def add_phase_info(gaf_path, tsv_path, out_path):
+    """This function adds phasing information to the gaf file using .tsv file of the WhatsHap
+    Haplotag... The information is added as two tags ("ps:Z" and "ht:Z") before the cigar string
+    """
+
+    logger.info("INFO: Adding phasing information...")
+
+    tsv_file = open(tsv_path, "r")
+
+    phase = {}
+    with timers("read_tsv"):
+        for line in tsv_file:
+            line_elements = line.rstrip().split("\t")
+
+            if line_elements[0] not in phase:
+                tmp = Node(line_elements[3], line_elements[1], line_elements[2])
+                phase[line_elements[0]] = tmp
+
+    gaf_out = FileWriter(out_path)
+
+    line_count = 0
+    missing_in_tsv = 0
+    phased = 0
+    gaf_file = GAF(gaf_path)
+    with timers("read_write_gaf"):
+        for gaf_line in gaf_file.read_file():
+            if line_count != 0:
+                gaf_out.write("\n")
+
+            line_count += 1
+
+            gaf_out.write(
+                "%s\t%s\t%s\t%s\t+\t%s\t%d\t%d\t%d\t%d\t%d\t%d"
+                % (
+                    gaf_line.query_name,
+                    gaf_line.query_length,
+                    gaf_line.query_start,
+                    gaf_line.query_end,
+                    gaf_line.path,
+                    gaf_line.path_length,
+                    gaf_line.path_start,
+                    gaf_line.path_end,
+                    gaf_line.residue_matches,
+                    gaf_line.alignment_block_length,
+                    gaf_line.mapping_quality,
+                )
+            )
+
+            in_tsv = True
+            if gaf_line.query_name not in phase:
+                missing_in_tsv += 1
+                in_tsv = False
+
+            if in_tsv and phase[gaf_line.query_name].haplotype != "none":
+                gaf_out.write(
+                    "\tps:Z:%s-%s\tht:Z:%s\t"
+                    % (
+                        phase[gaf_line.query_name].chr_name,
+                        phase[gaf_line.query_name].phase_set,
+                        phase[gaf_line.query_name].haplotype,
+                    )
+                )
+                phased += 1
+            else:
+                gaf_out.write("\tps:Z:none\tht:Z:none")
+
+            for k in gaf_line.tags.keys():
+                gaf_out.write("\t%s:%s" % (k, gaf_line.tags[k]))
+
+            gaf_out.write("\t%s" % gaf_line.cigar)
+
+    logger.info(
+        "INFO: Added phasing info (ps:Z and ht:Z) for %d reads out of %d GAF lines - (%d reads are missing in .tsv)"
+        % (phased, line_count, missing_in_tsv)
+    )
+    gaf_file.close()
+    tsv_file.close()
+    gaf_out.close()
+
+
+# fmt: off
+def add_arguments(parser):
+    arg = parser.add_argument
+
+    # Positional arguments
+    arg('gaf_file', metavar='GAF',
+        help='GAF file (can be bgzip-compressed)')
+    arg('tsv_file', metavar='TSV',
+        help='WhatsHap haplotag TSV file. Refer to https://whatshap.readthedocs.io/en/latest/guide.html#whatshap-haplotag')
+    arg('-o', '--output', default=None,
+        help='Output GAF file (bgzipped if the file ends with .gz). If omitted, output is directed to stdout.')
+# fmt: on
+
+
+def main(args):
+    run(**vars(args))