Pystan `get_mass_matrix` method

kon · July 19, 2019, 3:51pm

Is it possible to get or set the mass matrix from a fit object in Pystan?
More generally, is it possible to save the critical info from warmup and reuse it later for another chain? The idea being to run a long warmup chain, save the output and use it in the future to skip warmup by initializing a future chain with the saved output.

I came across this commit by @ahartikainen which seems to be the answer. I installed the development version of pystan but it seems like get_mass_matrix is not an available method for a fit object.

Thank you in advance for any help!

ahartikainen · July 19, 2019, 4:01pm

Hi, see

github.com

stan-dev/pystan/blob/develop/doc/hmc_euclidian_metric.rst

.. _hmc_euclidian_metric:

.. currentmodule:: pystan

======================
 HMC: Euclidian Metric
======================

Euclidian Metric (also known as mass matrix) is one of the tuned parameters for hmc algorithm. See
`hmc-algorithm-parameters <https://mc-stan.org/docs/2_19/reference-manual/hmc-algorithm-parameters.html>`_.

Two examples are shown where it is possible to use pre-tuned metric-matrix with other
pre-tuned parameters.

-------------------------------------
Example 1: (pseudo-)continue sampling
-------------------------------------

The first example shows an example how to (pseudo-)continue chains to get more draws to old fit.
The method is described as pseudo due to fact that continued sampling does not equal sampling

This file has been truncated. show original

and small example

github.com

ahartikainen/fit_check_fit_loop/blob/master/MoreSamples_round1/increase_draws_in_steps.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example how to iteratively get more draws (on PyStan)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pystan # current develop branch (2019-06-28)\n",
    "import arviz as az # current master branch (2019-06-28)\n",
    "\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "import numpy as np\n",

This file has been truncated. show original

kon · July 19, 2019, 7:46pm

That’s great - Thank you, @ahartikainen !
For future readers, the method that returns the mass matrix is
fit.get_inv_metric(as_dict=True)

davidh · July 30, 2019, 12:04am

Thank you so much @ahartikainen and Pystan team for implementing this, it looks incredibly useful. A couple questions about using this feature:

Is there any straightforward intuition for why continuing sampling “does not equal sampling done in one step”? Does this change how we should interpret inference results from this pseudo-sampling method? Or will it just not give you the same results because of the seed?
In the docs example for running the warmup first and then sampling in a different run, only one chain is used. Is it reasonable to do this? Would it make sense to compute some sort of r-hat on the inverse metric to see if you’ve run enough warmup steps?

ahartikainen · July 30, 2019, 3:16am

The random number generator used by Stan is not at the same state, when you restart your sampling

2a. For basic usage, no. It is something that can be done, if user has complicated model and wants to minimize the sampling time (this could happen for example in “production”). I think I should add some warning in there.

2b. (edit.) Maybe, sounds interesting.

davidh · July 30, 2019, 12:24pm

Thanks for the quick response!

The random number generator used by Stan is not at the same state, when you restart your sampling

So just to clarify, this means that the results won’t be exactly the same but inferences based on the samples won’t be any less valid than from samples collected in a single session?
In principle, could Stan return a seed at the end of sampling that would allow you to perfectly recover the state when you restart? I understand this might not be straightforward to implement though.

ahartikainen · July 30, 2019, 12:26pm

Yes.

I don’t know if it is possible to return / save the RNG state used by Stan.

cc @Bob_Carpenter

Topic		Replies	Views
How can we access/plot the warm up samples? PyStan	3	437	October 6, 2022
Sampling chains from the middle in cmdstanpy General	10	419	October 8, 2020
Sampler parameters which the typical user might need to set (hmc_nuts_diag_e_adapt only) Developers	18	1189	December 3, 2019
Checkpointing with CmdStanPy General	2	436	September 24, 2020
Fit.extract() takes a long time PyStan	9	2862	May 22, 2018

Pystan `get_mass_matrix` method

Related Topics