experiments.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Code for Reproducing Experimental Results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import libraries and functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pickle as pkl\n",
    "\n",
    "from utils import mae_quant_level\n",
    "from utils import mae_tail_synthetic, mae_varyk_real, mae_varyk_synthetic, mae_varyn_real, mae_varyn_synthetic, quant_synthetic\n",
    "from utils import supe_tail_synthetic, supe_varyn_real, supe_varyn_synthetic, supe_varyk_real, supe_varyk_synthetic\n",
    "from utils import plot_dist_est, plot_quant_error, plot_quant_level, plot_stat_bound"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prefix = 'results/mae'\n",
    "prefix_df = 'results/supe'\n",
    "\n",
    "os.system(f'mkdir -p {prefix}')\n",
    "os.system(f'mkdir -p {prefix_df}')\n",
    "os.system(f'mkdir -p results/real')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Synthetic data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Statistical error bound for frontier integral"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first vary the sample size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "supp = 1000\n",
    "nrange = np.logspace(4, 5, 12).astype(int)\n",
    "dist_pairs = [\n",
    "    [('zipf', 2), ('uniform', '')], [('zipf', 2), ('zipf', 2)],\n",
    "    [('zipf', 0), ('zipf', 0)], [('uniform', ''), ('dirichlet', '')]]\n",
    "b_varyn = mae_varyn_synthetic(supp, nrange, dist_pairs, prefix=prefix, save=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 9 (top) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'nvary-zipf2-uniform-supp1000',\n",
    "    'nvary-zipf2-zipf2-supp1000',\n",
    "    'nvary-zipf0-zipf0-supp1000',\n",
    "    'nvary-uniform-dirichlet-supp1000']\n",
    "b_varyn = [pd.read_pickle(f'{prefix}/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Sample size']*4\n",
    "titles = ['(a)', '(b)', '(c)', '(d)']\n",
    "fname = 'graphs/synthetic-bound-nvary.pdf'\n",
    "plot_stat_bound(b_varyn, xlabels, titles, const=50, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then vary the support size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 20000\n",
    "krange = np.logspace(1, 4, 12).astype(int)\n",
    "dist_pairs = [\n",
    "    [('zipf', 2), ('uniform', '')], [('zipf', 2), ('zipf', 2)],\n",
    "    [('zipf', 0), ('zipf', 0)], [('uniform', ''), ('dirichlet', '')]]\n",
    "b_varyk = mae_varyk_synthetic(krange, n, dist_pairs, prefix=prefix, save=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 10 (top) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'kvary-zipf2-uniform-size20000',\n",
    "    'kvary-zipf2-zipf2-size20000',\n",
    "    'kvary-zipf0-zipf0-size20000',\n",
    "    'kvary-uniform-dirichlet-size20000']\n",
    "b_varyk = [pd.read_pickle(f'{prefix}/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Support size']*4\n",
    "titles = ['(a)', '(b)', '(c)', '(d)']\n",
    "fname = 'graphs/synthetic-bound-kvary.pdf'\n",
    "plot_stat_bound(b_varyk, xlabels, titles, const=50, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we vary the tail decay index of the distribution $Q$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 10000\n",
    "supp = 1000\n",
    "orderq = np.linspace(0.5, 2, 12)\n",
    "dists = [('uniform', ''), ('dirichlet', ''), ('zipf', 1), ('zipf', 2)]\n",
    "b_tail = mae_tail_synthetic(supp, n, orderq, dists, prefix=prefix, save=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 11 (top) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'qvary-uniform-supp1000-size10000',\n",
    "    'qvary-dirichlet-supp1000-size10000',\n",
    "    'qvary-zipf1-supp1000-size10000',\n",
    "    'qvary-zipf2-supp1000-size10000']\n",
    "b_tail = [pd.read_pickle(f'{prefix}/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Tail decay']*4\n",
    "titles = ['(a)', '(b)', '(c)', '(d)']\n",
    "fname = 'graphs/synthetic-bound-qvary.pdf'\n",
    "plot_stat_bound(b_tail, xlabels, titles, const=100, fname=fname, save=False, log_scale=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 4 in the main text can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "const=100\n",
    "fnames = [\n",
    "    'nvary-zipf2-zipf2-supp1000',\n",
    "    'kvary-zipf2-zipf2-size20000',\n",
    "    'qvary-uniform-supp1000-size10000',\n",
    "    'qvary-zipf2-supp1000-size10000']\n",
    "dfs = [pd.read_pickle(f'{prefix}/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Sample size', 'Support size', 'Tail decay', 'Tail decay']\n",
    "titles = [\n",
    "    r'(a) $k=10^3$',\n",
    "    r'(b) $n=2\\times 10^4$',\n",
    "    r'(c) $k=10^3, n=10^4$',\n",
    "    r'(d) $k=10^3, n=10^4$']\n",
    "fname = 'graphs/synthetic-bound.pdf'\n",
    "plot_stat_bound(dfs, xlabels, titles, const=const, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Statistical error bound for divergence frontiers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first vary the sample size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "supp = 1000\n",
    "nrange = np.logspace(4, 5, 12).astype(int)\n",
    "dist_pairs = [\n",
    "    [('zipf', 2), ('uniform', '')], [('zipf', 2), ('zipf', 2)],\n",
    "    [('zipf', 0), ('zipf', 0)], [('uniform', ''), ('dirichlet', '')]]\n",
    "lambdas = np.linspace(0.01, 0.99, 100)\n",
    "b_varyn_df = supe_varyn_synthetic(supp, nrange, dist_pairs, lambdas, nrepeat=100, prefix=prefix_df, save=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 9 (bottom) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'nvary-zipf2-uniform-supp1000-df',\n",
    "    'nvary-zipf2-zipf2-supp1000-df',\n",
    "    'nvary-zipf0-zipf0-supp1000-df',\n",
    "    'nvary-uniform-dirichlet-supp1000-df']\n",
    "b_varyn_df = [pd.read_pickle(f'{prefix_df}/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Sample size']*4\n",
    "titles = ['(a)', '(b)', '(c)', '(d)']\n",
    "fname = 'graphs/synthetic-bound-nvary-df.pdf'\n",
    "plot_stat_bound(b_varyn_df, xlabels, titles, const=10, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then vary the support size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 20000\n",
    "krange = np.logspace(1, 4, 12).astype(int)\n",
    "dist_pairs = [\n",
    "    [('zipf', 2), ('uniform', '')], [('zipf', 2), ('zipf', 2)],\n",
    "    [('zipf', 0), ('zipf', 0)], [('uniform', ''), ('dirichlet', '')]]\n",
    "lambdas = np.linspace(0.01, 0.99, 100)\n",
    "b_varyk_df = supe_varyk_synthetic(krange, n, dist_pairs, lambdas, nrepeat=100, prefix=prefix_df, save=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 10 (bottom) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'kvary-zipf2-uniform-size20000-df',\n",
    "    'kvary-zipf2-zipf2-size20000-df',\n",
    "    'kvary-zipf0-zipf0-size20000-df',\n",
    "    'kvary-uniform-dirichlet-size20000-df']\n",
    "b_varyk_df = [pd.read_pickle(f'{prefix_df}/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Support size']*4\n",
    "titles = ['(a)', '(b)', '(c)', '(d)']\n",
    "fname = 'graphs/synthetic-bound-kvary-df.pdf'\n",
    "plot_stat_bound(b_varyk_df, xlabels, titles, const=10, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we vary the tail decay index of the distribution $Q$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 10000\n",
    "supp = 1000\n",
    "orderq = np.linspace(0.5, 2, 12)\n",
    "dists = [('uniform', ''), ('dirichlet', ''), ('zipf', 1), ('zipf', 2)]\n",
    "lambdas = np.linspace(0.01, 0.99, 100)\n",
    "b_tail_df = supe_tail_synthetic(supp, n, orderq, dists, lambdas, nrepeat=100, prefix=prefix_df, save=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 11 (bottom) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'qvary-uniform-supp1000-size10000-df',\n",
    "    'qvary-dirichlet-supp1000-size10000-df',\n",
    "    'qvary-zipf1-supp1000-size10000-df',\n",
    "    'qvary-zipf2-supp1000-size10000-df']\n",
    "b_tail_df = [pd.read_pickle(f'{prefix_df}/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Tail decay']*4\n",
    "titles = ['(a)', '(b)', '(c)', '(d)']\n",
    "fname = 'graphs/synthetic-bound-qvary-df.pdf'\n",
    "plot_stat_bound(b_tail_df, xlabels, titles, const=7, fname=fname, save=False, log_scale=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Distribution estimators for frontier integral"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first vary the sample size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "supp = 1000\n",
    "nrange = np.logspace(4, 5, 12).astype(int)\n",
    "dist_pairs = [\n",
    "    [('zipf', 1), ('step', '')], [('zipf', 0), ('dirichlet', '')],\n",
    "    [('zipf', 2), ('uniform', '')], [('zipf', 1), ('zipf', 1)]]\n",
    "est_varyn = mae_varyn_synthetic(supp, nrange, dist_pairs, prefix=prefix, save=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 12 (top) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'nvary-zipf1-step-supp1000',\n",
    "    'nvary-zipf0-dirichlet-supp1000',\n",
    "    'nvary-zipf2-uniform-supp1000',\n",
    "    'nvary-zipf1-zipf1-supp1000']\n",
    "est_varyn = [pd.read_pickle(f'{prefix}/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Sample size']*4\n",
    "titles = ['(a)', '(b)', '(c)', '(d)']\n",
    "fname = 'graphs/synthetic-smoothing-nvary.pdf'\n",
    "plot_dist_est(est_varyn, xlabels, titles, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then vary the support size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 20000\n",
    "krange = np.logspace(1, 4, 12).astype(int)\n",
    "dist_pairs = [\n",
    "    [('zipf', 1), ('step', '')], [('zipf', 0), ('dirichlet', '')],\n",
    "    [('zipf', 2), ('uniform', '')], [('zipf', 1), ('zipf', 1)]]\n",
    "est_varyk = mae_varyk_synthetic(krange, n, dist_pairs, prefix=prefix, save=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 13 (top) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'kvary-zipf1-step-size20000',\n",
    "    'kvary-zipf0-dirichlet-size20000',\n",
    "    'kvary-zipf2-uniform-size20000',\n",
    "    'kvary-zipf1-zipf1-size20000']\n",
    "est_varyk = [pd.read_pickle(f'{prefix}/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Support size']*4\n",
    "titles = ['(a)', '(b)', '(c)', '(d)']\n",
    "fname = 'graphs/synthetic-smoothing-kvary.pdf'\n",
    "plot_dist_est(est_varyk, xlabels, titles, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we vary the tail decay index of the distribution $Q$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 10000\n",
    "supp = 1000\n",
    "orderq = np.linspace(0.5, 2, 12)\n",
    "dists = [('uniform', ''), ('step', ''), ('zipf', 0), ('zipf', 2)]\n",
    "est_tail = mae_tail_synthetic(supp, n, orderq, dists, prefix=prefix, save=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 14 (top) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'qvary-uniform-supp1000-size10000',\n",
    "    'qvary-step-supp1000-size10000',\n",
    "    'qvary-zipf0-supp1000-size10000',\n",
    "    'qvary-zipf2-supp1000-size10000']\n",
    "est_tail = [pd.read_pickle(f'{prefix}/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Tail decay']*4\n",
    "titles = ['(a)', '(b)', '(c)', '(d)']\n",
    "fname = 'graphs/synthetic-smoothing-qvary.pdf'\n",
    "plot_dist_est(est_tail, xlabels, titles, fname=fname, save=False, log_scale=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 6 in the main text can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'nvary-zipf0-dirichlet-supp1000',\n",
    "    'kvary-zipf0-dirichlet-size20000',\n",
    "    'qvary-uniform-supp1000-size10000',\n",
    "    'qvary-zipf2-supp1000-size10000']\n",
    "dfs = [pd.read_pickle(f'{prefix}/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Sample size', 'Support size', 'Tail decay', 'Tail decay']\n",
    "titles = [\n",
    "    r'(a) $k=10^3$',\n",
    "    r'(b) $n=2\\times 10^4$',\n",
    "    r'(c) $k=10^3, n=10^4$',\n",
    "    r'(d) $k=10^3, n=10^4$']\n",
    "fname = 'graphs/synthetic-smoothing.pdf'\n",
    "plot_dist_est(dfs, xlabels, titles, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Distribution estimators for divergence frontiers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first vary the sample size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "supp = 1000\n",
    "nrange = np.logspace(4, 5, 12).astype(int)\n",
    "dist_pairs = [\n",
    "    [('zipf', 1), ('step', '')], [('zipf', 0), ('dirichlet', '')],\n",
    "    [('zipf', 2), ('uniform', '')], [('zipf', 1), ('zipf', 1)]]\n",
    "lambdas = np.linspace(0.01, 0.99, 100)\n",
    "est_varyn_df = supe_varyn_synthetic(supp, nrange, dist_pairs, lambdas, prefix=prefix_df, save=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 12 (bottom) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'nvary-zipf1-step-supp1000-df',\n",
    "    'nvary-zipf0-dirichlet-supp1000-df',\n",
    "    'nvary-zipf2-uniform-supp1000-df',\n",
    "    'nvary-zipf1-zipf1-supp1000-df']\n",
    "est_varyn_df = [pd.read_pickle(f'{prefix_df}/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Sample size']*4\n",
    "titles = ['(a)', '(b)', '(c)', '(d)']\n",
    "fname = 'graphs/synthetic-smoothing-nvary-df.pdf'\n",
    "plot_dist_est(est_varyn_df, xlabels, titles, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then vary the support size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 20000\n",
    "krange = np.logspace(1, 4, 12).astype(int)\n",
    "dist_pairs = [\n",
    "    [('zipf', 1), ('step', '')], [('zipf', 0), ('dirichlet', '')],\n",
    "    [('zipf', 2), ('uniform', '')], [('zipf', 1), ('zipf', 1)]]\n",
    "lambdas = np.linspace(0.01, 0.99, 100)\n",
    "est_varyk_df = supe_varyk_synthetic(krange, n, dist_pairs, lambdas, prefix=prefix_df, save=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 13 (bottom) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'kvary-zipf1-step-size20000-df',\n",
    "    'kvary-zipf0-dirichlet-size20000-df',\n",
    "    'kvary-zipf2-uniform-size20000-df',\n",
    "    'kvary-zipf1-zipf1-size20000-df']\n",
    "est_varyk_df = [pd.read_pickle(f'{prefix_df}/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Support size']*4\n",
    "titles = ['(a)', '(b)', '(c)', '(d)']\n",
    "fname = 'graphs/synthetic-smoothing-kvary-df.pdf'\n",
    "plot_dist_est(est_varyk_df, xlabels, titles, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we vary the tail decay index of the distribution $Q$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 10000\n",
    "supp = 1000\n",
    "orderq = np.linspace(0.5, 2, 12)\n",
    "dists = [('uniform', ''), ('step', ''), ('zipf', 0), ('zipf', 2)]\n",
    "lambdas = np.linspace(0.01, 0.99, 100)\n",
    "est_tail_df = supe_tail_synthetic(supp, n, orderq, dists, lambdas, prefix=prefix_df, save=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 14 (bottom) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'qvary-uniform-supp1000-size10000-df',\n",
    "    'qvary-step-supp1000-size10000-df',\n",
    "    'qvary-zipf0-supp1000-size10000-df',\n",
    "    'qvary-zipf2-supp1000-size10000-df']\n",
    "est_tail_df = [pd.read_pickle(f'{prefix_df}/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Tail decay']*4\n",
    "titles = ['(a)', '(b)', '(c)', '(d)']\n",
    "fname = 'graphs/synthetic-smoothing-qvary-df.pdf'\n",
    "plot_dist_est(est_tail_df, xlabels, titles, fname=fname, save=False, log_scale=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Quantization error"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "supp = 1000\n",
    "krange = np.logspace(1, 2.7, 12).astype(int)\n",
    "dist_pairs = [\n",
    "    [('uniform', ''), ('dirichlet', '')], [('zipf', 0), ('dirichlet', '')],\n",
    "    [('zipf', 2), ('step', '')], [('zipf', 1), ('zipf', 2)]]\n",
    "quant = quant_synthetic(supp, krange, dist_pairs, prefix=prefix, save=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 15 in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'quant-uniform-dirichlet-supp1000',\n",
    "    'quant-zipf0-dirichlet-supp1000',\n",
    "    'quant-zipf2-step-supp1000',\n",
    "    'quant-zipf1-zipf2-supp1000']\n",
    "quant = [pd.read_pickle(f'{prefix}/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Number of bins']*4\n",
    "titles = ['(a)', '(b)', '(c)', '(d)']\n",
    "fname = 'graphs/synthetic-quant.pdf'\n",
    "plot_quant_error(quant, xlabels, titles, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Quantization level"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nrange = np.logspace(2, 5, 10).astype(int)\n",
    "dim = 2\n",
    "dists = ['normal', 'normal', 't', 't']\n",
    "pars = [[(np.zeros(dim), np.identity(dim)),\n",
    "         (np.ones(dim), np.identity(dim))],\n",
    "        [(np.zeros(dim), np.identity(dim)),\n",
    "         (np.zeros(dim), 5*np.identity(dim))],\n",
    "        [(np.zeros(dim), np.identity(dim), 4),\n",
    "         (np.ones(dim), 5*np.identity(dim), 4)],\n",
    "        [(np.zeros(dim), np.identity(dim), 4),\n",
    "         (np.zeros(dim), 5*np.identity(dim), 4)]]\n",
    "true_fis = [0.2769, 0.3008, 0.2094, 0.3376]\n",
    "fnames = ['est-fi-normal-mean-dim2-nrates4-const5-kmeans',\n",
    "          'est-fi-t-mean-dim2-nrates4-const10-kmeans',\n",
    "          'est-fi-t-mean-dim2-nrates4-const5-kmeans',\n",
    "          'est-fi-t-var-dim2-nrates4-const10-kmeans']\n",
    "\n",
    "dfs = mae_quant_level(nrange, dists, pars, true_fis, nrepeat=10,\n",
    "                      prefix='results/mae', fnames=fnames, save=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 8 in the main text can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nrates = 4\n",
    "fnames = ['est-fi-normal-mean-dim2-nrates4-const5-kmeans',\n",
    "          'est-fi-normal-var-dim2-nrates4-const10-kmeans',\n",
    "          'est-fi-t-mean-dim2-nrates4-const5-kmeans',\n",
    "          'est-fi-t-var-dim2-nrates4-const10-kmeans']\n",
    "dfs = [np.loadtxt(f'results/mae/{file}.txt') for file in fnames]\n",
    "\n",
    "titles = ['(a)', '(b)', '(c)', '(d)']\n",
    "fname = 'graphs/synthetic-cont-quant-level.pdf'\n",
    "plot_quant_level(dfs, nrates, titles, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Real data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# download data\n",
    "os.system('wget https://www.stat.washington.edu/~liu16/divergence-frontier-bounds/parsed_outputs.zip')\n",
    "os.system('unzip parsed_outputs.zip')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Statistical error bound for frontier integral"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first vary the sample size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nrange = np.logspace(2, 4.4, 12).astype(int)\n",
    "nrepeat = 100\n",
    "\n",
    "# load vision data \n",
    "outs = pkl.load(open('parsed_outputs/cifar10_lattice.p', 'rb'))\n",
    "\n",
    "# 128 bins and 1024 bins\n",
    "o = outs[100000][128]\n",
    "p = o.p_hist\n",
    "q = o.q_hist\n",
    "print(p.shape, q.shape, np.linalg.norm(p), np.linalg.norm(q))\n",
    "varyn_vision_1 = mae_varyn_real(p, q, nrange, nrepeat)\n",
    "varyn_vision_1.to_pickle('results/real/nvary-vision-supp128.pkl')\n",
    "\n",
    "# end of training\n",
    "o = outs[100000][1024]\n",
    "p = o.p_hist\n",
    "q = o.q_hist\n",
    "print(p.shape, q.shape, np.linalg.norm(p), np.linalg.norm(q))\n",
    "varyn_vision_2 = mae_varyn_real(p, q, nrange, nrepeat)\n",
    "varyn_vision_2.to_pickle('results/real/nvary-vision-supp1024.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load nlp data \n",
    "outs = pkl.load(open('parsed_outputs/wikitext103.p', 'rb'))\n",
    "\n",
    "# start of training\n",
    "o = outs[2000][64]\n",
    "p = o.p_hist\n",
    "q = o.q_hist\n",
    "print(p.shape, q.shape, np.linalg.norm(p), np.linalg.norm(q))\n",
    "varyn_nlp_1 = mae_varyn_real(p, q, nrange, nrepeat)\n",
    "varyn_nlp_1.to_pickle('results/real/nvary-nlp-supp64.pkl')\n",
    "\n",
    "# end of training\n",
    "o = outs[2000][2048]\n",
    "p = o.p_hist\n",
    "q = o.q_hist\n",
    "print(p.shape, q.shape, np.linalg.norm(p), np.linalg.norm(q))\n",
    "varyn_nlp_2 = mae_varyn_real(p, q, nrange, nrepeat)\n",
    "varyn_nlp_2.to_pickle('results/real/nvary-nlp-supp2048.pkl')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 16 (top) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'nvary-vision-supp128',\n",
    "    'nvary-vision-supp1024',\n",
    "    'nvary-nlp-supp64',\n",
    "    'nvary-nlp-supp2048']\n",
    "dfs = [pd.read_pickle(f'results/real/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Sample size']*4\n",
    "titles = [r'(a) Images ($k=128$)', r'(b) Images ($k=1024$)',\n",
    "          r'(c) Text ($k=64$)', r'(d) Text ($k=2048$)']\n",
    "# dfs = [varyn_vision_1, varyn_vision_2, varyn_nlp_1, varyn_nlp_2]\n",
    "\n",
    "fname = 'graphs/real-bound-n.pdf'\n",
    "plot_stat_bound(dfs, xlabels, titles, const=15, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then vary the support size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load vision data \n",
    "outs = pkl.load(open('parsed_outputs/cifar10_lattice.p', 'rb'))\n",
    "\n",
    "o = outs[100000]\n",
    "varyk_vision_1 = mae_varyk_real(o, n=1000)\n",
    "varyk_vision_2 = mae_varyk_real(o, n=10000)\n",
    "varyk_vision_1.to_pickle('results/real/kvary-vision-size1000.pkl')\n",
    "varyk_vision_2.to_pickle('results/real/kvary-vision-size10000.pkl')\n",
    "\n",
    "# load text data\n",
    "outs = pkl.load(open('parsed_outputs/wikitext103.p', 'rb'))\n",
    "\n",
    "o = outs[2000]\n",
    "varyk_nlp_1 = mae_varyk_real(o, n=1000)\n",
    "varyk_nlp_2 = mae_varyk_real(o, n=10000)\n",
    "varyk_nlp_1.to_pickle('results/real/kvary-nlp-size1000.pkl')\n",
    "varyk_nlp_2.to_pickle('results/real/kvary-nlp-size10000.pkl')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 17 (top) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'kvary-vision-size1000',\n",
    "    'kvary-vision-size10000',\n",
    "    'kvary-nlp-size1000',\n",
    "    'kvary-nlp-size10000']\n",
    "dfs = [pd.read_pickle(f'results/real/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Support size']*4\n",
    "titles = [r'(a) Images ($n=1000$)', r'(b) Images ($n=10000$)',\n",
    "          r'(c) Text ($n=1000$)', r'(d) Text ($n=10000$)']\n",
    "# dfs = [varyk_vision_1, varyk_vision_2, varyk_nlp_1, varyk_nlp_2]\n",
    "\n",
    "fname = 'graphs/real-bound-k.pdf'\n",
    "plot_stat_bound(dfs, xlabels, titles, const=15, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Statistical error bound for divergence frontiers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first vary the sample size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nrange = np.logspace(2, 4.4, 12).astype(int)\n",
    "lambdas = np.linspace(0.01, 0.99, 100)\n",
    "nrepeat = 100\n",
    "\n",
    "# load vision data \n",
    "outs = pkl.load(open('parsed_outputs/cifar10_lattice.p', 'rb'))\n",
    "\n",
    "# 128 bins and 1024 bins\n",
    "o = outs[100000][128]\n",
    "p = o.p_hist\n",
    "q = o.q_hist\n",
    "print(p.shape, q.shape, np.linalg.norm(p), np.linalg.norm(q))\n",
    "varyn_vision_1 = supe_varyn_real(p, q, nrange, lambdas, nrepeat)\n",
    "varyn_vision_1.to_pickle('results/real/nvary-vision-supp128-df.pkl')\n",
    "\n",
    "# end of training\n",
    "o = outs[100000][1024]\n",
    "p = o.p_hist\n",
    "q = o.q_hist\n",
    "print(p.shape, q.shape, np.linalg.norm(p), np.linalg.norm(q))\n",
    "varyn_vision_2 = supe_varyn_real(p, q, nrange, lambdas, nrepeat)\n",
    "varyn_vision_2.to_pickle('results/real/nvary-vision-supp1024-df.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load nlp data\n",
    "outs = pkl.load(open('parsed_outputs/wikitext103.p', 'rb'))\n",
    "\n",
    "# start of training\n",
    "o = outs[2000][64]\n",
    "p = o.p_hist\n",
    "q = o.q_hist\n",
    "print(p.shape, q.shape, np.linalg.norm(p), np.linalg.norm(q))\n",
    "varyn_nlp_1 = supe_varyn_real(p, q, nrange, lambdas, nrepeat)\n",
    "varyn_nlp_1.to_pickle('results/real/nvary-nlp-supp64-df.pkl')\n",
    "\n",
    "# end of training\n",
    "o = outs[2000][2048]\n",
    "p = o.p_hist\n",
    "q = o.q_hist\n",
    "print(p.shape, q.shape, np.linalg.norm(p), np.linalg.norm(q))\n",
    "varyn_nlp_2 = supe_varyn_real(p, q, nrange, lambdas, nrepeat)\n",
    "varyn_nlp_2.to_pickle('results/real/nvary-nlp-supp2048-df.pkl')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 16 (bottom) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'nvary-vision-supp128-df',\n",
    "    'nvary-vision-supp1024-df',\n",
    "    'nvary-nlp-supp64-df',\n",
    "    'nvary-nlp-supp2048-df']\n",
    "dfs = [pd.read_pickle(f'results/real/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Sample size']*4\n",
    "titles = [r'(a) Images ($k=128$)', r'(b) Images ($k=1024$)',\n",
    "          r'(c) Text ($k=64$)', r'(d) Text ($k=2048$)']\n",
    "fname = 'graphs/real-bound-n-df.pdf'\n",
    "plot_stat_bound(dfs, xlabels, titles, const=7, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then vary the support size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lambdas = np.linspace(0.01, 0.99, 100)\n",
    "nrepeat = 100\n",
    "\n",
    "# load vision data\n",
    "outs = pkl.load(open('parsed_outputs/cifar10_lattice.p', 'rb'))\n",
    "\n",
    "o = outs[100000]\n",
    "varyk_vision_1 = supe_varyk_real(o, 1000, lambdas, nrepeat)\n",
    "varyk_vision_2 = supe_varyk_real(o, 10000, lambdas, nrepeat)\n",
    "varyk_vision_1.to_pickle('results/real/kvary-vision-size1000-df.pkl')\n",
    "varyk_vision_2.to_pickle('results/real/kvary-vision-size10000-df.pkl')\n",
    "\n",
    "# load text data\n",
    "outs = pkl.load(open('parsed_outputs/wikitext103.p', 'rb'))\n",
    "\n",
    "o = outs[2000]\n",
    "varyk_nlp_1 = supe_varyk_real(o, 1000, lambdas, nrepeat)\n",
    "varyk_nlp_2 = supe_varyk_real(o, 10000, lambdas, nrepeat)\n",
    "varyk_nlp_1.to_pickle('results/real/kvary-nlp-size1000-df.pkl')\n",
    "varyk_nlp_2.to_pickle('results/real/kvary-nlp-size10000-df.pkl')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 17 (bottom) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'kvary-vision-size1000-df',\n",
    "    'kvary-vision-size10000-df',\n",
    "    'kvary-nlp-size1000-df',\n",
    "    'kvary-nlp-size10000-df']\n",
    "dfs = [pd.read_pickle(f'results/real/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Support size']*4\n",
    "titles = [r'(a) Images ($n=1000$)', r'(b) Images ($n=10000$)',\n",
    "          r'(c) Text ($n=1000$)', r'(d) Text ($n=10000$)']\n",
    "fname = 'graphs/real-bound-k-df.pdf'\n",
    "plot_stat_bound(dfs, xlabels, titles, const=7, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Distribution estimators for frontier integral"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first vary the sample size."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 18 (top) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'nvary-vision-supp128',\n",
    "    'nvary-vision-supp1024',\n",
    "    'nvary-nlp-supp64',\n",
    "    'nvary-nlp-supp2048']\n",
    "dfs = [pd.read_pickle(f'results/real/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Sample size']*4\n",
    "titles = [r'(a) Images ($k=128$)', r'(b) Images ($k=1024$)',\n",
    "          r'(c) Text ($k=64$)', r'(d) Text ($k=2048$)']\n",
    "fname = 'graphs/real-smoothing-n.pdf'\n",
    "plot_dist_est(dfs, xlabels, titles, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then vary the support size."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 19 (top) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'kvary-vision-size1000',\n",
    "    'kvary-vision-size10000',\n",
    "    'kvary-nlp-size1000',\n",
    "    'kvary-nlp-size10000']\n",
    "dfs = [pd.read_pickle(f'results/real/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Support size']*4\n",
    "titles = [r'(a) Images ($n=1000$)', r'(b) Images ($n=10000$)',\n",
    "          r'(c) Text ($n=1000$)', r'(d) Text ($n=10000$)']\n",
    "# dfs = [varyk_vision_1, varyk_vision_2, varyk_nlp_1, varyk_nlp_2]\n",
    "\n",
    "fname = 'graphs/real-smoothing-k.pdf'\n",
    "plot_dist_est(dfs, xlabels, titles, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 7 in the main text can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'nvary-vision-supp128',\n",
    "    'nvary-nlp-supp2048',\n",
    "    'kvary-vision-size1000',\n",
    "    'kvary-nlp-size10000']\n",
    "dfs = [pd.read_pickle(f'results/real/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Sample size']*2 + ['Support size']*2\n",
    "titles = [r'(a) Images ($k=128$)', r'(b) Text ($k=2048$)',\n",
    "          r'(c) Images ($n=1000$)', r'(d) Text ($n=10000$)']\n",
    "fname = 'graphs/real-smoothing.pdf'\n",
    "plot_dist_est(dfs, xlabels, titles, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Distribution estimators for divergence frontiers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We first vary the sampel size."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 18 (bottom) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'nvary-vision-supp128-df',\n",
    "    'nvary-vision-supp1024-df',\n",
    "    'nvary-nlp-supp64-df',\n",
    "    'nvary-nlp-supp2048-df']\n",
    "dfs = [pd.read_pickle(f'results/real/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Sample size']*4\n",
    "titles = [r'(a) Images ($k=128$)', r'(b) Images ($k=1024$)',\n",
    "          r'(c) Text ($k=64$)', r'(d) Text ($k=2048$)']\n",
    "fname = 'graphs/real-smoothing-n-df.pdf'\n",
    "plot_dist_est(dfs, xlabels, titles, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We then vary the support size."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Figure 19 (bottom) in the appendix can be reproduced as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fnames = [\n",
    "    'kvary-vision-size1000-df',\n",
    "    'kvary-vision-size10000-df',\n",
    "    'kvary-nlp-size1000-df',\n",
    "    'kvary-nlp-size10000-df']\n",
    "dfs = [pd.read_pickle(f'results/real/{file}.pkl') for file in fnames]\n",
    "\n",
    "xlabels = ['Support size']*4\n",
    "titles = [r'(a) Images ($n=1000$)', r'(b) Images ($n=10000$)',\n",
    "          r'(c) Text ($n=1000$)', r'(d) Text ($n=10000$)']\n",
    "fname = 'graphs/real-smoothing-k-df.pdf'\n",
    "plot_dist_est(dfs, xlabels, titles, fname=fname, save=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### TODO: Correlations between K-means and Lattice quantization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On the image data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k_list = [8, 16, 32, 64, 128, 256, 512, 1024]\n",
    "\n",
    "with open('parsed_outputs/cifar10_lattice.p', 'rb') as f:\n",
    "    outs_new_lattice = pkl.load(f)\n",
    "o1 = outs_new_lattice[100000] ### from lattice\n",
    "o1.keys()\n",
    "arr1 = [o1[k].line_mauve for k in k_list]\n",
    "print(arr1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('parsed_outputs/cifar10_kmeans.p', 'rb') as f:\n",
    "    outs_new_lattice = pkl.load(f)\n",
    "o2 = outs_new_lattice[100000] ### from lattice\n",
    "o2.keys()\n",
    "arr2 = [o2[k].line_mauve for k in k_list]\n",
    "print(arr2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import spearmanr, pearsonr\n",
    "print(spearmanr(arr1, arr2))\n",
    "print(pearsonr(arr1, arr2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On the text data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k_list = [8, 16, 32, 64, 128, 256, 512, 1024]\n",
    "\n",
    "with open('parsed_outputs/wikitext103_lattice.p', 'rb') as f:\n",
    "    outs_new_lattice = pkl.load(f)\n",
    "o1 = outs_new_lattice[1800] ### from lattice\n",
    "arr1 = [o1[k].line_mauve for k in k_list]\n",
    "print(arr1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('parsed_outputs/wikitext103_kmeans_checkpoint-1800.pkl', 'rb') as f:\n",
    "    outs_new_kmeans = pkl.load(f)\n",
    "arr2 = [outs_new_kmeans[f'mauve_k={k}'].line_mauve for k in k_list]\n",
    "print(arr2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import spearmanr, pearsonr\n",
    "print(spearmanr(arr1, arr2))\n",
    "print(pearsonr(arr1, arr2))"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "b35703c28ef07c6dd3b1a9b7a6678bc943a659e52ddb160d6ab1f625af5ddba6"
  },
  "kernelspec": {
   "display_name": "divergence-frontier",
   "language": "python",
   "name": "divergence-frontier"
  },
  "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.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}