<html><head><meta name="color-scheme" content="light dark"></head><body><pre style="word-wrap: break-word; white-space: pre-wrap;">'''
Robert's attempt to make Christian's pretty 3D histograms. These methods could use
a lot of work still

'''



import numpy as np
import matplotlib.pyplot as pp
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.collections import PolyCollection
import matplotlib.cm
import scipy.stats

def hist3d_from_gkdes(xmin, xmax, gkdes, zs, title=None):
    '''data has to be in the right format'''
    
    zs = np.array(zs, dtype='float')
    xs = np.linspace(xmin, xmax, 100)
    max_y = 0
    
    if not len(gkdes) == len(zs):
        raise ValueError('Dimensions dont match')
    
    fig = pp.figure()
    ax = fig.gca(projection='3d')
    verts = []
    for gkde in gkdes:
        ys = gkde.evaluate(xs)
        # add values at the begining and end so that it comes down
        # to the axes
        ys = np.insert(ys, 0, 0)
        ys = np.insert(ys, len(ys), 0)
        xs = np.insert(xs, 0, xs[0])
        xs = np.insert(xs, len(xs), xs[-1])
        #ys = np.random.rand(len(xs))
        verts.append(zip(xs, ys))
        max_y = max(max_y, np.max(ys))
    
    colormap = matplotlib.cm.jet(zs / np.max(zs), alpha=0.6)
    poly = PolyCollection(verts, facecolors=colormap)
    poly.set_alpha(0.7)
    ax.add_collection3d(poly, zs=zs, zdir='y')
    ax.set_xlabel('X')
    ax.set_xlim3d(xmin, xmax)
    ax.set_ylabel('Y')
    ax.set_ylim3d(np.min(zs) - 1, np.max(zs) + 1)
    ax.set_zlabel('prob density')
    ax.set_zlim3d(0, max_y * 1.1)
    if title is not None:
        ax.set_title(title)

    pp.show()

def hist3d(A, xmin=None, xmax=None, z=None, title=None):
    ''' A is m rows by n columns of data where each row represnts
    n independent observatins from the same probability distribution
    
    z is a either None or a vector of length n given the label of that
    goes with each row in A
    
    plot 3D histogram of the data in A with the labels in y
    '''
    if xmin is None:
        xmin = np.min(A)
    if xmax is None:
        xmax = np.max(A)
    
    
    if z is None:
        z = np.arange(A.shape[0])
      
    gkdes = [None] * A.shape[0]
    for i in range(A.shape[0]):
        gkdes[i] = scipy.stats.gaussian_kde(A[i,:])
    hist3d_from_gkdes(xmin, xmax, gkdes, z, title)
  
  
if __name__ == '__main__':
    zs = [1, 1.25, 1.5, 2]
    gkdes = [None] * len(zs)
    xmin, xmax = float('inf'), float('-inf')
    for i, z in enumerate(zs):
        r = z + np.random.randn(1000)
        gkdes[i] = scipy.stats.gaussian_kde(r)
        xmax = max(xmax, np.max(r))
        xmin = min(xmin, np.min(r))
    
    hist3d_from_gkdes(xmin, xmax, gkdes, zs)</pre></body></html>