maphelios 0.8.3__py3-none-any.whl

maphelios/__init__.py

@@ -0,0 +1,5 @@
+from .comparison import Comparison
+from .insert import Insert
+from .mapping import Mapping
+
+__version__ = "0.8.3"

maphelios/__main__.py

@@ -0,0 +1,6 @@
+#!/usr/bin/env python3
+
+from .cli import main
+
+if __name__ == "__main__":
+    main()

maphelios/cli.py

@@ -0,0 +1,106 @@
+import argparse
+
+import matplotlib.pyplot as plt
+
+from .main import Pipeline
+from .plotting import plot_insert, plot_insert_dists, plot_genome
+
+
+def get_args(*args):
+    # Create the CLI using argparse
+    # TODO: maybe switch from argparse to click
+    parser = argparse.ArgumentParser(description="Generating plots")
+
+    subparsers = parser.add_subparsers(help="sub-command help")
+
+    parser_ncbi = subparsers.add_parser("ncbi", help="Download genome from NCBI")
+    parser_ncbi.add_argument("search_term", help="search term for BLAST")
+    parser_ncbi.add_argument("email", help="Your email address - needed by BLAST")
+    parser_ncbi.add_argument(
+        "--retmax",
+        default=200,
+        help="Maximum number of records to download from NCBI",
+    )
+
+    parser_genome = subparsers.add_parser("file", help="Specify a local genome.")
+    parser_genome.add_argument("path_to_genome")
+
+    parser.add_argument("seq_file", help="file with the sequences")
+    parser.add_argument("output_prefix", help="Output directory")
+
+    parser.add_argument(
+        "--fwd_suffix", default="_F", help="Suffix for forward seq. default='_F'"
+    )
+    parser.add_argument(
+        "--rev_suffix", default="_R", help="Suffix for reverse seq. default='_R'"
+    )
+    parser.add_argument(
+        "--insert_types", help="matched, unmatched or both inserts", default="both"
+    )
+    parser.add_argument(
+        "--filter",
+        default=0.5,
+        help="Threshold value for filtering out blast results with low coverage",
+    )
+
+    # Get all the command line arguments
+    return parser.parse_args(*args)
+
+
+def main():
+    args = get_args()
+    if "path_to_genome" in args:
+        search_term = None
+        email = None
+        retmax = None
+        path_to_genome = args.path_to_genome
+    else:
+        search_term = args.search_term
+        email = args.email
+        retmax = int(args.retmax)
+        path_to_genome = None
+
+    pipeline = Pipeline(
+        args.seq_file,
+        args.output_prefix,
+        genome_file=path_to_genome,
+        search_term=search_term,
+        email=email,
+        retmax=retmax,
+        fwd_suffix=args.fwd_suffix,
+        rev_suffix=args.rev_suffix,
+    )
+
+    # Create linear plots of inserts with annotations
+    for seq_id in pipeline.seq_ids:
+        inserts = pipeline.get_inserts(
+            seq_id, insert_types=args.insert_types, filter_threshold=args.filter
+        )
+        for i, insert in enumerate(inserts):
+            fig, axs = plt.subplots(2, 1, figsize=(10, 10), height_ratios=[3, 5])
+            axs = plot_insert(insert, pipeline.genome, axs=axs)
+            fig.savefig(pipeline.work_dir / f"{seq_id}_hit{i}.png", dpi=300)
+            plt.close()
+
+    # Get all inserts (used in both next plots)
+    inserts = pipeline.get_all_inserts(args.insert_types, filter_threshold=args.filter)
+
+    # Create a plot of genome / all contigs as circular plot with inserts
+    # layered on top
+    fig, ax = plt.subplots(figsize=(10, 30))
+    ax = plot_on_genome(inserts, pipeline.genome, ax=ax)
+    fig.savefig(pipeline.work_dir / "genome_plot.png", dpi=300)
+
+    # Plots of distributions of insert sizes
+    fig, axs = plt.subplots(1, 3, figsize=(12, 5))
+    axs = plot_insert_dists(inserts, axs=axs)
+    fig.savefig(pipeline.work_dir / "insert_length_dist.png", dpi=300)
+
+    df = pipeline.to_dataframe(
+        insert_types=args.insert_types, filter_threshold=args.filter
+    )
+    df.to_csv(pipeline.work_dir / "inserts.csv", index=False)
+
+
+if __name__ == "__main__":
+    main()

maphelios/comparison.py

@@ -0,0 +1,352 @@
+from collections import defaultdict
+from itertools import chain, cycle
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from dna_features_viewer import CircularGraphicRecord
+from matplotlib import color_sequences
+
+
+class Clusters(dict):
+    def __init__(self, clusters, genomes=None):
+        super(Clusters, self).__init__(clusters)
+
+        self.genomes = genomes or {g for g, clust_idx in self.keys()}
+        self.cluster_labels = list(self.keys())
+
+    def __contains__(self, item):
+        return item in self.cluster_labels
+
+    @property
+    def insert_ids(self):
+        insert_ids = {
+            genome: list(
+                chain.from_iterable(
+                    (
+                        insert_ids
+                        for (g, clust_idx), insert_ids in self.items()
+                        if g == genome
+                    )
+                )
+            )
+            for genome in self.genomes
+        }
+        return insert_ids
+
+    @property
+    def insert_labels(self):
+        seq_labels = defaultdict(dict)
+        for g, clust_idx in self.cluster_labels:
+            insert_ids = self.__getitem__((g, clust_idx))
+            # Set it such that only one label is given per cluster
+            for insert_id in insert_ids:
+                seq_labels[g][insert_id] = None
+            seq_labels[g][insert_ids[-1]] = f"Cluster {clust_idx}"
+        return seq_labels
+
+    @property
+    def other_repr(self):
+        other = defaultdict(list)
+        for (g, clust_idx), insert_ids in self.items():
+            other[g].append((clust_idx, insert_ids))
+        return other
+
+    def subset(self, selection):
+        return Clusters(
+            {
+                (g, clust_idx): self.__getitem__((g, clust_idx))
+                for g, clust_idx in selection
+            }
+        )
+
+
+class Comparison(dict):
+
+    def __init__(self, *args, **kwargs):
+        super(Comparison, self).__init__(*args, **kwargs)
+
+        self.insert_ids = {g: mapping.insert_ids for g, mapping in self.items()}
+
+    def get_inserts(
+        self,
+        genomes=None,
+        seq_ids=None,
+        insert_ids=None,
+        insert_type="both",
+        filter_threshold=None,
+        **kwargs,
+    ):
+        if genomes is None:
+            genomes = self.keys()
+
+        return [
+            self.__getitem__(g).get(
+                insert_ids=insert_ids,
+                seq_ids=seq_ids,
+                insert_type=insert_type,
+                filter_threshold=filter_threshold,
+            )
+            for g in genomes
+        ]
+
+    def get_inserts_df(
+        self, genomes=None, insert_type="both", filter_threshold=None, **kwargs
+    ):
+        if genomes is None:
+            genomes = self.keys()
+
+        inserts_dfs = [
+            self.__getitem__(genome)
+            .to_dataframe(insert_type=insert_type, filter_threshold=filter_threshold)
+            .assign(genome=genome)
+            for genome in genomes
+        ]
+
+        if len(inserts_dfs):
+            return pd.concat(inserts_dfs, ignore_index=True)
+
+    def get_insert_presence_df(
+        self, genomes=None, insert_type="both", filter_threshold=None, **kwargs
+    ):
+        if genomes is None:
+            genomes = self.keys()
+
+        dfs = [
+            pd.DataFrame(
+                self.__getitem__(g).get_seq_ids(insert_type, filter_threshold),
+                columns=["insert_ids"],
+            ).assign(genome=g)
+            for g in genomes
+        ]
+
+        all_seq_ids = sorted(
+            set(chain.from_iterable(self.__getitem__(g).seq_ids for g in genomes))
+        )
+
+        df_insert_presence = (
+            pd.concat(dfs, ignore_index=True)
+            .groupby(["insert_ids", "genome"], as_index=False)
+            .agg(num_inserts=pd.NamedAgg("insert_ids", "count"))
+            .pivot(index="insert_ids", columns="genome", values="num_inserts")
+            .reindex(columns=genomes, index=all_seq_ids)
+        )
+
+        return df_insert_presence
+
+    def get_clusters(
+        self,
+        genomes=None,
+        insert_type="both",
+        filter_threshold=None,
+        plain_dict=True,
+        **kwargs,
+    ):
+        if genomes is None:
+            genomes = self.keys()
+
+        if plain_dict:
+            return {
+                g: self.__getitem__(g).get_clusters(insert_type, filter_threshold)
+                for g in genomes
+            }
+        else:
+            return Clusters(
+                {
+                    (g, clust_idx): cluster
+                    for g in genomes
+                    for clust_idx, cluster in enumerate(
+                        self.__getitem__(g).get_clusters(insert_type, filter_threshold)
+                    )
+                },
+                genomes=genomes,
+            )
+
+    def get_labels(
+        self,
+        seq_ids,
+        insert_ids=None,
+        genomes=None,
+        insert_type="both",
+        filter_threshold=None,
+    ):
+        if genomes is None:
+            genomes = self.keys()
+
+        if insert_ids is None:
+            insert_ids = {g: None for g in genomes}
+
+        if seq_ids is not None or insert_ids is not None:
+            return {
+                g: self.__getitem__(g).get_labels(
+                    seq_ids, insert_ids.get(g), genomes, insert_type, filter_threshold
+                )
+                for g in genomes
+            }
+
+    def get_genes_df(
+        self,
+        seq_ids=None,
+        insert_ids=None,
+        insert_type="both",
+        filter_threshold=None,
+        buffer=4000,
+        **kwargs,
+    ):
+        if seq_ids is None:
+            seq_ids = {x: None for x in self.keys()}
+        elif isinstance(seq_ids, (list, tuple)):
+            seq_ids = {x: seq_ids for x in self.keys()}
+
+        if insert_ids is None:
+            insert_ids = {x: None for x in self.keys()}
+        elif isinstance(insert_ids, (list, tuple)):
+            insert_ids = {x: insert_ids for x in self.keys()}
+
+        genomes = set(seq_ids.keys()) | set(insert_ids.keys())
+
+        genes_dfs = [
+            self.__getitem__(g)
+            .genes_to_dataframe(
+                seq_ids=seq_ids.get(g),
+                insert_ids=insert_ids.get(g),
+                insert_type=insert_type,
+                filter_threshold=filter_threshold,
+                buffer=buffer,
+            )
+            .assign(genome=g)
+            for g in genomes
+        ]
+
+        if len(genes_dfs):
+            drop_by = ["start", "end", "type", "insert_idx", "seq_id", "genome"]
+            df_genes = pd.concat(genes_dfs, ignore_index=True).drop_duplicates(
+                subset=drop_by
+            )
+            df_genes.insert(0, "genome", df_genes.pop("genome"))
+            return df_genes
+
+    def get_graphic_features(
+        self,
+        seq_ids=None,
+        insert_ids=None,
+        genomes_order=None,
+        seq_labels=None,
+        show_contig_labels=True,
+        insert_type="both",
+        filter_threshold=None,
+        palette="tab10",
+        **kwargs,
+    ):
+        if seq_ids is None:
+            seq_ids = {x: None for x in self.keys()}
+        elif isinstance(seq_ids, (list, tuple)):
+            seq_ids = {x: seq_ids for x in self.keys()}
+
+        if insert_ids is None:
+            insert_ids = {x: None for x in self.keys()}
+        elif isinstance(insert_ids, (list, tuple)):
+            insert_ids = {x: insert_ids for x in self.keys()}
+
+        if genomes_order is None:
+            genomes_order = set(seq_ids.keys()) | set(insert_ids.keys())
+
+        if seq_labels is None:
+            seq_labels = {x: None for x in self.keys()}
+
+        palette = cycle(color_sequences[palette])
+
+        res = {
+            g: self.__getitem__(g).get_graphic_features(
+                seq_ids=seq_ids.get(g),
+                insert_ids=insert_ids.get(g),
+                insert_type=insert_type,
+                filter_threshold=filter_threshold,
+                contig_labels=show_contig_labels,
+                seq_labels=seq_labels.get(g),
+                col1=col,
+                col2=col,
+                linecolor=col,
+                **kwargs,
+            )
+            for g, col in zip(genomes_order, palette)
+        }
+
+        return res
+
+    def plot(
+        self,
+        seq_ids=None,
+        insert_ids=None,
+        genomes_order=None,
+        seq_labels=None,
+        insert_type="both",
+        filter_threshold=None,
+        feature_kwargs=None,
+        show_contig_labels=True,
+        show_titles=False,
+        palette="tab10",
+        facet_wrap=None,
+        fig=None,
+        **kwargs,
+    ):
+        feature_kwargs = {} if feature_kwargs is None else feature_kwargs
+
+        # Get the graphic feature objects
+        res = self.get_graphic_features(
+            seq_ids,
+            insert_ids,
+            genomes_order,
+            seq_labels,
+            show_contig_labels,
+            insert_type,
+            filter_threshold,
+            palette,
+            **feature_kwargs,
+        )
+
+        # Set all the plotting options based on input
+        if "figsize" not in kwargs:
+            kwargs["figsize"] = (10, 8)
+
+        genome_labels = res.keys()
+        num_genomes = len(genome_labels)
+
+        # Set the number of axes in the figure based on facet_wrap option
+        nrows, ncols = 1, 1
+        if facet_wrap is not None:
+            ncols = min(facet_wrap, num_genomes)
+            nrows = num_genomes // ncols
+
+        # Create figure and axes if not provided
+        if fig is None:
+            fig, axs = plt.subplots(nrows=nrows, ncols=ncols, **kwargs)
+        else:
+            axs = fig.get_axes()
+
+        # Check that axs is iterable even if a single axes object
+        if isinstance(axs, plt.Axes):
+            axs = np.array([axs])
+
+        # Make and label all the plots
+        if facet_wrap is None:
+            features = list(chain.from_iterable([x[1] for x in res.values()]))
+            rec = CircularGraphicRecord(max([x[0] for x in res.values()]), features)
+            _ = rec.plot(axs[0], annotate_inline=False)
+
+            if show_titles:
+                axs[0].set_title(", ".join(list(genome_labels)))
+
+        else:
+            for (label, (seq_len, features)), ax in zip(res.items(), axs.flatten()):
+                rec = CircularGraphicRecord(seq_len, features)
+                _ = rec.plot(ax, annotate_inline=False)
+
+                if show_titles:
+                    ax.set_title(label)
+
+            # Remove redundant axis
+            for ax in axs[num_genomes:]:
+                ax.axis("off")
+
+        return fig

maphelios/genome_funcs.py

@@ -0,0 +1,136 @@
+import subprocess
+import warnings
+from hashlib import sha1
+from pathlib import Path
+
+from BCBio import GFF
+from Bio import Entrez, SearchIO, SeqIO
+
+
+def _download_genome(search_term, retmax=100):
+    # Set the email address for NCBI queries
+    # Entrez.email = email
+
+    # Get record ids from NCBI
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        with Entrez.esearch(db="nucleotide", term=search_term, retmax=retmax) as handle:
+            res = Entrez.read(handle)
+
+    # Download these records from NCBI rettype="gbwithparts" is used to
+    # download the sequences with record features
+    with Entrez.efetch(
+        db="nucleotide",
+        id=",".join(res["IdList"]),
+        rettype="gbwithparts",
+        retmode="text",
+    ) as handle:
+        data_gb = SeqIO.to_dict(SeqIO.parse(handle, "gb"))
+
+    return data_gb
+
+
+def download_genome(search_term, retmax=100, output_path=None):
+
+    if output_path is None:
+        data_gb = _download_genome(search_term, retmax)
+
+    else:
+        output_path = Path(output_path)
+
+        # Check output_path is given
+        if Path(output_path).exists():
+            # Read in the file
+            data_gb = SeqIO.to_dict(SeqIO.parse(output_path, "genbank"))
+
+        # Otherwise retrieve the records from NCBI
+        else:
+            data_gb = _download_genome(search_term, retmax)
+            SeqIO.write(data_gb.values(), output_path, "genbank")
+
+    return data_gb
+
+
+def get_genome_file(work_dir, search_term, retmax):
+    work_dir = Path(work_dir)
+    # Hash the search term to use as filename in cache dir
+    search_hashed = sha1((search_term + str(retmax)).encode()).hexdigest()
+    return work_dir / f"db_{search_hashed}.gbk"
+
+
+def get_genome(genome_file, genome_fasta, search_term=None, retmax=None):
+    # Get genome
+    if search_term is not None:
+        genome = download_genome(search_term, retmax, genome_file)
+    else:
+        if genome_file.suffix == ".gff":
+            genome = SeqIO.to_dict(GFF.parse(genome_file))
+        elif genome_file.suffix == ".gbk":
+            genome = SeqIO.to_dict(SeqIO.parse(genome_file, "genbank"))
+        elif genome_file.suffix in (".fa", ".fasta"):
+            genome = SeqIO.to_dict(SeqIO.parse(genome_file, "fasta"))
+        else:
+            raise RuntimeError("Wrong format expected either `.gbk` or `.gff` file.")
+
+    defined_seqs = [x for x in genome.values() if x.seq.defined]
+
+    if not len(defined_seqs):
+        message = f"Genome file {genome_file} does not contain any sequences!"
+        if search_term is not None:
+            message = (
+                f"Search term '{search_term}' did not yield a genome with sequences"
+            )
+        raise RuntimeError(message)
+
+    # Save in fasta format (only acceptable format for makeblastdb)
+    SeqIO.write(defined_seqs, genome_fasta, "fasta")
+
+    return genome
+
+
+def _correct_hit_id(x):
+    """Helper function to get rid red| or emb| added by BLAST to contig ID"""
+    if "|" in x.id:
+        x.id = x.id.split("|")[1]
+    return x
+
+
+def run_blast(seq_file, db_file, blast_output, blast_options=None):
+    # TODO: check how to catch errors from blast
+    # make blast database
+    run1 = subprocess.run(
+        f"makeblastdb -in '{db_file}' -parse_seqids -dbtype nucl",
+        shell=True,
+        check=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    if run1.returncode != 0:
+        raise OSError(f"Failed to run makeblastdb: {run1.stderr.decode()}")
+
+    if blast_options is None:
+        blast_options = ""
+
+    # align input sequences with query using blastn with default options
+    run2 = subprocess.run(
+        "blastn "
+        f"-query '{seq_file}' -db '{db_file}' "
+        f"-out '{blast_output}' -outfmt 5 {blast_options}",
+        shell=True,
+        check=True,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    )
+    if run2.returncode != 0:
+        raise OSError(f"Failed to run blastn: {run2.stderr.decode()}")
+
+    # Return a dictionary correcting the hit IDs which get an unnecessary
+    # prefix from BLAST
+    return {
+        x.id: x.hit_map(_correct_hit_id)
+        for x in SearchIO.parse(blast_output, "blast-xml")
+    }
+
+
+def run_minimap2(seq_file, db_file, output_file):
+    pass