AFT_toy.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "47793b2b-3c77-4e34-9f19-bba8d1465d1b",
   "metadata": {},
   "source": [
    "# Annealed Flow Transport Monte Carlo"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d5ce5d2-51d2-4850-ba20-2f0bdbeeb140",
   "metadata": {},
   "source": [
    "I'm going to try to write up a toy example of annealed flow transport monte carlo..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 181,
   "id": "a278cbd7-a082-42cb-9689-3b53d7b7b49d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Sequence, Callable, Dict, Tuple, Optional\n",
    "import jax\n",
    "import flax.linen as nn\n",
    "import jax.numpy as jnp\n",
    "from functools import partial\n",
    "from jax import lax, ops, vmap, jit, grad, random\n",
    "from jax.scipy.special import logsumexp\n",
    "\n",
    "from jax.config import config\n",
    "from typing import Callable\n",
    "\n",
    "config.update(\"jax_enable_x64\", True)\n",
    "\n",
    "Conf = Params = Array = Seed = jnp.array\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9352b0e6-d28c-410f-b4d0-3cdb55cc9ad0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cf47625-c949-42a9-ac58-0182f1ada847",
   "metadata": {},
   "source": [
    "## defining simple distributions to anneal from/to"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8ab9efef-daaa-4ade-9c4b-dc188eba7987",
   "metadata": {},
   "outputs": [],
   "source": [
    "def unnormalized_Normal_logp(x: Conf,\n",
    "                             mu: Conf, \n",
    "                             cov: Conf) -> Conf:\n",
    "    \"\"\"\n",
    "    compute an unnormalized gaussian logp\n",
    "    arguments\n",
    "        x : jnp.array(Dx)\n",
    "            position\n",
    "        mu : jnp.array(Dx)\n",
    "            mean vector\n",
    "        cov : jnp.array(Dx)\n",
    "            covariance vector\n",
    "    returns\n",
    "        out : float\n",
    "            unnormalized gaussian logp\n",
    "    \"\"\"\n",
    "    delta = x-mu\n",
    "    return -0.5*(delta/cov).dot(delta)\n",
    "\n",
    "def Normal_energy(x: Conf,\n",
    "                  mu: Conf, \n",
    "                  cov: Conf) -> Conf:\n",
    "    \n",
    "    return -unnormalized_Normal_logp(x, mu, cov)\n",
    "\n",
    "def unnormalized_gmm_logp(x: Conf, \n",
    "                          mus: Conf, \n",
    "                          covs: Conf, \n",
    "                          lws: Conf) -> Conf:\n",
    "    \"\"\"\n",
    "    return unnormalized gaussian mixture model logp\n",
    "    \"\"\"\n",
    "    dim = len(x)\n",
    "    def mapper(entry):\n",
    "        _mu, _cov = entry[:dim], entry[dim:]\n",
    "        return unnormalized_Normal_logp(x, _mu, _cov)\n",
    "    \n",
    "    unnorm_logps = lax.map(mapper, jnp.hstack((mus, covs)))\n",
    "    weighted_ps = jnp.exp(lws + unnorm_logps).sum()\n",
    "    return jnp.log(weighted_ps)\n",
    "\n",
    "def gmm_energy(x: Conf, \n",
    "               mus: Conf, \n",
    "               covs: Conf, \n",
    "               lws: Conf) -> Conf:\n",
    "    return -unnormalized_gmm_logp(x, mus, covs, lws)\n",
    "    \n",
    "\n",
    "#samplers\n",
    "def sample_normal(seed: Conf, \n",
    "                  N: int, \n",
    "                  mu: Conf, \n",
    "                  cov: Conf) -> Conf:\n",
    "    \"\"\"\n",
    "    sample a normal distribution\n",
    "    \"\"\"\n",
    "    dim = len(mu)\n",
    "    return random.normal(seed, shape=(N,dim)) * jnp.sqrt(cov) + mu\n",
    "\n",
    "def sample_gmm(seed: Conf, \n",
    "               mus: Conf, \n",
    "               covs: Conf, \n",
    "               lws: Conf) -> Conf:\n",
    "    from jax.scipy.special import logsumexp\n",
    "    num_mixtures = len(lws)\n",
    "    dim = len(mu)\n",
    "    seed1, seed2 = random.split(seed)\n",
    "    mixture_idx = random.choice(seed1, len(lws), p=jnp.exp(lws - logsumexp(lws)))\n",
    "    return random.normal(seed2, shape=(dim,)) * jnp.sqrt(covs[mixture_idx]) + mus[mixture_idx]\n",
    "\n",
    "def kinetic_energy(m : Array, v: Array) -> float:\n",
    "    return 0.5 * v.dot(v) / m\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "be1e30ad-ad39-4afa-af03-426ba2a1a51b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def free_energy(works):\n",
    "    \"\"\"\n",
    "    compute the free energy from a work array\n",
    "    \"\"\"\n",
    "    from jax.scipy.special import logsumexp\n",
    "    N = len(works)\n",
    "    w_min = jnp.min(works)\n",
    "    return w_min - logsumexp(-works + w_min) + jnp.log(N)\n",
    "\n",
    "def ESS(works):\n",
    "    log_weights = -works\n",
    "    Ws = jnp.exp(log_weights - logsumexp(log_weights))\n",
    "    ESS = 1. / jnp.sum(Ws**2) / len(works)\n",
    "    return ESS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d8693bb0-cc66-49ef-b443-82ce98d31ea4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def BAOAB_kernel(X: Conf,\n",
    "                 seed: Seed,\n",
    "                 energy_parameters: Params,\n",
    "                 energy_fn : Callable[[Conf, Params], float],\n",
    "                 dt: float,\n",
    "                 gamma: float,\n",
    "                 mass: Params) -> Conf:\n",
    "    \"\"\"\n",
    "    a metropolized BAOAB kernel\n",
    "    \"\"\"\n",
    "    twice_dim = X.shape[0] #get the dimension of x, v together\n",
    "    single_dim = twice_dim // 2 #get the dimension of x, v separately\n",
    "    noise_seed, MH_seed = random.split(seed) # split the random seed\n",
    "    noise = random.normal(noise_seed, shape=(single_dim,)) # get the noise as white noise\n",
    "    a, b = jnp.exp(-gamma * dt), jnp.sqrt(1. - jnp.exp(-2. * gamma * dt)) # get the a and b parameters\n",
    "    r0, v0 = X[:single_dim], X[single_dim :] #pull x0, v0\n",
    "    e0 = energy_fn(r0, energy_parameters) + kinetic_energy(mass, v0) # get start total reduced energy...\n",
    "    #print(f\"v0: {v0}\")\n",
    "    # do splitting procedure\n",
    "    v1 = v0 + -grad(energy_fn)(r0, energy_parameters) * dt/mass #V\n",
    "    #print(f\"v1: {v1}\")\n",
    "    r1 = r0 + v1 * dt #R\n",
    "    k_b = kinetic_energy(mass, v1) #get t before O\n",
    "    v2 = -a * v1 + b * jnp.sqrt(1. / mass) * noise # O\n",
    "    #print(f\"v2: {v2}\")\n",
    "    k_a = kinetic_energy(mass, v2) #get t after O\n",
    "    r2 = r1 + v2 * dt #R\n",
    "    v3 = v2 + -grad(energy_fn)(r2, energy_parameters) * dt/ mass #V\n",
    "    #print(f\"v3: {v3}\")\n",
    "    \n",
    "    ef = energy_fn(r2, energy_parameters) + kinetic_energy(mass, v3) # get final energy\n",
    "    \n",
    "    # compute log work\n",
    "    lw = ef - e0 - (k_a - k_b)\n",
    "    log_acceptance_prob = jnp.min(Array([0., -lw]))\n",
    "    lu = jnp.log(random.uniform(MH_seed))\n",
    "    accept = (lu <= log_acceptance_prob)\n",
    "    #print(f\"accept: {accept}\")\n",
    "    out_X = lax.cond(accept, lambda x: jnp.concatenate([r2, v3]), lambda x: jnp.concatenate([r0, -v0]), None)\n",
    "    return out_X"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0cbc9f7-a829-4ae5-a6e1-db668d81c650",
   "metadata": {},
   "source": [
    "write a little test?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7b533744-5a6b-49aa-acc9-671046eabd94",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_AIS_sampler(energy_fn: Callable[[Conf, Params], float], # energy function\n",
    "                     energy_parameters: Params,  # energy parameters\n",
    "                     forward_kernel: Callable[[Conf, Seed, Params], Conf], # forward kernel \n",
    "                     t0_sampler: Callable[Seed, Conf] # sampler at \n",
    "                    ) -> Callable[Seed, Tuple[float, Conf, Conf]]:\n",
    "    \"\"\"simple function that generates a AIS sampler with the following criteria.\n",
    "    - this should be vmappable across N particles\n",
    "    \"\"\"\n",
    "    len_energy_parameters = energy_parameters.shape[0]\n",
    "    ts = jnp.arange(len_energy_parameters)[1:]\n",
    "    #print(f\"ts returning: {ts}\")\n",
    "    def execute_iteration(state: Tuple[Conf, Array, Seed], \n",
    "                          t: float) -> Tuple[Tuple[Conf, Array, Array], Tuple[Conf, Array]]:\n",
    "        x_tm1, energy_tally, seed = state # separate the state\n",
    "        _full_dim = x_tm1.shape[0] # get the dim(x) + dim(v)\n",
    "        _particle_dim = _full_dim // 2 # get dim(x), dim(v) (same)\n",
    "        run_seed, next_seed = random.split(seed) # split the random seed\n",
    "        energy_diff = energy_fn(x_tm1[:_particle_dim], energy_parameters[t]) - energy_fn(x_tm1[: _particle_dim], energy_parameters[t-1]) # compute energy difference\n",
    "        x_t = forward_kernel(x_tm1, run_seed, energy_parameters[t]) # propagate the particle forward\n",
    "        new_energy_tally = energy_tally + energy_diff # update the energy tally\n",
    "        new_state = (x_t, new_energy_tally, next_seed) # wrap the new state in a tuple\n",
    "        return new_state, new_state[:2] # return the new state\n",
    "        \n",
    "    \n",
    "    def AIS_sampler(seed: Seed) -> Tuple[float, Conf, Conf]: \n",
    "        init_seed, seed = random.split(seed) # split the seed\n",
    "        x0 = t0_sampler(init_seed) # sample [x,v] @ t0\n",
    "        init = (x0, 0., seed) # initialize x0 with zero work\n",
    "        (x_T, cumulative_work, final_seed), (traj, work_traj) = lax.scan(execute_iteration, init, ts)\n",
    "        out = (cumulative_work, jnp.vstack([x0[jnp.newaxis, ...], traj]), work_traj)\n",
    "        return out  \n",
    "    \n",
    "    # jit the new sampler\n",
    "    jAIS_sampler = jit(AIS_sampler)\n",
    "    \n",
    "    return jAIS_sampler"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6743fb0-d0a0-4997-b29c-cb58c4b74c37",
   "metadata": {},
   "source": [
    "1D test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 407,
   "id": "5a9889d8-c955-4363-8768-5da43428cc62",
   "metadata": {},
   "outputs": [],
   "source": [
    "def energy_fn(x, param):\n",
    "    _param = param[0]\n",
    "    out = Normal_energy(x, jnp.array([0.]), _param)\n",
    "    return out "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 408,
   "id": "e099eeeb-6a2b-4bfb-859a-9fa033848dc1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _t0_sampler(seed):\n",
    "    xseed, vseed = random.split(seed)\n",
    "    x = sample_normal(xseed, \n",
    "                  1, \n",
    "                  jnp.zeros((1,)), \n",
    "                  jnp.ones((1,))).flatten()\n",
    "    v = random.normal(vseed, shape=(1,))\n",
    "    return jnp.concatenate([x, v])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 409,
   "id": "3d3343d2-6426-4475-82b6-a0beeb9cb196",
   "metadata": {},
   "outputs": [],
   "source": [
    "_forward_kernel = partial(BAOAB_kernel, energy_fn = energy_fn, dt=1e-1, gamma=1., mass=1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 410,
   "id": "9b930c2f-f99f-4d86-a758-5f48062cf742",
   "metadata": {},
   "outputs": [],
   "source": [
    "energy_parameters = jnp.linspace(1., 2., 10)[..., jnp.newaxis]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 411,
   "id": "12b01f37-c4e2-42d9-ab2a-e375f20eacea",
   "metadata": {},
   "outputs": [],
   "source": [
    "ais_sampler = make_AIS_sampler(energy_fn, energy_parameters, _forward_kernel, _t0_sampler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 412,
   "id": "79f1302a-a0c9-495a-8a2b-d1ed79163a7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "vais_sampler = vmap(ais_sampler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 416,
   "id": "11704c8f-2aed-47b1-aff4-03adfc34b948",
   "metadata": {},
   "outputs": [],
   "source": [
    "key = random.PRNGKey(33)\n",
    "all_keys = random.split(key, num=100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0cdcb81-b610-4add-a71e-f21c6c2f3fe9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 417,
   "id": "c476beac-63f4-4e83-8db2-b5577904f646",
   "metadata": {},
   "outputs": [],
   "source": [
    "(lws, trajs, lws_trajs) = vais_sampler(all_keys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 425,
   "id": "01045110-1024-4b67-ad42-6db3335feb3f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([ 2.,  0.,  2.,  3.,  3.,  3.,  8.,  6.,  7., 66.]),\n",
       " array([-1.68757860e+00, -1.51883749e+00, -1.35009639e+00, -1.18135529e+00,\n",
       "        -1.01261419e+00, -8.43873084e-01, -6.75131981e-01, -5.06390878e-01,\n",
       "        -3.37649776e-01, -1.68908673e-01, -1.67570249e-04]),\n",
       " <BarContainer object of 10 artists>)"
      ]
     },
     "execution_count": 425,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXAAAAD4CAYAAAD1jb0+AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8rg+JYAAAACXBIWXMAAAsTAAALEwEAmpwYAAAOHElEQVR4nO3dYYxlZ13H8e/PLg1awHbt7GakhBGzqTQmtM2k1tSQ6FJSW8MuL2raRJ2YJhsSMG2iMau+wXetiURNDMlK0VERrUCzG0BgHWmICVZmS1taF1hKSqkddocqUnwBKfx9MWfLMHt37pmZe+/cR76f5Oac89xz7vPv6dNfzzxzz5lUFZKk9vzIbhcgSdoeA1ySGmWAS1KjDHBJapQBLkmN2jPJzq688sqam5ubZJeS1LxTp059vapmNrZPNMDn5uZYXl6eZJeS1LwkXxnU7hSKJDXKAJekRhngktQoA1ySGmWAS1KjDHBJapQBLkmNMsAlqVEGuCQ1aqJ3YkrSbpo7+pFd6/vpe28b+Wd6BS5JjTLAJalRBrgkNcoAl6RGGeCS1CgDXJIaZYBLUqMMcElqlAEuSY0ywCWpUQa4JDXKAJekRhngktQoA1ySGtUrwJNcnuQDST6f5HSSn0+yN8nJJGe65RXjLlaS9H19r8D/FPhYVf0M8AbgNHAUWKqqA8BSty1JmpChAZ7kVcAbgfsBquo7VfUN4BCw2O22CBweT4mSpEH6XIG/DlgF/jLJZ5O8J8llwP6qWgHolvvGWKckaYM+Ab4HuB54d1VdB/wvW5guSXIkyXKS5dXV1W2WKUnaqE+APws8W1UPd9sfYC3QzyaZBeiW5wYdXFXHqmq+quZnZmZGUbMkiR4BXlVfA76a5Oqu6SDwH8AJYKFrWwCOj6VCSdJAff8q/W8B70tyKfBl4DdZC/8HktwFPAPcPp4SJUmD9ArwqnoUmB/w1sGRViNJ6s07MSWpUQa4JDXKAJekRhngktQoA1ySGmWAS1KjDHBJapQBLkmNMsAlqVEGuCQ1ygCXpEYZ4JLUKANckhplgEtSowxwSWqUAS5JjTLAJalRBrgkNcoAl6RGGeCS1CgDXJIaZYBLUqMMcElq1J4+OyV5GngB+C7wYlXNJ9kL/AMwBzwN/GpV/fd4ypQkbbSVK/BfrKprq2q+2z4KLFXVAWCp25YkTchOplAOAYvd+iJweMfVSJJ66xvgBXwiyakkR7q2/VW1AtAt9w06MMmRJMtJlldXV3desSQJ6DkHDtxUVc8l2QecTPL5vh1U1THgGMD8/Hxto0ZJ0gC9rsCr6rlueQ54ELgBOJtkFqBbnhtXkZKkCw0N8CSXJXnl+XXgzcATwAlgodttATg+riIlSRfqM4WyH3gwyfn9/66qPpbkM8ADSe4CngFuH1+ZkqSNhgZ4VX0ZeMOA9ueBg+MoSpI0nHdiSlKjDHBJapQBLkmNMsAlqVEGuCQ1ygCXpEYZ4JLUKANckhplgEtSowxwSWqUAS5JjTLAJalRBrgkNcoAl6RGGeCS1CgDXJIaZYBLUqMMcElqlAEuSY0ywCWpUQa4JDXKAJekRvUO8CSXJPlskg9323uTnExyplteMb4yJUkbbeUK/G7g9Lrto8BSVR0AlrptSdKE9ArwJFcBtwHvWdd8CFjs1heBwyOtTJK0qb5X4H8C/C7wvXVt+6tqBaBb7ht0YJIjSZaTLK+uru6kVknSOkMDPMmvAOeq6tR2OqiqY1U1X1XzMzMz2/kISdIAe3rscxPwliS3Ai8HXpXkb4GzSWaraiXJLHBunIVKkn7Q0Cvwqvq9qrqqquaAO4B/qapfA04AC91uC8DxsVUpSbrATr4Hfi9wc5IzwM3dtiRpQvpMobykqh4CHurWnwcOjr4kSVIf3okpSY0ywCWpUQa4JDXKAJekRhngktQoA1ySGmWAS1KjDHBJapQBLkmNMsAlqVEGuCQ1ygCXpEYZ4JLUKANckhplgEtSowxwSWqUAS5JjTLAJalRBrgkNcoAl6RGGeCS1CgDXJIaZYBLUqOGBniSlyf59ySPJXkyyR927XuTnExyplteMf5yJUnn9bkC/zbwS1X1BuBa4JYkNwJHgaWqOgAsdduSpAkZGuC15lvd5su6VwGHgMWufRE4PI4CJUmD9ZoDT3JJkkeBc8DJqnoY2F9VKwDdct9Fjj2SZDnJ8urq6ojKliT1CvCq+m5VXQtcBdyQ5Gf7dlBVx6pqvqrmZ2ZmtlmmJGmjLX0Lpaq+ATwE3AKcTTIL0C3Pjbo4SdLF9fkWykySy7v1HwXeBHweOAEsdLstAMfHVKMkaYA9PfaZBRaTXMJa4D9QVR9O8mnggSR3Ac8At4+xTknSBkMDvKoeB64b0P48cHAcRUmShvNOTElqlAEuSY0ywCWpUQa4JDXKAJekRhngktQoA1ySGmWAS1KjDHBJapQBLkmNMsAlqVEGuCQ1ygCXpEYZ4JLUKANckhplgEtSowxwSWqUAS5JjTLAJalRBrgkNcoAl6RGGeCS1KihAZ7kNUk+meR0kieT3N21701yMsmZbnnF+MuVJJ3X5wr8ReC3q+r1wI3A25NcAxwFlqrqALDUbUuSJmRogFfVSlU90q2/AJwGXg0cAha73RaBw2OqUZI0wJbmwJPMAdcBDwP7q2oF1kIe2HeRY44kWU6yvLq6usNyJUnn9Q7wJK8APgjcU1Xf7HtcVR2rqvmqmp+ZmdlOjZKkAXoFeJKXsRbe76uqD3XNZ5PMdu/PAufGU6IkaZA+30IJcD9wuqrete6tE8BCt74AHB99eZKki9nTY5+bgF8HPpfk0a7t94F7gQeS3AU8A9w+lgolSQMNDfCq+lcgF3n74GjLkST15Z2YktQoA1ySGmWAS1KjDHBJapQBLkmNMsAlqVEGuCQ1ygCXpEYZ4JLUKANckhplgEtSowxwSWqUAS5JjTLAJalRBrgkNcoAl6RGGeCS1CgDXJIaZYBLUqMMcElqlAEuSY0ywCWpUUMDPMl7k5xL8sS6tr1JTiY50y2vGG+ZkqSN+lyB/xVwy4a2o8BSVR0AlrptSdIEDQ3wqvoU8F8bmg8Bi936InB4tGVJkobZ7hz4/qpaAeiW+y62Y5IjSZaTLK+urm6zO0nSRmP/JWZVHauq+aqan5mZGXd3kvRDY7sBfjbJLEC3PDe6kiRJfWw3wE8AC936AnB8NOVIkvrq8zXC9wOfBq5O8mySu4B7gZuTnAFu7rYlSRO0Z9gOVXXnRd46OOJaJElb4J2YktSooVfgksZr7uhHdqXfp++9bVf61eh4BS5JjTLAJalRTqFImrjdmjb6/8YrcElqlAEuSY1yCkX6IeU0Rvu8ApekRhngktQoA1ySGmWAS1KjDHBJapQBLkmNMsAlqVEGuCQ1yht5dAFv8JDa4BW4JDXKAJekRhngktSoZubAd3Nedrf+9JRz0ZI24xW4JDXKAJekRu0owJPckuQLSb6U5OioipIkDbftAE9yCfDnwC8D1wB3JrlmVIVJkja3kyvwG4AvVdWXq+o7wN8Dh0ZTliRpmJ18C+XVwFfXbT8L/NzGnZIcAY50m99K8oV1b18JfH0HNUxE7vuBzSZq3sCaJ8OaJ6PFmsl9O6r7tYMadxLgGdBWFzRUHQOODfyAZLmq5ndQw8RZ82RY82RY8+SMo+6dTKE8C7xm3fZVwHM7K0eS1NdOAvwzwIEkP5XkUuAO4MRoypIkDbPtKZSqejHJO4CPA5cA762qJ7f4MQOnVqacNU+GNU+GNU/OyOtO1QXT1pKkBngnpiQ1ygCXpEaNPcCT3J7kySTfSzLwKzRJrk7y6LrXN5Pc0733ziT/ue69W6eh5m6/p5N8rqtreV373iQnk5zplldMQ81JXpPkk0lOd/veve69aT7PAx/ZsEvneWifUziee52nKRvPfc7zVIznYY8UyZo/695/PMn1fY8dqqrG+gJeD1wNPATM99j/EuBrwGu77XcCvzPuOrdTM/A0cOWA9j8CjnbrR4H7pqFmYBa4vlt/JfBF4JppPs/deHgKeB1wKfDYupp34zxvqc8pGc+9ap6y8Ty0z2kYz5uNz3X73Ar8E2v3ztwIPNz32GGvsV+BV9XpqvrC8D1fchB4qqq+Mq6ahtlGzRsdAha79UXg8I6LGqJPzVW1UlWPdOsvAKdZu6N2V/Q8z5s9smHi53kbfe76eGbn52kqz/OUjOc+jxQ5BPx1rfk34PIksz2P3dQ0zoHfAbx/Q9s7uh893juJH9+2oIBPJDmVtUcGnLe/qlZgbZAB+3aluk0kmQOuAx5e1zyN53nQIxvO/0e6G+d5q31Ow3juW/M0ject9bmL43mz8Tlsnz7HbmokAZ7kn5M8MeC1pf+bZO2GoLcA/7iu+d3ATwPXAivAH09RzTdV1fWsPZHx7UneOIraLmaE5/kVwAeBe6rqm13ztJ7nXo9sGCXHs+N5K90PaNs4Pi+2z47H9kj+pFpVvWkUn8PawHmkqs6u++yX1pP8BfDhUXQ0ipqr6rlueS7Jg6z9SPQp4GyS2apa6X5UOrfTvrp+dlxzkpexNtjfV1UfWvfZ03qeN3tkw8TPc5Kt9DkV47lvzdM0nvvWPOnxPECfR4pcbJ9Lexy7qWmbQrmTDT9udv/yznsr8MREK7qIJJcleeX5deDNfL+2E8BCt74AHJ98hRdKEuB+4HRVvWvDe1N5ntn8kQ27cZ630ue0jOehNU/heO5T8zSM5z6PFDkB/Mbal1FyI/A/3bTQzh9HMoHf0r6Vtf8DfRs4C3y8a/9J4KPr9vsx4Hngxzcc/zfA54DHu3+42WmombXfHD/WvZ4E/mDd8T8BLAFnuuXeKan5F1j7Ee1x4NHudes0n+f6/m/xv8jab+x3+zwP7HPKx/PQmqdwPPepeSrG86DxCbwNeFu3Htb++M1TXU3zmx27lZe30ktSo6ZtCkWS1JMBLkmNMsAlqVEGuCQ1ygCXpEYZ4JLUKANckhr1fxzet6s3L/bqAAAAAElFTkSuQmCC\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.hist(lws)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 427,
   "id": "574cc79c-73c3-4dc4-be8e-845f47b295fd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(-0.35377424, dtype=float64)"
      ]
     },
     "execution_count": 427,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "free_energy(lws)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 428,
   "id": "a32d6422-dddf-4a0f-bb9c-f876eb3bd791",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(0.76650916, dtype=float64)"
      ]
     },
     "execution_count": 428,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ESS(lws)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5642aa11-4fcd-44c2-b38e-c8939078f988",
   "metadata": {},
   "source": [
    "neato, now let's do this in 2 dimensions with a trivial free energy difference..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c5300ab-3873-4673-bc85-95eb83f56482",
   "metadata": {},
   "source": [
    "# free energy in 2D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "id": "654f90b3-2c5e-4174-83c3-610f18a2d100",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _energy_fn(x, param):\n",
    "    out = Normal_energy(x, param, jnp.ones(2))\n",
    "    return out "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "id": "6b0eab41-82d2-465c-95bc-7b268c6328a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _t0_sampler(seed):\n",
    "    xseed, vseed = random.split(seed)\n",
    "    x = sample_normal(xseed, \n",
    "                  1, \n",
    "                  jnp.zeros((2,)), \n",
    "                  jnp.ones((2,))).flatten()\n",
    "    v = random.normal(vseed, shape=(2,))\n",
    "    return jnp.concatenate([x, v])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "id": "1c405a26-f525-47a8-b4e4-a5f7b0a11140",
   "metadata": {},
   "outputs": [],
   "source": [
    "_forward_kernel = partial(BAOAB_kernel, energy_fn = _energy_fn, dt=1e-1, gamma=1., mass=1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "id": "d2957a95-b3b7-4ada-aa02-04549f001f0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "energy_parameters = jnp.vstack((jnp.linspace(0., 10., 30000), jnp.linspace(0., 10., 30000))).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "id": "bf11f62c-6ae6-4206-8037-f5bf3806c928",
   "metadata": {},
   "outputs": [],
   "source": [
    "ais_sampler = make_AIS_sampler(_energy_fn, energy_parameters, _forward_kernel, _t0_sampler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "id": "495b553f-be87-4b50-8e03-66c34f5fa132",
   "metadata": {},
   "outputs": [],
   "source": [
    "vais_sampler = vmap(ais_sampler)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "id": "3b1a9a7a-a1ec-44f5-85cf-c5f1871471e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "key = random.PRNGKey(350)\n",
    "all_keys = random.split(key, num=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "id": "8fac4120-3515-4530-afe0-7a7a0640653f",
   "metadata": {},
   "outputs": [],
   "source": [
    "(lws, trajs, lws_trajs) = vais_sampler(all_keys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "id": "8d6dfd60-8768-4b7d-a030-0b83134c092f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([ 5.,  4.,  8., 12., 13.,  3.,  3.,  1.,  0.,  1.]),\n",
       " array([ 1.22216554,  2.74093758,  4.25970962,  5.77848166,  7.2972537 ,\n",
       "         8.81602574, 10.33479778, 11.85356982, 13.37234186, 14.8911139 ,\n",
       "        16.40988594]),\n",
       " <BarContainer object of 10 artists>)"
      ]
     },
     "execution_count": 124,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXAAAAD4CAYAAAD1jb0+AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8rg+JYAAAACXBIWXMAAAsTAAALEwEAmpwYAAANMUlEQVR4nO3db4xldX3H8fenjEQXMWAYrLJsRxqktYQWMm1RUpuCJFuXgA/6ACJmW0kmaVpEY6tLSOqzZluN1cRGswFcEjeQZsVKJFo2qCVNkHZ3Qf4titEtLKI7hrRabYqbfvtgLulwd+fP3nN2zvyW9yvZzL1nzsz57vx577nn3nM2VYUkqT2/NPQAkqTJGHBJapQBl6RGGXBJapQBl6RGTa3lxs4666yamZlZy01KUvP27dv346qaHl++pgGfmZlh7969a7lJSWpekn8/1nIPoUhSowy4JDXKgEtSowy4JDXKgEtSowy4JDXKgEtSowy4JDXKgEtSo9b0TExpJTPb7h1kuwe3bxlku1IX7oFLUqMMuCQ1yoBLUqMMuCQ1yoBLUqMMuCQ1yoBLUqMMuCQ1yoBLUqMMuCQ1yoBLUqNWDHiS25McTvL4omUfS/JUkkeTfDHJGSd0SknSUVazB74T2Dy2bA9wYVVdBHwHuLnnuSRJK1gx4FX1APDC2LL7qurI6O43gY0nYDZJ0jL6OAb+PuArPXweSdJx6HQ98CS3AEeAXcusMwfMAWzatKnL5rRGhromt6TjM/EeeJKtwFXAe6qqllqvqnZU1WxVzU5PT0+6OUnSmIn2wJNsBj4C/H5V/bzfkSRJq7GalxHeCTwIXJDkUJIbgE8DpwN7kjyS5LMneE5J0pgV98Cr6rpjLL7tBMwiSToOnokpSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY1aMeBJbk9yOMnji5a9PsmeJE+P3p55YseUJI1bzR74TmDz2LJtwP1VdT5w/+i+JGkNrRjwqnoAeGFs8TXAHaPbdwDv7ncsSdJKJj0G/oaqeh5g9Pbs/kaSJK3GCX8SM8lckr1J9s7Pz5/ozUnSK8akAf9RkjcCjN4eXmrFqtpRVbNVNTs9PT3h5iRJ4yYN+D3A1tHtrcCX+hlHkrRaq3kZ4Z3Ag8AFSQ4luQHYDlyZ5GngytF9SdIamlpphaq6bol3XdHzLJKk4+CZmJLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUKAMuSY0y4JLUqE4BT/LBJE8keTzJnUle3ddgkqTlTRzwJOcA7wdmq+pC4BTg2r4GkyQtr+shlCngNUmmgA3AD7qPJElajalJP7CqnkvyceAZ4L+B+6rqvvH1kswBcwCbNm2adHOvSDPb7h16BEnrWJdDKGcC1wBvBt4EnJbk+vH1qmpHVc1W1ez09PTkk0qSXqbLIZR3At+vqvmq+gVwN/D2fsaSJK2kS8CfAS5NsiFJgCuAA/2MJUlaycQBr6qHgN3AfuCx0efa0dNckqQVTPwkJkBVfRT4aE+zSJKOg2diSlKjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjOgU8yRlJdid5KsmBJG/razBJ0vKmOn78p4CvVtUfJTkV2NDDTJKkVZg44EleB7wD+GOAqnoReLGfsSRJK+lyCOU8YB74XJKHk9ya5LTxlZLMJdmbZO/8/HyHzUmSFusS8CngEuAzVXUx8DNg2/hKVbWjqmaranZ6errD5iRJi3UJ+CHgUFU9NLq/m4WgS5LWwMQBr6ofAs8muWC06ArgyV6mkiStqOurUG4Edo1egfI94E+6jyRJWo1OAa+qR4DZfkaRJB0Pz8SUpEYZcElqlAGXpEYZcElqlAGXpEYZcElqlAGXpEYZcElqlAGXpEYZcElqVNdroayZmW33Drbtg9u3DLZtSVqKe+CS1CgDLkmNMuCS1CgDLkmNMuCS1CgDLkmNMuCS1CgDLkmNMuCS1CgDLkmNMuCS1CgDLkmN6hzwJKckeTjJl/sYSJK0On3sgd8EHOjh80iSjkOngCfZCGwBbu1nHEnSanW9HvgngQ8Dpy+1QpI5YA5g06ZNHTc3jCGvRS5JS5l4DzzJVcDhqtq33HpVtaOqZqtqdnp6etLNSZLGdDmEchlwdZKDwF3A5Uk+38tUkqQVTRzwqrq5qjZW1QxwLfC1qrq+t8kkScvydeCS1Khe/lPjqvoG8I0+PpckaXXcA5ekRhlwSWqUAZekRhlwSWqUAZekRhlwSWqUAZekRhlwSWqUAZekRhlwSWpUL6fSS617JV7z/eD2LUOPoI7cA5ekRhlwSWqUAZekRhlwSWqUAZekRhlwSWqUAZekRhlwSWqUAZekRhlwSWqUAZekRhlwSWrUxAFPcm6Sryc5kOSJJDf1OZgkaXldrkZ4BPhQVe1PcjqwL8meqnqyp9kkScuYeA+8qp6vqv2j2z8FDgDn9DWYJGl5vRwDTzIDXAw8dIz3zSXZm2Tv/Px8H5uTJNFDwJO8FvgC8IGq+sn4+6tqR1XNVtXs9PR0181JkkY6BTzJq1iI966qurufkSRJq9HlVSgBbgMOVNUn+htJkrQaXfbALwPeC1ye5JHRn3f1NJckaQUTv4ywqv4FSI+zSJKOg2diSlKjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNcqAS1KjDLgkNarL/8gjqWEz2+4dbNsHt28ZZLsn29/ZPXBJapQBl6RGGXBJapQBl6RGGXBJapQBl6RGGXBJapQBl6RGGXBJapQBl6RGGXBJapQBl6RGdQp4ks1Jvp3ku0m29TWUJGllEwc8ySnA3wN/CLwVuC7JW/saTJK0vC574L8DfLeqvldVLwJ3Adf0M5YkaSVdrgd+DvDsovuHgN8dXynJHDA3uvtfSb7dYZtdnAX8eKBtr4bzdeN83azpfPmb4/6Q5r9+E/ydF/uVYy3sEvAcY1kdtaBqB7Cjw3Z6kWRvVc0OPcdSnK8b5+vG+boZar4uh1AOAecuur8R+EG3cSRJq9Ul4P8GnJ/kzUlOBa4F7ulnLEnSSiY+hFJVR5L8OfBPwCnA7VX1RG+T9W/wwzgrcL5unK8b5+tmkPlSddRha0lSAzwTU5IaZcAlqVEndcCTnJvk60kOJHkiyU1Dz3QsSU5J8nCSLw89y7gkZyTZneSp0dfxbUPPtFiSD46+t48nuTPJq9fBTLcnOZzk8UXLXp9kT5KnR2/PXGfzfWz0PX40yReTnLGe5lv0vr9IUknOGmK20QzHnC/JjaNLizyR5G/XYpaTOuDAEeBDVfXrwKXAn63T0/1vAg4MPcQSPgV8tap+DfhN1tGcSc4B3g/MVtWFLDyZfu2wUwGwE9g8tmwbcH9VnQ/cP7o/lJ0cPd8e4MKqugj4DnDzWg+1yE6Ono8k5wJXAs+s9UBjdjI2X5I/YOFM9Iuq6jeAj6/FICd1wKvq+araP7r9Uxbic86wU71cko3AFuDWoWcZl+R1wDuA2wCq6sWq+o9BhzraFPCaJFPABtbBuQhV9QDwwtjia4A7RrfvAN69ljMtdqz5quq+qjoyuvtNFs7rGMQSXz+AvwM+zDFOGFxLS8z3p8D2qvqf0TqH12KWkzrgiyWZAS4GHhp4lHGfZOGH8n8HnuNYzgPmgc+NDvHcmuS0oYd6SVU9x8KezjPA88B/VtV9w061pDdU1fOwsGMBnD3wPMt5H/CVoYdYLMnVwHNV9a2hZ1nCW4DfS/JQkn9O8ttrsdFXRMCTvBb4AvCBqvrJ0PO8JMlVwOGq2jf0LEuYAi4BPlNVFwM/Y9iH/i8zOo58DfBm4E3AaUmuH3aqtiW5hYVDj7uGnuUlSTYAtwB/NfQsy5gCzmThUO1fAv+Q5FiXG+nVSR/wJK9iId67quruoecZcxlwdZKDLFzN8fIknx92pJc5BByqqpcetexmIejrxTuB71fVfFX9ArgbePvAMy3lR0neCDB6uyYPsY9Hkq3AVcB7an2dIPKrLPwj/a3R78pGYH+SXx50qpc7BNxdC/6VhUfUJ/yJ1pM64KN/AW8DDlTVJ4aeZ1xV3VxVG6tqhoUn375WVetmD7Kqfgg8m+SC0aIrgCcHHGncM8ClSTaMvtdXsI6eZB1zD7B1dHsr8KUBZzlKks3AR4Crq+rnQ8+zWFU9VlVnV9XM6HflEHDJ6OdzvfhH4HKAJG8BTmUNrp54UgechT3c97KwZ/vI6M+7hh6qMTcCu5I8CvwW8NfDjvP/Ro8MdgP7gcdY+Hke/JTrJHcCDwIXJDmU5AZgO3BlkqdZeCXF9nU236eB04E9o9+Tz66z+daNJea7HThv9NLCu4Cta/EoxlPpJalRJ/seuCSdtAy4JDXKgEtSowy4JDXKgEtSowy4JDXKgEtSo/4P7twGNdXPc/IAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.hist(lws)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "id": "82ce4a80-70ba-4442-b232-449f40cf7613",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(50, 30000, 4)"
      ]
     },
     "execution_count": 125,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trajs.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "id": "23d9e3d0-06e4-43aa-989d-03f4d255026b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x7f24241cfcd0>]"
      ]
     },
     "execution_count": 127,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXMAAAD4CAYAAAAeugY9AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8rg+JYAAAACXBIWXMAAAsTAAALEwEAmpwYAAAxZElEQVR4nO3dd3gVVfoH8O+bQgKhQ2ihhN57pIqoFBFYC1bslR+r2HbXFUVXLCjr2tsqK3ZFFLuIFEWRTiD00Am9hJYESEg7vz9uycy9c/vcmu/neXyYOXfKGa68OTlzzntEKQUiIopuceGuABERBY7BnIgoBjCYExHFAAZzIqIYwGBORBQDEsJx0/r166v09PRw3JqIKGqtXr36mFIq1eizsATz9PR0ZGZmhuPWRERRS0T2uPqM3SxERDGAwZyIKAYwmBMRxQAGcyKiGMBgTkQUAxjMiYhiAIM5EVEMYDAnIvJAKYXxn6xGQVFJuKviEoM5EZEHD89aj182HUbXyfPCXRWXGMyJiNxYufsEZq3eH+5qeMRgTkTkxrXvLgt3FbzCYE5EFAMYzImIYgCDORGRl5ITIzdkRm7NiIgiTFFJebir4BKDORFRDGAwJyKKAQzmREQulJbpu1XSalcNU008YzAnInJh38lC3f6BU4VYv/9UeCrjAYM5EZELF734u1PZZW8uCX1FvMBgTkQUA7wO5iLyvogcFZGNmrL/iMgWEVkvIt+KSO2g1JKIKIwiua/cxpeW+YcARjiUzQfQRSnVDcA2AI+aVC8ioohx5/ktw10Fj7wO5kqpRQBOOJTNU0qVWneXA2hqYt2IiMJGKWXfvn1gOu67uE0Ya+OZmX3mdwCY4+pDERknIpkikpmbm2vibYmIzPeVJu2tiGBA6/r2/eLScrzwyxacPldqdGpYmBLMRWQSgFIAn7k6Rik1TSmVoZTKSE1NNeO2RERB8Wv2Efxz1npdWf/W9VAjKQFd02rhq9X78PbvO/Hi3K1hqqGzgIO5iNwKYDSAG5X29xIioih150eZhuVtG1bHhgN5+GhpDgDgQ+ufkSAhkJNFZASARwAMVkqdNadKRESR45nLO9u31+w9BQDYduR0mGrjmi9DE2cAWAagvYjsF5E7AbwJoAaA+SKyVkTeCVI9iYhCYvH2Y7r94Z0bhakmvvG6Za6UGmtQPN3EuhARhd1N01fo9guLy8JUE99wBigRkRtVq8S7/fzP7ZExOo/BnIjIjQY1ktx+fvP0lSGqiXsM5kREboiIx2N2HA3/C1EGcyIiFz66o49uf9NTlxgeN33xbpSUhXdJOQZzIiIDz13ZFYPb6Sc4piRVjBlJSqgInzNW7sUT321EODGYExE5qJmcgLF9mhl+ljN1FHKmjsKDQ9vpyr9YtS8UVXOJwZyIyMFN/Vp47CtPiPPclx5KDOZERFa7ci0vMmd60cp+c+GOYFfHJwzmRERWtq6S42eKPR4bSRkTAQZzIiK789LrAgAmXOQ5d/mTf+kU7Or4hMGciMjq7o8t2RJb1k/xeGx6Pc/HhBKDORGRg8ISz/lYGtZMdioLZxZwBnMiIgAbD+TZtx3Hlxtp36gGvhrfX1fW8tGfTa+XtxjMiahSU0qhz5QFGP3GYntZs7rVvDr3vPS62PKM4zr34cFgTkSVWmm5wtGCc36fn5zoPqtiqDCYE1GlVloWeD93er2Klny4+s0ZzImoUuv+1LyAr3FL/3T7dq5DK//U2WJc9uZi/G/RroDv444vy8a9LyJHRWSjpqyuiMwXke3WP+sEp5pERMFR7JDtcNmjF/t8DW1bXDuZaP7mI+jx9Hys35+HKT9n+1tFr/jSMv8QgGNP/0QAvyql2gL41bpPRBS1khJ87wO/JqOpfXvPibMoL1e48u0l+C7rgJlVc8uXNUAXiUi6Q/HlAC60bn8E4HcAj5hRMSKiYFu0zXnJN3/yZ9VMTrRvv7pgOzYfzEfW3lPIwqkAauebQPvMGyqlDgGA9c8GgVeJiCg0jFYISoj3Lyz++6quAICBrevhP3O3BlQvf4TsBaiIjBORTBHJzM2NjAVQiahye/qnzbr958d0RfUkrzssdHq3sOR1mb3hUMD18od/ta5wREQaK6UOiUhjAEddHaiUmgZgGgBkZGSEb84rEZGDa3o3xYPD2iGtdlW/r1G1iqWvfc/xs2ZVyyeBtsx/AHCrdftWAN8HeD0iIlPtOX4GG/bnuT3m4RHtAwrkAFDVi8lDn6/YG9A93PFlaOIMAMsAtBeR/SJyJ4CpAIaJyHYAw6z7REQRY/B/fsdf3lyMQ3mFLo9pUMM5aZavqlXxHMwf+3ZDwPdxxZfRLGNdfDTEpLoQEQVN/+d/w7f3DEDP5sGZDhPuaf2cAUpElcbuY2fs22XlsfXqjsGciCqNeM0g8jPFwV327bqMZvbtUV0bAwD6t6oXtPsFOpqFiAgnzxTjxNlitE6t7vTZ4bwiNKyZ5HG1+1D4dPkeXN4jDekTZ9vLzksPTrfL5Ms6o2qVeDStUxV3nt8Sv/3rKNo1dP77MQtb5kQUsJ7PzMeQl/5wKt94IA/9nv8VM1Z6Xu0+FFblnHQqM7O7RbuoRVJCHCZf1hl3DWoFEUH15AQUm5Ch0RUGcyICAOQVliB94myMs66D6Y/0ibN1KWBtMyxX7D4ecP381aKefqGJbUcKdPt9Tez6+OiOPsiZOgo5U0chziEvQGKcoNQhqZeZGMyJCABw2wcrAQDzNh8J6Do/rq+YAWlr9caHsYulqKQMA1pXBOw1e/St8xrJoeltTkyIQwmDOREFW9beU6Zc5/4ZWfbtv3+1DoD+xaMnHZ6Yg2+z9ptSl2d/2owj+efQtkFFX7XjD6uL2ocmpVRCnKAkiCNoGMyJyHSbDuZhr2Za+9YjBfhti+cW/9er96OopBwPzVwXcB125p7Ge4t3AwCSNRN6ftuizzrSsXHNgO/lXX3OYPb64OVtYTAnItONen0xThUW2/fX78/DHR8698Wv2XsSBUUl9v0X5m4xrQ43vbfCvp3sR47yYAnWsnIM5kSVTLfJczHm7SWmXjOjhfPwvhW7TjiVFZWU2bcLi8sw5u2lGP3GYnvZkXz/F1Z2dCivyL69eo/zKJZw+XP7saBcl8GcqJLJLyrFmr2ndIE1r7DE8NjsQ/nYf9JzFkDtUmk20/50XvPyOc3SaSXllpeBrrIMBtKCdazPlsMFLo4MvWOnzfuBpcVgTlRJnThTjO1HCnDFW0uQtbei5Vq7WsWqOZe+9ifO//dC+/6OowW44X/LUVhcprvW2eIyXNGjia7McWFjAPh42R779hmDHwBaLR/92bsHMTDDITuhUQB9+drueGhoO7/v4a+lO4MzTJPBnKgSyTtb0QJftvM4xn2yGmv3ncJtH6yylxeXWlrMjgEbAIa+vAhLdx7HHw7LrZ0tLkU1LxZ1eO36Hrr72/R9bgGW7nDufpi36TCen5Pt88Se95fs1u0PbOM8lnxMr6Z4YGhbn64biOeutKxEdHXvph6O9A+DOVElMmHGGvv2379ap0s8BQBtG1THWWsQf3BmFlwZ/+lqLN9VEYzPnCtDihcpYM9qfkA0qlWRdvZI/jncoHlhaTPuk9V4949daP3YzygrV/hi5V6vxmpr+8sBYMoVXT2eE2w39G2OnKmj0C9I+VkYzIkqEU9Lom23zti866NMzN3kfijh9dOWA7BMDCosKUO1Kp5b5o9+s8EejH9cd9CbKtvNXLUPE7/ZgLaT5rjtT9eOjgGAnKmjkF4/xad7RSMGc6JK5Nrzmnk+CMCCbO9ngRZaX6SmJMVj3ZPDPR5/ztqN82WmbxODtAs7PDxrvcvjuk6eZ1j+/m0ZaFQzGeMHt8ZX4/v7dO9owGBOVImUe+h7rlU10e3nWpdbX3iePGMZT376XJlX56/ffwoAAlqmbdZq32eIXtyhIZY/NgQTL+2A89Lr+n3vSMVgTlSJeOpvHmnNu+1o7qbDTmW2fCt3WxNzvf7rdq/qYGuZm/0isKikDA98oe/nT3dIshXLTAnmIvKQiGwSkY0iMkNEAl9Qj4hMV+ImBWuv5rUx4eI2hp8tMxhO903WAZSVK/sY7iEdLDlOXriqG+4f4nqUiG1K+8vzt3ldb0dGQwpv/2AVvl+r74f/+q8D/L5HtAk4mItIGoD7AWQopboAiAdwfaDXJSLzlZa7bplvOVzgsuvjw6U5OFpQ5FQ+4fOK0TH/N7g1AEu//N+GuR6/Pahtfbd1HNqxAR6+pL3bY86W6MeonzxTjGW7nH/g1Kue5PY6scSsbpYEAFVFJAFANQC+vaYmopCwZUZ8Y2xPp89swwZXPz7U8Nw+U351KpuzsaL7Ja2O/gfBy9d2x9wHL8Du50ciZ+ooe/mK3Sew6WCeyzoOaF0f91zY2r5v1FXiOAa+5zPznY7p3qy2y3vEooAT+SqlDojIiwD2AigEME8p5fQ6WUTGARgHAM2bNw/0tkTkB9sMzL4tXb8ArFc9yR58Oz7xi320iidNaul7V8f0Mu4T/3zFXnzuMENTq1vTWrol5nIMpvufNZjQ5GhWDI5YcceMbpY6AC4H0BJAEwApInKT43FKqWlKqQylVEZqaqrjx0QUZL9oWtGJ8d790/elz9nTGp+3DUh3+3mv5rVxYftUZHgYadIqNcVwdqpWer1qXj9jrDDjaYcC2K2UylVKlQD4BkDleetAFCXGf7ravp0QL1g1aSgecPOiEgA6NXHO9f3Kdd3x6Z19fb7/PRe1dvlZu4bV8c09A/Hh7X2cPhvTKw2AZShkztRRqFYlHmeL3ed1uTBEC05EEjPWS9oLoJ+IVIOlm2UIAP8XESSioEuMj0ON5EQ8NKwdTpwpxifL92DSyI5enTuwTX00qKHvUultkALXkeM5ADCicyO8PrYn3C1E9NI13TGicyMM7dgQABAfF4ejDkm8uqTVxMYD+QAsvwE85OYFbKwyo898hYjMArAGQCmALADTAr0uEQVPcmJFHpX7hrTBwVOFXs8ONQrKj17awa96DOnYAFUS3HcQiAiGd25k31+375TTMamaUSuTL+vsV12inSmdSkqpJ5VSHZRSXZRSNyulgpOwl4j88vOGiuXKhnbUd0E0qJGM6bed53L2pqt8LtpMhJ76uV25JsO7HyCelJQpVEmIw5wHBplyvWhUud4QEFVS93y2RrPn/eLKALDxqUtQL6UKAOgmAz0b5EyEOVNH6YY0OtIm1CopK0ePZrVDtp5nJGIwJ6pkujet5fM539xjGdPwfxe0spfV9iGPSzBstc48zT6UjxW7T6BKJRu94qhyPz1RjPrnrHV4xcV0+XsvMp6y706LeinImToKKZoul2pJvi+S/Pldvo+CccW2NNylr/0JANh+NHKWhgsHBnOiGPRl5n689ut2HDxVCADo0KgGAOC9WzIQ527oiA/8aQkPaON+Kr83bCkH6jtM1TdzMehoZMbQRCKKUAOm/qbbH9qpoWnXFhG8c1MvdEnzvdsmEM9e0QW3f7gKo99YjHdu6h3Se0cyBnOiGOPrepmBGNHFOGVuMGmHVWonQlV27GYhijGzVu8LdxXc6tm8NoCKxS185SFrQKXFYE4UYz5YkhPuKrg1yNpvPkIzEcgXh/IKDcun35rhd51iAbtZiGKMbbGISDXh4rZIr5+CEV38C+aNahrnXB/S0bz3AdGIwZyokghkzU0zVUmIc5ke1xv9WsXe+p1mYDcLUSXx9OWxkbPEU6rdyorBnChG3dhXvwiMbcx5LEpOZCjj3wBRjJp8WWcsevgi1KlmmXYfSzm+HX9QrZxkvNRdZcI+c6IYVCUhDonxcWherxrWPDEM+UWlLrMiRqMpV3ZFn5Z18cAXazHj7n6omRw7z+YvBnOiGPH27zvwwi9bAQDFpeX2chGJqUBuc3mPNFzeIy3c1YgY7GYhihG2QE6VE4M5UQwa26e554MopjCYE8WAcod8LF3SKu8iDZWVKcFcRGqLyCwR2SIi2SLS34zrEpF3FmQf0e3nFlTudLCVkVkt89cA/KKU6gCgO4Bsk65LRAZ25Z7Gv77fiPJyhdV7TmDcJ/rsga8u2B6mmlG4BDyaRURqArgAwG0AoJQqBlAc6HWJyLWLX/oDAHDdec3wybI9Tp+PH9w61FWiMDOjZd4KQC6AD0QkS0TeE5EUx4NEZJyIZIpIZm5urgm3JaJRry9GflGpU/nFHWJnghB5x4xgngCgF4D/KqV6AjgDYKLjQUqpaUqpDKVURmpqqgm3JYoeh/IKcfy0537sopIy/P3LdTiSX6QrLy4tx4hXF2Hx9mNO5/y25ahTWfO61fyvLEUlMyYN7QewXym1wro/CwbBnKgy6/+8Zfm2nKmj3B43b/MRfL1mP4pKy/DWDb3s5QdPFWLL4QLcNH0FLmrvuTGUGM9kVJVNwC1zpdRhAPtEpL21aAiAzYFelyhW7Dtx1r596qz710n3z8gCADiG4jhNpsCFWz13UyYmcNRxZWPWN34fgM9EZD2AHgCeM+m6RFFv86F8+/aDM9d6dY42zeu2IwWY9N0Gr+838dIOzFVSCZmSm0UptRZA5V6ziSqt8nKFuDjX3RraRSFyjp3x6pqni0rs28NfWeTVOSM6N8Lfh7dD24Y1vDqeYgt/FyMKQH5RCVo99jOmLdrp8pi3Fu6wb+ccP4tduaedjvlgyW6MfuNP+76tK+WAmxzkPZrV1u2/c3NvBvJKjMGcKADfZx0AADz38xaXx8zZeFi3bxsjrvXUj5ux8UC+U/nAqb8FWEOqLBjMiQLQol7FlIq7PsoM6b3jNV071ZOYzbqyYzAn8sOZc6WYv/kItOmtFmQfwXt/7tIdd7bYeUIPAJw4UzGqRTvaRWvLYeeWus2UK7vgih5N7Pv/u4WvrCo7BnMiPzzz02bc/XEmbn1/pa782dn6tERni8sMz7/9g4rz8gpLDI+Zv+mIYTkA3Ni3BUo1mRLdBX6qHPi7GZEf3L2Y1HJMTWuzbn+effvjZTmGx7w0f5tuf9HDFyG/qAQlZZZVhBI03SxlLu5DlQeDOZGXco6dwYUv/g5AP9zQnY9cBOoBrevZt7/M3O/53gYzR6/JaIYnvt8EAOjUhPnLKzsGcyIvLdxakQPFm5b54bwivLXQMmQxrXZV3Tliwmz75MR47JhyKZbsPI4BresHfkGKauwzJ/LSUz/6lqVC+5JzQOt62PncSHuLvqTMuVukQyPfx4gnxMdhcDsmriMGcyJTvHpdD6cybet7SMcGiI8T/N/gVgCAGtahhEpVBPXv7h2I2wemo2taLd11vr93oPkVppjDYE5kgit6puGhoe0AVLz01Abzwe0s+cVv6Z+Ojo1r2nOvnCsttx+TnBiPJ//SGd/eM0B3baazJW8wmBN5QduCdrTuyeEAgMQES4DeYZ2urx1hkqTJYpicGIdzpZYhi+dKKoK5TUK8/p9l1SrxftaaKhMG8yillEJpmXMgoODQ9n87qlXVkqHQFu9tibFKNf3i2kRcyQnxKLSOP3c1xnxU18YAgJ/uOx/JiQzm5BlHs0SpZ2dnY/ri3dj13Ei3GfvIHJ8u32tYvuu5kfbtYk2XSWlZOa7671IAwF+6N9Gds2zXcQDAun2nUGCw5BsAvHlDT7yhevK7Ja+xZR6lpi/eDQDIjqGZfzuOFuDJ7zeGuxpOCovL8MqCbU7lH9/RRxdsX/t1u317zH+X2mdoju3TzPC6l7+1BDdNX2H4mYgwkJNP2DKPcvfNyMJvf78w3NUwxdCXLd0T9w9pi3rVk8Jcmwrzs/XT6jc+dQkS4sRt98d6zQzPxHi2mSj4+H9ZlNuV691iB+XlCle8tcTraeiBKCwuQ/rE2UifOBvfrz3g8/m9n10QhFr5b3XOCd1+9aQEw0BuNDwR0E+7d8VV653IWwzmlcT9X2Rh7b5TIcmPfe/na+zbM1ftC/r9gm3rkQKvjruiZ5pheZHDiJVvHIYeAsDVvRnMKTCmBXMRiReRLBH5yaxrkmvp9Xwbe/zT+kNBqomz37ZUTHtfuvO4x+PLy5XTaJGTbkaPhNryXRUt86EdG7o99oPbznMqW7P3pG6/V/M6TseUuxn6SOQNM1vmDwDI9ngU+a2opAzbrK3Erk1r+32dnQbLlgWTuyGUZeUKrR77Gb2ema8rzz19LtjV8su7N/d2+7l28WYbo+DtqITDTClApgRzEWkKYBSA98y4Hhn725drMfyVRTh9rhQlpf7/4/98hfEwu2BpM2kOjhYUGX72msEoEQCI1IEcnurVrWktp7KaVT2PM+jXsp7HY4jcMatl/iqAfwJwGWFEZJyIZIpIZm5urkm3rVyWWbssco6dwbr9p/y+zvTFu3H6XGlIc2D3mfKrYbmrbhhXizponTlXitPnjMdpB4t4SHfYuYlzMDcq69Oyrm6fwxApUAEHcxEZDeCoUmq1u+OUUtOUUhlKqYzUVGZ580ecNZCMfmMxDuUZt3SNHMpzHsHS5cm5mPTtBtPqZuNqCTRXMvecNCy/7M0lLs+Z8PkazNlwCJ2fnIsuT8716X7BVqdaolfHrdx9wvNBRD4wY5z5QACXichIAMkAaorIp0qpm0y4Nmkc9/OloONoCpsvVu3D1Ku6BVIlJ66Ceb9WdQ3LvTVz1V689+duHM4vQkFRache6Lqbxm9E23If0bkRbujb3PC4KglxuhmjRIEKOJgrpR4F8CgAiMiFAP7BQB49SsvKnRI7BeLjZXsMy21ZAwFLq/TjZTl4Y2xPr6/7yNfe/RZxrrQMSQnm5TJZqBmZ46s3b+jp8u928SMXYdofu9AyNcXrVYuI3OE48xjgLqMfYBkF44q3QdJbv2w6bFiubYVe++4y/LT+kN+/abjS+5n5aP/4L9ju5bhwb6Qk+f6D4fkxXdGhUQ23PyQb1EjG46M74ca+LXBh+wYujyPylqnBXCn1u1JqtJnXJM/Oefh1/bC1f/3tG3s5fbZ2n3GftdlsQ++0wxRHv77Yvv3M5Z3t2/Fxgu4Go0Lcycw5Yf/h8PCs9YFUVWfXMe9m2GqN7dMcvzx4gWl1IPIGW+ZRJLWGcb4ST1P6b/9wFQBgweYjTp8dzQ/OeG7HSU3FZeUoLStHm0lz7GWH8yte4t7cP92+XVausG5/nv0FradRNz9vOISXNSvZ57tIK+tVvSfORsazljHvWw7n44Vftto/u6V/C7+vSxRsDOZRpFHNZMPyWas9r+4OAK0bVHcqKwjS0L53b87AmF4V09unLdrl1TDCns1r27c/s46H9xSc7/lsjW7xh8a1jf+e3Nl6uADpE2cDAI6dtrTwR7z6p/3zG/s2x9OXd/H5ukShwmAeRTYcyNPtd25SEwDw/pLdOJpfZF+9pri03HBUyd2DWgW/klatUlPw4tXdcf+QtvayHk/Pd3MG8MfDF+LTO/s6lf/3j50e77dwa8Xchaa1fV9m7ZJXF+n2DzsM/byqd1Ofr0kUSgzmEe7gqUJc9+4ynDqrf1nYKjUFo7tVLHrQ57lf0f7xX7Dn+Bm0e3wOBr2wEGeLLS3h6zKaoXGtZFRJCO7X/V1WRYbEhDhLPu6/DWvn8Tzb4g0t6qUgJSkBVR0yEk5btEu3b1sM2ZWZmb4l9zJ6gfzSvK26/e4BpE8gCgUG8wj33993YsXuE/bFKAAgo0Ud/PfG3mhax3lI24LsiqF0tvUli8vK7Tm1f5gQvJXeH5y51r7taaak1kvXdNft20519Y7g/Lb1fa6bO8UGeVG+cui6iucMTYpwDOYhUFpWbl+x3Ve2wKYdsTLrrwPQvlENXNqlkdPxz/y02b5dUlYOpRRmbzhkzynSNc15lMiny43HhvtrwkVtfDre8TcGW0M5t+Cc0/uAT+/sizkbjYc/+quomJN3KPoxmIdAm0lzcOdHq/w617ZiTW6B86gTT5N9CkvK8Of2YyguLUfOcUsfulGL+fHv/F+qbeOBPJwtLtWNI39/yW43Z+g1MGh9l2m6Pf7x1Tr79t2DWprWKr/wPwuRPnE2pi/ejRW7PafpJYp0DOYhsnBrLjJzfM/HsXbfKQDAt1nGK/asnzzc5bkvzdvmdsKQN/7yxmJ86CI4nzhTjNFvLMbTP27WdVU4Jska1a2xbl+bv8SoN8bVNPcHh1r6322jZDJaeE4t64rth9szP23G3E3OQzaJog2DeZBpu1eufmeZ39exdam8c5M+n3bNZNeJnX5Yd9D+w0Br/eTh+Pqv/XUvGl3lHN9wIA+Tf9ysK9t+pAB9pizAcWvO8fmbj+D9xa5b42/doJ+sdNegVtg+5VK0qFcNz17R1eV5jlKsLz6rW/8c3a0x3rqhF+4e1NLpWF/yg3eyjgoCgDYGwzcfGdHB62sRhQuDeZC5ygroK1s/sdHiB69d38PleW//7jysr2ZyInq3qIvsZ0bYy9pMmoMdR/XT4F0F+GGvLMLRgnP2ceDHzxTrJu1s0VzX5pM7++Cr8f3x3i0ZuOfC1kiMj8MfD1+EYZ2cV+7p29J9Ui5b+tjeLepiVLfGeGBoOwzv1FD3cvfASX2myKMFRUifOBsX/mchXpyrH6mSZV0JaEiHBpj/0AWoVqXih9yKx4Zg/ODQDekk8heDeZAZBcR9J85iweYjmG8wI9N2zsFThfYgo1WrqnNL/PIeaciZOspwYQSbh4Z6HiI49GX9WGvtEL+lO48B0Od5qZdSxekaH9/Rx3Cx40FtU3Feel0M7dTQ40iXz+5yHmuuNbpbE6z713B0tT5v9aQETLslA900wwdttygrV5j8wyZk7T0FwNK98ubCHbrr2TIwPjG6E0RE102UWj3Jp5E5ROFiRgpccsNobPegFxbat3c/P9IpWEz4PMtlwqrhBi1ZG9vLUiOD2/ueQ37StxUvRjcfzEdJmcKt76+0lxkt7WYUyH3lTRbHWi7yhr8+tifun5Fl72Z5/LsNmLFyHz5cmuPxmkbfFReNoGjBYB5knnJWz1q9H9dk6FdmdxXIASAp0b9fpgKNSc/Odl7e1WiVoHAvTGyb1m/L4T5jpfcTiGok858DRS92swTZybP6vCJH8vXTxB+etR57jnufmc/fXN0C42jesKZ+aKCndLpaO446LwydF0CSK62Z4/rp9r3t6bD9ZmBLbeCL6h5mlhJFMgbzIHvyB/0Y7ikGLVxPKWy1kt20zG0pbv81upPTZ66C4ez7B+n2vVl70x2zFiZulVoxquS2Aen44x8XeXWerWXuzw8V9o1TNGMwD6LtRwrsGfhsflh30Om437cexfuLd+P6acuglMIgNxNjqrjpTx7ZtTFypo7CHec7D9XrrBl+p1W/ur5lbvvNwd8Zq676sn2lnco/+bLOaF7Pu+RZp4ss+Wh+3+r/ouFX9kxDRos6yJk6yu9rEIUag3kQDXulYnSIu9wez/28BU//tBnLd51AcVk5/tx+TPf5HQMrgrO3rcef7jvfvp0zdZTb87KfrhhKuNeabdGXKfNfje/v9bG+yHpiGLY9e6lP59jS3368bA8KPfyWkfXEMMPyV67rgVl/HeDTfYnCjcE8RG61Lr5gNMFFy7F7YMqVXdC4lu/5uRv5cE5Vzbjqb9ZYZpq+MHeL4bFj+zRzWqQ4IUgjPuqkVPE506N28epXFmxzc6Tl+kSxIuBgLiLNRGShiGSLyCYRecCMisWaxHhLwPPUJ71qd8XY8lWThuLGvi2g4HuXh69DBAe0tvR1D+loWY9yz/GKfOi9NdPmu6TVwnNX6mdttqiXgpnj+mFuBCyVph37vsxgtI3NJZ0tQzyvYZ5yihFmvL4vBfB3pdQaEakBYLWIzFdKbfZ0YmViGwNumzXpyr2fr7FvO6aAvW1Autf3S/axRXv3oFZYuvM4mtfV9003rpWMm/o1x2rrTNb8QufVguqmVEHfVua8+AyUY/1dedGadvffV3VDpyY1cYtm2TqiaBRwy1wpdUgptca6XQAgG0Ca+7MCk1dYgoIic4bABcvCLUd1+4UOCa8uaOd+Eo/28/aNLC8vuxikr3XFNvHm2gzvWp627gzHcfGNaiXrhjVuOeycTiCSaCf5uHtPUcOa0yYuTnD7wJbMV05Rz9SBtSKSDqAngBUGn40DMA4Amjdv7vixT7o/NQ+A8ezJSGFbRNkmrU5VXdKrq3s3xaJtrkdcjOpakat8cLtUzH/oAsMkUO5seWaE29EvWrbFK5bsOKZrZV+X0QwjujQCZlr2k/0c5x5Kn97ZFzdNX+E0UmfxIxdhV+4ZHDhV6OJMouhl2gtQEakO4GsADyqlnJpvSqlpSqkMpVRGaqrvU8uNGA3zi1SOXSS2PmpXnvxhk26/bcMaPv/gSk6M93o6+jHr1PzXf9uBVxdsQ4r1peh15zXT9b8PaKOv94sOqwRFgu7NLL/BLMi25L65rHsTLPzHhWhapxouaJeKsX0Ca0wQRSJTgrmIJMISyD9TSn1jxjWNnCst07Vmsw8VuDk6cjStUxUt66foyupXT9KNY3bMFKgdlREK2gRery7YjjPFZRjTM83pB4ht3dExPS09aUOtL0wjieMs2bopVZz+/olijRmjWQTAdADZSqmXA6+Sa+0f/wW3aBI9vePFqu3hNn5wa3x370DDDINaj47sqNsP9SiLgW2cJyodzKvojrClqrX1Lb9wdTcse/Ri1K4WecP7bCOHbNgSp8rAjD7zgQBuBrBBRNZayx5TSv1swrU9Ukr53W9+7PQ5p35Vs13du6lX99Dm0AaAZ67oEqwqeW35roqVkd6+sZduWGVCfBwa13JeUDoSOP7/0K6hb+8aiKKRGaNZFiulRCnVTSnVw/pfSAI5ALy/JMev8zJzTiDj2QV4/dft5lYIwK5cSwKq5MQ4w5eWH93Rx6ksxSHJkxmpZM2UGB9nmEs9Umlz2ETqS3IiM0XVDNDnxzgvMfbndv9ycCzeYZkyr10hxywXv/QHAOd+74/u6IO7B7XEYM2ww8/v6ot+reqiYY0kn2c7kmu2v/v61SOvG4goGKIqehi1DP1NqKSdNp8fojHrg9ulYtIofUbDAW3q44tx/ZEQH4dZ1hwnHRrVCEl9PJl+a0a4qxAwx0RnRLEqqoL54bwizwd56QNN90y3yfNMu24g4qzdAdsN8oSHgmO/fYt6HAFCFC2iKpjfPjA93FUIqmbWqeh3eUjGFSw/TDhft1/HpHS24eTtpCmiaBdVS6uICJ6+vDP+9f0mVKsSbx9dEciIlmBq39C37pJaVRPDmkO7TYPq+PdVXdElrRZqJieiXpBH+gTT53f3xQ3/W4G6zIxIlURUBXMAuKV/Om7pn47ycoVWj1kGzXyz5gCuisDsd7P+Gpw838F03XmxMSY71fqDyJ+Mk0TRKGp/B9VOU//7V+vCWBPXbMmcKPRsuWb8XDCJKOpEbTD3x0/rD2LN3pOeD6Sol1ojCYnxgn9e0j7cVSEKiajrZtFKiBOU+tD0mvB5FgBwbcdKICUpAdunjAx3NYhCJqpb5o4ry7ujVEXQv/G95U6fO+bx9pdtIeTh1lwmREShENXBvL0Pk2vOaYL1kh3Oy4mdKjRncsm8zd4vhExEZJaoDua+2OFiIs751myB2tV0AjH+U8uyb/M2HzHlekRE3oiZYP6jm4UqSsrKse2Ic+7z/1zdDX/p3th+jJnSakdmRkEiik0xE8zvm2F5uZm19ySOW1fNsWk7aQ7+9qXz8MXeLerYh7CZHcxfuLqbqdcjInInZoI5ABSVlOHKt5fiyreX2styC845HWcbot6oVrLLhYwD5WlZOCIiM0V9MK+RXDG6ssMTvwAA9p44ay87b8oCp3NsoxmrxMfZc3dc+fZS5Bacw6/Z/vd1a38gRGJ6ASKKXVEfzFc8NsTvc+PjxN4yP32uFDdPX4E7P8r0u5V+zTtLPR9ERBQEUR/Mq1VxPe9pn6aFbkREsCv3jH1/y2HLS9Jy5d8c8HX78/w6j4goUKYEcxEZISJbRWSHiEw045pmWLn7hMdjjhr0qZv9MpSIKNgCDuYiEg/gLQCXAugEYKyIdHJ/VmhsOOC5pdyghnOa1+xDzsMYfREfx/5yIgotM1rmfQDsUErtUkoVA/gCwOUmXNdvVRPjcceHq/Dh0hxduXa44Ds39QYA9GvlPOrk2neX+XzPU2crZpB2Tavl8/lERIEwI5inAdin2d9vLdMRkXEikikimbm5/q3b6cpnd/XV7TetUxW/bTmqK9v13Ehcm9HMvm9b6LdTk5q6ldz99Y+v1tu3X7u+R8DXIyLyhRnB3KhPwekNolJqmlIqQymVkZqaanCK/wZap+TbGK2hact/3r1pLd0+ACx+5GKn49fuO4WvMvfpyrIP5WPyD5sM67Art+KeXDuTiELNjBS4+wE00+w3BeB6bn2Qje7WGD+tP6Qrsw0/BICSMsvPmcS4irL6BsujXfHWEgDANdbW/LdZ+/HQTMss0nOlZbh9YEu00ywLt+vYGadrEBGFihkt81UA2opISxGpAuB6AD+YcF2/JCXEO5X9MGGgffuFq7thSIcGXmdcLCqxrDNqC+QAMGPlPgx/ZZHh7FIionAIOJgrpUoBTAAwF0A2gC+VUsZ9EUE0tGMDAMDXa/Y7fdahUU37dpe0Wph+23m61ro7D3yR5fKzJTuOYeHWo8g+lG8vs3XjEBGFkikrDSmlfgbwsxnX8te7N2egrFxh3CeZ+H2r7y9Y7x7UEjtzzzi9OJ27yfX0/gdnrnUqe/+283y+NxFRoKJ+BqiNbWr++Y4vQ6dc6tX5k0Z1wls39HL5eZ/0ul5dJyUpqlfiI6IoFTPB3CY5saLPfEyvNHuKW29UreLc326zMsfzbFIA9sRdREShFHORRxtMp44JfU7xOM7+JKIwiLlgvkCTwtbbl5xaj4zo4FT2xHcbA6oTEVGwxVww33Qw3/NBbowf3Mqp7JPlewK6JhFRsMVcML+qd9OAzve0qMTP9w/C9/cONPysaR2u+0lE4RFzwXxUV8sCzW0bVA/K9Ts1qYnuzWobfrb/ZGFQ7klE5EnMBfPEeEvL2miKvrd+mDAQvzw4yO0xabWdW+HPj+nq9z2JiAIRc8G8VWp1PD+mK968oaff1+jWtDY6NKqJoR0bujzmy/H9ncquCbCLh4jIXzEXzAFgbJ/mqBdAy9zmmSs66/b/NbpizY0mtZLt24PbWbJAJnCMORGFCacrutG4lr4rpWPjihwv2hel79zU2+9FoImIzMBg7qWW9VPQu0UdXdmwTg2RX1iCqlXi3c4eJSIKNgZzLy38x4VOZf+7JSP0FSEiMsBOXiKiGMBgTkQUA9jN4sEDQ9qiSe1kzwcSEYURg7kHDw1rF+4qEBF5xG4WIqIYEFAwF5H/iMgWEVkvIt+KSG2T6kVERD4ItGU+H0AXpVQ3ANsAPBp4lYiIyFcBBXOl1DylVKl1dzkAJichIgoDM/vM7wAwx9WHIjJORDJFJDM3N9fE2xIRkcfRLCKyAEAjg48mKaW+tx4zCUApgM9cXUcpNQ3ANADIyMhQftWWiIgMeQzmSqmh7j4XkVsBjAYwRCnFIE1EFAYBjTMXkREAHgEwWCl11pwqERGRrySQxrSI7ACQBOC4tWi5Umq8F+flAvB3leT6AI75eW6k4bNEnlh5DoDPEqkCeZYWSqlUow8CCubhICKZSqmYSFfIZ4k8sfIcAJ8lUgXrWTgDlIgoBjCYExHFgGgM5tPCXQET8VkiT6w8B8BniVRBeZao6zMnIiJn0dgyJyIiBwzmREQxIKqCuYiMEJGtIrJDRCaGuz5GRCRHRDaIyFoRybSW1RWR+SKy3fpnHc3xj1qfZ6uIXKIp7229zg4ReV1EJAR1f19EjorIRk2ZaXUXkSQRmWktXyEi6SF+lskicsD63awVkZGR/iwi0kxEFopItohsEpEHrOVR9724eZao+l5EJFlEVorIOutzPGUtD+93opSKiv8AxAPYCaAVgCoA1gHoFO56GdQzB0B9h7IXAEy0bk8E8G/rdifrcyQBaGl9vnjrZysB9AcgsCQwuzQEdb8AQC8AG4NRdwD3AHjHun09gJkhfpbJAP5hcGzEPguAxgB6WbdrwJJqulM0fi9uniWqvhfrPatbtxMBrADQL9zfSVCDg8l/gf0BzNXsPwrg0XDXy6CeOXAO5lsBNNb8D73V6BkAzLU+Z2MAWzTlYwG8G6L6p0MfAE2ru+0Y63YCLLPgJITP4ipoRPyzaOrwPYBh0fy9GDxL1H4vAKoBWAOgb7i/k2jqZkkDsE+zv99aFmkUgHkislpExlnLGiqlDgGA9c8G1nJXz5Rm3XYsDwcz624/R1ny4OcBqBe0mhubIJaVsd7X/BocFc9i/VW7Jywtwaj+XhyeBYiy70VE4kVkLYCjAOYrpcL+nURTMDfqM47EcZUDlVK9AFwK4F4RucDNsa6eKRqe1Z+6h/u5/gugNYAeAA4BeMlaHvHPIiLVAXwN4EGlVL67Qw3KIv1Zou57UUqVKaV6wLIgTx8R6eLm8JA8RzQF8/0Ammn2mwI4GKa6uKSUOmj98yiAbwH0AXBERBoDgPXPo9bDXT3TfuhXbQrns5pZd/s5IpIAoBaAE0GruQOl1BHrP8JyAP+D5bvR1csqop5FRBJhCX6fKaW+sRZH5fdi9CzR+r1Y634KwO8ARiDM30k0BfNVANqKSEsRqQLLS4EfwlwnHRFJEZEatm0AwwFshKWet1oPuxWWvkJYy6+3vrluCaAtgJXWX9EKRKSf9e32LZpzQs3MumuvdTWA35S1UzAUbP/QrK6E5bux1Ssin8V63+kAspVSL2s+irrvxdWzRNv3IiKpYl28XkSqAhgKYAvC/Z0E+yWHyS8bRsLyBnwnLCsdhb1ODvVrBctb63UANtnqCEtf168Atlv/rKs5Z5L1ebZCM2IFQAYs/1PvBPAmQvNCagYsv+aWwNIyuNPMugNIBvAVgB2wvMVvFeJn+QTABgDrrf9YGkf6swA4H5Zfr9cDWGv9b2Q0fi9uniWqvhcA3QBkWeu7EcC/rOVh/U44nZ+IKAZEUzcLERG5wGBORBQDGMyJiGIAgzkRUQxgMCciigEM5kREMYDBnIgoBvw/JeJ1Ya2wWS0AAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(trajs[20,:,0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f18e3a84-9254-49c5-b163-dd22058773a8",
   "metadata": {},
   "source": [
    "ok, this seems to be working."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ac8fe13-a4bc-4a94-be71-799a44edd94e",
   "metadata": {},
   "source": [
    "## write a deterministic, learnable sampler\n",
    "the first implementation will be a singleton, 3-step process which transforms momentum and position with a jacobian of 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "id": "b75333ad-7c2f-4d38-a3f8-5d8ab44683d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Sequence, Callable\n",
    "import flax.linen as nn\n",
    "import flax\n",
    "nnParams = flax.core.frozen_dict.FrozenDict\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "id": "ac9124f8-300f-4f36-945e-bf9628b4ced6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 150,
   "id": "63b140cd-3708-46bf-930b-ec1164a9aece",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_mod_leapfrog(seed : Seed, \n",
    "                      features: Sequence) -> Tuple[nnParams, Callable[[Conf, nnParams], Conf]]:\n",
    "    \"\"\"\n",
    "    make a modified leapfrog function\n",
    "    \"\"\"\n",
    "    _dim = features[0] # features of positions x\n",
    "    model = tanhMLP(features=features) # define the model\n",
    "    test_x = jnp.ones(_dim) # make a test x \n",
    "    params = model.init(seed, test_x) # initialize the nnParams (FrozenDict)\n",
    "    test_y = model.apply(params, test_x) # apply the model the first time.\n",
    "    #print(f\"input of shape {test_x.shape} returned out of shape {test_y.shape}\")\n",
    "    \n",
    "    def leapfrog(X_in : Conf, dt: float, in_params: nnParams) -> Conf:\n",
    "        x_1, v_1 = X_in[:_dim], X_in[_dim:] #pull apart x, v\n",
    "        v_one_and_half = v_1 + dt * model.apply(in_params, x_1)\n",
    "        x_2 = x_1 + dt * v_one_and_half\n",
    "        v_2 = v_one_and_half + dt * model.apply(in_params, x_2)\n",
    "        return jnp.concatenate([x_2, v_2])\n",
    "    \n",
    "    return params, leapfrog\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65035b9e-bb07-4774-9e7c-ddc4f1aa7a52",
   "metadata": {},
   "source": [
    "can we test this and make sure the jacobian is actually 1?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 151,
   "id": "362a74db-f90e-4121-9068-9fa224b77ebe",
   "metadata": {},
   "outputs": [],
   "source": [
    "nn_params, deterministic_leapfrog = make_mod_leapfrog(random.PRNGKey(3462), [2,2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 154,
   "id": "90c47c43-caf4-43dd-988a-4551e9339eb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "X = random.normal(random.PRNGKey(2462), shape=(4,))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 155,
   "id": "6c88f983-781a-4977-834a-510fa016c5fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_out = deterministic_leapfrog(X, 1e-1, nn_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 157,
   "id": "c66bde8f-afe9-41ce-be5f-77eda9165661",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([-0.96499554,  0.59350874, -1.38793319,  0.55717537], dtype=float64)"
      ]
     },
     "execution_count": 157,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 156,
   "id": "fe9e1421-0112-4b46-861f-db699657f20c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([-1.1180315 ,  0.64786676, -1.68356186,  0.52882456], dtype=float64)"
      ]
     },
     "execution_count": 156,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 158,
   "id": "2060a535-7557-4bcd-ba79-77bbf46a2ba7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax import jacfwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 159,
   "id": "00db1b97-abdf-4f31-ab4d-cde503d59f4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "J = jacfwd(deterministic_leapfrog)(X, 1e-1, nn_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 162,
   "id": "9a021325-dbc6-4dbf-ab3a-8248c82544be",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([[1.00829751e+00, 9.24930722e-04, 1.00000000e-01,\n",
       "              0.00000000e+00],\n",
       "             [9.55510605e-04, 1.00028251e+00, 0.00000000e+00,\n",
       "              1.00000000e-01],\n",
       "             [1.48013949e-01, 1.72562925e-02, 1.00644961e+00,\n",
       "              7.94508611e-04],\n",
       "             [1.70953111e-02, 5.11553162e-03, 7.47599173e-04,\n",
       "              1.00022829e+00]], dtype=float64)"
      ]
     },
     "execution_count": 162,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "J"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 164,
   "id": "2d711d35-b2ba-47ce-b52f-9d4e3f577021",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f23aebfb970>"
      ]
     },
     "execution_count": 164,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAQcAAAD8CAYAAAB6iWHJAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8rg+JYAAAACXBIWXMAAAsTAAALEwEAmpwYAAANN0lEQVR4nO3dcchd9X3H8fdnWUolOrIRnWmSaWFh0hWm7iFVhOFcHRqE9A8p8Y8qMnio2GJh/lE2cOy//VWYWHSBSg2UdoKtC1264jqHCnU1SZNMzdyCKxgMi7OaGI0r0e/+uEd5ePw9Scw999z7JO8XXJ5z7vnlfn+X5Pnk3HPOPd9UFZK02K9NewKSZpPhIKnJcJDUZDhIajIcJDUZDpKafn2cP5zkt4C/By4HfgF8sareaIz7BfAW8B5wsqrmxqkrafLG3XP4OvCTqtoI/KRbX8ofV9WVBoO0PIwbDluAR7rlR4AvjPl6kmZExrlCMsmbVbV6wfobVfWbjXH/DbwBFPB3VbXtFK85D8wDrFq16g+vuOKKs57frPr57t3TnsLEXDTtCUzIW9OewIS8D1RVWttOGw5J/hm4tLHpL4FHzjAcPlVVrya5BHgC+GpVPXW6ic/NzdWuXbtON2zZWZXm38U54YZpT2BC/mXaE5iQd4H3lgiH0x6QrKrPL7Utyf8kWVtVh5OsBY4s8Rqvdj+PJPkBsAk4bThImp5xjznsAO7olu8A/mHxgCSrklz0wTLwp8DzY9aVNGHjhsPfADcm+S/gxm6dJJ9KsrMb89vAM0n2AT8D/rGq/mnMupImbKzrHKrqdeBPGs+/Cmzull8G/mCcOpKG5xWSkpoMB0lNhoOkJsNBUpPhIKnJcJDUZDhIajIcJDUZDpKaDAdJTYaDpCbDQVKT4SCpyXCQ1GQ4SGoyHCQ1GQ6SmgwHSU29hEOSm5K8lORgko90vcrI/d32/Umu7qOupMkZOxySrAC+CdwMfAa4LclnFg27GdjYPeaBB8etK2my+thz2AQcrKqXq+pXwPcYtclbaAuwvUaeBVZ3fS4kzag+wmEd8MqC9UPdcx93jKQZ0kc4tFppLe6xdyZjRgOT+SS7kux67bXXxp6cpLPTRzgcAjYsWF8PvHoWYwCoqm1VNVdVcxdffHEP05N0NvoIh+eAjUk+neQTwFZGbfIW2gHc3p21uAY4WlWHe6gtaULG6ngFUFUnk3wF+DGwAni4ql5I8uVu+0PATkYdsA4C7wB3jltX0mSNHQ4AVbWTUQAsfO6hBcsF3N1HLUnD8ApJSU2Gg6Qmw0FSk+EgqclwkNRkOEhqMhwkNRkOkpoMB0lNhoOkJsNBUpPhIKnJcJDUZDhIajIcJDUZDpKaDAdJTYaDpCbDQVLTUL0yr09yNMne7nFfH3UlTc7YN5hd0CvzRkb9KZ5LsqOqXlw09OmqumXcepKG0cfdpz/slQmQ5INemYvD4WP7+e7drEqrWdby9nadmPYUJmZVLpj2FCbihmlPYEKePsW2oXplAlybZF+SHyX5/aVebGE7vGa/PEmD6GPP4Uz6YO4BLquq40k2A48DG1svVlXbgG0AKxLzQZqSQXplVtWxqjreLe8EViZZ00NtSRMySK/MJJcmo4MHSTZ1dV/vobakCRmqV+atwF1JTgIngK1dizxJMyqz/Du6IqlPTnsSE+DZiuXnXD5b8WZV85SgV0hKajIcJDUZDpKaDAdJTYaDpCbDQVKT4SCpyXCQ1GQ4SGoyHCQ1GQ6SmgwHSU2Gg6Qmw0FSk+EgqclwkNRkOEhqMhwkNfXVDu/hJEeSPL/E9iS5v2uXtz/J1X3UlTQ5fe05fBu46RTbb2bUp2IjMA882FNdSRPSSzhU1VPAL08xZAuwvUaeBVYnWdtHbUmTMdQxhzNtmWc7PGlG9NEO70ycScu80ZO2w5NmwlB7DqdtmSdptgwVDjuA27uzFtcAR6vq8EC1JZ2FXj5WJPkucD2wJskh4K+AlfBhO7ydwGbgIPAOcGcfdSVNTi/hUFW3nWZ7AXf3UUvSMLxCUlKT4SCpyXCQ1GQ4SGoyHCQ1GQ6SmgwHSU2Gg6Qmw0FSk+EgqclwkNRkOEhqMhwkNRkOkpoMB0lNhoOkJsNBUpPhIKlpqHZ41yc5mmRv97ivj7qSJqevvhXfBh4Atp9izNNVdUtP9SRN2FDt8CQtM0N1vAK4Nsk+Rs1s7q2qF1qDkswzarbLhcAdw81vMJfkgmlPYWLernOzSdmqtJq2LX/vnmLbUOGwB7isqo4n2Qw8zqjj9kcsbId3ie3wpKkZ5GxFVR2rquPd8k5gZZI1Q9SWdHYGCYcklyaj/bIkm7q6rw9RW9LZGaod3q3AXUlOAieArV0XLEkzKrP8O3pJUl+c9iQm4NFpT2CCjszwv6dxnMsHJN+rar45r5CU1GQ4SGoyHCQ1GQ6SmgwHSU2Gg6Qmw0FSk+EgqclwkNRkOEhqMhwkNRkOkpoMB0lNhoOkJsNBUpPhIKnJcJDUZDhIaho7HJJsSPJkkgNJXkhyT2NMktyf5GCS/UmuHreupMnq4wazJ4E/r6o9SS4Cdid5oqpeXDDmZkZ9KjYCnwMe7H5KmlFj7zlU1eGq2tMtvwUcANYtGrYF2F4jzwKrk6wdt7akyen1mEOSy4GrgH9btGkd8MqC9UN8NEA+eI35JLuS7DrR5+QkfSy9hUOSC4HHgK9V1bHFmxt/pHkP86raVlVzVTV37naUlGZfL+GQZCWjYPhOVX2/MeQQsGHB+npGDXUlzag+zlYE+BZwoKq+scSwHcDt3VmLa4CjVXV43NqSJqePsxXXAV8C/j3J3u65vwB+Bz5sh7cT2AwcBN4B7uyhrqQJGjscquoZ2scUFo4p4O5xa0kajldISmoyHCQ1GQ6SmgwHSU2Gg6Qmw0FSk+EgqclwkNRkOEhqMhwkNRkOkpoMB0lNhoOkJsNBUpPhIKnJcJDUZDhIajIcJDUN1Q7v+iRHk+ztHveNW1fSZA3VDg/g6aq6pYd6kgYwVDs8SctMH3sOHzpFOzyAa5PsY9TM5t6qemGJ15gH5mGUXI/2OcEZ8e60JzBBq3LKG5EvW2/XudmccW7uuiW39RYOp2mHtwe4rKqOJ9kMPM6o4/ZHVNU2YBvAyqTZMk/S5A3SDq+qjlXV8W55J7AyyZo+akuajEHa4SW5tBtHkk1d3dfHrS1pcoZqh3crcFeSk8AJYGvXBUvSjBqqHd4DwAPj1pI0HK+QlNRkOEhqMhwkNRkOkpoMB0lNhoOkJsNBUpPhIKnJcJDUZDhIajIcJDUZDpKaDAdJTYaDpCbDQVKT4SCpyXCQ1GQ4SGrq4wazn0zysyT7unZ4f90YkyT3JzmYZH+Sq8etK2my+rjB7P8BN3Q9KVYCzyT5UVU9u2DMzYz6VGwEPgc82P2UNKP6aIdXH/SkAFZ2j8V3lt4CbO/GPgusTrJ23NqSJqevpjYrutvSHwGeqKrF7fDWAa8sWD+E/TSlmdZLOFTVe1V1JbAe2JTks4uGtG5d3+xbkWQ+ya4ku97vY3KSzkqvZyuq6k3gX4GbFm06BGxYsL6eUUPd1mtsq6q5qprzVIo0PX2crbg4yepu+QLg88B/LBq2A7i9O2txDXC0qg6PW1vS5PRxtmIt8EiSFYzC5tGq+mGSL8OH7fB2ApuBg8A7wJ091JU0QX20w9sPXNV4/qEFywXcPW4tScPxY72kJsNBUpPhIKnJcJDUZDhIajIcJDUZDpKaDAdJTYaDpCbDQVKT4SCpyXCQ1GQ4SGoyHCQ1GQ6SmgwHSU2Gg6Qmw0FSk+EgqWmoXpnXJzmaZG/3uG/cupIma6hemQBPV9UtPdSTNIA+7j5dwOl6ZUpaZvrYc6DrWbEb+F3gm41emQDXJtnHqNPVvVX1whKvNQ/Md6vHX4OX+pjjGVgD/O9AtYbk++rBqF/TYIZ8b5cttSGj//j70XW++gHw1ap6fsHzvwG833302Az8bVVt7K1wD5Lsqqq5ac+jb76v5WdW3tsgvTKr6lhVHe+WdwIrk6zps7akfg3SKzPJpUnSLW/q6r4+bm1JkzNUr8xbgbuSnAROAFurz88z/dg27QlMiO9r+ZmJ99brMQdJ5w6vkJTUZDhIajrvwyHJTUleSnIwydenPZ++JHk4yZEkz59+9PKRZEOSJ5Mc6C7Xv2fac+rDmXwNYfA5nc/HHLqDqP8J3AgcAp4DbquqF6c6sR4k+SNGV65ur6rPTns+fUmyFlhbVXuSXMTo4rsvLPe/s+5s3qqFX0MA7ml8DWEw5/uewybgYFW9XFW/Ar4HbJnynHpRVU8Bv5z2PPpWVYerak+3/BZwAFg33VmNr0Zm6msI53s4rANeWbB+iHPgH9r5IsnlwFVA63L9ZSfJiiR7gSPAE0t8DWEw53s4pPHc+fs5axlJciHwGPC1qjo27fn0oareq6orgfXApiRT/Th4vofDIWDDgvX1jL4YphnWfSZ/DPhOVX1/2vPp21JfQxja+R4OzwEbk3w6ySeArcCOKc9Jp9AduPsWcKCqvjHt+fTlTL6GMLTzOhyq6iTwFeDHjA5sPbrUV8mXmyTfBX4K/F6SQ0n+bNpz6sl1wJeAGxbcWWzztCfVg7XAk0n2M/pP64mq+uE0J3Ren8qUtLTzes9B0tIMB0lNhoOkJsNBUpPhIKnJcJDUZDhIavp/bOAKhKI3ezQAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.imshow(J, cmap='hot')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 165,
   "id": "635711e0-a455-4866-9d7a-6a0058f7a360",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(1., dtype=float64)"
      ]
     },
     "execution_count": 165,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jnp.linalg.det(J)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7e17eb0-41f5-4227-8db0-95e44838b5cc",
   "metadata": {},
   "source": [
    "neat-o..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0ff2eb0-eb45-4001-a7d9-9db9f430f895",
   "metadata": {},
   "source": [
    "it should also be the case that if we define \n",
    "$$\n",
    "s(x, v) = (x, -v)\n",
    "$$\n",
    "it should be the case that \n",
    "\n",
    "$$\n",
    "s \\circ \\Phi \\circ s = \\Phi^{-1}\n",
    "$$\n",
    "let's see if we can numerically show this..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "id": "e1bb3f0d-9214-435d-a567-b26ec2cd8df5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def s(X: Conf) -> Conf:\n",
    "    _dim = X.shape[0] // 2\n",
    "    return jnp.concatenate([X[:_dim], -X[_dim:]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 178,
   "id": "e48679a6-94cf-420f-8dfb-36bf28c784df",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([-0.96499554,  0.59350874, -1.38793319,  0.55717537], dtype=float64)"
      ]
     },
     "execution_count": 178,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b41fbcd-3caa-4640-8a76-fccd25dcffba",
   "metadata": {},
   "source": [
    "specifically, \n",
    "$$\n",
    "\\Phi^{-1} \\circ \\Phi(X) = s \\circ \\Phi \\circ s \\circ \\Phi(X)\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 180,
   "id": "b1abd634-abdd-47ad-8e74-830572d87fd8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray([-0.96499554,  0.59350874, -1.38793319,  0.55717537], dtype=float64)"
      ]
     },
     "execution_count": 180,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "s(\n",
    "    deterministic_leapfrog(s(deterministic_leapfrog(X, 1e-1, nn_params)), 1e-1, nn_params)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dd76efe-ad4c-400d-b97e-8b13447d9b94",
   "metadata": {},
   "source": [
    "and that does it. we can now define the inverse map!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0158b56-8de7-4c7a-b856-0e62917bed0a",
   "metadata": {},
   "source": [
    "## Annealed Flow Transport\n",
    "Let's now make a modified AIS sampler that will run annealed flow transport...\n",
    "But first, let's try to build _just_ the trainable block where we do a deterministic map between two distributions...the stochastic kernel is already build and just requires us to append it on the end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 491,
   "id": "f818a426-af95-441e-a8aa-2ea97a3602bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_leapfrog_NF(_t0_sampler: Callable[Seed, Conf], \n",
    "                     _t0_energy_fn: Callable[[Conf, Params], Conf],\n",
    "                     _t1_energy_fn: Callable[[Conf, Params], Conf],\n",
    "                     dt : float,\n",
    "                     features: Optional[Sequence]=[2,2],\n",
    "                     seed: Optional[Seed]=random.PRNGKey(13)\n",
    "                    ) -> Tuple[nnParams, Callable[[Conf, nnParams, Seed, float], Conf]]:\n",
    "    \n",
    "    params, leapfrog_nf = make_mod_leapfrog(seed, features)\n",
    "    _dim = features[0]\n",
    "    \n",
    "    def Importance_Sampler(seed: Seed, \n",
    "                           nn_params_as_tuple: Tuple[nnParams], \n",
    "                           t0_energy_params: Params, \n",
    "                           t1_energy_params: Params) -> Tuple[Conf, Conf, float]:\n",
    "        X_in = _t0_sampler(seed)\n",
    "        X_out = leapfrog_nf(X_in, dt, nn_params_as_tuple[0])\n",
    "        \n",
    "        t0_u, t0_k = _t0_energy_fn(X_in[:_dim], t0_energy_params), kinetic_energy(1., X_in[_dim:])\n",
    "        t1_u, t1_k = _t1_energy_fn(X_out[:_dim], t1_energy_params), kinetic_energy(1., X_out[_dim:])\n",
    "        \n",
    "        work = t1_u + t1_k - t0_u - t0_k\n",
    "        \n",
    "        return X_in, X_out, work\n",
    "    \n",
    "    return params, Importance_Sampler  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 492,
   "id": "c560fb86-224d-404a-a4b9-647093394d77",
   "metadata": {},
   "outputs": [],
   "source": [
    "is_nn_params, isampler = make_leapfrog_NF(_t0_sampler, _energy_fn, _energy_fn, 1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 493,
   "id": "ed33f9d6-0d8f-4e50-8ffa-7df2c01114e8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(DeviceArray([ 0.96907176,  0.51501533, -1.43994373, -1.6511834 ], dtype=float64),\n",
       " DeviceArray([-0.1030787 , -1.23066942, -0.41721773, -2.18985484], dtype=float64),\n",
       " DeviceArray(2.57900905, dtype=float64))"
      ]
     },
     "execution_count": 493,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "isampler(random.PRNGKey(34155), (is_nn_params,), jnp.zeros(2), jnp.ones(2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 494,
   "id": "645e3901-146d-4c04-8024-b07f3ef6577e",
   "metadata": {},
   "outputs": [],
   "source": [
    "visampler = vmap(partial(isampler, t0_energy_params = jnp.zeros(2), t1_energy_params = jnp.ones(2)), in_axes = (0,None))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 495,
   "id": "5f3c9e32-f350-4899-a74c-58a5da041eed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(DeviceArray([[-0.13315785,  1.01386513, -0.6212633 ,  0.60865124],\n",
       "              [ 0.14958146,  0.43117586,  0.05241661,  0.05378837],\n",
       "              [-0.20707079, -0.55351003,  0.70936913, -0.66941612],\n",
       "              [-0.86510991, -1.36719591, -1.32867807, -0.09460279],\n",
       "              [ 0.14172441, -0.01347149, -0.84180898, -0.1111099 ],\n",
       "              [ 0.53151389, -0.90371657,  0.46633612,  0.28262607],\n",
       "              [ 0.574772  ,  1.10342848,  0.16865889,  0.83267297],\n",
       "              [-0.36430806, -0.83109545, -1.67787892, -0.75173715],\n",
       "              [ 0.55795223,  0.4525778 ,  0.35620414, -1.81124353],\n",
       "              [-0.58820621,  1.52723167, -0.21199303, -1.32970617]],            dtype=float64),\n",
       " DeviceArray([[-1.41461745,  2.05706838, -1.87873927,  1.59475793],\n",
       "              [-0.16087358,  0.65367767, -0.9466206 ,  0.5700123 ],\n",
       "              [ 0.93157583, -1.43000784,  1.75411696, -1.40996385],\n",
       "              [-1.62229365, -1.80458228, -0.2986438 , -0.74485972],\n",
       "              [-0.52942469, -0.18617663, -0.99167394, -0.07049572],\n",
       "              [ 1.65730698, -1.08574131,  1.73617314, -0.7211537 ],\n",
       "              [ 0.17391125,  2.25957034, -1.01357087,  1.69156383],\n",
       "              [-1.51430875, -1.86071866, -0.61031067, -1.38260964],\n",
       "              [ 0.93583491, -1.32989618,  0.99719163, -2.31176757],\n",
       "              [-1.42109755,  0.7245177 , -1.46663659, -0.28907386]],            dtype=float64),\n",
       " DeviceArray([5.60930358, 1.23732164, 4.83698509, 5.49706745, 1.99664333,\n",
       "              3.4600858 , 1.9439734 , 6.29286094, 3.92378597, 1.84035386],            dtype=float64))"
      ]
     },
     "execution_count": 495,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "visampler(random.split(_s, num=10), (is_nn_params,))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 496,
   "id": "6752d3df-b990-442c-bc57-1937b7306aaa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def loss(in_nn_params: nnParams, run_keys: Seed) -> float:\n",
    "    _, _, works = visampler(run_keys, (in_nn_params,))\n",
    "    return jnp.mean(works)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 497,
   "id": "beeb5897-29f7-4876-b3e6-0dc7de15af4c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DeviceArray(3.66383811, dtype=float64)"
      ]
     },
     "execution_count": 497,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss(is_nn_params, random.split(_s, num=10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 498,
   "id": "5e1f9bc9-4b95-4635-bfd0-c7526fdaa3c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax.experimental import optimizers\n",
    "from jax import value_and_grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 499,
   "id": "950e1dc5-49bf-4343-b85d-431e39eaaa3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "opt_init, opt_update, get_params = optimizers.adam(step_size=1e-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 500,
   "id": "3d6c53e7-d9f0-4dbc-9db4-b9f6af676dc6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def step(i, opt_state, seeds):\n",
    "    params = get_params(opt_state)\n",
    "    _val, g = value_and_grad(loss)(params, seeds)\n",
    "    return _val, opt_update(i, g, opt_state)\n",
    "\n",
    "#step = jit(step)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 501,
   "id": "87d18b3b-9b61-4010-b0b3-9c83eb5b667b",
   "metadata": {},
   "outputs": [],
   "source": [
    "iters = int(1e2)\n",
    "opt_state = opt_init(is_nn_params)\n",
    "parent_seed = random.PRNGKey(781)\n",
    "mean_works = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 502,
   "id": "87b0ac04-5f52-40c3-8f4d-1f38617f1c45",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 503,
   "id": "fee0c88a-4ebe-4e60-a6ee-fae60ab1a58a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [00:10<00:00,  9.45it/s]\n"
     ]
    }
   ],
   "source": [
    "for i in tqdm.trange(iters):\n",
    "    #parent_seed = random.PRNGKey(781)\n",
    "    parent_seed, child_seed = random.split(parent_seed)\n",
    "    run_seeds = random.split(child_seed, num=100)\n",
    "    _val, opt_state = step(i, opt_state, run_seeds)\n",
    "    \n",
    "    mean_works.append(_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 506,
   "id": "d7866f4d-a348-4bbd-9eb4-4655a053f0c5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0, 0.5, 'mean work (fwd)')"
      ]
     },
     "execution_count": 506,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYIAAAEGCAYAAABo25JHAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjQuMiwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy8rg+JYAAAACXBIWXMAAAsTAAALEwEAmpwYAAA7j0lEQVR4nO3dd3hc5ZX48e/RqHfLktwkW5J7wQWMKzbGELDpEJIAgRAgOBASkmVDFpJNgGRTdhOSQPgFh0Ag9NBjencB4yL33ptcVG1Vq5/fH/dqrDKyRrZGsjXn8zx6PHPvnZlzjZmjt51XVBVjjDHBK6SrAzDGGNO1LBEYY0yQs0RgjDFBzhKBMcYEOUsExhgT5EK7OoD2Sk5O1oyMjK4OwxhjTisrVqwoUNUUX+dOu0SQkZFBdnZ2V4dhjDGnFRHZ09o56xoyxpggF/BEICIeEVklIm/7OCci8oiIbBeRtSJyZqDjMcYY01RntAh+CGxq5dxsYLD7Mwd4rBPiMcYY00hAE4GIpAGXAE+0cskVwDPqWAIkikifQMZkjDGmqUC3CP4M/ASob+V8P2Bfo+c57rEmRGSOiGSLSHZ+fn6HB2mMMcEsYIlARC4F8lR1xfEu83GsRRU8VX1cVcer6viUFJ+zn4wxxpygQLYIpgKXi8hu4CVgpog81+yaHCC90fM04EAAYzLGGNNMwBKBqt6nqmmqmgFcC3yqqjc0u2we8C139tAkoFhVDwYini2HSvnDB1soKq8OxNsbY8xpq9PXEYjI7SJyu/v0XWAnsB34O/C9QH3uzvwyHv1sO7kllYH6CGOMOS11yspiVZ0PzHcfz210XIE7OyOG6AjnViuq6zrj44wx5rQRNCuLo8M9AFRU13ZxJMYYc2oJwkRgLQJjjGksiBJBQ9eQtQiMMaaxoEkEMdYiMMYYn4ImEUQ1JIIqSwTGGNNY0CSCY11DlgiMMaaxoEkEnhAhIjTExgiMMaaZoEkEADERodYiMMaYZoIqEUSFeSi3FoExxjQRVIkgJsLDUWsRGGNME0GVCKLCQym3RGCMMU0EVSKICfdw1LqGjDGmiaBKBNHhHsptHYExxjQRZIkglKM1lgiMMaaxIEsEHsqrrGvIGGMaC7JEYOsIjDGmuSBLBB4qqmtx9sMxxhgDAUwEIhIpIstEZI2IbBCRB31cM0NEikVktfvzi0DFAxAd4aFeoaq2PpAfY4wxp5VAblVZBcxU1TIRCQM+F5H3VHVJs+sWqeqlAYzDKzrsWCnqSPexMcYEu4C1CNRR5j4Nc3+6tE+mYd9iGzA2xphjAjpGICIeEVkN5AEfqepSH5dNdruP3hORka28zxwRyRaR7Pz8/BOOp2G7SptCaowxxwQ0EahqnaqOBdKACSIyqtklK4EBqjoG+AvwZivv87iqjlfV8SkpKSccT0y4tQiMMaa5Tpk1pKpHgPnArGbHSxq6j1T1XSBMRJIDFUfDLmVWeM4YY44J5KyhFBFJdB9HARcAm5td01tExH08wY2nMFAxeVsElgiMMcYrkLOG+gD/FBEPzhf8y6r6tojcDqCqc4FrgDtEpBY4ClyrAZzk79232ArPGWOMV8ASgaquBcb5OD630eNHgUcDFUNzMRHHpo8aY4xxBNnKYtvA3hhjmguyROC2CGzWkDHGeAVVIgjzhBDuCaHC1hEYY4xXUCUCcAaMrUVgjDHHBF0iiAn32BiBMcY0EnSJIMoSgTHGNBF0iSAmItTWERhjTCNBlwiiwjy2stgYYxoJukQQExFqtYaMMaaRoEsEUeEeyq1ryBhjvIIuEcSEe6xFYIwxjQRdIogOD7X9CIwxppEgTAQe26HMGGMaCcpEUFOnVNfWd3UoxhhzSgjCROBUILVxAmOMcQRhInAqkNrMIWOMcQRfIoho2JPAEoExxkAwJoIw26XMGGMaC+Tm9ZEiskxE1ojIBhF50Mc1IiKPiMh2EVkrImcGKp4G0e52leVVlgiMMQYCu3l9FTBTVctEJAz4XETeU9Ulja6ZDQx2fyYCj7l/Box3sLjGuoaMMQYC2CJQR5n7NMz90WaXXQE84167BEgUkT6BigmclcVgLQJjjGkQ0DECEfGIyGogD/hIVZc2u6QfsK/R8xz3WPP3mSMi2SKSnZ+ff1IxRbmJwKaPGmOMI6CJQFXrVHUskAZMEJFRzS4RXy/z8T6Pq+p4VR2fkpJyUjHFuF1DNn3UGGMcnTJrSFWPAPOBWc1O5QDpjZ6nAQcCGUtDi8BmDRljjCOQs4ZSRCTRfRwFXABsbnbZPOBb7uyhSUCxqh4MVEwAEaEheELE1hEYY4wrkLOG+gD/FBEPTsJ5WVXfFpHbAVR1LvAucDGwHagAbg5gPACICNFhtm+xMcY0CFgiUNW1wDgfx+c2eqzAnYGKoTXRER4qbNaQMcYAQbiyGJy1BBVWitoYY4CgTQQeKmxzGmOMAYI5EdgYgTHGAEGbCEJt1pAxxrj8GiwWkVRgKtAXOAqsB7JV9bTc5is63MOBI9YiMMYYaCMRiMh5wL1AErAKp1REJHAlMFBEXgUeUtWSAMfZoZwWgSUCY4yBtlsEFwO3qere5idEJBS4FPgK8FoAYgsYZ4zAuoaMMQbaSASqes9xztUCb3Z0QJ0hOsJDubUIjDEGaLtr6O7jnVfVP3ZsOJ0jOiyU6tp6auvqCfUE5Xi5McZ4tdU1FOf+ORQ4G6c2EMBlwMJABRVoMe4uZRU1dcRbIjDGBLm2uoYeBBCRD4EzVbXUff4A8ErAowuQxnsSxEeGdXE0xhjTtfz9dbg/UN3oeTWQ0eHRdBLvngS2utgYY/wuOvcssExE3sDZOOYq4J8BiyrAbE8CY4w5xq9EoKq/FpH3gGnuoZtVdVXgwgosaxEYY8wx/q4s/iWwCHhCVcsDG1LgxUU6t11micAYY/weI9gNXAdki8gyEXlIRK4IXFiBFR/lDBCXVNZ0cSTGGNP1/EoEqvoPVb0FOA94Dvia++dpKd5tEZQctRaBMcb4lQhE5AkRWQw8htOddA3Qo43XpIvIZyKySUQ2iMgPfVwzQ0SKRWS1+/OLE7mJ9opzp4yWHLUWgTHG+DtrqCfgAY4ARUCBW2LieGqB/1TVlSISB6wQkY9UdWOz6xap6qXtCfpkhYeGEBkWYl1DxhiD/7OGrgIQkeHARcBnIuJR1bTjvOYgcNB9XCoim4B+QPNE0CXiI8MorbSuIWOM8XfW0KU4U0en43QJfYozi8gvIpKBs5H9Uh+nJ4vIGuAA8GNV3eDj9XOAOQD9+/f392OPKz4qzFoExhhD20XnMlV1FzAbp7bQw6p6oD0fICKxOGWqf+Rj34KVwABVLRORi3GqmQ5u/h6q+jjwOMD48eO1PZ/fmvjIUBssNsYY2h4sftX9c5iq/usEkkAYThJ4XlVfb35eVUtUtcx9/C4QJiLJ7fmMExUXaS0CY4yBtruGQkTkfmCIr5LUxytDLSICPAlsau06EekN5KqqisgEnMRU6Hf0JyE+Koy9RRWd8VHGGHNKaysRXIuzLWUox0pS+2sqcCOwTkRWu8d+ilPADlWdizMN9Q4RqcXZC/laVe2Qrp+2OF1D1iIwxpi2ylBvAf5XRNaq6nvteWNV/RyQNq55FHi0Pe/bURoGi1UVp/FijDHB6bhjBCJyg4iEtJYERGSgiJwTmNACKz4yjJo6pbKmvqtDMcaYLtVW11BPYJWIrABWAPlAJDAIOBcoAO4NaIQBEh/llpmorPGWpTbGmGDUVtfQwyLyKDATp89/NE5f/ibgRlXdG/gQA6OhzERpZQ294iO7OBpjjOk6bS4oU9U64CP3p9toKDxXbGsJjDFBLmh3brdS1MYY4wjeRGAVSI0xBvC/DHWEj2NJHR9O5zk2WGxdQ8aY4OZvi+B1t1wEACLSh9N8zCC+0WCxMcYEM38TwZvAKyLicSuJfgDcF6igOkNEaAjhnhArPGeMCXr+7kfwdxEJx0kIGcB3VXVxAOMKOBEhPirUBouNMUGvrTLUjQvNCZAOrAYmicik4xWdOx3ER4bZYLExJui11SJoXmjujVaOn5biomyXMmOMaWtl8YMi4gF+p6r3dFJMnSY+0rqGjDGmzcFid2XxmZ0QS6ezriFjjPFzsBhYLSLzgFeA8oaDvnYdO504g8XWNWSMCW7+JoIknJ3DZjY6psDpnQisRWCMMX5PH7050IF0hfioMKpq66mqrSMi1EpRG2OCk78lJtJE5A0RyRORXBF5TUTSAh1coDVUILWZQ8aYYObvyuKngHlAX6Af8JZ7rFUiki4in4nIJhHZICI/9HGNiMgjIrJdRNaKSKcOSsdZ4TljjPE7EaSo6lOqWuv+PA2ktPGaWuA/VXU4MAm4U0RGNLtmNjDY/ZkDPOZ/6CfPCs8ZY4z/iaDA3b/Y4/7cgDN43CpVPaiqK93HpTi7mvVrdtkVwDPqWAIkugXtOoWVojbGGP8TwS3A14FDwEHgGveYX9xCdeOApc1O9QP2NXqeQ8tkgYjMEZFsEcnOz8/392Pb1LA5jY0RGGOCmb/TR/NU9fIT+QARiQVeA36kqiXNT/t4ibY4oPo48DjA+PHjW5w/Ud4Wga0uNsYEMX8TwXoRyQUWAQuBL1S1uK0XuXsYvAY838risxycQnYN0oADfsZ00rxjBNY1ZIwJYn51DanqIOA6YB1wKbBGRFYf7zUiIsCTwKbjVCmdB3zLnT00CShW1YP+Bn+yosI8eELEWgTGmKDmV4vAXTMwFZgGjAE2AJ+38bKpwI3AukZJ46dAfwBVnQu8C1wMbAcqgE5duCYiTuE525zGGBPE/O0a2gssB36jqrf78wJV/RzfYwCNr1HgTj9jCIj4qDDbrtIYE9T8nTU0DngGuF5EvhSRZ0Tk1gDG1WniI8NsHYExJqj5W2tojYjsAHbgdA/dAEzHGQM4rcVHhdpgsTEmqPk7RpANRACLccYGpqvqnkAG1lniIsLILy3r6jCMMabL+DtGMFtVO24l1ykkPirUFpQZY4Kav9NHu2USANuTwBhj/B0s7rbio8Ior66jtq4egKPVdV0ckTHGdC5LBI32JFiys5Axv/yQ99d32po2Y4zpcv6OESAiU4CMxq9R1WcCEFOnaig8t+lgCd9/cRXVtfWsySlm1qhOK4JqjDFdyt9ZQ88CA4HVQEPfieKsLTitNWxO84MXV1FTV09KXAS78su7OCpjjOk8/rYIxgMj3JXA3UpD11BRRTX/+PbZvLB0LzsLbDqpMSZ4+DtGsB7oHchAukrfxChE4N5ZwzhvaCpZyTHsLqygrr7b5TxjjPHJ3xZBMrBRRJYBVQ0HT3SPglNJelI0q39+IQnRThdRZnIM1bX1HDhylPSk6C6OzhhjAs/fRPBAIIPoag1JAJxEALCroNwSgTEmKPhba2hBoAM5VWSlxAKwM7+M6UNSujgaY4wJPL/GCERkkogsF5EyEakWkToRab7tZLeQHBtOXEQouwps5pAxJjj4O1j8KM4OZduAKOA77rFuR0TITIlhpyUCY0yQ8HtlsapuBzyqWqeqTwEzAhZVF8tMjmGnrSUwxgQJfxNBhYiEA6tF5P9E5D+AmOO9QET+ISJ5IrK+lfMzRKRYRFa7P79oZ+wBk5Ucy4Hio1TWWN0hY0z3528iuNG99vtAOZAOfLWN1zwNzGrjmkWqOtb9+aWfsQRcZkoMqrCnsKKrQzHGmIDzd9bQHhGJAvqo6oN+vmahiGScTHBdJcs7hbSMob3jujgaY4wJLH9nDV2GU2fofff5WBGZ1wGfP1lE1ojIeyIy8jifP0dEskUkOz8/8FsjZLiJYIeNExhjgoC/XUMPABOAIwCquhqnEunJWAkMUNUxwF+AN1u7UFUfV9Xxqjo+JSXwc/tjI0LpFR9hU0iNMUHB30RQq6rFHfnBqlqiqmXu43eBMBFJ7sjPOBmZyTGWCIwxQcHvonMicj3gEZHBIvIXnI3sT5iI9BYRcR9PcGMpPJn37EiZybGWCIwxQcHfRPADYCROwbkXgRLgR8d7gYi8CHwJDBWRHBG5VURuF5Hb3UuuwUkwa4BHgGtPpTLXWckxFJVXc6SiuqtDMcaYgPJ31lAF8DP3xy+qel0b5x/lFF6dnJXiDBivzSlm86ES/rV8H9+bMYivnpXWxZEZY0zH8neHsvHAT2m5VeXowITV9RqqkN701DJUIURg/tZ8SwTGmG7H3zLUzwP3AOuA+sCFc+pIT4pm/IAe9E6IZM70LB76cCs7823nMmNM9+NvIshX1Y5YN3DaCPOE8OodU7zPM5NjWL67CFXFHeM2xphuwd9EcL+IPAF8QtMdyl4PSFSnoIEpMVRU15FbUkXvhMiuDscYYzqMv4ngZmAYEMaxriEFgiYRZCa7G9YUlFkiMMZ0K/4mgjGqekZAIznFNcwi2plfzpSBp8y6N2OMOWn+riNYIiIjAhrJKa53fCSRYSF+LzI7Wl1HSWVNgKMyxpiT528iOAdnL4ItIrJWRNaJyNpABnaqCQkRMpNj/Z459MC8Ddz4xNIAR2WMMSfP366htvYVCApZyTFsOOBfyaWNB0vYdLCEmrp6wjx+bwRnjDGdzu/9CAIdyOkgKyWG9zccorq2nvDQ43+57y2qoLZe2VNYwaDU2E6K0Bhj2s9+VW2HzOQY6uqVvUXH37msuKKG4qPO+MAOW4RmjDnFWSJoh6wU5zf7tgaM9x0+lii251kiMMac2iwRtENmz4YppMf/cm9oMYjADksExphTnCWCdkiIDqNnTHjbLQI3EYxOS7SuIWPMKc8SQTtlpcSws429jPcWVdAjOoxx6YnsyC/nFNpmwRhjWrBE0E6ZyTHsbKNFsLeogvSkaAamxFBWVcuhkspOis4YY9rPEkE7ZaXEUlBWddxVw/saEoE7bXRHnm15aYw5dQUsEYjIP0QkT0TWt3JeROQREdnurlY+M1CxdKSGDWt2tdI9VFev5Bw+Sv+kaO/6ge15pT6vLauqZfGOgsAEaowxfgpki+Bpjr8ieTYw2P2ZAzwWwFg6zMCG4nMFvgeBDxYfpbZe6Z8UTUpsBHGRoexoJWn89t1NXP/3pWw8UBKweI0xpi0BSwSquhAoOs4lVwDPqGMJkCgifQIVT0dJT4omRGh1wLhh6mj/pGhEhEGpsT7XEuSWVPJKdg4Azy45sYXbi7blU1ZVe0KvNcaYBl05RtAP2NfoeY577JQWEephRN94/r36AJU1dS3O72uUCAAGpsSy3ccU0r8v3EmdKucMSubNVfu9K5H9tX5/MTc+uYzfv7/5BO7CGGOO6cpE4Gu/R5/zLEVkjohki0h2fn5+gMNq209nD2dvUQWPzd/R4ty+oqN4QoQ+7uY1g1JjyS+tavJFf7i8mueX7uXyMX35r1nDOFpTx+src9oVw7NfOq2If2Xvo6i8+iTuxhgT7LoyEeQA6Y2epwEHfF2oqo+r6nhVHZ+SktIpwR3PlEHJXD6mL48t2MHuZlNJ9xZV0C8xilC34uggtyxF44VlT32xi6M1ddwxYyBnpCUwNj2RZ5fs8Xu9QXFFDf9es59JWUlU1tR7k4IxxpyIrkwE84BvubOHJgHFqnqwC+Npl/++ZDjhnhDun7ehyRf43qIKb7cQ0GgKqZMISitreHrxbi4a2YshveIAuHHSAHbml7N4RyGqykcbc/nxK2ta7S56ZcU+Kmvq+fmlI5g5LJVnvtzts5vKGGP8Ecjpoy8CXwJDRSRHRG4VkdtF5Hb3kneBncB24O/A9wIVSyCkxkdy91eGsGBrPu+vP+Q97qwhiPI+T+8RRbgnhO35ZZRW1nDv6+soqazlzvMGea+5ZHQfkmLCefjjbVz39yXc9kw2r67I4Z21LfNifb3y/NK9nDWgByP7JjBnehaF5dW8uqJ9XUt7CsvZU2jrG4wxgZ01dJ2q9lHVMFVNU9UnVXWuqs51z6uq3qmqA1X1DFXNDlQsgfKtyQMY1juO/3lnE5U1dZRV1VJYXk16oxZBqCeEjORoFmzJZ/bDi3hv3UF+fOEQRqcleq+JDPPw9fHpLNtdxNbcMn51xUgG9Izm/Q2HWnzm59sL2FVQzo2TBgAwMTOJMWkJ/H3RTurq2+5aqqyp4w8fbOH8hxbwvedXnvxfgjHmtGcri09CqCeEX1w6gv1HjvKPL3a1mDHUYFBqLJsPlRIiwiu3T+H7Mwe3eK87zxvIb68+g/n3zODGyRnMGtmbxdsLWnQPPbtkDz1jwpl9Rm8ARITvnjuQPYUVfOgjcTT2+bYCZj+8iEc/207P2HC255VR70fy6AiLtuVzpMIGtY05FVkiOElTBiVzwfBe/PWzHazaewRomQhumZrJXTMH8e4Pp3HWgB4+3ycuMozrJvQnPjIMgItG9aa2Xvl0c673mr2FFXyyKZdvnJ1ORKjHe/yikb0Z0DOauQt2+BxwXr67iGsf/5IbnlxKXb3y3K0Tuev8wVTV1nOwE+og7Sks58Ynl/H4wp0B/yxjTPtZIugA9108jMqaOn733iagZSIYn5HE3RcOJTbC3y2iYWxaIr3iI/hg/bFE8OePtxLmCeGmKRlNrvWECLdNy2JNTjFf7iz0HldV7n55NV+b+yXb88q5/7IRfPgf0zlncLJ3b4Xms54C4Z11zljHyr2HA/5Zxpj2s0TQAQamxHLDpAGUVNYSFxlKQlTYSb9nSIhw0cjezN+ax9HqOrbmlvLG6v18e0oGveIjW1x/zVlpJMeGM3fBsd+631i1n9dX7mfO9CwW/eQ8bp6aSWSY05LIaKiZ1CwRbDhQzLNf7j7p+BtrGPRem1Ps1ziGMaZzWSLoID88fzDxkaHe0hId4aKRvamsqWfB1nz++OFWYsJDuf3cgT6vjQzzcPPUTBZuzWfDgWLyS6v45dsbOWtAD+6dNYyocE+T63vHRxIRGtKiRfDkol38/N8bKCirane8eaWVTPj1xyzadmzR366CcjYcKOGMfglUuAnNGHNqsUTQQXrEhPPETWdz/2UjO+w9J2QmkRgdxqOfbeP9DYf4zrRMesSEt3r9DZMGEBsRyt8W7OSBtzZQUVXH/371DEJCWiamkBBhQM9odjebQrrxoFMA74vt7a+KuulgKXmlVfzq7Y3e3/zfWeusEbzv4mEArN53pN3va4wJLEsEHWhCZhITMpM67P3CPCFcMLwX6/eX0CM6jFvPyTzu9QlRYVw/sT9vrT3AO2sP8oOZgxiUGtfq9Rk9Y5p0DVXV1nkL5C3c2v5EkHPYmTW1NbfMWzLj7bUHOWtADyZn9aRHdBir3QF1Y8ypwxLBKW72KGea6O3nDiQusu2xh1umZhIaIgzrHcd3W+lGapCZHMO+oqPe396355VRW6/ERYSyaFt+qyUv1uYcYeZD8yls1n2Uc/goYR5hdFoCf/xoKxsOFLP5UCmXju6DiDA2PZFV+zpuwFhVbRtQYzqAJYJT3Mxhqfzj2+PbbA006J0QySu3T+GZWycQHnr8/7wZyTFU19Vz4MhRAO++CNdP6k9eaRVbc33vufDqihx25pezdn9xk+M5h4/SNzGK+2YP52BxJXc8txIRuPgMp7r42PQebMtzVliD80V+/d+X8JdPtvl1b809+fkupv3fZ9TU1Z/Q640xDksEpzgRYeawXt4idv4Ym55IalzLmUXNZfRsOnNo08FSIsNCvKuWGw/6NlBVPtmUB8CeZgPNOYcrSOsRxeSBPZkxNIW9RRWcnZHkneU0tn8iqrAux0kg87fms3hHIc8s2dPu2USVNXXMXbCDnMNHWdcsIYHTzXW6UVXW7DtirRzT6SwRBLGGbTcbag5tPFjM0N7xpPVwttlcuK3lOMHmQ6Xsd1sQuwsrmpzLOXyUtERnDcVPLhpGmEe4etyxLSbGumU1VrkDxo8v2IknRMgvrWLprkLaY97qAxSUOSuVv9zR9LWfbMplzIMfescsWtPePSDaa/W+I8xd0LJUeWteyc7hiv/3BfO3dH2pdRNcLBEEsV7xEUSFedhVUIGqsulgKSP6xAMwbXAyS3cWtqhq+skmZ4Fbn4TIJjOOKmvqyC+tIq2HU3BvRN94Ft97Pt84+1il8YToMLJSYli19wjr3MVvd80cTHS4h7d9FNhrjaryxOc7GdY7jiG9Ylmys2kicDYNqufDDbk+X78tt5Q7nlvBmAc/5O21LSuff7Y5j9wOWHH9x4+28rv3NrPNjymz5VW1/P7DLQB8ujnvpD/bmPawRBDERI5NIT1QXEnx0RpG9HFmGU0fnEJVbT3Zu5sO7n68KY/RaQmc2b9HkzUIDa2EtEaVV1PiIlqsqRiX3oPV+47wt4U7iI0I5eZzMjh/eC/eX3+IWj/7+hdtK2BrbhnfmZbF5KyeZO8+THWt89qaunrmb3G+SD/Z3DQRVNbU8eNX1nDhnxeyaFsBCVFh3u1CG+wrquDmp5fzhw+2HDeGuno9bp2movJq7xTcV/3YdOhvC3eSX1pFRs9oFvrokjMmkCwRBLnM5Bh2F5SzyR0oHtHXaRFMzEoizCNNxgnyS6tYk3OEC4b3IiM5mpzDR70DtTmH3UTQI5rjGds/kYKyKt5Zd5DrJzq1lS4b3Yei8moWu108qsqv39nIQx/6/jJ+4vNdpMRFcNmYPkwe2JOjNXWszTkCQPbuw5RU1jIoNZalO4soqTzW/fPmqv28uiKHm6dksvAn53H9xP58vr2gyQ5vb6zaD8DHm3JbTUyqyg1PLOWO51e0ep8fbDhEXb2SlRLDGyv3HzfJHSqu5PGFO7hkdB9unprJnsKKTin9cSLq6pWNB0p4dskeHnxrA4dtd7xuwRJBkMtIjmFvUYV3wHVobycRRIeHMn5AEgu2HksEn23OQxXOH57KgJ4x1NYr+90E0NAf39A11Jpx6YkAeET4tlsz6dyhKcRFhPLWGqeb5rkle/j7ol08vXh3i0HkLYdKWbg1n5smDyAi1MPEzJ7AsXGCTzblEu4J4b8vGU5tvbKwUfwvLtvLkF6x/PzS4STFhHPp6D7U1SvvrXe6pVSV11fmEBsRyuGKGpbv9j3V9aONuXy5s5CPNuaS10oX0jtrD5KZHMM9Fw4lr7SKz4+zQO+hD7dQXw/3zhrG9CHODnynYqtga24pE379MRc/soifv7mep77Yzb9X7+/qsEwHsEQQ5DLdL/SPN+WS0TO6SWG8C0b0YvOhUh58awN17jV9EyIZ0SfeO9DcME7QsIagrdlKQ3vHERcZyhVj+9E30UkaEaEevjKyFx9sOMSXOwp58K2NpMZFUFpZy4YDTWcEPb90DxGhIVw/0ZnZ1CMmnGG941jiDjZ/sjmPSQN7Mm1wConRYXzqznBav7+YNTnFXDehv7e7akSfeLJSYnh7jZMIVuw5zO7CCv5r1lDCQ0P4wEdZ77p65fcfbCE1LoJ6hXlrWo4xFJZVsXhHAZec0YeZw1NJjA5rdeOgxTsKeHVlDt+emkF6UjQZPaNJT4pqksBOFY/N30FVbT1//PoYFt5zHulJUXyxo32D/ObUZIkgyDUUn9twoITh7kBxg29PyeCWqZk89cVubnsmm8+3FzBzeKp3bAGOVS9tWEPg8VHOorEwTwjv/GAav7qyaSmOy0b3paSylpueWka/HlE8/52JQNMZQarKxxtzOXdICkmNSm1MHuiME2w6WMKugnIuGJ6KJ0Q4b2gqn23Jo65eeWn5XiJCQ7h6XJr3dSLCZaP7smRXIXkllby2MofocA9Xn5nG9MHJfLjhUIupnK+vzGFbXhkPXD6SM/ol8O/VLRPBe+sPUa/OznMRoR6uGNOXDzfmtpil9OnmXG5+ajlZyTHeHetEhOmDU/hyR6F33ONkvbB0L09/seuk3iOvpJK31x7gmrPSuPrMNPr3jGbqwGSW7Cz0e2zHnLosEQS5jORjffojmiUCT4jwi8tG8KsrR7Fgaz4V1XWcP7wXACmxEcSEe7xTSBvWEPijf89oosObluSeOiiZhKgwPCLMveEsBveKIyslpklZ7S25pRwormTmsNQmr52c1ZOq2noe+nArgPf8+cNTOVxRw+fbC3hz1QEuGd2HhOimq7MvG9MHVXht5X7eXnOQWaN6ExMRyoUje3OguLLJGoXKmjr+/PE2RqclMHtUb64c1491+4u9ZTkavLP2IFkpMQzr7Qy8f/WsNKpr65vMUPr36v3MeWYFQ3rF8crtU5pUrJ0+JIXy6jpW7Dn5Vdi7C8q5f956fvveZoorTny67PNL91Jbr97uPHD24iitrPW5jqOzlFTWdMjfU7ALaCIQkVkiskVEtovIvT7OzxCRYhFZ7f78IpDxmJYavtCBFi2CBjdOGsDTN5/NNyf2Z8pAp0/eaRXEeLuG9jdaQ3AiwkNDeOybZ/LcdyZ645ic1ZPlu4q8A9IN0yrPa5YIJmb2RMQZ4B3WO847YD19SAqhIcLP3lhHWVUt10/o3+JzB6XGMax3HA9/spXSqlquOdNpMVwwvBeeEGnSPfTckj3sP3KU/5o1zGlNjOlDiNCknzyvtJKluwq5dHRfbxfUGf0SGNIrlicX7eLuf63m0r8s4kf/Ws1ZA3rwwm0Tm7RuAKYM7EloiDQZn/GHr5Ibv31vEyJCVW09r69q377WDapq63h+6R7OG5rqbUE2xAl4B/kD6aONudzwxFJeWLqX0soaqmvrefqLXcz4/Xy++tjiFntd5Byu4Bt/+5Id+b5Xxzd4+otdbV7TXm+vPcA1jy0+rVpKgdy83gP8P2A2MAK4TkRG+Lh0kaqOdX9+Gah4jG8i4v2fu2HGkC/TBqfw66vOaLIzWsOMo8qaOvIarSE4UVMGJTfZwW3ywJ6UV9d5f+P8dFMeo/rFt9iPISE6zNuaucBtsQDER4YxMSuJnMNHGdIrttXd4S4b05fKmnr6JkQyKcv5ckuKCWdCRhIfuGsRXl+Zw/+9v4Vpg5OZOigZgNS4SKYOSuaNVfu9X8BvrNxPvcKlo/t4319EuHHSAHYWlLN4RyE9osP5wXmD+OctE3zWj4qLDOPMAT38HieoravnpWV7mfq7T/nOP7Mpq6oFnG61Dzbk8sPzBzMmLYEXlu49oVXLb685SEFZNTdPzWhyPDk2gmG949pdqXZPYTm/fGsjb/hITCWVNT4H4J/6YheLdxTw0zfWcfavP+a8P8zngbc2MrRXHLERoTz35Z4m1z+xaBdLdxXxm3c2tRrHzvwyHnhro88SJ++vP8SaE6yUO3fBDrL3HD6tKu0GskUwAdiuqjtVtRp4CbgigJ9nTtDAlFh6RIfRJ6HtshSNDejpTCHd6+7V3HgNQUdo+FL+ckchh8urWbn3MDOHpvq8drJ77fnDm54/f5iTGBoPEjd32ei+hIjThdO4ZPdFI3uxPa+Me15Zw90vr+HMAYk8cu24Jq+9cmw/cg4fZf7WfO57fS2/fW8zZ2f0YEivplVfb5ycweZfzWLJT8/n2VsncveFQ72bBPly7pAUNh4sIb/0+PtCfLYljwv/vJB7X19HQnQ487fm842/fcnB4qP86u2N9EuM4tZzMrluQn+25ZW1uxtFVXlq8S4GpcZyjpsAG5s6KJnsPYdbLDz0ZfOhEu58YSXn/WE+//hiFz9+ZS3Zu4u85wvLqrjkkUV87W9fNklYJZU1LNtVxJzpA3nzzqlcNS6N/knRPPXts3nhtolcfWY/3l530DsNuPhoDS9n76NHdBifbM5rsfK8QUML85PNeU3GY0ora/jhS6u459U17U6cmw6WsH6/MxX7VBzwb00gE0E/YF+j5znuseYmi8gaEXlPRHwW8xeROSKSLSLZ+fmnz1/u6eKei4byxE3j272hTkayM+NoqduP39YagvZKjo3wrhxesDWfeoWZjX7jb+ymKRncc9FQxrhlLBp89cw0vntuFl8fn+7zdeCMWcz7/jneAdsGF450Kr++siKH6yb059lbJ7bYD+KiUb2JDAvhlqeX86/l+/ju9CyevXWiz8853hd/c+e5Ce/ul1dzpML3XP3Syhq++6yzluHxG8/i3bvO4YmbxjsD5g8tYOPBEu6dPYzIMA+XjelLbEQoLyzb63cM4Gwvun5/Cd+ekuHz38c5g5Kp9rHwsLmDxUe55rEvWbglnznTB/Lx3eeS1iOK77+wisKyKipr6pjz7Ar2FR1lT2FFk66ez7cVUFuvnD88lbHpifz26jN4cc4kzhvmTFy4YdIAqmvreSXb+bp5adleKqrreOKms+mbEMlv39vkc/HfZ1vyCPMIpZW1Tcai3l9/iKraerbmlrWaRFrz2oocwjzC4NRYFvgo0XKqCmQi8PWt0vy/xkpggKqOAf4CvOnrjVT1cVUdr6rjU1JSOjZKQ3pSNGcNaP8+Cg1F6xrmyJ9s15AvDSuHP9x4iOTYcEb3S/B5XXpSNHeeN6jFJjwJ0WHcN3s4MW3sFz2qX0KLL+q+iVHcNXMQv75qFL+5ahRhPgr/xUaEct2E/gzvHc9rd0zhvouHt+sLvzUj+sbzv189g6U7i7j80S/YfKikxTWfbcmnurae//3qaC4c2RsRZ6bUy9+dTHREKBMzk7xdVDERoVwxti/vrD3YrkHj99YdItwTwpXjfP0O5+zBERoifLHj+F96D87bSE1dPW/fdQ73zh7GoNRY/vrNMymqqOZH/1rNT15dy4o9h/m/a0YTERrCvEazsT7ZlEdidJh3DUpzQ3rFMSEzieeX7qW6tp5/Lt7NpKwkzhrQg/+8cChrc4p5e13TEiZlVbUs21XENycOIDrc02Qs6N+rD5CeFEVSTDhPL97t318Uzqr2N1fvZ+awVC4Z3Ye1OUdOmwV3gUwEOUDjX8PSgCZz7VS1RFXL3MfvAmEi0rL9aU5JDTOOFu8o9GsNwYmYlOWsHH5v/SFmDE31udtaIN194VC+OXHAcVtL9182knd/OI1x/X2PQZyob5zdnxfnTKKypo6r/7qYVc0GRD/c4CTHM5t97qh+CSz6yXn885YJTeK+fmL/dg8af+quy4htJZHGRIQyNj2RxccZJ/h4Yy7vbzjEXecPZkDPY4PNI/sm8ODlI1m0rYB5aw5wz0VD+fr4dGYOS+WddQepraunvl6ZvyWPc4ekHLcC7w2TBrC3qIKfvrGOA8WV3HpOFgBXjuvH8D7x/P6DzU0q0n6+LZ+aOmXWqN6cNzSVDzfkUlev5JZU8sWOAq4al8a1Z6fz8aZc9hUdK1644UCxt3pucwu25FNQVs01Z6UzfUgKqhx3IeGpJJCJYDkwWEQyRSQcuBaY1/gCEekt7r9UEZngxmMrVE4TDTOOSitr/VpDcCImun3/qrSYNhoMzhrQg7d+cA7R4aH85dPt3uNVtXXM35Lvnd3UXGSYp0XLZGTfBMakJTB3wQ7v3hPHs6ugnJ0F5Zzfxt/71EHJrNtfzLbcUj7bnMcTi3ayZGchqkpFdS33z9vA4NRYbpuW1eK1156dzvdmDOSOGQP53gxnI6XLx/SloKyaL3cWsibnCIXl1W3+t581sjfJseG8uiKHjJ7R3pg9IcJPLx7GvqKjPPXFbu/1n27OIy4ylLMG9OCiUb0pKKti1d7DvLXmAKpw5di+3DDJ+QXguSXOQPQX2wu4+q+Lufnp5T5nBL26IoeeMeHMGJrCmLREEqLC2j1O8NaaA8x8aD6fd3K30vHbyydBVWtF5PvAB4AH+IeqbhCR293zc4FrgDtEpBY4ClyrVoz9tNEwhXTjwZKAdAuBM3tnWO84tueVMW1wcDYWe8VH8s2J/Xn4k23sKignMzmGxTsKKauq5SJ3HMNfv7n6DG59OpurH/uC/7tmDJeP6dvqtQ2DqW19CU8dlMzDn2zjK39a2OR4/yRnlfT+I0d55fbJPjdKEhF+MmtYk2PnDUslNiKUeasP0DshEk+IcO6Q43cJh4eGcO3Z/Xn0s+3cPDWzSctx2uAULhjei0c+2cYVY/vSKy6Sz7bkM31ICmGeEM4bmkK4x1lJvnhHIWPSEshKiQXgwhG9eGn5PiYN7Mkdz60gJiKUgrIqFmzN966pAafI4Cebc/nW5AxvF+I5g5JZ6O7019b42+Hyan7+7/XeKrwPfbSFqYN6tnvc7kQFdB2Bqr6rqkNUdaCq/to9NtdNAqjqo6o6UlXHqOokVV0cyHhMx2voHuqXGJhEAHDHjIHcdf5gv7bq7K6+ObE/YR7hmS93A063UEy4h8nuXH5/jeybwFs/OIcz+iVw14uruPvl1fxz8W4+3dyybtKnm3MZnBpLetLxJwGMH9CD/5o1jF9eMZKXvzuZZT87nz99YwxpPaL4Ynsh103oz9kZ/o9BRYZ5uHBkL97fcIgPNhzirP49SIwOb/N1t56TyV3nD25S+rzB/ZeNoK5e+Z93NrHhgDMbq2EGWlxkGFMH9eTl7Bw2HChpMh5y05QMio/WcPNTy0nvEc17P5xGz5jwFlVr31i1n5o65Zqzjq1cnz4kmdyS1nf6a5BXWsmshxfy/vpD/PjCIfzi0hGs2nuE7E5cKBewFoEJDg0Dxh09Y6ixK8b6HqgMJqnxkVx8Rh9eyc7hP74yhI825jJjWOoJDUynxEXw/Hcm8Zt3N/HS8r28vtJZEBcd7uGdu6aRmRxDaWUNS3cWceu0trdIDQkR7pjRdH/sq8alcdW4NArLqpqsmvbX5WP68vrK/ZRWlnHv7GFtvwCn7tTdXxni81x6UjTfmzGIP328leKKGkRgxtBjrYyLRvbmsy35eEKES0cfayVNzExiTHoiZZU1PH/bRFLjIrlyXD+e+XI3ReXVJMWEU1pZw2Pzt3N2Ro8mizK9BQS35jO0d9PpxI29uHQfuSVVvPG9KYzr34OK6loe+XQbf1uws10J9GRYiQlzUo4lgsC1CIzj21MyKKuq5WdvrKegrJoLR/ieSuuP8NAQHrh8JJt+OYtlPzufF2+bhCdEuO/1tagqixqmbA478c8A6Bkb0a5tVhtMHZTsXXHd1hiFv757bhb9k6L5fHsBY9IS6Rkb4T13wYhehIjzuSlxx46LCC/dNokPfjTdOxnia+PTqKlT3nRLlv91/g4Kyqr570uarpftkxDF4NTY41aSratX/rV8L9MGJ3snG0SHh/KtSQP4eFNui/IlgWKJwJyUsf0TCfeEMKqVaZ2m44zr34Mx6Ym8teYAYR5pUWrjRIg4s70mD+zJzy4ezpKdRby0fB+fbMojISqMM/snnnzgJyDME8LXx6czok88g1JjO+Q9I8M8PHC582XdPLkkx0bwp2+M5b8vGd7idVHhnibJbFjveM7ol8ArK3LYV1TBk5/v4upx/RjjY3rr9CEpLN1VRLm72ru5BVvzOFBc2aL8ybemZBARGsKTn+9s722eEEsE5qQM6RXHpl/NarGS1gTGzW7Rt8kDk4nv4DGTb5ydzuSsnvzmnU18sjm3zSmbgfZfs4byzl3ndOiA6cxhvXjhOxN9dnldMbaf3/+OvzY+jU0HS/je8ysJEbhn1lCf180e1Zvq2nreW9+ypDnAC0v3kRwbwQXNWnfJsRF89aw0Xlu5n225pcfdDa8jWCIwJy0Q00aNbxef0YfpQ1K4afKADn9vEeG3V59BTX09RypqWpTr6GwiEpBZM1MGJbeofttel4/pS7gnhHX7i7n93IH0SfDdNXrWgB5kJcfwcva+FucOFh/l0825fH18ms/FirdNy6K+XvnKnxYy4v73mfXnhTy7ZE+L6zqCJQJjTiPhoSE8c8uEJlMXO1JGcgz3zhpGQlRYm1M2g1lidDiXjO5Dv8Qo5kxvuT6igYjwtfHpLNtVxK5m24++vDyHeoVrz25ZFRecoo5v33UO/3PlKG6YOIB+iVGEewLzS5ecbtP2x48fr9nZ2V0dhjHdWl29WkuvDVW1dVTX1rc5rTmvpJLJv/uU28/N4p6LnBlQdfXKtP/9lIGpsa3WpupoIrJCVcf7OmctAmNMC5YE2hYR6vFrbUtqfCQzhqTw6ooc7x7ccxfs4EBxJd+c6Ls10NksERhjTIB9bXw6uSVVLNyaz0vL9vL7D7Zw+Zi+XDiifSvDA8UWlBljTIDNHJZKz5hw/uedjewqKGfG0BT+8LUxnV5EsTXWIjDGmAALDw3hqnH92JFfztj0RP76zTN91l7qKtYiMMaYTjBnehYej3DHuQNPevpqRzu1ojHGmG4qNT6S+2a3XLl8Kjh12ibGGGO6hCUCY4wJcpYIjDEmyFkiMMaYIBfQRCAis0Rki4hsF5F7fZwXEXnEPb9WRM4MZDzGGGNaClgiEBEP8P+A2cAI4DoRGdHsstnAYPdnDvBYoOIxxhjjWyBbBBOA7aq6U1WrgZeAK5pdcwXwjDqWAIki0ieAMRljjGkmkImgH9C4CHeOe6y91xhjjAmgQC4o81VEo3nNa3+uQUTm4HQdAZSJyJYTjCkZKDjB157OgvG+g/GeITjvOxjvGdp/363uZhTIRJADpDd6ngYcOIFrUNXHgcdPNiARyW6tHnd3Foz3HYz3DMF538F4z9Cx9x3IrqHlwGARyRSRcOBaYF6za+YB33JnD00CilX1YABjMsYY00zAWgSqWisi3wc+ADzAP1R1g4jc7p6fC7wLXAxsByqAmwMVjzHGGN8CWnROVd/F+bJvfGxuo8cK3BnIGJo56e6l01Qw3ncw3jME530H4z1DB973abdnsTHGmI5lJSaMMSbIWSIwxpggFzSJoK26R92BiKSLyGcisklENojID93jSSLykYhsc//s0dWxdjQR8YjIKhF5230eDPecKCKvishm97/55CC57/9w/32vF5EXRSSyu923iPxDRPJEZH2jY63eo4jc5363bRGRi9r7eUGRCPyse9Qd1AL/qarDgUnAne593gt8oqqDgU/c593ND4FNjZ4Hwz0/DLyvqsOAMTj3363vW0T6AXcB41V1FM6MxGvpfvf9NDCr2TGf9+j+P34tMNJ9zV/d7zy/BUUiwL+6R6c9VT2oqivdx6U4Xwz9cO71n+5l/wSu7JIAA0RE0oBLgCcaHe7u9xwPTAeeBFDValU9Qje/b1coECUioUA0ziLUbnXfqroQKGp2uLV7vAJ4SVWrVHUXznT8Ce35vGBJBEFX00hEMoBxwFKgV8NCPffP1C4MLRD+DPwEqG90rLvfcxaQDzzldok9ISIxdPP7VtX9wB+AvcBBnEWoH9LN79vV2j2e9PdbsCQCv2oadRciEgu8BvxIVUu6Op5AEpFLgTxVXdHVsXSyUOBM4DFVHQeUc/p3h7TJ7Re/AsgE+gIxInJD10bV5U76+y1YEoFfNY26AxEJw0kCz6vq6+7h3Iby3u6feV0VXwBMBS4Xkd04XX4zReQ5uvc9g/NvOkdVl7rPX8VJDN39vi8AdqlqvqrWAK8DU+j+9w2t3+NJf78FSyLwp+7RaU9EBKfPeJOq/rHRqXnATe7jm4B/d3ZsgaKq96lqmqpm4Px3/VRVb6Ab3zOAqh4C9onIUPfQ+cBGuvl943QJTRKRaPff+/k4Y2Hd/b6h9XucB1wrIhEikomz0deydr2zqgbFD05No63ADuBnXR1PgO7xHJwm4VpgtftzMdATZ5bBNvfPpK6ONUD3PwN4233c7e8ZGAtku/+93wR6BMl9PwhsBtYDzwIR3e2+gRdxxkBqcH7jv/V49wj8zP1u2wLMbu/nWYkJY4wJcsHSNWSMMaYVlgiMMSbIWSIwxpggZ4nAGGOCnCUCY4wJcpYITLflVuf83gm+9l0RSezgkE6YiMwXkaDboN10DksEpjtLBHwmgraqM6rqxeoUcTOm27NEYLqz3wEDRWS1iPxeRGa4+zW8AKwDEJE3RWSFW99+TsMLRWS3iCSLSIZb6//v7jUfikhU8w8SkRQReU1Elrs/U93jD4jIsyLyqVtH/jb3uLgxrReRdSLyjUbv9RP32BoR+V2jj/maiCwTka0iMi1Af2cmCAV083pjuti9wChVHQsgIjNwyvOOUqdcL8AtqlrkfrkvF5HXVLWw2fsMBq5T1dtE5GXgq8Bzza55GPiTqn4uIv2BD4Dh7rnROPtDxACrROQdYDLOyuAxQLL72QvdY1cCE1W1QkSSGn1GqKpOEJGLgftx6u4Yc9IsEZhgs6xREgC4S0Such+n43zpN08Eu1R1tft4BZDh430vAEY45W8AiBeROPfxv1X1KHBURD7DSUbnAC+qah1OMbEFwNnAucBTqloBoKqNa9I3FBFsLQZjToglAhNsyhseuC2EC4DJ7m/f84FIH6+pavS4DmjRNYTTzTrZ/cL3chND8zouiu/SwbjHW6v70hBHHfb/rulANkZgurNSIO445xOAw24SGIbTfXOiPgS+3/BERMY2OneFu69uT5zCeMuBhcA3xNlrOQVnt7Fl7vvcIiLR7vs07hoyJiAsEZhuy+3r/8IdkP29j0veB0JFZC3wK2DJSXzcXcB4EVkrIhuB2xudWwa8477/r1T1APAGTtXQNcCnwE9U9ZCqvo9TVjhbRFYDPz6JmIzxi1UfNSaAROQBoExV/9DVsRjTGmsRGGNMkLMWgTHGBDlrERhjTJCzRGCMMUHOEoExxgQ5SwTGGBPkLBEYY0yQ+//Ir5c7XFW+aQAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(mean_works)\n",
    "plt.xlabel(f\"train epoch\")\n",
    "plt.ylabel(f\"mean work (fwd)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "516e069f-4c0e-40e8-a8ca-09c0d59fbd14",
   "metadata": {},
   "source": [
    "above, I conducted the NF part of [AFT](https://arxiv.org/pdf/2102.07501.pdf), considered the case where there is a _single_ NF between my prior/posterior (both Normal distributions in 2D centerd at (0,0), (1,1), respectively). <br>\n",
    "I demonstrated that we can train the NF when it is defined by a modified leapfrog integrator as described in  eq 14 from [Nonreversible MCMC from conditional invertible transforms](https://arxiv.org/pdf/2012.15550.pdf) (thanks to [Josh Fass](https://github.com/maxentile) for pointing out the neat volume-preserving leapfrog integrator). <br>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6a10281-0c11-4d76-af4a-96b978c71db2",
   "metadata": {},
   "source": [
    "next, let's try to implement this in a multi-step (i.e. N>1 bridging distributions) and see if we can't train this sequentially. hopefully we can observe an improvement over vanilla AIS. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "83246d4a-3bdf-48be-a036-220d8c9aeef3",
   "metadata": {},
   "source": [
    "afterwards, maybe we can apply to a molecular system?!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83a34043-7208-4399-9505-800236d0fa21",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}