scalingcollapse.py

import numpy as np
import matplotlib.pyplot as pl
from matplotlib import cm
from matplotlib.colors import Normalize
from matplotlib.widgets import Slider

def rescale(xs, ys, params, values):
    """
    An example scaling_func that can be provided to interactive_collapse.  This
    scales the zombie data by the variables alphac, sigma, tau for curves at
    various values of alpha (the bite to kill ratio packed into `values`)
    """
    # unpack the scaling parameters from the params list
    alphac, sigma, tau = params

    outx, outy = [], []
    for alpha,x,y in zip(values, xs, ys):
        outx.append(x/abs(alpha-alphac)**(-1./sigma))
        outy.append(x**(-tau)*y*abs(alpha-alphac)**(-1./sigma))

    return np.array(outx), np.array(outy)

def interactive_collapse(xs, ys, scaling_func, values, params,
        param_names=None, param_ranges=None, fig=None):
    """
    Interactive scaling collapse utility that allows you to change the scaling
    parameters to achieve a good scaling collapse.

    Parameters:
    -----------
    xs : ndarray [N, X]
        x-values for N one dimensional curves of length X

    ys : ndarray [N, X]
        y-values for N one dimensional curves of length X corresponding to the xs

    scaling_func : function
        A function that takes the arguments (xs, ys, params, values) and returns
        a set of (xs', ys') that have been scaled appropriately.

    values : ndarray [N], [N, Y]
        Parameter values at which the data in (xs, ys) was taken. For example,
        temperature or system size (or both) for each x,y curve.

    params : ndarray [Y']
        Parameters which to optimize for the collapse such as exponents,
        critical point value, or other non-universal constants

    param_names : list [Y'] optional
        Names of the parameters which the user will manipulate

    param_ranges : list [Y', 2] optional
        range for the particular parameters in question

    fig : `pl.figure`
        Use an existing figure for the scaling collapse
    """
    # create the plot itself, clear and resize
    if fig is None:
        fig = pl.figure()
    pl.subplots_adjust(left=0.15, bottom=0.15, right=0.6)
    pl.clf()

    lines = []
    out = scaling_func(xs, ys, params, values)
    colors = cm.copper(Normalize()(np.arange(len(xs))))
    for a, c, x, y in zip(values, colors, out[0], out[1]):
        lines.append(pl.loglog(x, y, 'o-', label=str(a), c=c)[0])

    # create defaults for values that don't exist
    nparams = len(params)
    if param_names is None:
        param_names = ['param-{}'.format(i) for i in xrange(nparams)]

    if param_ranges is None:
        param_ranges = [(p-1,p+1) for p in params]

    # function to be called when any slider is changed
    def update(val):
        params = np.array([i.val for i in slider_wgt])
        print ' '.join("{}={}".format(k,v) for (k,v) in zip(param_names, params))

        out = scaling_func(xs, ys, params, values)
        for line, x, y in zip(lines, out[0], out[1]):
            line.set_xdata(x)
            line.set_ydata(y)

            fig.canvas.draw_idle()

    # create the widgets to use for the update
    slider_axs = []
    slider_wgt = []
    for i, (param, name, prange) in enumerate(zip(params, param_names, param_ranges)):
        axs = pl.axes([0.65, 0.2+0.6*float(i+1)/nparams, 0.25, 0.03], axisbg='white')
        wgt = Slider(axs, name, prange[0], prange[1], valinit=param)
        wgt.on_changed(update)
        slider_axs.append(axs)
        slider_wgt.append(wgt)