--- a
+++ b/bpnet/plot/profiles.py
@@ -0,0 +1,398 @@
+from tqdm import tqdm
+import seaborn as sns
+import pandas as pd
+import os
+import matplotlib.pyplot as plt
+import numpy as np
+from bpnet.plot.utils import simple_yaxis_format, strip_axis, spaced_xticks
+from bpnet.modisco.utils import bootstrap_mean, nan_like, ic_scale
+from bpnet.plot.utils import show_figure
+
+
+# TODO - make it as a bar-plot with two standard colors:
+# #B23F49 (pos), #045CA8 (neg)
+def plot_stranded_profile(profile, ax=None, ymax=None, profile_std=None, flip_neg=True, set_ylim=True):
+    """Plot the stranded profile
+    """
+    if ax is None:
+        ax = plt.gca()
+
+    if profile.ndim == 1:
+        # also compatible with single dim
+        profile = profile[:, np.newaxis]
+    assert profile.ndim == 2
+    assert profile.shape[1] <= 2
+    labels = ['pos', 'neg']
+
+    # determine ymax if not specified
+    if ymax is None:
+        if profile_std is not None:
+            ymax = (profile.max() - 2 * profile_std).max()
+        else:
+            ymax = profile.max()
+
+    if set_ylim:
+        if flip_neg:
+            ax.set_ylim([-ymax, ymax])
+        else:
+            ax.set_ylim([0, ymax])
+
+    ax.axhline(y=0, linewidth=1, linestyle='--', color='black')
+    # strip_axis(ax)
+
+    xvec = np.arange(1, len(profile) + 1)
+
+    for i in range(profile.shape[1]):
+        sign = 1 if not flip_neg or i == 0 else -1
+        ax.plot(xvec, sign * profile[:, i], label=labels[i])
+
+        # plot also the ribbons
+        if profile_std is not None:
+            ax.fill_between(xvec,
+                            sign * profile[:, i] - 2 * profile_std[:, i],
+                            sign * profile[:, i] + 2 * profile_std[:, i],
+                            alpha=0.1)
+    # return ax
+
+
+def multiple_plot_stranded_profile(d_profile, figsize_tmpl=(4, 3), normalize=False):
+    fig, axes = plt.subplots(1, len(d_profile),
+                             figsize=(figsize_tmpl[0] * len(d_profile), figsize_tmpl[1]),
+                             sharey=True)
+    if len(d_profile)==1: #If only one task, then can't zip axes
+        ax = axes
+        task = [*d_profile][0]
+        arr = d_profile[task].mean(axis=0)
+        if normalize:
+            arr = arr / arr.max()
+        plot_stranded_profile(arr, ax=ax, set_ylim=False)
+        ax.set_title(task)
+        ax.set_ylabel("Avg. counts")
+        ax.set_xlabel("Position")
+        fig.subplots_adjust(wspace=0)  # no space between plots
+        return fig
+    else:
+        for i, (task, ax) in enumerate(zip(d_profile, axes)):
+            arr = d_profile[task].mean(axis=0)
+            if normalize:
+                arr = arr / arr.max()
+            plot_stranded_profile(arr, ax=ax, set_ylim=False)
+            ax.set_title(task)
+            if i == 0:
+                ax.set_ylabel("Avg. counts")
+                ax.set_xlabel("Position")
+        fig.subplots_adjust(wspace=0)  # no space between plots
+        return fig
+
+
+def aggregate_profiles(profile_arr, n_bootstrap=None, only_idx=None):
+    if only_idx is not None:
+        return profile_arr[only_idx], None
+
+    if n_bootstrap is not None:
+        return bootstrap_mean(profile_arr, n=n_bootstrap)
+    else:
+        return profile_arr.mean(axis=0), None
+
+
+def extract_signal(x, seqlets, rc_fn=lambda x: x[::-1, ::-1]):
+    def optional_rc(x, is_rc):
+        if is_rc:
+            return rc_fn(x)
+        else:
+            return x
+    return np.stack([optional_rc(x[s['example'], s['start']:s['end']], s['rc'])
+                     for s in seqlets])
+
+
+def plot_profiles(seqlets_by_pattern,
+                  x,
+                  tracks,
+                  contribution_scores={},
+                  figsize=(20, 2),
+                  start_vec=None,
+                  width=20,
+                  legend=True,
+                  rotate_y=90,
+                  seq_height=1,
+                  ymax=None,  # determine y-max
+                  n_limit=35,
+                  n_bootstrap=None,
+                  flip_neg=False,
+                  patterns=None,
+                  fpath_template=None,
+                  only_idx=None,
+                  mkdir=False,
+                  rc_fn=lambda x: x[::-1, ::-1]):
+    """
+    Plot the sequence profiles
+    Args:
+      x: one-hot-encoded sequence
+      tracks: dictionary of profile tracks
+      contribution_scores: optional dictionary of contribution scores
+
+    """
+    import matplotlib.pyplot as plt
+    from concise.utils.plot import seqlogo_fig, seqlogo
+
+    # Setup start-vec
+    if start_vec is not None:
+        if not isinstance(start_vec, list):
+            start_vec = [start_vec] * len(patterns)
+    else:
+        start_vec = [0] * len(patterns)
+        width = len(x)
+
+    if patterns is None:
+        patterns = list(seqlets_by_pattern)
+    # aggregated profiles
+    d_signal_patterns = {pattern:
+                         {k: aggregate_profiles(
+                             extract_signal(y, seqlets_by_pattern[pattern])[:, start_vec[ip]:(start_vec[ip] + width)],
+                             n_bootstrap=n_bootstrap, only_idx=only_idx)
+                          for k, y in tracks.items()}
+                         for ip, pattern in enumerate(patterns)}
+    if ymax is None:
+        # infer ymax
+        def take_max(x, dx):
+            if dx is None:
+                return x.max()
+            else:
+                # HACK - hard-coded 2
+                return (x + 2 * dx).max()
+
+        ymax = [max([take_max(*d_signal_patterns[pattern][k])
+                     for pattern in patterns])
+                for k in tracks]  # loop through all the tracks
+    if not isinstance(ymax, list):
+        ymax = [ymax] * len(tracks)
+
+    figs = []
+    for i, pattern in enumerate(tqdm(patterns)):
+        j = i
+        # --------------
+        # extract signal
+        seqs = extract_signal(x, seqlets_by_pattern[pattern])[:, start_vec[i]:(start_vec[i] + width)]
+        ext_contribution_scores = {s: extract_signal(contrib, seqlets_by_pattern[pattern])[:, start_vec[i]:(start_vec[i] + width)]
+                                   for s, contrib in contribution_scores.items()}
+        d_signal = d_signal_patterns[pattern]
+        # --------------
+        if only_idx is None:
+            sequence = ic_scale(seqs.mean(axis=0))
+        else:
+            sequence = seqs[only_idx]
+
+        n = len(seqs)
+        if n < n_limit:
+            continue
+        fig, ax = plt.subplots(1 + len(contribution_scores) + len(tracks),
+                               1, sharex=True,
+                               figsize=figsize,
+                               gridspec_kw={'height_ratios': [1] * len(tracks) + [seq_height] * (1 + len(contribution_scores))})
+
+        # signal
+        ax[0].set_title(f"{pattern} ({n})")
+        for i, (k, signal) in enumerate(d_signal.items()):
+            signal_mean, signal_std = d_signal_patterns[pattern][k]
+            plot_stranded_profile(signal_mean, ax=ax[i], ymax=ymax[i],
+                                  profile_std=signal_std, flip_neg=flip_neg)
+            simple_yaxis_format(ax[i])
+            strip_axis(ax[i])
+            ax[i].set_ylabel(f"{k}", rotation=rotate_y, ha='right', labelpad=5)
+
+            if legend:
+                ax[i].legend()
+
+        # -----------
+        # contribution scores (seqlogo)
+        # -----------
+        # average the contribution scores
+        if only_idx is None:
+            norm_contribution_scores = {k: v.mean(axis=0)
+                                        for k, v in ext_contribution_scores.items()}
+        else:
+            norm_contribution_scores = {k: v[only_idx]
+                                        for k, v in ext_contribution_scores.items()}
+
+        max_scale = max([np.maximum(v, 0).sum(axis=-1).max() for v in norm_contribution_scores.values()])
+        min_scale = min([np.minimum(v, 0).sum(axis=-1).min() for v in norm_contribution_scores.values()])
+        for k, (contrib_score_name, logo) in enumerate(norm_contribution_scores.items()):
+            ax_id = len(tracks) + k
+
+            # Trim the pattern if necessary
+            # plot
+            ax[ax_id].set_ylim([min_scale, max_scale])
+            ax[ax_id].axhline(y=0, linewidth=1, linestyle='--', color='grey')
+            seqlogo(logo, ax=ax[ax_id])
+
+            # style
+            simple_yaxis_format(ax[ax_id])
+            strip_axis(ax[ax_id])
+            # ax[ax_id].set_ylabel(contrib_score_name)
+            ax[ax_id].set_ylabel(contrib_score_name, rotation=rotate_y, ha='right', labelpad=5)  # va='bottom',
+
+        # -----------
+        # information content (seqlogo)
+        # -----------
+        # plot
+        seqlogo(sequence, ax=ax[-1])
+
+        # style
+        simple_yaxis_format(ax[-1])
+        strip_axis(ax[-1])
+        ax[-1].set_ylabel("Inf. content", rotation=rotate_y, ha='right', labelpad=5)
+        ax[-1].set_xticks(list(range(0, len(sequence) + 1, 5)))
+
+        figs.append(fig)
+        # save to file
+        if fpath_template is not None:
+            pname = pattern.replace("/", ".")
+            basepath = fpath_template.format(pname=pname, pattern=pattern)
+            if mkdir:
+                os.makedirs(os.path.dirname(basepath), exist_ok=True)
+            plt.savefig(basepath + '.png', dpi=600)
+            plt.savefig(basepath + '.pdf', dpi=600)
+            plt.close(fig)    # close the figure
+            show_figure(fig)
+            plt.show()
+    return figs
+
+
+def plot_profiles_single(seqlet,
+                         x,
+                         tracks,
+                         contribution_scores={},
+                         figsize=(20, 2),
+                         legend=True,
+                         rotate_y=90,
+                         seq_height=1,
+                         flip_neg=False,
+                         rc_fn=lambda x: x[::-1, ::-1]):
+    """
+    Plot the sequence profiles
+    Args:
+      x: one-hot-encoded sequence
+      tracks: dictionary of profile tracks
+      contribution_scores: optional dictionary of contribution scores
+
+    """
+    import matplotlib.pyplot as plt
+    from concise.utils.plot import seqlogo_fig, seqlogo
+
+    # --------------
+    # extract signal
+    seq = seqlet.extract(x)
+    ext_contribution_scores = {s: seqlet.extract(contrib) for s, contrib in contribution_scores.items()}
+
+    fig, ax = plt.subplots(1 + len(contribution_scores) + len(tracks),
+                           1, sharex=True,
+                           figsize=figsize,
+                           gridspec_kw={'height_ratios': [1] * len(tracks) + [seq_height] * (1 + len(contribution_scores))})
+
+    # signal
+    for i, (k, signal) in enumerate(tracks.items()):
+        plot_stranded_profile(seqlet.extract(signal), ax=ax[i],
+                              flip_neg=flip_neg)
+        simple_yaxis_format(ax[i])
+        strip_axis(ax[i])
+        ax[i].set_ylabel(f"{k}", rotation=rotate_y, ha='right', labelpad=5)
+
+        if legend:
+            ax[i].legend()
+
+    # -----------
+    # contribution scores (seqlogo)
+    # -----------
+    max_scale = max([np.maximum(v, 0).sum(axis=-1).max() for v in ext_contribution_scores.values()])
+    min_scale = min([np.minimum(v, 0).sum(axis=-1).min() for v in ext_contribution_scores.values()])
+    for k, (contrib_score_name, logo) in enumerate(ext_contribution_scores.items()):
+        ax_id = len(tracks) + k
+        # plot
+        ax[ax_id].set_ylim([min_scale, max_scale])
+        ax[ax_id].axhline(y=0, linewidth=1, linestyle='--', color='grey')
+        seqlogo(logo, ax=ax[ax_id])
+
+        # style
+        simple_yaxis_format(ax[ax_id])
+        strip_axis(ax[ax_id])
+        # ax[ax_id].set_ylabel(contrib_score_name)
+        ax[ax_id].set_ylabel(contrib_score_name, rotation=rotate_y, ha='right', labelpad=5)  # va='bottom',
+
+    # -----------
+    # information content (seqlogo)
+    # -----------
+    # plot
+    seqlogo(seq, ax=ax[-1])
+
+    # style
+    simple_yaxis_format(ax[-1])
+    strip_axis(ax[-1])
+    ax[-1].set_ylabel("Inf. content", rotation=rotate_y, ha='right', labelpad=5)
+    ax[-1].set_xticks(list(range(0, len(seq) + 1, 5)))
+    return fig
+
+
+def hist_position(dfp, tasks):
+    """Make the positional histogram
+
+    Args:
+      dfp: pd.DataFrame with columns: peak_id, and center
+      tasks: list of tasks for which to plot the different peak_id columns
+    """
+    fig, axes = plt.subplots(1, len(tasks), figsize=(5 * len(tasks), 2),
+                             sharey=True, sharex=True)
+    if len(tasks) == 1:
+        axes = [axes]
+    for i, (task, ax) in enumerate(zip(tasks, axes)):
+        ax.hist(dfp[dfp.peak_id == task].center, bins=100)
+        ax.set_title(task)
+        ax.set_xlabel("Position")
+        ax.set_xlim([0, 1000])
+        if i == 0:
+            ax.set_ylabel("Frequency")
+    plt.subplots_adjust(wspace=0)
+    return fig
+
+
+def bar_seqlets_per_example(dfp, tasks):
+    """Make the positional histogram
+
+    Args:
+      dfp: pd.DataFrame with columns: peak_id, and center
+      tasks: list of tasks for which to plot the different peak_id columns
+    """
+    fig, axes = plt.subplots(1, len(tasks), figsize=(5 * len(tasks), 2),
+                             sharey=True, sharex=True)
+    if len(tasks) == 1:
+        axes = [axes]
+    for i, (task, ax) in enumerate(zip(tasks, axes)):
+        dfpp = dfp[dfp.peak_id == task]
+        ax.set_title(task)
+        ax.set_xlabel("Frequency")
+        if i == 0:
+            ax.set_ylabel("# per example")
+        if not len(dfpp):
+            continue
+        dfpp.groupby("example_idx").\
+            size().value_counts().plot(kind="barh", ax=ax)
+    plt.subplots_adjust(wspace=0)
+    return fig
+
+
+def box_counts(total_counts, pattern_idx):
+    """Make a box-plot with total counts in the region
+
+    Args:
+      total_counts: dict per task
+      pattern_idx: array with example_idx of the pattern
+    """
+    dfs = pd.concat([total_counts.melt().assign(subset="all peaks"),
+                     total_counts.iloc[pattern_idx].melt().assign(subset="contains pattern")])
+    dfs.value = np.log10(1 + dfs.value)
+
+    fig, ax = plt.subplots(figsize=(5, 5))
+    sns.boxplot("variable", "value", hue="subset", data=dfs, ax=ax)
+    ax.set_xlabel("Task")
+    ax.set_ylabel("log10(1+counts)")
+    ax.set_title("Total number of counts in the region")
+    return fig