import bnpy

from matplotlib import pylab
%pylab inline

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

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

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

bnpy.viz.BarsViz.plotExampleBarsDocs(Data, nDocToPlot=16, **imshowArgs);

hmodel, RInfo = bnpy.run(Data, 'FiniteTopicModel', 'Mult', 'VB',
                         K=5, alpha=0.1, lam=0.1, initname='kmeansplusplus',
                         nLap=500, printEvery=25, nCoordAscentItersLP=25,
                         nTask=2, jobname='demobarsVB-Kinit=5-kmeans')

hmodel, RInfo = bnpy.run(Data, 'FiniteTopicModel', 'Mult', 'VB',
                         K=10, alpha=0.1, lam=0.1, initname='kmeansplusplus',
                         nLap=500, printEvery=25, nCoordAscentItersLP=25,
                         nTask=2, jobname='demobarsVB-Kinit=10-kmeans')

hmodel, RInfo = bnpy.run(Data, 'FiniteTopicModel', 'Mult', 'VB',
                         K=20, alpha=0.1, lam=0.1, initname='kmeansplusplus',
                         nLap=500, printEvery=25, nCoordAscentItersLP=25,
                         nTask=2, jobname='demobarsVB-Kinit=20-kmeans')

from matplotlib import pylab
%pylab inline

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

bnpy.viz.PlotComps.plotCompsForTask('BarsK10V900/demobarsVB-Kinit=5-kmeans/1/', **imshowArgs)

bnpy.viz.PlotComps.plotCompsForTask('BarsK10V900/demobarsVB-Kinit=10-kmeans/1/', **imshowArgs)

bnpy.viz.PlotComps.plotCompsForTask('BarsK10V900/demobarsVB-Kinit=20-kmeans/1/', **imshowArgs)