libhts.py

# Copyright (C) 2020 Blaise Li
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.
from math import floor, ceil, sqrt, log
from functools import reduce
# from re import sub
import warnings
import numpy as np
import pandas as pd
# To compute bins in bigwig data
from scipy.stats import binned_statistic
# To compute correlation coefficient, and compute linear regression
from scipy.stats.stats import pearsonr, linregress
# To compute geometric mean
from scipy.stats.mstats import gmean
import matplotlib as mpl
import matplotlib.pyplot as plt
# TODO: set this at the "correct" place
# https://stackoverflow.com/a/42768093/1878788
#from matplotlib.backends.backend_pgf import FigureCanvasPgf
#mpl.backend_bases.register_backend('pdf', FigureCanvasPgf)
#TEX_PARAMS = {
#    "text.usetex": True,            # use LaTeX to write all text
#    "pgf.rcfonts": False,           # Ignore Matplotlibrc
#    "pgf.texsystem": "lualatex",  # hoping to avoid memory issues
#    "pgf.preamble": [
#        r'\usepackage{color}'     # xcolor for colours
#    ]
#}
#mpl.rcParams.update(TEX_PARAMS)
import seaborn as sns
# from rpy2.robjects import r, pandas2ri, Formula, StrVector
# as_df = r("as.data.frame")
# from rpy2.rinterface import RRuntimeError
# from rpy2.robjects.packages import importr
# deseq2 = importr("DESeq2")
from pybedtools import BedTool
import pyBigWig
import networkx as nx
from libworkflows import texscape


def formatwarning(
        message, category, filename, lineno, line):  # pylint: disable=W0613
    """Used to format warning messages."""
    return "%s:%s: %s: %s\n" % (filename, lineno, category.__name__, message)


warnings.formatwarning = formatwarning


# This might represent any type of genomic interval.
class Exon():
    """Object representing an exon."""
    __slots__ = ("chrom", "start", "end")
    def __init__(self, chrom, start, end):
        self.chrom = chrom
        self.start = start
        self.end = end

    def overlap(self, other):
        """
        Tell whether *self* and *other* overlap.
        """
        if self.chrom != other.chrom:
            return False
        return (self.start <= other.start < self.end) or (other.start <= self.start < other.end)

    def merge(self, other):
        """
        Create a new Exon object by merging *self* with *other*.
        """
        # Not necessary: can be indirectly linked
        #assert overlap(self, other)
        return Exon(self.chrom, min(self.start, other.start), max(self.end, other.end))

    def __len__(self):
        return self.end - self.start

OVERLAP = Exon.overlap
MERGE = Exon.merge

class Gene():
    """This object contains information obtained from a gtf file."""
    __slots__ = ("gene_id", "exons", "union_exon_length")
    def __init__(self, gene_id):
        self.gene_id = gene_id
        #self.transcripts = {}
        self.exons = nx.Graph()
        self.union_exon_length = None

    #def add_transcript(self, feature):
    #    the_id = feature.attrs["transcript_id"]
    #    assert the_id not in self.transcripts
    #    self.transcripts[the_id] = feature

    def add_exon(self, feature):
        """
        Add one Exon object to the exon graph based in the information in gtf
        information *feature*.
        """
        #the_id = feature.attrs["exon_id"]
        #assert the_id not in self.exons
        #self.exons[the_id] = feature
        exon = Exon(feature.chrom, feature.start, feature.end)
        if exon not in self.exons:
            self.exons.add_node(exon)

    # The merging cannot be done on the full BedTool because we dont want
    # to merge together exons not belonging to the same gene.
    def set_union_exon_length(self):
        """The exons are used to make a BedTool, which enables convenient merging of
        overlapping features. The sum of the lengths of the merged exons is returned."""
        if len(self.exons) == 1:
            # No need to merge when there is only one exon
            self.union_exon_length = len(next(iter(self.exons.nodes())))
        else:
            # Too slow
            #self.union_exon_length = sum(map(
            #    len, BedTool(self.exons.values()).merge().features()))
            #self.union_exon_length = 0
            # We group nodes that overlap, and merge them
            #overlapping_exons = nx.quotient_graph(self.exons, OVERLAP)
            #for node in overlapping_exons.nodes():
            #    mex = reduce(MERGE, node)
            #    self.union_exon_length += len(mex)
            self.union_exon_length = sum((len(reduce(
                MERGE, node)) for node in nx.quotient_graph(
                    self.exons, OVERLAP).nodes()))


def gtf_2_genes_exon_lengths(gtf_filename):
    """Returns a pandas DataFrame where union exon lengths are associated to gene IDs."""
    gtf_file = open(gtf_filename, "r")
    gtf = BedTool(gtf_file)
    genes = {}
    for feature in gtf.features():
        feat_type = feature[2]
        if feat_type != "exon":
            continue
        attrs = feature.attrs
        gene_id = attrs["gene_id"]
        if gene_id not in genes:
            genes[gene_id] = Gene(gene_id)
        gene = genes[gene_id]
        try:
            gene.add_exon(feature)
        except AssertionError:
            # A given exon may be registered for several transcripts, hence several gtf entries
            already = gene.exons[feature.attrs["exon_id"]]
            assert already.attrs["transcript_id"] != feature.attrs["transcript_id"]
            assert (already.start, already.end) == (feature.start, feature.end)
    for gene in genes.values():
        gene.set_union_exon_length()
    return pd.DataFrame(pd.Series(
        {gene.gene_id : gene.union_exon_length for gene in genes.values()},
        name=("union_exon_len")).rename_axis("gene"))


def repeat_bed_2_lengths(repeat_bed):
    """Computes the lengths of repetitive elements in a bed file, grouped by families.
    This assumes that the elements have their names composed of the family name,
    then a colon, then a number. For instance:
    Simple_repeat|Simple_repeat|(TTTTTTG)n:1
    Simple_repeat|Simple_repeat|(TTTTTTG)n:2
    Simple_repeat|Simple_repeat|(TTTTTTG)n:3
    Simple_repeat|Simple_repeat|(TTTTTTG)n:4
    -> Simple_repeat|Simple_repeat|(TTTTTTG)n
    Returns a DataFrame associating the summed lengths to the family names.
    """
    # usecols=[1, 2, 3]: start, end, id
    # index_col=2: id (relative to the selected columns)
    start_ends = pd.read_table(repeat_bed, usecols=[1, 2, 3], header=None, index_col=2)
    # bed lengths
    lens = start_ends[2] - start_ends[1]
    lens.name = "union_exon_len"
    repeat_families = [":".join(name.split(":")[:-1]) for name in start_ends.index]
    # The reads assigned to a repeated element can come
    # from the summed length of all the members of the family
    # We call this "gene" for convenience and compatibility
    return pd.DataFrame(lens).assign(gene=repeat_families).groupby("gene").sum()


def spikein_gtf_2_lengths(spikein_gtf):
    """Computes the lengths of spike-ins, grouped by families.
    Returns a DataFrame associating the summed lengths to the spike-in names.
    """
    spikein_ends = {}
    with open(spikein_gtf) as gtf_file:
        for line in gtf_file:
            fields = line.strip().split("\t")
            spikein_ends[fields[0]] = int(fields[4])
    return pd.DataFrame(pd.Series(
        spikein_ends,
        name=("union_exon_len")).rename_axis("gene"))


def id_list_gtf2bed(
        identifiers, gtf_filename,
        feature_type="transcript", id_kwd="gene_id"):
    """
    Extract bed coordinates of an iterable of identifiers from a gtf file.

    *identifiers* is the iterable of identifiers.
    *gtf_filename* is the name of the gtf file.
    *feature_type* is the type of feature to be considered
    in the gtf file (third columns).
    *id_kwd* is the keyword under which the feature ID is expected to be found
    in the feature annotations in the gtf_file. These feature IDs will be
    matched against the elements in *identifiers*.
    """
    if identifiers:
        ids = set(identifiers)

        def feature_filter(feature):
            return feature[2] == feature_type and feature[id_kwd] in ids
        gtf = BedTool(gtf_filename)
        return gtf.filter(feature_filter)

    # https://stackoverflow.com/a/13243870/1878788
    def empty_bed_generator():
        return
        yield  # pylint: disable=W0101
    return empty_bed_generator()


def make_empty_bigwig(filename, chrom_sizes):
    """Writes *filename* so that it is an empty bigwig file.
    *chrom_sizes is a dictionary giving chromosome sizes* given
    chomosome names.
    """
    bw_out = pyBigWig.open(filename, "w")
    bw_out.addHeader(list(chrom_sizes.items()))
    # for (chrom, chrom_len) in bw_out.chroms().items():
    for (chrom, chrom_len) in chrom_sizes.items():
        bw_out.addEntries(
            chrom, 0,
            values=np.nan_to_num(np.zeros(chrom_len)[0::10]),
            span=10, step=10)
    bw_out.close()


# Possible improvements: use an itearble of (from_region, to_region) pairs,
# or two zippable iterables
def paste_bigwig_region(
        from_fname, to_fname, from_region, to_region,
        dest_fname, nanmean=False):
    """

    Take values from a region *from_region* in a bigwig file *from_fname* and
    paste them into a region *to_region* in another bigwig file *to_fname*, and
    write the results in a third bigwig file *dest_fname*.

    A region should be specified as a (chrom, start, stop) triplet, where start
    is zero-based and stop is 1-based (like BED coordinates or Python slices).

    *dest_fname* will have the same chromosomes as *to_fname*.
    The values in *to_region* are substituted.

    Current limitation: The regions should have the same length.
    Option *nanmean* is likely not working correctly.
    """
    (from_chrom, from_start, from_stop) = from_region
    from_len = from_stop - from_start
    (to_chrom, to_start, to_stop) = to_region
    assert to_stop - to_start == from_len, (
        "Regions should have the same lengths.")
    from_bw = pyBigWig.open(from_fname)
    to_bw = pyBigWig.open(to_fname)
    chrom_sizes = list(to_bw.chroms().items())
    dest_bw = pyBigWig.open(dest_fname, "w")
    dest_bw.addHeader(chrom_sizes)
    for (chrom, chrom_len) in chrom_sizes:
        nb_bins = ceil(chrom_len / 10)
        # values = to_bw.values(chrom, 0, chrom_len)
        # Original values, plus some nans for binned_statistics to properly set the bin boundaries
        values = np.pad(
            to_bw.values(chrom, 0, chrom_len),
            # pad zero on the left, and what is needed to complete the bin on the right
            (0, 10 * nb_bins - chrom_len),
            mode="constant",
            constant_values=np.nan)
        if chrom == to_chrom:
            # Replace the values in the desired region
            values[to_start:to_stop] = from_bw.values(
                from_chrom, from_start, from_stop)
        if nanmean:
            bin_means = binned_statistic(
                range(0, 10 * nb_bins), values,
                statistic=np.nanmean, bins=nb_bins).statistic
        else:
            bin_means = binned_statistic(
                # range(0, chrom_len), np.nan_to_num(values),
                range(0, 10 * nb_bins), values,
                statistic="mean", bins=nb_bins).statistic
        dest_bw.addEntries(
            chrom, 0,
            # Mean for each bin of size 10
            values=bin_means,
            # values=np.nan_to_num(np.zeros(chrom_len)[0::10]),
            span=10, step=10)
    dest_bw.close()
    to_bw.close()
    from_bw.close()


#################
# Bowtie2 stuff #
#################
def zero(value):  # pylint: disable=W0613
    """Constant zero."""
    return 0


def identity(value):
    """Identity function."""
    return value


BOWTIE2_FUNCTION_SELECTOR = {
    "C": zero,
    "L": identity,
    "S": sqrt,
    "G": log}


def make_seeding_function(seeding_string):
    """Generates a function that computes the seeding pattern given a
    string representing bowtie2 seeding settings (-L and -i options).

    >>> make_seeding_function("-L 6 -i S,1,0.8")(18)
    [[0, 6], [4, 10], [8, 14], [12, 18]]
    """
    [opt1, val1, opt2, val2] = seeding_string.split()
    if opt1 == "-L":
        assert opt2 == "-i"
        seed_len = int(val1)
        interval_string = val2
    else:
        assert opt2 == "-L"
        seed_len = int(val2)
        interval_string = val1
    [func_type, constant, coeff] = interval_string.split(",")
    constant = float(constant)
    coeff = float(coeff)
    func_type = BOWTIE2_FUNCTION_SELECTOR[func_type]

    def seeding_function(read_len):
        interval = floor(constant + (coeff * func_type(read_len)))
        seeds = []
        seed = [0, seed_len]
        while seed[1] <= read_len:
            seeds.append(seed)
            next_seed_start = seed[0] + interval
            seed = [next_seed_start, next_seed_start + seed_len]
        return seeds
    return seeding_function


def aligner2min_mapq(aligner, wildcards):
    """
    Option to filter on MAPQ value in featureCounts.

    What minimal MAPQ value should a read have to be considered uniquely mapped?
    See <https://sequencing.qcfail.com/articles/mapq-values-are-really-useful-but-their-implementation-is-a-mess/>.
    """  # pylint: disable=C0301
    mapping_type = None
    try:
        mapping_type = wildcards.mapping_type
    except AttributeError:
        pass
    if mapping_type is None:
        try:
            mapping_type = wildcards.mapped_type
        except AttributeError:
            pass
    if mapping_type is None:
        try:
            mapping_type = wildcards.read_type
        except AttributeError:
            pass
    if mapping_type is None or mapping_type.startswith("unique_"):
        if aligner == "hisat2":
            return "-Q 60"
        if aligner == "bowtie2":
            return "-Q 23"
        raise NotImplementedError(f"{aligner} not handled (yet?)")
    return ""


# Not sure this is a good idea...
# def masked_gmean(a, axis=0, dtype=None):
#     """Modified from stats.py."""
#     # Converts the data into a masked array
#     ma = np.ma.masked_invalid(a)
#     # Apply gmean
#     if not isinstance(ma, np.ndarray):
#         # if not an ndarray object attempt to convert it
#         log_a = np.log(np.array(ma, dtype=dtype))
#     elif dtype:
#         # Must change the default dtype allowing array type
#         if isinstance(ma, np.ma.MaskedArray):
#             log_a = np.log(np.ma.asarray(ma, dtype=dtype))
#         else:
#             log_a = np.log(np.asarray(ma, dtype=dtype))
#     else:
#         log_a = np.log(ma)
#     return np.exp(log_a.mean(axis=axis))


def median_ratio_to_pseudo_ref_size_factors(counts_data):
    """Adapted from DESeq paper (doi:10.1186/gb-2010-11-10-r106)
    All libraries are used to define a pseudo-reference, which has
    the geometric mean across libraries for a given gene in *counts_data*.
    For a given library, the median across genes of the ratios to the
    pseudo-reference is used as size factor."""
    # Add pseudo-count to compute the geometric mean, then remove it
    # pseudo_ref = (counts_data + 1).apply(gmean, axis=1) - 1
    # Ignore lines with zeroes instead (may be bad for IP: many zeroes expected):
    pseudo_ref = (counts_data[counts_data.prod(axis=1) > 0]).apply(gmean, axis=1)
    # Ignore lines with only zeroes
    # pseudo_ref = (counts_data[counts_data.sum(axis=1) > 0]).apply(masked_gmean, axis=1)

    def median_ratio_to_pseudo_ref(col):
        return (col / pseudo_ref).median()
    # size_factors = counts_data.apply(median_ratio_to_pseudo_ref, axis=0)
    median_ratios = counts_data[counts_data.prod(axis=1) > 0].apply(
        median_ratio_to_pseudo_ref, axis=0)
    # Not sure fillna(0) is appropriate
    if any(median_ratios.isna()):
        msg = "Could not compute median ratios to pseudo reference.\n"
        warnings.warn(msg)
        return median_ratios.fillna(1)
    return median_ratios


def size_factor_correlations(counts_data, summaries, normalizer):
    """Is there a correlation, across libraries, between normalized values and size factors?
    The size factor type *normalizer* is either computed or taken from *summaries*.
    The normalized data are computed by dividing *counts_data* by this size factor."""
    if normalizer == "median_ratio_to_pseudo_ref":
        size_factors = median_ratio_to_pseudo_ref_size_factors(counts_data)
    else:
        size_factors = summaries.loc[normalizer]
    # by_norm = counts_data / size_factors

    def compute_pearsonr_with_size_factor(row):
        try:
            return pearsonr(row, size_factors)[0]
        except ValueError as err:
            if str(err) == "array must not contain infs or NaNs":
                msg = (
                    f"Cannot compute Pearson correlation with size factors "
                    f"when normalizing using {normalizer}")
                warnings.warn(msg)
                return np.nan
            else:
                raise
    # return by_norm.apply(compute_pearsonr_with_size_factor, axis=1)
    return (counts_data / size_factors).apply(compute_pearsonr_with_size_factor, axis=1)


def plot_norm_correlations(correlations):
    """
    Make violin plots to represent data in *correlations*.
    """
    # fig, (ax1, ax2) = plt.subplots(nrows=2, ncols=1, sharex=True)
    # correlations.plot.kde(ax=ax1)
    # sns.violinplot(data=correlations, orient="h", ax=ax2)
    # ax2.set_xlabel("Pearson correlation coefficient")
    usetex = mpl.rcParams.get("text.usetex", False)
    if usetex:
        correlations.columns = [texscape(colname) for colname in correlations.columns]
    axis = sns.violinplot(data=correlations, cut=0)
    axis.set_ylabel("Pearson correlation coefficient")


def plot_counts_distribution(data, xlabel):
    """
    Plot a kernel density estimate of the distribution of counts in *data*.
    """
    # TODO: try to plot with semilog x axis
    # axis = data.plot.kde(legend=None)
    # axis.set_xlabel(xlabel)
    # axis.legend(ncol=len(REPS))
    usetex = mpl.rcParams.get("text.usetex", False)
    if usetex:
        xlabel = texscape(xlabel)
        data.columns = [texscape(colname) for colname in data.columns]
    try:
        axis = data.plot.kde()
    # except ValueError as e:
    except ValueError:
        msg = "".join([
            "There seems to be a problem with the data.\n",
            "The data matrix has %d lines and %d columns.\n" % (len(data), len(data.columns))])
        warnings.warn(msg)
        raise
    axis.set_xlabel(xlabel)


def plot_histo(outfile, data, title=None):
    """
    Plot a histogram of *data* in file *outfile*.
    """
    fig = plt.figure(figsize=(15, 7))
    axis = fig.add_subplot(111)
    axis.set_xlim([data.index[0] - 0.5, data.index[-1] + 0.5])
    # axis.set_ylim([0, 100])
    bar_width = 0.8
    # letter2legend = dict(zip("ACGT", "ACGT"))
    usetex = mpl.rcParams.get("text.usetex", False)
    if usetex:
        data.columns = [texscape(colname) for colname in data.columns]
        title = texscape(title)
    for (read_len, count) in data.iterrows():
        plt.bar(
            read_len,
            count,
            align="center",
            width=bar_width)
            # color=letter2colour[letter],
            # label=letter2legend[letter])
    axis.legend()
    axis.set_xticks(data.index)
    axis.set_xticklabels(data.index)
    axis.set_xlabel("read length")
    axis.set_ylabel("number of reads")
    plt.setp(axis.get_xticklabels(), rotation=90)
    if title is not None:
        plt.title(title)
    ## debug
    try:
        plt.savefig(outfile)
    # except RuntimeError as e:
    except RuntimeError:
        print(data.index)
        print(title)
        raise
    ##


def plot_boxplots(data, ylabel):
    """
    Plot boxplots of data in *data* using *ylabel* as y-axis label.
    """
    fig = plt.figure(figsize=(6, 12))
    axis = fig.add_subplot(111)
    usetex = mpl.rcParams.get("text.usetex", False)
    if usetex:
        ylabel = texscape(ylabel)
        data.columns = [texscape(colname) for colname in data.columns]
    data.plot.box(ax=axis)
    axis.set_ylabel(ylabel)
    for label in axis.get_xticklabels():
        label.set_rotation(90)
    plt.tight_layout()


############
# DE stuff #
############
# Cutoffs in log fold change
LFC_CUTOFFS = [0.5, 1, 2]
UP_STATUSES = [f"up{cutoff}" for cutoff in LFC_CUTOFFS]
DOWN_STATUSES = [f"down{cutoff}" for cutoff in LFC_CUTOFFS]


def status_setter(lfc_cutoffs=None, fold_type="log2FoldChange"):
    """*fold_type* can also be "lfcMLE", which is based on uncorrected values.
    This may not be good for genes with low expression levels."""
    if lfc_cutoffs is None:
        lfc_cutoffs = LFC_CUTOFFS
    def set_status(row):
        """Determines the up- or down-regulation status corresponding to a given
        row of a deseq2 results table."""
        if row["padj"] < 0.05:
            # if row["log2FoldChange"] > 0:
            lfc = row[fold_type]
            if lfc > 0:
                # Start from the highest cutoff,
                # and decrease until below lfc
                for cutoff in sorted(lfc_cutoffs, reverse=True):
                    if lfc > cutoff:
                        return f"up{cutoff}"
                return "up"
            for cutoff in sorted(lfc_cutoffs, reverse=True):
                if lfc < -cutoff:
                    return f"down{cutoff}"
            return "down"
        return "NS"
    return set_status


# res = res.assign(is_DE=res.apply(set_de_status, axis=1))
def set_de_status(row):
    """Determines whether a gene is differentially expressed (DE) of not (NS)
    based on the adjusted p-value in row of a deseq2 results table."""
    if row["padj"] < 0.05:
        return "DE"
    return "NS"
DE2COLOUR = {
    # black
    "DE": "k",
    # pale grey
    "NS": "0.85"}


def plot_lfc_distribution(res, contrast, fold_type=None):
    """*fold_type* is "log2FoldChange" by default.
    It can also be "lfcMLE", which is based on uncorrected values.
    This may not be good for genes with low expression levels."""
    # lfc = res.lfcMLE.dropna()
    if fold_type is None:
        fold_type = "log2FoldChange"
    lfc = getattr(res, fold_type).dropna()
    usetex = mpl.rcParams.get("text.usetex", False)
    if usetex:
        lfc.name = texscape(contrast)
    else:
        lfc.name = contrast
        # lfc.columns = [texscape(colname) for colname in lfc.columns]
    axis = sns.kdeplot(lfc)
    axis.set_xlabel(fold_type)
    axis.set_ylabel("frequency")


def make_status2colour(down_statuses, up_statuses):
    """
    Generate a dictionary associating colours to statuses.
    """
    statuses = list(reversed(down_statuses)) + ["down", "NS", "up"] + up_statuses
    return dict(zip(statuses, sns.color_palette("coolwarm", len(statuses))))


STATUS2COLOUR = make_status2colour(DOWN_STATUSES, UP_STATUSES)


# TODO: use other labelling than logfold or gene lists, i.e. biotype
def plot_MA(res,
            grouping=None,
            group2colour=None,
            mean_range=None,
            lfc_range=None,
            fold_type=None):
    """*fold_type* is "log2FoldChange" by default.
    It can also be "lfcMLE", which is based on uncorrected values.
    This may not be good for genes with low expression levels."""
    # if not len(res):
    # if not res:
    if res.empty:
        raise ValueError("No data to plot.")
    fig, axis = plt.subplots()
    # Make a column indicating whether the gene is DE or NS
    data = res.assign(is_DE=res.apply(set_de_status, axis=1))
    x_column = "baseMean"
    data = data.assign(logx=np.log10(data[x_column]))
    if fold_type is None:
        y_column = "log2FoldChange"
    else:
        y_column = fold_type
    usetex = mpl.rcParams.get("text.usetex", False)

    def scatter_group(group, label, colour, size=1):
        """Plots the data in *group* on the scatterplot."""
        if usetex:
            label = texscape(label)
        group.plot.scatter(
            # x=x_column,
            x="logx",
            y=y_column,
            s=size,
            # logx=True,
            c=colour,
            label=label, ax=axis)
    if usetex:
        data.columns = [texscape(colname) for colname in data.columns]
        y_column = texscape(y_column)
        de_status_column = "is\_DE"  # pylint: disable=W1401
    else:
        de_status_column = "is_DE"
    # First plot the data in grey and black
    for (de_status, group) in data.groupby(de_status_column):
        label = f"{de_status} ({len(group)})"
        colour = DE2COLOUR[de_status]
        scatter_group(group, label, colour, size=2)
    if grouping is not None:
        if isinstance(grouping, str):
            # Overlay colours based on the "grouping" column
            if group2colour is None:
                group2colour = STATUS2COLOUR
            for status, group in data.groupby(grouping):
                label = f"{status} ({len(group)})"
                colour = group2colour[status]
                scatter_group(group, label, colour)
        else:
            (status, colour) = group2colour
            row_indices = data.index.intersection(grouping)
            try:
                label = f"{status} ({len(row_indices)})"
                scatter_group(data.loc[row_indices], label, colour)
            except ValueError as err:
                if str(err) != "scatter requires x column to be numeric":
                    print(data.loc[row_indices])
                    raise
                warnings.warn(f"Nothing to plot for {status}\n")
    axis.axhline(y=1, linewidth=0.5, color="0.5", linestyle="dashed")
    axis.axhline(y=-1, linewidth=0.5, color="0.5", linestyle="dashed")
    # TODO: check data basemean range
    if mean_range is not None:
        # axis.set_xlim(mean_range)
        axis.set_xlim(np.log10(mean_range))
    if lfc_range is not None:
        (lfc_min, lfc_max) = lfc_range
        lfc_here_min = getattr(data, y_column).min()
        lfc_here_max = getattr(data, y_column).max()
        if (lfc_here_min < lfc_min) or (lfc_here_max > lfc_max):
            warnings.warn(
                f"Cannot plot {y_column} data "
                f"([{lfc_here_min}, {lfc_here_max}]) in requested range "
                f"([{lfc_min}, {lfc_max}])\n")
        else:
            axis.set_ylim(lfc_range)
    # https://stackoverflow.com/a/24867320/1878788
    x_ticks = np.arange(
        floor(np.ma.masked_invalid(data["logx"]).min()),
        ceil(np.ma.masked_invalid(data["logx"]).max()),
        1)
    x_ticklabels = [r"$10^{{{}}}$".format(tick) for tick in x_ticks]
    plt.xticks(x_ticks, x_ticklabels)
    axis.set_xlabel(x_column)


def plot_scatter(data,
                 x_column,
                 y_column,
                 regression=False,
                 grouping=None,
                 group2colour=None,
                 x_range=None,
                 y_range=None,
                 axes_style=None):
    """
    Plot a scatterplot using data from *data*, using columns
    """
    # No rows
    # if not len(data):
    # Does it work like that too?
    # if not data:
    if data.empty:
        raise ValueError("No data to plot.")
    # fig, axis = plt.subplots()
    _, axis = plt.subplots()
    # axis.set_adjustable('box')
    # First plot the data in grey
    data.plot.scatter(
        x=x_column, y=y_column,
        s=2, c="black", alpha=0.15, edgecolors='none',
        ax=axis)
    if regression:
        linreg = linregress(data[[x_column, y_column]].dropna())
        a = linreg.slope  # pylint: disable=C0103
        b = linreg.intercept  # pylint: disable=C0103

        def fit(x):
            return (a * x) + b
        min_x = data[[x_column]].min()[0]
        max_x = data[[x_column]].max()[0]
        min_y = fit(min_x)
        max_y = fit(max_x)
        xfit, yfit = (min_x, max_x), (min_y, max_y)
        axis.plot(
            xfit, yfit,
            linewidth=0.5, color="0.5", linestyle="dashed")
    # Overlay colour points
    if grouping is not None:
        if isinstance(grouping, str):
            # Determine colours based on the "grouping" column
            if group2colour is None:
                statuses = data[grouping].unique()
                group2colour = dict(zip(
                    statuses,
                    sns.color_palette("colorblind", len(statuses))))
            for (status, group) in data.groupby(grouping):
                group.plot.scatter(
                    x=x_column, y=y_column, s=1, c=group2colour[status],
                    label=f"{status} ({len(group)})", ax=axis)
        else:
            # Apply a colour to a list of genes
            (status, colour) = group2colour
            row_indices = data.index.intersection(grouping)
            try:
                data.loc[row_indices].plot.scatter(
                    x=x_column, y=y_column, s=1, c=colour,
                    label=f"{status} ({len(row_indices)})", ax=axis)
            except ValueError as err:
                if str(err) != "scatter requires x column to be numeric":
                    print(data.loc[row_indices])
                    raise
                warnings.warn(f"Nothing to plot for {status}\n")
    if axes_style is None:
        axes_style = {"linewidth": 0.5, "color": "0.5", "linestyle": "dashed"}
    axis.axhline(y=0, **axes_style)
    axis.axvline(x=0, **axes_style)
    # axis.axhline(y=0, linewidth=0.5, color="0.5", linestyle="dashed")
    # axis.axvline(x=0, linewidth=0.5, color="0.5", linestyle="dashed")
    # Set axis limits
    if x_range is not None:
        (x_min, x_max) = x_range
        x_here_min = getattr(data, x_column).min()
        x_here_max = getattr(data, x_column).max()
        if (x_here_min < x_min) or (x_here_max > x_max):
            warnings.warn(
                f"Cannot plot {x_column} data "
                f"([{x_here_min}, {x_here_max}]) in requested range "
                f"([{x_min}, {x_max}])\n")
        else:
            axis.set_xlim(x_range)
    if y_range is not None:
        (y_min, y_max) = y_range
        y_here_min = getattr(data, y_column).min()
        y_here_max = getattr(data, y_column).max()
        if (y_here_min < y_min) or (y_here_max > y_max):
            warnings.warn(
                f"Cannot plot {y_column} data ([{y_here_min}, {y_here_max}]) "
                f"in requested range ([{y_min}, {y_max}])\n")
        else:
            axis.set_ylim(y_range)
    return axis


def plot_paired_scatters(data, columns=None, hue=None, log_log=False):
    """Alternative to pairplot, in order to avoid histograms on the diagonal."""
    if columns is None:
        columns = data.columns
    usetex = mpl.rcParams.get("text.usetex", False)
    if usetex:
        data.columns = [texscape(colname) for colname in data.columns]
        columns = [texscape(colname) for colname in columns]
    grid = sns.PairGrid(data, vars=columns, hue=hue, size=8)
    # grid.map_offdiag(plt.scatter, marker=".")
    grid.map_lower(plt.scatter, marker=".")
    if log_log:
        for axis in grid.axes.ravel():
            axis.set_xscale('log')
            axis.set_yscale('log')
    grid.add_legend()