DEMetropolis(Z): tune_drop_fraction#

The implementation of DEMetropolisZ in PyMC3 uses a different tuning scheme than described by ter Braak & Vrugt, 2008. In our tuning scheme, the first tune_drop_fraction * 100 % of the history from the tuning phase is dropped when the tune iterations end and sampling begins.

In this notebook, a D-dimenstional multivariate normal target densities is sampled with DEMetropolisZ at different tune_drop_fraction settings to show why the setting was introduced.

import time

import arviz as az
import ipywidgets
import numpy as np
import pandas as pd
import pymc3 as pm

from matplotlib import cm, gridspec
from matplotlib import pyplot as plt

print(f"Running on PyMC3 v{pm.__version__}")

Running on PyMC3 v3.9.0

%config InlineBackend.figure_format = 'retina'
az.style.use("arviz-darkgrid")

Setting up the Benchmark#

We use a multivariate normal target density with some correlation in the first few dimensions.

def get_mvnormal_model(D: int) -> pm.Model:
    true_mu = np.zeros(D)
    true_cov = np.eye(D)
    true_cov[:5, :5] = np.array(
        [
            [1, 0.5, 0, 0, 0],
            [0.5, 2, 2, 0, 0],
            [0, 2, 3, 0, 0],
            [0, 0, 0, 4, 4],
            [0, 0, 0, 4, 5],
        ]
    )

    with pm.Model() as pmodel:
        x = pm.MvNormal("x", mu=true_mu, cov=true_cov, shape=(D,))

    true_samples = x.random(size=1000)
    truth_id = az.data.convert_to_inference_data(true_samples[np.newaxis, :], group="random")
    return pmodel, truth_id

The problem will be 10-dimensional and we run 5 independent repetitions.

D = 10
N_tune = 10000
N_draws = 10000
N_runs = 5
pmodel, truth_id = get_mvnormal_model(D)
pmodel.logp(pmodel.test_point)

C:\Users\osthege\AppData\Local\Continuum\miniconda3\envs\pm3-dev2\lib\site-packages\arviz\data\inference_data.py:99: UserWarning: random group is not defined in the InferenceData scheme
  "{} group is not defined in the InferenceData scheme".format(key), UserWarning

array(-9.99410429)

df_results = pd.DataFrame(columns="drop_fraction,r,ess,t,idata".split(",")).set_index(
    "drop_fraction,r".split(",")
)

for drop_fraction in (0, 0.5, 0.9, 1):
    for r in range(N_runs):
        with pmodel:
            t_start = time.time()
            step = pm.DEMetropolisZ(tune="lambda", tune_drop_fraction=drop_fraction)
            idata = pm.sample(
                cores=6,
                tune=N_tune,
                draws=N_draws,
                chains=1,
                step=step,
                start={"x": [7.0] * D},
                discard_tuned_samples=False,
                return_inferencedata=True,
                # the replicates (r) have different seeds, but they are comparable across
                # the drop_fractions. The tuning will be identical, they'll divergen in sampling.
                random_seed=2020 + r,
            )
            t = time.time() - t_start
            df_results.loc[(drop_fraction, r), "ess"] = float(az.ess(idata).x.mean())
            df_results.loc[(drop_fraction, r), "t"] = t
            df_results.loc[(drop_fraction, r), "idata"] = idata

Sequential sampling (1 chains in 1 job)
DEMetropolisZ: [x]