import bnpy

from matplotlib import pylab
%pylab inline

imshowArgs = dict(interpolation='nearest', 
                  cmap='bone_r', 
                  vmin=0.0, 
                  vmax=10./900,
                  )

import MixBarsK10V900
Data = MixBarsK10V900.get_data(nDocTotal=1000, nWordsPerDoc=100)
Data.name = 'MixBarsK10V900'

bnpy.viz.BarsViz.showTopicsAsSquareImages(MixBarsK10V900.Defaults['topics'], **imshowArgs)

bnpy.viz.BarsViz.plotExampleBarsDocs(Data, **imshowArgs)

hmodel, RInfo = bnpy.run(Data, 'DPMixtureModel', 'Mult', 'moVB',
                         nLap=20, K=1, lam=0.1, moves='birth,merge', 
                         nTask=2, jobname='demobirthmerge-Kinit=1')

hmodel, RInfo = bnpy.run(Data, 'DPMixtureModel', 'Mult', 'moVB',
                         nLap=20, K=4, lam=0.1, moves='birth,merge', 
                         nTask=2, jobname='demobirthmerge-Kinit=4')

hmodel, RInfo = bnpy.run(Data, 'DPMixtureModel', 'Mult', 'moVB',
                         nLap=20, K=16, lam=0.1, moves='birth,merge', 
                         nTask=2, jobname='demobirthmerge-Kinit=16')

from matplotlib import pylab
%pylab inline

bnpy.viz.PlotELBO.plotJobsThatMatchKeywords('MixBarsK10V900/demobirthmerge-*');
pylab.legend(loc='lower right');

bnpy.viz.PlotELBO.plotJobsThatMatchKeywords('MixBarsK10V900/demobirthmerge-*', yvar='K');

bnpy.viz.PlotComps.plotCompsForTask('MixBarsK10V900/demobirthmerge-Kinit=16/1/', **imshowArgs)

bnpy.viz.PlotComps.plotCompsForTask('MixBarsK10V900/demobirthmerge-Kinit=16/1/', lap=0, **imshowArgs)

bnpy.viz.PlotComps.plotCompsForTask('MixBarsK10V900/demobirthmerge-Kinit=16/1/', lap=2, **imshowArgs)