It would be nice to compute the log probability of the model during each fitting iteration of fitting and plot it so that people have a sense of whether the model has converged or not. Here are a few notes on implementation:
model_likelihood returns a dictionary that maps each variable name to the conditional probability of that variable. We could store all of these values separately, but it might be simpler to just store their sum (the total joint log prob)
In terms of where/when to call model_likelihood, one option would be every resample iteration (and perhaps store the result in the model dictionary). Another option would be just when updating the history during model fitting (i.e. here)
When plotting, there's going to be a huge increase in the first few iterations and then a slow increase after that. If the y-axis includes the full range, it may be hard to see then the slow increase becomes a plateau. Therefore, I propose that we superimpose two copies of the plot, one with the full-yaxis and one with a zoomed in y-axis (and perhaps use color coded axis ticks on the left and right sides of the subplot). Another option would be to just show the zoomed-in version.
It would be nice to compute the log probability of the model during each fitting iteration of fitting and plot it so that people have a sense of whether the model has converged or not. Here are a few notes on implementation:
jax_moseq.models.keypoint_slds.model_likelihoodmodel_likelihoodreturns a dictionary that maps each variable name to the conditional probability of that variable. We could store all of these values separately, but it might be simpler to just store their sum (the total joint log prob)model_likelihood, one option would be every resample iteration (and perhaps store the result in themodeldictionary). Another option would be just when updating the history during model fitting (i.e. here)plot_progressfunction