From e69367118582e943fd05d13008b79c7cb675e61d Mon Sep 17 00:00:00 2001 From: Haibao Tang Date: Sat, 27 Apr 2024 23:31:25 -0700 Subject: [PATCH] Fix type --- jcvi/apps/base.py | 4 ++-- jcvi/assembly/geneticmap.py | 8 ++++++-- jcvi/assembly/hic.py | 8 +++----- jcvi/graphics/base.py | 24 ++++++++++++------------ 4 files changed, 23 insertions(+), 21 deletions(-) diff --git a/jcvi/apps/base.py b/jcvi/apps/base.py index ddac6699..0d101bf9 100644 --- a/jcvi/apps/base.py +++ b/jcvi/apps/base.py @@ -25,7 +25,7 @@ from socket import gethostname from subprocess import PIPE, call, check_output from optparse import OptionParser as OptionP, OptionGroup, SUPPRESS_HELP -from typing import Any, Collection, List, Optional, Union +from typing import Any, Collection, List, Optional, Tuple, Union from natsort import natsorted from rich.console import Console @@ -39,7 +39,7 @@ os.environ["LC_ALL"] = "C" JCVIHELP = "JCVI utility libraries {} [{}]\n".format(__version__, __copyright__) -TextCollection = Union[str, Iterable[str]] +TextCollection = Union[str, List[str], Tuple[str]] def debug(level=logging.DEBUG): diff --git a/jcvi/assembly/geneticmap.py b/jcvi/assembly/geneticmap.py index 2baeb9e2..612bb000 100644 --- a/jcvi/assembly/geneticmap.py +++ b/jcvi/assembly/geneticmap.py @@ -12,6 +12,7 @@ from functools import lru_cache from itertools import combinations, groupby from random import sample +from typing import List import numpy as np import seaborn as sns @@ -20,7 +21,7 @@ from ..algorithms.matrix import symmetrize from ..formats.base import BaseFile, LineFile, must_open, read_block from ..formats.bed import Bed, fastaFromBed -from ..graphics.base import plt, savefig, Rectangle, draw_cmap +from ..graphics.base import Rectangle, draw_cmap, plt, plot_heatmap, savefig MSTheader = """population_type {0} @@ -332,7 +333,10 @@ def dotplot(args): @lru_cache(maxsize=None) -def calc_ldscore(a, b): +def calc_ldscore(a: List[str], b: List[str]) -> float: + """ + Calculate Linkage disequilibrium (r2) between two genotypes. + """ assert len(a) == len(b), "{0}\n{1}".format(a, b) # Assumes markers as A/B c = Counter(zip(a, b)) diff --git a/jcvi/assembly/hic.py b/jcvi/assembly/hic.py index 021224a7..d5bb099e 100644 --- a/jcvi/assembly/hic.py +++ b/jcvi/assembly/hic.py @@ -46,7 +46,7 @@ savefig, ) from ..graphics.dotplot import dotplot -from ..utils.cbook import gene_name, human_size +from ..utils.cbook import gene_name # Map orientations to ints FF = {"+": 1, "-": -1, "?": 1} @@ -782,9 +782,7 @@ def heatmap(args): breaks = sorted(breaks)[1:] if contig or opts.nobreaks: breaks = [] - binsize = human_size(resolution, precision=0) - title = f"Resolution = {binsize} per bin" - plot_heatmap(ax, B, breaks, iopts, groups=new_groups, title=title) + plot_heatmap(ax, B, breaks, groups=new_groups, binsize=resolution) # Title pf = npyfile.rsplit(".", 1)[0] @@ -1639,7 +1637,7 @@ def movieframe(args): ax2_root = fig.add_axes([0.5, 0, 0.5, 1]) # dot plot canvas # Left axis: heatmap - plot_heatmap(ax1, M, breaks, iopts) + plot_heatmap(ax1, M, breaks, binsize=BINSIZE) # Right axis: synteny qbed, sbed, qorder, sorder, is_self = check_beds(anchorsfile, p, opts, sorted=False) diff --git a/jcvi/graphics/base.py b/jcvi/graphics/base.py index 755b3299..9e00a620 100644 --- a/jcvi/graphics/base.py +++ b/jcvi/graphics/base.py @@ -11,7 +11,6 @@ logging.getLogger("numexpr").setLevel(logging.WARNING) logging.getLogger("PIL").setLevel(logging.INFO) - from functools import partial import numpy as np @@ -39,6 +38,7 @@ from ..apps.base import datadir, glob, listify, logger, sample_N, which from ..formats.base import LineFile +from ..utils.cbook import human_size CHARS = { @@ -333,7 +333,7 @@ def savefig(figname, dpi=150, iopts=None, cleanup=True): # human readable size (Kb, Mb, Gb) -def human_readable(x, pos, base=False): +def human_readable(x: Union[str, int], _, base=False): x = str(int(x)) if x.endswith("000000000"): x = x[:-9] + "G" @@ -504,10 +504,9 @@ def plot_heatmap( ax, M: np.array, breaks: List[int], - iopts: ImageOptions, groups: List[Tuple[int, int, List[Tuple[int, str]], str]] = [], plot_breaks: bool = False, - title: str = "", + binsize: Optional[int] = None, ): """Plot heatmap illustrating the contact probabilities in Hi-C data. @@ -518,7 +517,7 @@ def plot_heatmap( iopts (OptionParser options): Graphical options passed in from commandline groups (List, optional): [(start, end, [(position, seqid)], color)]. Defaults to []. plot_breaks (bool): Whether to plot white breaks. Defaults to False. - title (str): Title of the heatmap at bottom + binsize (int, optional): Resolution of the heatmap. """ cmap = sns.cubehelix_palette(rot=0.5, as_cmap=True) ax.imshow(M, cmap=cmap, interpolation="none") @@ -550,13 +549,14 @@ def simplify_seqid(seqid): ax.set_xticklabels(ax.get_xticks(), family="Helvetica", color="gray") ax.set_yticklabels(ax.get_yticks(), family="Helvetica", color="gray", rotation=90) ax.tick_params(left=True, bottom=True, labelleft=True, labelbottom=True) - formatter = ticker.FuncFormatter( - lambda x, pos: human_readable(int(x) * binsize, pos, base=True) - ) - ax.xaxis.set_major_formatter(formatter) - ax.yaxis.set_major_formatter(formatter) - binlabel = "Resolution = {} per bin".format(human_size(binsize, precision=0)) - ax.set_xlabel(binlabel) + if binsize is not None: + formatter = ticker.FuncFormatter( + lambda x, pos: human_readable(int(x) * binsize, pos, base=True) + ) + ax.xaxis.set_major_formatter(formatter) + ax.yaxis.set_major_formatter(formatter) + title = f"Resolution = {human_size(binsize, precision=0)} per bin" + ax.set_xlabel(title) def discrete_rainbow(N=7, cmap=cm.Set1, usepreset=True, shuffle=False, plot=False):