The Cracked Bassoon


Drift-diffusion plots

Tags: python | reaction-times | 📊 visualization

The drift-diffusion model (DDM), sometimes simply called the diffusion model, is one of the most popular models of decision making in psychology and cognitive neuroscience (for a fairly recent review, see Ratcliff et al., 2016). The DDM is usually applied to experiments that require subjects to quickly choose between two response alternatives on each trial. The model is able to predict the proportions of responses and the distributions of their associated reaction times (RTs).

DDM papers often contain figures similar the one below. Some time ago, I needed to make several such “DDM plots” for one of my own papers (Fig. 1 in Mathias et al., 2017), so I wrote a Python script to automate the process. This script (with some minor modifications and extra comments) is presented at the end of this post.

A barebones DDM plot. The plot is divided into the three panels. The top panel shows a kernel density estimate (KDE) of the degenerate distribution of RTs associated with one of the response options. These are actual trials simulated from a given set of DDM parameters via the HDDM package (Wiecki et al., 2013). The middle panel illustrates the paths of the decision variable on two trials. The lower panel is the degenerate distribution of RTs associated with the other response options.

The first thing the script does is simulate a bunch of trials under the DDM for a given set of parameters. This is done using HDDM, a Python package for fitting DDMs using Bayesian hierarchical inference (Wiecki et al., 2013). HDDM is slightly tricky to install these days. However, if you are interested in making DDM plots in Python, there is a good chance you’re already using or want to use HDDM. For this post, I installed HDDM into a new conda environment using the following commands:

conda create --name py36 python=3.6
conda activate py36
conda install conda-build
conda install pymc3
conda install pandas patsy matplotlib scipy tqdm
pip install cython
pip install hddm

Thanks to Esin Turkakin for posting this solution to the HDDM mailing list!

If you are interested in learning more about the DDM, you may wish to read a mini-review I wrote on the topic for a psychoacoustics audience (Mathias, 2016). Or for a thorough review, see Wagenmakers (2009).

Here is the code:

"""Create a DDM figure.

"""
import hddm
import numpy as np
from scipy.stats import norm
import matplotlib.pyplot as plt
from matplotlib.gridspec import GridSpec
from tqdm import tqdm


def setupfig():
    """Tweak for the target journal.

    """
    fig = plt.figure()
    gs = GridSpec(3, 1, height_ratios=[1, 2, 1], hspace=0)
    return fig, gs


def delabel(ax):
    """Strip labels.

    """
    ax.xaxis.set_ticklabels([])
    ax.yaxis.set_ticklabels([])
    ax.xaxis.set_ticks([])
    ax.yaxis.set_ticks([])


def kde(ax, x, mx, c):
    """Plot a KDE for reaction times.

    """
    x = x[x <= mx]
    bandwidth = 0.8 * x.std() * x.size ** (-1 / 5.0)
    support = np.linspace(0, mx, 500)
    kernels = []

    for x_i in tqdm(x):

        kernel = norm(x_i, bandwidth).pdf(support)
        kernels.append(kernel)
        density = np.sum(kernels, axis=0)

    my = np.max(density)
    ax.plot(support, density, c=c)
    ax.fill_between(support, 0, density, alpha=0.5, facecolor=c)
    ax.set_ylim(0, my * 1.05)
    delabel(ax)

    return my


def traces(ax, n, mx, **params):
    """Draw example diffusions.

    """
    x = np.linspace(0, mx, 101)
    delta = x[1]
    nd_samples = np.round(params["t"] / delta).astype(int)
    d_samples = len(x) - nd_samples
    y0 = np.zeros(nd_samples) * np.nan
    y0[-1] = 0

    for i in range(n):

        y1 = np.cumsum(norm.rvs(params["v"] * delta, np.sqrt(delta), size=d_samples))
        y = params["a"] * params["z"] + np.concatenate([y0, y1])

        try:

            idx1 = np.where(y > params["a"])[0][0] + 1

        except:

            idx1 = np.inf

        try:

            idx2 = np.where(y < 0)[0][0] + 1

        except:

            idx2 = np.inf

        if idx1 < idx2:

            y[idx1:] = np.nan
            ax.plot(x, y, c="C0", zorder=-12, alpha=0.25)

        if idx2 < idx1:

            y[idx2:] = np.nan
            ax.plot(x, y, c="C3", zorder=-11, alpha=0.25)

        ax.set_ylim(0, params["a"])
        ax.set_xlim(0, mx)
        delabel(ax)


def ddmfig(**params):
    """Draw a DDM plot with the given parameter values.

    """
    mx = 3.5  # max x-value; adjust this if simulated RTs are slower/faster
    size = 1500  # increase this number for better KDEs
    ntraces = 2  # increase this for more example diffusions

    # set up fig
    fig, gs = setupfig()

    # traces
    ax = plt.subplot(gs[1])
    traces(ax, ntraces, mx, **params)

    # data for kdes
    df, _ = hddm.generate.gen_rand_data(params, subjs=1, size=size)

    # top KDE
    ax = plt.subplot(gs[0])
    my = kde(ax, df[df.response == 1].rt.values, mx, "C0")

    # bottom KDE
    ax = plt.subplot(gs[2])
    kde(ax, df[df.response == 0].rt.values, mx, "C3")
    ax.set_ylim(0, my * 1.05)
    ax.invert_yaxis()

    # remove whitespace around fig
    plt.tight_layout(0)


def main():

    np.random.seed(25)
    ddmfig(v=0.7, a=1.5, t=0.6, z=0.5)
    plt.savefig("assets/scripts/ddm-fig.svg", bbox_inches=0, transparent=True)


if __name__ == "__main__":

    main()


References

Mathias, S. R. (2016). Unified analysis of accuracy and reaction times via models of decision making. 10.1121/2.0000219

Mathias, S. R., Knowles, E. E. M., Barrett, J., Leach, O., Buccheri, S., Beetham, T., Blangero, J., Poldrack, R. A., & Glahn, D. C. (2017). The processing-speed impairment in psychosis is more than just accelerated aging. Schizophrenia Bulletin, 43(4), 814–823. 10.1093/schbul/sbw168

Ratcliff, R., Smith, P. L., Brown, S. D., & McKoon, G. (2016). Diffusion decision model: Current issues and history. Trends in Cognitive Sciences, 20(4), 260–281. 10.1016/j.tics.2016.01.007

Wagenmakers, E. (2009). Methodological and empirical developments for the Ratcliff diffusion model of response times and accuracy. European Journal of Cognitive Psychology, 21(5), 641–671. 10.1080/09541440802205067

Wiecki, T. V., Sofer, I., & Frank, M. J. (2013). HDDM: Hierarchical Bayesian estimation of the drift-diffusion model in Python. Frontiers in Neuroinformatics, 7. 10.3389/fninf.2013.00014

Version history