craft_manywell_train_eval.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ace069b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from annealed_flow_transport.train import run_experiment\n",
    "import numpy as np\n",
    "import jax.numpy as jnp\n",
    "import jax\n",
    "import matplotlib.pyplot as plt\n",
    "from annealed_flow_transport.many_well_plotting import plot, plot_marginal_pair\n",
    "from annealed_flow_transport.densities import ManyWell\n",
    "from annealed_flow_transport.resampling import log_effective_sample_size\n",
    "\n",
    "import matplotlib as mpl\n",
    "\n",
    "mpl.rcParams['figure.dpi'] = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24899873",
   "metadata": {},
   "outputs": [],
   "source": [
    "from configs.many_well_original import get_config as many_well_original_get_config\n",
    "from configs.many_well import get_config as many_well_get_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd24956f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_experiment_config(seed, experiment_type = \"original\", long=True):\n",
    "    # Set long to False to run just a short version of the experiment.\n",
    "    if experiment_type == \"original\":\n",
    "        exp_config = many_well_original_get_config()\n",
    "        exp_config.n_samples_plotting = 100\n",
    "    else:\n",
    "        assert experiment_type == \"custom\"\n",
    "        exp_config = many_well_get_config()\n",
    "    if not long: # Run for short\n",
    "        exp_config.craft_num_iters = 200\n",
    "        exp_config.report_step = int(exp_config.craft_num_iters // 6)\n",
    "    exp_config.seed = seed\n",
    "    exp_config.params_filename = f\"checkpoint_blob_{experiment_type}_L{long}_seed{exp_config.seed}\"\n",
    "    \n",
    "    n_mill_target_eval = exp_config.craft_batch_size*exp_config.craft_num_iters*exp_config.num_temps*exp_config.mcmc_config.hmc_num_leapfrog_steps/1e6\n",
    "    print(f\"experiment using {n_mill_target_eval} target evaluations\")\n",
    "    return exp_config"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73c7d4b8",
   "metadata": {},
   "source": [
    "# Train using the config from the original CRAFT repo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be78266f",
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_config = get_experiment_config(0, \"original\")\n",
    "results0 = run_experiment(exp_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9705176b",
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_config = get_experiment_config(1, \"original\")\n",
    "results1 = run_experiment(exp_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c8e185b",
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_config = get_experiment_config(2, \"original\")\n",
    "results2 = run_experiment(exp_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c64e59d3",
   "metadata": {},
   "source": [
    "# Train using config similar to FAB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1908a8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_config = get_experiment_config(0, \"custom\")\n",
    "results0 = run_experiment(exp_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6912fd1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_config = get_experiment_config(1, \"custom\")\n",
    "results1 = run_experiment(exp_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2367ab3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "exp_config = get_experiment_config(2, \"custom\")\n",
    "results2 = run_experiment(exp_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f1b7cff",
   "metadata": {},
   "source": [
    "# Evaluation\n",
    "Load checkpoints and evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89b0399b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from evaluation import evaluate_many_well, load_checkpoint, make_forward_pass_func, \\\n",
    "    make_get_ess, make_get_resample_info, get_flow_init_params\n",
    "import scipy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e4eb060",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_density = ManyWell(exp_config.final_config.density, (32,))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5abb074",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_batch_size = 1000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fcf5406",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate_Z_estimation(seed, experiment_type = \"original\", long=True, \n",
    "                          eval_batch_size=1000, n_eval=50, flow_identity=False):\n",
    "    exp_config = get_experiment_config(seed, experiment_type, long)\n",
    "    if flow_identity:\n",
    "        transition_params = get_flow_init_params(exp_config)\n",
    "    else:\n",
    "        transition_params = load_checkpoint(exp_config.params_filename)\n",
    "    forward_pass_function = make_forward_pass_func(exp_config, transition_params=transition_params,\n",
    "                                                  eval_batch_size=eval_batch_size)\n",
    "    key = jax.random.PRNGKey(0)\n",
    "    abs_errors = []\n",
    "    for i in range(n_eval):\n",
    "        key, subkey = jax.random.split(key)\n",
    "        particle_state = forward_pass_function(subkey)\n",
    "        log_Z_estimate = particle_state.log_normalizer_estimate\n",
    "        relative_error = jnp.exp(log_Z_estimate - target_density.log_Z) - 1\n",
    "        abs_error = jnp.abs(relative_error)\n",
    "        abs_errors.append(abs_error)\n",
    "    return jnp.stack(abs_errors)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eafdf44c",
   "metadata": {},
   "source": [
    "#### Original"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9b45640",
   "metadata": {},
   "outputs": [],
   "source": [
    "# standard\n",
    "abs_errors0 = evaluate_Z_estimation(0, \"original\")\n",
    "jnp.mean(abs_errors0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "413cd069",
   "metadata": {},
   "outputs": [],
   "source": [
    "# standard\n",
    "abs_errors1 = evaluate_Z_estimation(1, \"original\")\n",
    "jnp.mean(abs_errors1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff8948ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "abs_errors2 = evaluate_Z_estimation(2, \"original\")\n",
    "jnp.mean(abs_errors2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45f7aeb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# for table\n",
    "vals = [float(jnp.mean(abs_errors0)), float(jnp.mean(abs_errors1)), jnp.mean(abs_errors2)]\n",
    "sum(vals) / 3, scipy.stats.sem(vals, ddof=0) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bd13774",
   "metadata": {},
   "source": [
    "#### Similar to FAB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d818248",
   "metadata": {},
   "outputs": [],
   "source": [
    "abs_errors0 = evaluate_Z_estimation(0, \"custom\")\n",
    "jnp.mean(abs_errors0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa7bc499",
   "metadata": {},
   "outputs": [],
   "source": [
    "abs_errors1 = evaluate_Z_estimation(1, \"custom\")\n",
    "jnp.mean(abs_errors1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34bb97fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "abs_errors2 = evaluate_Z_estimation(2, \"custom\")\n",
    "jnp.mean(abs_errors2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7c40666",
   "metadata": {},
   "outputs": [],
   "source": [
    "# for table\n",
    "vals = [float(jnp.mean(abs_errors0)), float(jnp.mean(abs_errors1)), jnp.mean(abs_errors2)]\n",
    "sum(vals) / 3, scipy.stats.sem(vals, ddof=0) "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}