5

I am trying to make a nice free energy surface (heat map) using Seaborn's kdeplot. I am very close but can not figure out a way to change the color bar scale. The color bar scale is important since it is supposed to represent the difference in energy at different coordinates on the map. I need to know how to scale the values of the color bar by -(0.5961573)*log(x), where x is the values of the color bar. I may also then need to normalize the color bar from there so that the max value is 0.

Here is what I currently have:

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.pyplot as plt
import seaborn as sns


rs=[]
dihes=[]
with open(sys.argv[1], 'r') as f:
    for line in f:
        time,r,dihe = line.split()
        rs.append(float(r))
        dihes.append(float(dihe))
    sns.set_style("white")
    sns.kdeplot(rs, dihes, n_levels=25, cbar=True, cmap="Purples_d")
    plt.show()

This gets me:

sample output

The arrays rs and dihes are simple one dimensional arrays.

Any suggestions on how to scale the color bar (z-axis) would be very helpful!

Trenton McKinney
  • 56,955
  • 33
  • 144
  • 158
Chemistry for Computers
  • 87
  • 1
  • 9

3 Answers3

1

One way to do it is to create the graph manually and then modify the labels directly. This involves a couple more lines of code. You may have to tweak the formatting a bit but something like this should get you on the right track.

The following is adapted from this answer and this answer.

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import seaborn as sns
from scipy import stats

rs=[]
dihes=[]
with open(sys.argv[1], 'r') as f:
    for line in f:
        time,r,dihe = line.split()
        rs.append(float(r))
        dihes.append(float(dihe))

    x = rs
    y = dihes
    kde = stats.gaussian_kde([x, y])

    xx, yy = np.mgrid[min(x):max(x):(max(x)-min(x))/100, min(y):max(y):(max(y)-min(y))/100]
    density = kde(np.c_[xx.flat, yy.flat].T).reshape(xx.shape)

    sns.set_style("white")
    fig, ax = plt.subplots()
    cset = ax.contour(xx, yy, density, 25, cmap="Purples_r")
    cb = fig.colorbar(cset)
    cb.ax.set_yticklabels(map(lambda x: -0.5961573*np.log(float(x.get_text())), 
                              cb.ax.get_yticklabels()))
Leo
  • 1,773
  • 12
  • 19
1

A bit late to the party, but I ended up putting together this context manager which switches plotted density values to a logarithmic scale:

import contextlib
import seaborn as sns

@contextlib.contextmanager
def plot_kde_as_log(base=np.exp(1), support_threshold=1e-4):
    """Context manager to render density estimates on a logarithmic scale.

    Usage:

        with plot_kde_as_log():
            sns.jointplot(x='x', y='y', data=df, kind='kde')
    """
    old_stats = sns.distributions._has_statsmodels
    old_univar = sns.distributions._scipy_univariate_kde
    old_bivar = sns.distributions._scipy_bivariate_kde

    sns.distributions._has_statsmodels = False
    def log_clip_fn(v):
        v = np.log(np.clip(v, support_threshold, np.inf))
        v -= np.log(support_threshold)
        v /= np.log(base)
        return v
    def new_univar(*args, **kwargs):
        x, y = old_univar(*args, **kwargs)
        y = log_clip_fn(y)
        return x, y
    def new_bivar(*args, **kwargs):
        x, y, z = old_bivar(*args, **kwargs)
        z = log_clip_fn(z)
        return x, y, z

    sns.distributions._scipy_univariate_kde = new_univar
    sns.distributions._scipy_bivariate_kde = new_bivar
    try:
        yield
    finally:
        sns.distributions._has_statsmodels = old_stats
        sns.distributions._scipy_univariate_kde = old_univar
        sns.distributions._scipy_bivariate_kde = old_bivar

The benefit of this approach is that it keeps all of the styling and other options of sns.jointplot() without any additional effort.

Walt W
  • 3,261
  • 3
  • 30
  • 37
0

I updated Walt W's context manager to work with newer versions of seaborn

@contextlib.contextmanager
def plot_kde_as_log(base=np.exp(1), support_threshold=1e-4):
    """Context manager to render density estimates on a logarithmic scale.

    Usage:

        with plot_kde_as_log():
            sns.jointplot(x='x', y='y', data=df, kind='kde')
    """
    old_call = sns._statistics.KDE.__call__

    def log_clip_fn(v):
        v = np.log(np.clip(v, support_threshold, np.inf))
        v -= np.log(support_threshold)
        v /= np.log(base)
        return v
    def new_call(*args, **kwargs):
        density, support = old_call(*args, **kwargs)
        density = log_clip_fn(density)
        return density, support

    sns._statistics.KDE.__call__ = new_call
    try:
        yield
    finally:
        sns._statistics.KDE.__call__ = old_call
Dalton Worsnup
  • 1