Source code for schrodinger.application.matsci.plotwidgets

"""
Contains widgets and functionality that are useful for plotting.

Copyright Schrodinger, LLC. All rights reserved.
"""

import numpy as np
import matplotlib
import matplotlib.pyplot as plt


def greyColorMap(total_iterations, reverse=False):
    """
    Matplotlib color map customized for a frame's figure
        with various shades of grey.

    :type total_iterations: int
    :param total_iterations: The total number of iterations

    :type reverse: bool
    :param reverse: True if brightness of plot needs to increase with
        increase in iteration number
    """

    if not total_iterations:
        return

    base_cmap = plt.cm.Greys_r if reverse else plt.cm.Greys

    # Some color sequences start in white(when reverse is False), which is
    # hard to see on a white background. The ends of these sequences can be
    # excluded using the two `color_bonud_*` values. They should stay
    # within [0, 1]. Raise the lower bound to exclude the bottom of the
    # sequence(when reverse is False), and lower the upper bound to exclude
    # the top of the sequence(when reverse is True).
    # The logic is similar when reverse is True.
    color_bound_min, color_bound_high = (0.0, 0.8) if reverse else (0.2, 1.0)
    colors = base_cmap(np.linspace(color_bound_min, color_bound_high))
    cmap = matplotlib.colors.LinearSegmentedColormap.from_list(
        name=f'truncated_{base_cmap.name}', colors=colors)

    # We want to color the profile lines in a gradient to make it easier to
    # tell which line is which. Labeling these lines with a legend will be
    # difficult, because there could be too many lines to fit in a legend.
    # This is why we'd rather use a colorbar to label the colors.
    # Matplotlib isn't exactly built to make color bars with lines, so we
    # do some StackOverflow magic here to make it work.
    iter_nums = [i + 1 for i in range(total_iterations)]
    norm = matplotlib.colors.Normalize(vmin=min(iter_nums), vmax=max(iter_nums))
    cmap = matplotlib.cm.ScalarMappable(norm=norm, cmap=cmap)
    cmap.set_array([])

    return cmap