nRhat-convergence.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Section 6: Numerical experiments\n",
    "Numerical experiments for $\\mathfrak n \\widehat R$ on several test models. Runs and saves results."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup: import libraries"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To install `fun_mc`, run the following commands (in virtual environnment).\n",
    "```\n",
    "!rm -Rf probability\n",
    "!rm -Rf fun_mc\n",
    "!rm -Rf inference_gym\n",
    "!git clone https://github.com/tensorflow/probability.git\n",
    "!mv probability/spinoffs/fun_mc/fun_mc .\n",
    "!mv probability/spinoffs/inference_gym/inference_gym .\n",
    "!pip install tf-nightly tfp-nightly jax jaxlib\n",
    "```\n",
    "and\n",
    "```\n",
    "!pip install immutabledict\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from matplotlib.pyplot import *\n",
    "font = {'family' : 'normal',\n",
    "        'weight' : 'bold',\n",
    "        'size'   : 14}\n",
    "matplotlib.rc('font', **font)\n",
    "\n",
    "# Use this to silence check type warning messages.\n",
    "import logging\n",
    "logging.disable(logging.WARNING)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "from jax import random\n",
    "from jax import numpy as jnp\n",
    "\n",
    "# import tfp models and datasets\n",
    "from inference_gym import using_jax as gym\n",
    "\n",
    "from fun_mc import using_jax as fun_mcmc\n",
    "\n",
    "from tensorflow_probability.python.internal import prefer_static as ps\n",
    "from tensorflow_probability.python.internal import unnest\n",
    "import tensorflow_probability as _tfp\n",
    "tfp = _tfp.substrates.jax\n",
    "tfd = tfp.distributions\n",
    "tfb = tfp.bijectors\n",
    "\n",
    "tfp_np = _tfp.substrates.numpy\n",
    "tfd_np = tfp_np.distributions\n",
    "\n",
    "\n",
    "# import arviz as az\n",
    "from tensorflow_probability.python.internal.unnest import get_innermost  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Directory to save results (adjust to your setting!)\n",
    "deliv_dir = \"/mnt/home/cmargossian/Code/nested-rhat/deliv/\"\n",
    "data_dir = \"/mnt/home/cmargossian/Code/nested-rhat/data/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"utility.py\") as f: exec(f.read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tuning for the numerical experiments\n",
    "num_seed = 10\n",
    "adapt_warmup = False\n",
    "if (not adapt_warmup): nRhat_upper = 1\n",
    "\n",
    "num_chains = 2048      # total number of chains\n",
    "num_super_chains = 16  # K, options: 2, 8, 16, 64, 256, 1024\n",
    "num_sub_chains = num_chains // num_super_chains  # M\n",
    "num_samples = 1        # length of sampling phase\n",
    "\n",
    "num_warmup = 10  # 1000\n",
    "total_samples = num_warmup + num_samples + 1\n",
    "\n",
    "max_warmup = 1000\n",
    "warmup_window = 100\n",
    "\n",
    "# Start with a series of small windows for initialized chains\n",
    "# and then switch to a wider warmup window.\n",
    "window_array = np.append(np.repeat(10, 10),\n",
    "                      np.repeat(warmup_window, max_warmup // warmup_window - 1))\n",
    "\n",
    "naive_super_chains = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num_samples: 1\n",
      "num_super_chains: 16\n",
      "(total) num_chains: 2048\n",
      "num_seed: 10\n",
      "nRhat_upper: 1\n",
      "naive_super_chains: True\n"
     ]
    }
   ],
   "source": [
    "# check tuning parameters of experiment\n",
    "print(\"num_samples:\", num_samples)\n",
    "print(\"num_super_chains:\", num_super_chains)\n",
    "print(\"(total) num_chains:\", num_chains)\n",
    "print(\"num_seed:\", num_seed)\n",
    "print(\"nRhat_upper:\", nRhat_upper)\n",
    "print(\"naive_super_chains:\", naive_super_chains)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run experiments"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rosenbrock distribution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "metadata": {},
   "outputs": [],
   "source": [
    "target = gym.targets.VectorModel(gym.targets.Banana(),\n",
    "                                 flatten_sample_transformations=True)\n",
    "num_dimensions = target.event_shape[0]\n",
    "init_step_size = 1.\n",
    "\n",
    "def target_log_prob_fn(x):\n",
    "  \"\"\"Unnormalized, unconstrained target density.\n",
    "\n",
    "  This is a thin wrapper that applies the default bijectors so that we can\n",
    "  ignore any constraints.\n",
    "  \"\"\"\n",
    "  y = target.default_event_space_bijector(x)\n",
    "  fldj = target.default_event_space_bijector.forward_log_det_jacobian(x)\n",
    "  return target.unnormalized_log_prob(y) + fldj\n",
    "\n",
    "# NOTE: Avoid initials centered around the true mean.\n",
    "offset = 2\n",
    "def initialize (shape, key = random.PRNGKey(37272709)):\n",
    "  return 10 * random.normal(key, shape + (num_dimensions,)) + offset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get some estimates of the mean and variance.\n",
    "try:\n",
    "  mean_est = target.sample_transformations['identity'].ground_truth_mean\n",
    "except:\n",
    "  print('no ground truth mean')\n",
    "  mean_est = (result.all_states[num_warmup:, :]).mean(0).mean(0)\n",
    "try:\n",
    "  var_est = target.sample_transformations['identity'].ground_truth_standard_deviation**2\n",
    "except:\n",
    "  print('no ground truth std dev')\n",
    "  var_est = ((result.all_states[num_warmup:, :]**2).mean(0).mean(0) -\n",
    "             mean_est**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define MCMC kernel\n",
    "# NOTE: to compute classic Rhat, need at least 2 iterations per chain.\n",
    "# To insure this, total_samples is incremented by 1.\n",
    "\n",
    "\n",
    "kernel = construct_kernel(target_log_prob_fn = target_log_prob_fn,\n",
    "                          init_step_size = init_step_size,\n",
    "                          num_warmup = num_warmup)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "metadata": {},
   "outputs": [],
   "source": [
    "# seed = jax.random.split(jax.random.PRNGKey(1), num_seed)\n",
    "# initial_state = initialize((num_chains,), key = seed[0] + 1954)\n",
    "# # initial_state = np.repeat(initial_state, num_chains // num_super_chains,\n",
    "# #                                 axis = 0)\n",
    "\n",
    "# initial_state\n",
    "num_seed = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEED : [696157669 674520459]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<string>:212: DeprecationWarning: the `interpolation=` argument to quantile was renamed to `method=`, which has additional options.\n",
      "Users of the modes 'nearest', 'lower', 'higher', or 'midpoint' are encouraged to review the method they used. (Deprecated NumPy 1.22)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEED : [1771252691  383428323]\n",
      "SEED : [3312214627   69373237]\n",
      "SEED : [1094216230 2034336350]\n",
      "SEED : [1788847948 1554193616]\n",
      "SEED : [338758650 111605681]\n",
      "SEED : [4073466969 2556348314]\n",
      "SEED : [ 167784864 1021574613]\n",
      "SEED : [3458295564 1820549682]\n",
      "SEED : [3023805236 1296165206]\n"
     ]
    }
   ],
   "source": [
    "kernel_cold, kernel_warm = adaptive_kernels(target_log_prob_fn, init_step_size, num_warmup)\n",
    "index_param = np.array([0, 1])  # only two dimensions for this distribution\n",
    "\n",
    "mc_mean_list, warmup_length,\\\n",
    "squared_error_list, nrhat_list = run_forge_chain(num_seed = num_seed,\n",
    "                                          kernel_cold = kernel_cold,\n",
    "                                          kernel_warm = kernel_warm,\n",
    "                                          initialize = initialize,\n",
    "                                          num_super_chains = num_super_chains,\n",
    "                                          num_warmup = window_array,\n",
    "                                          num_samples = num_samples,\n",
    "                                          target_rhat = nRhat_upper,\n",
    "                                          max_num_steps = window_array.shape[0],\n",
    "                                          index_param = index_param,\n",
    "                                          mean_benchmark = mean_est,\n",
    "                                          var_benchmark = var_est,\n",
    "                                          naive_super_chains = naive_super_chains)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Save output into npy files\n",
    "model_name = \"rosenbrock\"\n",
    "exp_parm = \"_K\" + str(num_super_chains) + \"_M\" + str(num_sub_chains) + \"_N\" + str(num_samples)\n",
    "\n",
    "if naive_super_chains:\n",
    "    exp_parm = \"_naive\" + exp_parm\n",
    "\n",
    "np.save(deliv_dir + model_name + exp_parm + \"_nrhat\", nrhat_list)\n",
    "np.save(deliv_dir + model_name + exp_parm + \"_squared_error\", squared_error_list)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### German Credit Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "metadata": {},
   "outputs": [],
   "source": [
    "target = gym.targets.VectorModel(gym.targets.GermanCreditNumericLogisticRegression(),\n",
    "                                  flatten_sample_transformations=True)\n",
    "num_dimensions = target.event_shape[0]\n",
    "init_step_size = 0.02\n",
    "\n",
    "def target_log_prob_fn(x):\n",
    "  \"\"\"Unnormalized, unconstrained target density.\n",
    "\n",
    "  This is a thin wrapper that applies the default bijectors so that we can\n",
    "  ignore any constraints.\n",
    "  \"\"\"\n",
    "  y = target.default_event_space_bijector(x)\n",
    "  fldj = target.default_event_space_bijector.forward_log_det_jacobian(x)\n",
    "  return target.unnormalized_log_prob(y) + fldj\n",
    "\n",
    "offset = 0.5\n",
    "def initialize (shape, key = random.PRNGKey(37272709)):\n",
    "  return 3 * random.normal(key, shape + (num_dimensions,)) + offset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get some estimates of the mean and variance.\n",
    "try:\n",
    "  mean_est = target.sample_transformations['identity'].ground_truth_mean\n",
    "except:\n",
    "  print('no ground truth mean')\n",
    "  mean_est = (result.all_states[num_warmup:, :]).mean(0).mean(0)\n",
    "try:\n",
    "  var_est = target.sample_transformations['identity'].ground_truth_standard_deviation**2\n",
    "except:\n",
    "  print('no ground truth std dev')\n",
    "  var_est = ((result.all_states[num_warmup:, :]**2).mean(0).mean(0) -\n",
    "             mean_est**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEED : [696157669 674520459]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-03-19 11:24:37.695920: W tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: OUT_OF_RANGE: End of sequence\n",
      "<string>:212: DeprecationWarning: the `interpolation=` argument to quantile was renamed to `method=`, which has additional options.\n",
      "Users of the modes 'nearest', 'lower', 'higher', or 'midpoint' are encouraged to review the method they used. (Deprecated NumPy 1.22)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEED : [1771252691  383428323]\n",
      "SEED : [3312214627   69373237]\n",
      "SEED : [1094216230 2034336350]\n",
      "SEED : [1788847948 1554193616]\n",
      "SEED : [338758650 111605681]\n",
      "SEED : [4073466969 2556348314]\n",
      "SEED : [ 167784864 1021574613]\n",
      "SEED : [3458295564 1820549682]\n",
      "SEED : [3023805236 1296165206]\n"
     ]
    }
   ],
   "source": [
    "kernel_cold, kernel_warm = adaptive_kernels(target_log_prob_fn, init_step_size, num_warmup)\n",
    "index_param = np.arange(0, 25)\n",
    "\n",
    "mc_mean_list, warmup_length,\\\n",
    "squared_error_list, nrhat_list = run_forge_chain(num_seed = num_seed,\n",
    "                                          kernel_cold = kernel_cold,\n",
    "                                          kernel_warm = kernel_warm,\n",
    "                                          initialize = initialize,\n",
    "                                          num_super_chains = num_super_chains,\n",
    "                                          num_warmup = window_array,\n",
    "                                          num_samples = num_samples,\n",
    "                                          target_rhat = nRhat_upper,\n",
    "                                          max_num_steps = window_array.shape[0],\n",
    "                                          index_param = index_param,\n",
    "                                          mean_benchmark = mean_est,\n",
    "                                          var_benchmark = var_est,\n",
    "                                          naive_super_chains = naive_super_chains)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save output into npy files\n",
    "model_name = \"german\"\n",
    "exp_parm = \"_K\" + str(num_super_chains) + \"_M\" + str(num_sub_chains) + \"_N\" + str(num_samples)\n",
    "\n",
    "if naive_super_chains:\n",
    "    exp_parm = \"_naive\" + exp_parm\n",
    "\n",
    "np.save(deliv_dir + model_name + exp_parm + \"_nrhat\", nrhat_list)\n",
    "np.save(deliv_dir + model_name + exp_parm + \"_squared_error\", squared_error_list)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Eight Schools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: inference gym stores the centered parameterization\n",
    "target_raw = gym.targets.EightSchools()  # store raw to examine doc.\n",
    "target = gym.targets.VectorModel(target_raw,\n",
    "                                  flatten_sample_transformations = True)\n",
    "num_dimensions = target.event_shape[0]\n",
    "init_step_size = 1\n",
    "\n",
    "# Using underdispersed initis can show case problems with our diagnostics.\n",
    "# underdispered = False\n",
    "# Options: underdispersed, overdispersed, prior\n",
    "init_type = \"prior\"\n",
    "if init_type == \"underdispersed\":\n",
    "  offset = 0.0\n",
    "  def initialize (shape, key = random.PRNGKey(37272709)):\n",
    "    return 1 * random.normal(key, shape + (num_dimensions,)) + offset\n",
    "elif init_type == \"overdispersed\":\n",
    "  offset = 0.0\n",
    "  def initialize (shape, key = random.PRNGKey(37272709)):\n",
    "    return 100 * random.normal(key, shape + (num_dimensions,)) + offset\n",
    "elif init_type == \"prior\":\n",
    "  def initialize (shape, key = random.PRNGKey(37272709)):\n",
    "    prior_scale = jnp.append(jnp.array([10., 1.]), jnp.repeat(1., 8))\n",
    "    prior_offset = jnp.append(jnp.array([0., 5.]), jnp.repeat(0., 8))\n",
    "    return prior_scale * random.normal(key, shape + (num_dimensions,)) + prior_offset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_schools = 8\n",
    "y = np.array([28, 8, -3, 7, -1, 1, 18, 12], dtype = np.float32)\n",
    "sigma = np.array([15, 10, 16, 11, 9, 11, 10, 18], dtype = np.float32)\n",
    "\n",
    "# NOTE: the reinterpreted batch dimension specifies the dimension of\n",
    "# each indepdent variable, here the school.\n",
    "model = tfd.JointDistributionSequential([\n",
    "    tfd.Normal(loc = 0., scale = 10., name = \"mu\"),\n",
    "    tfd.Normal(loc = 5., scale = 1., name = \"log_tau\"),\n",
    "    tfd.Independent(tfd.Normal(loc = jnp.zeros(num_schools),\n",
    "                               scale = jnp.ones(num_schools),\n",
    "                               name = \"eta\"),\n",
    "                    reinterpreted_batch_ndims = 1),\n",
    "    lambda eta, log_tau, mu: (\n",
    "        tfd.Independent(tfd.Normal(loc = (mu[..., jnp.newaxis] +\n",
    "                                        jnp.exp(log_tau[..., jnp.newaxis]) *\n",
    "                                        eta),\n",
    "                                   scale = sigma),\n",
    "                        name = \"y\",\n",
    "                        reinterpreted_batch_ndims = 1))\n",
    "  ])\n",
    "\n",
    "def target_log_prob_fn(x):\n",
    "  mu = x[:, 0]\n",
    "  log_tau = x[:, 1]\n",
    "  eta = x[:, 2:10]\n",
    "  return model.log_prob((mu, log_tau, eta, y))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use results from running 128 chains with 1000 + 5000 iterations each,\n",
    "# for non-centered parameterization.\n",
    "mean_est = np.array([5.8006573 ,  2.4502006 ,  0.6532423 ,  0.09639207,\n",
    "             -0.23725411,  0.04723661, -0.33556408, -0.19666635,\n",
    "              0.5390533 ,  0.14633301])\n",
    "\n",
    "var_est = np.array([29.60382   ,  0.26338503,  0.6383733 ,  0.4928926 ,\n",
    "              0.65307987,  0.52441144,  0.46658015,  0.5248887 ,\n",
    "              0.49544162,  0.690975])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEED : [696157669 674520459]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<string>:212: DeprecationWarning: the `interpolation=` argument to quantile was renamed to `method=`, which has additional options.\n",
      "Users of the modes 'nearest', 'lower', 'higher', or 'midpoint' are encouraged to review the method they used. (Deprecated NumPy 1.22)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEED : [1771252691  383428323]\n",
      "SEED : [3312214627   69373237]\n",
      "SEED : [1094216230 2034336350]\n",
      "SEED : [1788847948 1554193616]\n",
      "SEED : [338758650 111605681]\n",
      "SEED : [4073466969 2556348314]\n",
      "SEED : [ 167784864 1021574613]\n",
      "SEED : [3458295564 1820549682]\n",
      "SEED : [3023805236 1296165206]\n"
     ]
    }
   ],
   "source": [
    "kernel_cold, kernel_warm = adaptive_kernels(target_log_prob_fn, init_step_size, num_warmup)\n",
    "index_param = np.arange(0, 10)\n",
    "\n",
    "mc_mean_list, warmup_length,\\\n",
    "squared_error_list, nrhat_list = run_forge_chain(num_seed = num_seed,\n",
    "                                          kernel_cold = kernel_cold,\n",
    "                                          kernel_warm = kernel_warm,\n",
    "                                          initialize = initialize,\n",
    "                                          num_super_chains = num_super_chains,\n",
    "                                          num_warmup = window_array,\n",
    "                                          num_samples = num_samples,\n",
    "                                          target_rhat = nRhat_upper,\n",
    "                                          max_num_steps = window_array.shape[0],\n",
    "                                          index_param = index_param,\n",
    "                                          mean_benchmark = mean_est,\n",
    "                                          var_benchmark = var_est,\n",
    "                                          naive_super_chains = naive_super_chains)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save output into npy files\n",
    "model_name = \"schools\"\n",
    "exp_parm = \"_K\" + str(num_super_chains) + \"_M\" + str(num_sub_chains) + \"_N\" + str(num_samples)\n",
    "\n",
    "if naive_super_chains:\n",
    "    exp_parm = \"_naive\" + exp_parm\n",
    "\n",
    "np.save(deliv_dir + model_name + exp_parm + \"_nrhat\", nrhat_list)\n",
    "np.save(deliv_dir + model_name + exp_parm + \"_squared_error\", squared_error_list)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Pharmacokinetic"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### load simulated data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"pk_model.py\") as f: exec(f.read())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_patients = 20\n",
    "y_obs = jnp.array(np.load(data_dir + \"pk_y_obs.npy\"))\n",
    "mean_est = jnp.array(np.load(data_dir + \"pk_npatients_\" + str(n_patients) + \"_mean_est.npy\"))\n",
    "var_est = jnp.array(np.load(data_dir + \"pk_npatients_\" + str(n_patients) + \"_var_est.npy\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 149,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEED : [696157669 674520459]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<string>:212: DeprecationWarning: the `interpolation=` argument to quantile was renamed to `method=`, which has additional options.\n",
      "Users of the modes 'nearest', 'lower', 'higher', or 'midpoint' are encouraged to review the method they used. (Deprecated NumPy 1.22)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEED : [1771252691  383428323]\n",
      "SEED : [3312214627   69373237]\n",
      "SEED : [1094216230 2034336350]\n",
      "SEED : [1788847948 1554193616]\n",
      "SEED : [338758650 111605681]\n",
      "SEED : [4073466969 2556348314]\n",
      "SEED : [ 167784864 1021574613]\n",
      "SEED : [3458295564 1820549682]\n",
      "SEED : [3023805236 1296165206]\n"
     ]
    }
   ],
   "source": [
    "init_step_size = 0.001\n",
    "kernel_cold, kernel_warm = adaptive_kernels(pop_target_log_prob_fn_flat,\n",
    "                                            init_step_size, num_warmup)\n",
    "index_param = np.arange(0, 45)\n",
    "\n",
    "mc_mean_list, warmup_length,\\\n",
    "squared_error_list, nrhat_list = run_forge_chain(num_seed = num_seed,\n",
    "                                          kernel_cold = kernel_cold,\n",
    "                                          kernel_warm = kernel_warm,\n",
    "                                          initialize = initialize_flat,\n",
    "                                          num_super_chains = num_super_chains,\n",
    "                                          num_warmup = window_array,\n",
    "                                          num_samples = num_samples,\n",
    "                                          target_rhat = nRhat_upper,\n",
    "                                          max_num_steps = window_array.shape[0],\n",
    "                                          index_param = index_param,\n",
    "                                          mean_benchmark = mean_est,\n",
    "                                          var_benchmark = var_est,\n",
    "                                          naive_super_chains = naive_super_chains)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save output into npy files\n",
    "model_name = \"pk\"\n",
    "exp_parm = \"_K\" + str(num_super_chains) + \"_M\" + str(num_sub_chains) + \"_N\" + str(num_samples)\n",
    "\n",
    "if naive_super_chains:\n",
    "    exp_parm = \"_naive\" + exp_parm\n",
    "\n",
    "np.save(deliv_dir + model_name + exp_parm + \"_nrhat\", nrhat_list)\n",
    "np.save(deliv_dir + model_name + exp_parm + \"_squared_error\", squared_error_list)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Item Response Theory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "metadata": {},
   "outputs": [],
   "source": [
    "target = gym.targets.VectorModel(gym.targets.SyntheticItemResponseTheory(),\n",
    "                                 flatten_sample_transformations=True)\n",
    "num_dimensions = target.event_shape[0]\n",
    "init_step_size = 1.\n",
    "\n",
    "def target_log_prob_fn(x):\n",
    "  \"\"\"Unnormalized, unconstrained target density.\n",
    "\n",
    "  This is a thin wrapper that applies the default bijectors so that we can\n",
    "  ignore any constraints.\n",
    "  \"\"\"\n",
    "  y = target.default_event_space_bijector(x)\n",
    "  fldj = target.default_event_space_bijector.forward_log_det_jacobian(x)\n",
    "  return target.unnormalized_log_prob(y) + fldj\n",
    "\n",
    "offset = 0\n",
    "def initialize (shape, key = random.PRNGKey(37272709)):\n",
    "  return 10 * random.normal(key, shape + (num_dimensions,)) + offset\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 152,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get some estimates of the mean and variance.\n",
    "try:\n",
    "  mean_est = target.sample_transformations['identity'].ground_truth_mean\n",
    "except:\n",
    "  print('no ground truth mean')\n",
    "  mean_est = (result.all_states[num_warmup:, :]).mean(0).mean(0)\n",
    "try:\n",
    "  var_est = target.sample_transformations['identity'].ground_truth_standard_deviation**2\n",
    "except:\n",
    "  print('no ground truth std dev')\n",
    "  var_est = ((result.all_states[num_warmup:, :]**2).mean(0).mean(0) -\n",
    "             mean_est**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 153,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEED : [696157669 674520459]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<string>:212: DeprecationWarning: the `interpolation=` argument to quantile was renamed to `method=`, which has additional options.\n",
      "Users of the modes 'nearest', 'lower', 'higher', or 'midpoint' are encouraged to review the method they used. (Deprecated NumPy 1.22)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEED : [1771252691  383428323]\n",
      "SEED : [3312214627   69373237]\n",
      "SEED : [1094216230 2034336350]\n",
      "SEED : [1788847948 1554193616]\n",
      "SEED : [338758650 111605681]\n",
      "SEED : [4073466969 2556348314]\n",
      "SEED : [ 167784864 1021574613]\n",
      "SEED : [3458295564 1820549682]\n",
      "SEED : [3023805236 1296165206]\n"
     ]
    }
   ],
   "source": [
    "kernel_cold, kernel_warm = adaptive_kernels(target_log_prob_fn, init_step_size, num_warmup)\n",
    "index_param = np.arange(0, 501)\n",
    "\n",
    "mc_mean_list, warmup_length,\\\n",
    "squared_error_list, nrhat_list = run_forge_chain(num_seed = num_seed,\n",
    "                                          kernel_cold = kernel_cold,\n",
    "                                          kernel_warm = kernel_warm,\n",
    "                                          initialize = initialize,\n",
    "                                          num_super_chains = num_super_chains,\n",
    "                                          num_warmup = window_array,\n",
    "                                          num_samples = num_samples,\n",
    "                                          target_rhat = nRhat_upper,\n",
    "                                          max_num_steps = window_array.shape[0],\n",
    "                                          index_param = index_param,\n",
    "                                          mean_benchmark = mean_est,\n",
    "                                          var_benchmark = var_est,\n",
    "                                          naive_super_chains = naive_super_chains)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save output into npy files\n",
    "model_name = \"itr\"\n",
    "exp_parm = \"_K\" + str(num_super_chains) + \"_M\" + str(num_sub_chains) + \"_N\" + str(num_samples)\n",
    "\n",
    "if naive_super_chains:\n",
    "    exp_parm = \"_naive\" + exp_parm\n",
    "\n",
    "np.save(deliv_dir + model_name + exp_parm + \"_nrhat\", nrhat_list)\n",
    "np.save(deliv_dir + model_name + exp_parm + \"_squared_error\", squared_error_list)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Bimodal Gaussian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_dimensions = 100\n",
    "dist = tfd.MixtureSameFamily(\n",
    "    mixture_distribution=tfd.Categorical(probs=[0.3, 0.7]), \n",
    "    components_distribution=tfd.MultivariateNormalDiag(\n",
    "      loc=[jnp.repeat(-5., num_dimensions), jnp.repeat(5., num_dimensions)],\n",
    "      scale_diag=jnp.repeat(1., num_dimensions)))\n",
    "\n",
    "def target_log_prob_fn(x):\n",
    "  return dist.log_prob(x)\n",
    "\n",
    "offset = 0\n",
    "def initialize (shape, key = random.PRNGKey(37272709)):\n",
    "  return 10 * random.normal(key, shape + (num_dimensions,)) + offset\n",
    "\n",
    "mean_est = jnp.repeat(2, num_dimensions)\n",
    "var_est = jnp.repeat(22, num_dimensions)\n",
    "\n",
    "init_step_size = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEED : [696157669 674520459]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<string>:212: DeprecationWarning: the `interpolation=` argument to quantile was renamed to `method=`, which has additional options.\n",
      "Users of the modes 'nearest', 'lower', 'higher', or 'midpoint' are encouraged to review the method they used. (Deprecated NumPy 1.22)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SEED : [1771252691  383428323]\n",
      "SEED : [3312214627   69373237]\n",
      "SEED : [1094216230 2034336350]\n",
      "SEED : [1788847948 1554193616]\n",
      "SEED : [338758650 111605681]\n",
      "SEED : [4073466969 2556348314]\n",
      "SEED : [ 167784864 1021574613]\n",
      "SEED : [3458295564 1820549682]\n",
      "SEED : [3023805236 1296165206]\n"
     ]
    }
   ],
   "source": [
    "kernel_cold, kernel_warm = adaptive_kernels(target_log_prob_fn, init_step_size, num_warmup)\n",
    "index_param = np.arange(0, num_dimensions)\n",
    "\n",
    "mc_mean_list, warmup_length,\\\n",
    "squared_error_list, nrhat_list = run_forge_chain(num_seed = num_seed,\n",
    "                                          kernel_cold = kernel_cold,\n",
    "                                          kernel_warm = kernel_warm,\n",
    "                                          initialize = initialize,\n",
    "                                          num_super_chains = num_super_chains,\n",
    "                                          num_warmup = window_array,\n",
    "                                          num_samples = num_samples,\n",
    "                                          target_rhat = nRhat_upper,\n",
    "                                          max_num_steps = window_array.shape[0],\n",
    "                                          index_param = index_param,\n",
    "                                          mean_benchmark = mean_est,\n",
    "                                          var_benchmark = var_est,\n",
    "                                          naive_super_chains = naive_super_chains)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31mThe Kernel crashed while executing code in the the current cell or a previous cell. Please review the code in the cell(s) to identify a possible cause of the failure. Click <a href='https://aka.ms/vscodeJupyterKernelCrash'>here</a> for more info. View Jupyter <a href='command:jupyter.viewOutput'>log</a> for further details."
     ]
    }
   ],
   "source": [
    "model_name = \"bimodal\"\n",
    "exp_parm = \"_K\" + str(num_super_chains) + \"_M\" + str(num_sub_chains) + \"_N\" + str(num_samples)\n",
    "\n",
    "if naive_super_chains:\n",
    "    exp_parm = \"_naive\" + exp_parm\n",
    "\n",
    "np.save(deliv_dir + model_name + exp_parm + \"_nrhat\", nrhat_list)\n",
    "np.save(deliv_dir + model_name + exp_parm + \"_squared_error\", squared_error_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}