Sunday, 6 January 2019

Animate gaussian distribution and scatter

The code below plots a bivariate gaussian distribution from a single frame of xy coordinates. I have recorded the entire code to display how this is done. But I really only need to determine how to animate the plot. As in iterate the code over each frame.

import numpy as np
import pandas as pd
from mpl_toolkits.axes_grid1 import make_axes_locatable
import matplotlib.pyplot as plt
import scipy.stats as sts
import matplotlib.colors as mcolors
import matplotlib.animation as animation

def datalimits(*data, pad=.15):
    dmin,dmax = min(d.min() for d in data), max(d.max() for d in data)
    spad = pad*(dmax - dmin)
    return dmin - spad, dmax + spad

def rot(theta):
    theta = np.deg2rad(theta)
    return np.array([
        [np.cos(theta), -np.sin(theta)],
        [np.sin(theta), np.cos(theta)]
    ])

def getcov(radius=1, scale=1, theta=0):
    cov = np.array([
        [radius*(scale + 1), 0],
        [0, radius/(scale + 1)]
    ])

    r = rot(theta)
    return r @ cov @ r.T

cim = plt.imread("https://i.stack.imgur.com/4q2Ev.png")
cim = cim[cim.shape[0]//5, 8:700, :]
cim_10 = cim[cim.shape[0] // 9 * np.arange(10)]  # array of 10 colors

cmap = mcolors.ListedColormap(cim_10)


def mvpdf(x, y, xlim, ylim, radius=1, velocity=0, scale=0, theta=0):
    """Creates a grid of data that represents the PDF of a multivariate gaussian.

    x, y: The center of the returned PDF
    (xy)lim: The extent of the returned PDF
    radius: The PDF will be dilated by this factor
    scale: The PDF be stretched by a factor of (scale + 1) in the x direction, and squashed by a factor of 1/(scale + 1) in the y direction
    theta: The PDF will be rotated by this many degrees

    returns: X, Y, PDF. X and Y hold the coordinates of the PDF.
    """
    # create the coordinate grids
    X,Y = np.meshgrid(np.linspace(*xlim), np.linspace(*ylim))

    # stack them into the format expected by the multivariate pdf
    XY = np.stack([X, Y], 2)

    # displace xy by half the velocity
    x,y = rot(theta) @ (velocity/2, 0) + (x, y)

    # get the covariance matrix with the appropriate transforms
    cov = getcov(radius=radius, scale=scale, theta=theta)

    # generate the data grid that represents the PDF
    PDF = sts.multivariate_normal([x, y], cov).pdf(XY)

    return X, Y, PDF

def mvpdfs(xs, ys, xlim, ylim, radius=None, velocity=None, scale=None, theta=None):
    PDFs = []
    for i,(x,y) in enumerate(zip(xs,ys)):
        kwargs = {
            'radius': radius[i] if radius is not None else 1,
            'velocity': velocity[i] if velocity is not None else 0,
            'scale': scale[i] if scale is not None else 0,
            'theta': theta[i] if theta is not None else 0,
            'xlim': xlim,
            'ylim': ylim
        }
        X, Y, PDF = mvpdf(x, y, **kwargs)
        PDFs.append(PDF)

    return X, Y, np.sum(PDFs, axis=0)

fig, ax = plt.subplots(figsize = (10,4))

def plotmvs(df, xlim=None, ylim=None, fig=fig, ax=ax):

    """Plot an xy point with an appropriately tranformed 2D gaussian around it.
    Also plots other related data like the reference point.
    """
    if xlim is None: xlim = datalimits(df['X'])
    if ylim is None: ylim = datalimits(df['Y'])

    if fig is None:
        fig = plt.figure(figsize=(8,8))
        ax = fig.gca()
    elif ax is None:
        ax = fig.gca()

    PDFs = []
    for (group,gdf),color in zip(df.groupby('group'), ('red', 'blue')):
        # plot the xy points of each group
        ax.plot(*gdf[['X','Y']].values.T, '.', c=color, alpha = 0.5)

        # fetch the PDFs of the 2D gaussian for each group
        kwargs = {
            'radius': gdf['Radius'].values if 'Radius' in gdf else None,
            'velocity': gdf['Velocity'].values if 'Velocity' in gdf else None,
            'scale': gdf['Scaling'].values if 'Scaling' in gdf else None,
            'theta': gdf['Rotation'].values if 'Rotation' in gdf else None,
            'xlim': xlim,
            'ylim': ylim
        }
        X, Y, PDF = mvpdfs(gdf['X'].values, gdf['Y'].values, **kwargs)
        PDFs.append(PDF)

    # create the PDF for all points from the difference of the sums of the 2D Gaussians from group A and group B
    PDF = PDFs[0] - PDFs[1]

    # normalize PDF by shifting and scaling, so that the smallest value is 0 and the largest is 1
    normPDF = PDF - PDF.min()
    normPDF = normPDF/normPDF.max()

    # plot and label the contour lines of the 2D gaussian
    cs = ax.contour(X, Y, normPDF, levels=5, colors='white', alpha=.1)
    ax.clabel(cs, colors = 'black', fontsize=8)

    # plot the filled contours of the 2D gaussian. Set levels high for smooth contours
    cfs = ax.contourf(X, Y, normPDF, levels=100, cmap=cmap)

    # create the colorbar and ensure that it goes from 0 -> 1
    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.1)
    cbar = fig.colorbar(cfs, ax=ax, cax=cax)
    cbar.set_ticks([0,.1,.2,.3,.4,.5,.6,.7,.8,.9,1])

    ax.grid(False)

    return fig, ax


time = [1,2,3,4,5,6,7,8,9,10]
d = ({      
    'A1_Y' : [10,20,15,20,25,40,50,60,61,65],                 
    'A1_X' : [15,10,15,20,25,25,30,40,60,61], 
    'A2_Y' : [10,13,17,10,20,24,29,30,33,40],                 
    'A2_X' : [10,13,15,17,18,19,20,21,26,30],
    'A3_Y' : [11,12,15,17,19,20,22,25,27,30],                 
    'A3_X' : [15,18,20,21,22,28,30,32,35,40], 
    'A4_Y' : [15,20,15,20,25,40,50,60,61,65],   
    'A4_X' : [16,20,15,20,25,40,50,60,61,65],                 
    'B1_Y' : [18,20,15,20,25,40,50,60,61,65],                 
    'B1_X' : [17,20,15,20,25,40,50,60,61,65], 
    'B2_Y' : [13,20,15,20,25,40,50,60,61,65],                 
    'B2_X' : [12,20,15,20,25,40,50,60,61,65],
    'B3_Y' : [15,20,15,20,25,40,50,60,61,65],                 
    'B3_X' : [18,20,15,20,25,40,50,60,61,65], 
    'B4_Y' : [19,20,15,20,25,40,50,60,61,65],   
    'B4_X' : [20,20,15,20,25,40,50,60,61,65],                                                                                   
     })        


# a list of tuples of the form ((time, group_id, point_id, value_label), value)
 tuples = [((t, k.split('_')[0][0], int(k.split('_')[0][1:]), k.split('_')[1]), v[i]) for k,v in d.items() for i,t in enumerate(time)]

df = pd.Series(dict(tuples)).unstack(-1)
df.index.names = ['time', 'group', 'id']

for time,tdf in df.groupby('time'):
    plotmvs(tdf)

Below is my attempt to animate all the [A_'s]. I'm also wondering if it's possible to animate the gaussian?

def animate(i) :
    tdf.set_offsets([[tdf.iloc[0:,1][0+i][0], tdf.iloc[0:,0][0+i][0]], [tdf.iloc[0:,1][0+i][1], tdf.iloc[0:,0][0+i][1]], [tdf.iloc[0:,1][0+i][2], tdf.iloc[0:,0][0+i][2]], [tdf.iloc[0:,1][0+i][3], tdf.iloc[0:,0][0+i][3]], [tdf.iloc[0:,1][0+i][4], tdf.iloc[0:,0][0+i][4]]])


ani = animation.FuncAnimation(fig, animate, np.arange(0,10),# init_func = init,
                          interval = 10, blit = False)



from Animate gaussian distribution and scatter

No comments:

Post a Comment