Visualise the two parameter Indian Buffet Process

ibp_vis

import sys
import time
from pylab import *
import scipy.stats as stats

"""
Script to visualise draws from a two parameter Indian Buffet Process.

usage example:

$ python ibp.py 5 100 10 2
Showing 5 draws from an IBP(10.0, 2.0) for 100 samples

Parameter a affects the total number of dishes tried by diners. 
        Large a promotes dietary diversity.

Parameter b affects the sparsity of the matrix and overall level of dietary diversity. 
        Diners with a large b are like hipsters who only want novel dishes.

"""

def ibp(n,a,b,v=True):
  """Draws from a 2 parameter Indian Buffet Process IBP(a,b) for n diners. 
  Ref: Knowles & Ghahramani (2007) 'Infinite sparse factor analysis...' for further details.
  """
  ab = float(a*b)
  new_dishes = stats.poisson.rvs(array([ab/(b+i-1) for i in range(1,n+1)]))
  K = sum(new_dishes)
  oK = cumsum(new_dishes)
  Z = zeros((n,K))
  Z[0,0:oK[0]] = ones((1,new_dishes[0]))
  M = Z[0,:]
  for i in range(2,n+1):
    if not new_dishes[i-1] == 0:
      on = ones((1,new_dishes[i-1]))
      Z[i-1, oK[i-2]:oK[i-1]] = on 
    p = M[0:oK[i-2]]/(b + i - 1) 
    old_dishes = stats.binom.rvs(1,p)
    Z[i-1,0:oK[i-2]] = old_dishes
    M = M + Z[i-1,:]  
  return Z

def vis_ibp(R,N,a,b,t=.4):
  for r in range(R):
    Z = ibp(N,a,b)
    imshow(Z,cmap="gray",interpolation="none")
    show()
    time.sleep(t)

if __name__ == "__main__":
  R,N,a,b = int(sys.argv[1]), int(sys.argv[2]), float(sys.argv[3]), float(sys.argv[4])
  t=None
  if len(sys.argv)>5:
    t = float(sys.argv[5])
  print "Showing %s draws from an IBP(%s, %s) for %s samples" % (R,a,b,N)
  if not t:
    vis_ibp(R,N,a,b)
  else:
    vis_ibp(R,N,a,b,t=t)