{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "715ecee1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import math\n",
    "\n",
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.autograd as autograd\n",
    "\n",
    "from torch import distributions as D\n",
    "\n",
    "def set_pub():\n",
    "    plt.rcParams.update({\n",
    "        #\"font.weight\": \"bold\",  # bold fonts\n",
    "        \"xtick.labelsize\": 13,\n",
    "        \"ytick.labelsize\": 13,\n",
    "        \"axes.labelsize\": 13,\n",
    "        #\"tick.labelsize\": 15,   # large tick labels\n",
    "        #\"lines.linewidth\": 1,   # thick lines\n",
    "        #\"lines.color\": \"k\",     # black lines\n",
    "        \"grid.color\": \"0.5\",    # gray gridlines\n",
    "        \"grid.linestyle\": \"-\",  # solid gridlines\n",
    "        \"grid.linewidth\": 0.5,  # thin gridlines\n",
    "        \"savefig.dpi\": 300,     # higher resolution output.\n",
    "    })"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "369b27e6",
   "metadata": {},
   "source": [
    "### Definition of the Energy Discrepancy Loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5ab4a192",
   "metadata": {},
   "outputs": [],
   "source": [
    "def energy_discrepancy(energy_net, samples, m_particles=4, t_noise=1.0, w_stable=1.0, seed = None):\n",
    "    device = samples.device\n",
    "    \n",
    "    if seed != None:\n",
    "        torch.manual_seed(seed)\n",
    "\n",
    "    browniani = torch.randn_like(samples).to(device) * np.sqrt(t_noise)\n",
    "    brownianij = torch.randn(samples.size(0), m_particles, *samples.shape[1:]).to(device) * np.sqrt(t_noise)\n",
    "\n",
    "    pert_data = samples.unsqueeze(1) + browniani.unsqueeze(1) + brownianij\n",
    "\n",
    "    pos_energy = energy_net(samples)\n",
    "    neg_energy = energy_net(pert_data.view(-1, *samples.shape[1:])).view(samples.size(0), -1)\n",
    "    val = pos_energy - neg_energy\n",
    "    if w_stable != 0:\n",
    "        val = torch.cat([val, np.log(w_stable) * torch.ones_like(val[:, :1])], dim=-1)\n",
    "\n",
    "    loss = val.logsumexp(dim=-1).mean()\n",
    "    if t<1.:\n",
    "        loss = loss/t\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b46b28f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SmallMLP(nn.Module):\n",
    "    def __init__(self, n_dims, n_hid=8, layer=nn.Linear):\n",
    "        super(SmallMLP, self).__init__()\n",
    "        self._built = False\n",
    "        self.net = nn.Sequential(\n",
    "            layer(n_dims, n_hid),\n",
    "            nn.SiLU(),\n",
    "            layer(n_hid, n_hid),\n",
    "            nn.SiLU(),\n",
    "            layer(n_hid, 1)\n",
    "        )\n",
    "        self.normalized = False\n",
    "\n",
    "    def forward(self, u):\n",
    "        out = self.net(u)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "baf89d03",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "def train(data, model = None, n=4096, m=4, t=1., w=1., n_epochs=20, batch_size=64):    \n",
    "        \n",
    "    opt = torch.optim.Adam(model.parameters(), lr=0.01)\n",
    "    \n",
    "    dataloader = torch.utils.data.DataLoader(data, batch_size=batch_size,\n",
    "                                          shuffle=True)\n",
    "    \n",
    "    loss_history = [0.]\n",
    "    for i in tqdm(range(n_epochs)):\n",
    "        batch_loss = 0.\n",
    "        for batch in dataloader:\n",
    "            opt.zero_grad()\n",
    "            loss_val = energy_discrepancy(energy_net = model, samples = data, m_particles=m, t_noise=t, w_stable=w)\n",
    "            loss_val.backward()\n",
    "            opt.step()\n",
    "            batch_loss += loss_val.item()\n",
    "            \n",
    "        batch_loss = batch_loss/len(dataloader)\n",
    "        loss_history.append(batch_loss)\n",
    "    \n",
    "    return model, loss_history"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "5d646866",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mse_log_prob(distribution, energy_net, device=None):\n",
    "    samples = distribution.sample((10000,))\n",
    "\n",
    "    if device is not None:\n",
    "        samples = samples.to(device)\n",
    "\n",
    "    log_ground_truth = -distribution.log_prob(samples)\n",
    "    \n",
    "\n",
    "    log_prob = -energy_net(samples.unsqueeze(-1))\n",
    "    print(log_prob.shape)\n",
    "    logZ = torch.logsumexp(log_prob.squeeze(-1)-log_ground_truth, dim=(0)) - np.log(10000)\n",
    "    print(logZ.shape)\n",
    "    log_prob_ebm = log_prob - logZ\n",
    "    print(log_prob_ebm.shape)\n",
    "\n",
    "    return torch.mean((log_prob_ebm - log_ground_truth) ** 2).cpu().item(), logZ.cpu().item()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1c212d7",
   "metadata": {},
   "source": [
    "### Energy Estimation for various choices of w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "67d7da0e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████| 50/50 [00:08<00:00,  6.04it/s]\n",
      "100%|███████████████████████████████████████████| 50/50 [00:08<00:00,  5.90it/s]\n",
      "100%|███████████████████████████████████████████| 50/50 [00:08<00:00,  5.98it/s]\n",
      "100%|███████████████████████████████████████████| 50/50 [00:08<00:00,  5.89it/s]\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 46320 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "set_pub()\n",
    "m = 4\n",
    "n = 4096\n",
    "t = 1.\n",
    "ws = [0., 0.1, 1., 4.]\n",
    "x = torch.arange(-3, 3, 0.1)\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "data = torch.randn(n, 1)\n",
    "y_true = x**2/2\n",
    "plt.plot(x, y_true, '--', label = f'ground_truth', color = 'grey')\n",
    "\n",
    "for w in ws:\n",
    "    model = SmallMLP(n_dims = 1, n_hid = 2)\n",
    "    model, loss_history = train(data, model = model, n=n, m=m, t=t, w=w, n_epochs=50, batch_size=64)\n",
    "    minimum = model(x.unsqueeze(1)).min()\n",
    "    y = model(x.unsqueeze(1)).detach()-minimum.detach()\n",
    "    plt.plot(x, y, label = f'w = {w}')\n",
    "    \n",
    "plt.ylim(-0.1, 3.)\n",
    "plt.xlabel(\"x\")\n",
    "plt.ylabel(\"estimated energy\")\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.spines['right'].set_visible(False)\n",
    "plt.grid(linestyle='--', linewidth=0.5, alpha=0.5)\n",
    "plt.grid(True)\n",
    "plt.legend(frameon = False)\n",
    "plt.savefig('Images/Tuningwform=4.pdf', dpi=300)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bf6b333",
   "metadata": {},
   "source": [
    "### Comparison of Nearsightedness of Energy Discrepancy, Maximum Likelihood Estimation, and Score Matching"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "504baea5",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Mixture_model():\n",
    "    def __init__(self, distribution_a, distribution_b, weights = None):\n",
    "        self.weights = weights\n",
    "        self.distribution_a = distribution_a\n",
    "        self.distribution_b = distribution_b\n",
    "        \n",
    "    def set_parameter(self, value):\n",
    "        self.weights = value\n",
    "        \n",
    "    def sample(self, n_samples):\n",
    "        alpha = self.weights\n",
    "        b = torch.distributions.Bernoulli(torch.tensor([alpha]))\n",
    "\n",
    "        data = torch.zeros(n_samples, 1)\n",
    "\n",
    "        index = b.sample((n_samples,))\n",
    "\n",
    "        for i in range(n_samples):\n",
    "            if index[i]==1:\n",
    "                data[i, :]=self.distribution_a.sample()\n",
    "            else:\n",
    "                data[i, :]=self.distribution_b.sample()\n",
    "        return data\n",
    "\n",
    "    def energy(self, x):\n",
    "        alpha = self.weights\n",
    "\n",
    "        log_prob_a = distribution_a.log_prob(x) + np.log(alpha)\n",
    "        log_prob_b = distribution_b.log_prob(x) + np.log(1-alpha)\n",
    "\n",
    "        log_probs = torch.cat((log_prob_a.unsqueeze(1), log_prob_b.unsqueeze(1)), 1)\n",
    "\n",
    "        energy = -torch.logsumexp(log_probs, 1)\n",
    "        return energy\n",
    "    \n",
    "    def negative_loglikelihood(self, data):\n",
    "        return self.energy(data).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8678093f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x129bba190>]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 16128 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "rho = 0.2\n",
    "distribution_a = D.Normal(loc = torch.tensor([-5.]), scale = torch.tensor([1.0]))\n",
    "distribution_b = D.Normal(loc = torch.tensor([5.]), scale = torch.tensor([1.0]))\n",
    "\n",
    "gauss_mixture = Mixture_model(distribution_a, distribution_b, rho)\n",
    "\n",
    "x = torch.arange(-10., 10., 0.1)\n",
    "plt.plot(x, np.exp(-gauss_mixture.energy(x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c15a124a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import minimize_scalar\n",
    "\n",
    "def energy_discrepancy_opt_fn(rho, *args):\n",
    "    data, seed, m, t, w = args\n",
    "    gauss_mixture.set_parameter(rho)\n",
    "    val = energy_discrepancy(energy_net = gauss_mixture.energy, samples = data, m_particles=m, t_noise=t, w_stable=w, seed = seed)\n",
    "    return val\n",
    "\n",
    "def mle_opt_fn(rho, *args):\n",
    "    data, = args\n",
    "    gauss_mixture.set_parameter(rho)\n",
    "    return gauss_mixture.negative_loglikelihood(data)\n",
    "\n",
    "def score(energy, samples):\n",
    "    with torch.enable_grad():\n",
    "        samples = samples.detach()\n",
    "        samples.requires_grad_(True)\n",
    "        log_probs = energy(samples).sum()\n",
    "        return autograd.grad(log_probs, samples)[0]\n",
    "\n",
    "def fish_divergence(rho, *args):\n",
    "    gauss_mixture.set_parameter(0.2)\n",
    "    energy_truth = gauss_mixture.energy\n",
    "    gauss_mixture.set_parameter(rho)\n",
    "    #samples = selsample(num_samples)\n",
    "    return 0.5 * ((score(gauss_mixture.energy, data) - score(energy_truth, data))**2).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "43c47e08",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 48888 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "set_pub()\n",
    "\n",
    "gauss_mixture.set_parameter(0.2)\n",
    "data = gauss_mixture.sample(4096)\n",
    "\n",
    "rhos = np.arange(0.01, 0.99, 0.01)\n",
    "\n",
    "seed = np.random.randint(100000)\n",
    "\n",
    "ts = [1, 2, 4, 8, 32]\n",
    "fig, ax = plt.subplots()\n",
    "ax.plot(rhos, [fish_divergence(rho, data) for rho in rhos], label = 'SM', linestyle = '--')\n",
    "for t in ts[1:]:\n",
    "    losses = []\n",
    "    for rho in rhos:\n",
    "        m = 4\n",
    "        w = 1.\n",
    "        gauss_mixture.set_parameter(rho)\n",
    "        val = energy_discrepancy(energy_net = gauss_mixture.energy, samples = data, m_particles=m, t_noise=t, w_stable=w, seed = seed)\n",
    "        losses.append(val)\n",
    "\n",
    "    losses = np.array(losses)\n",
    "\n",
    "    min_val = min(losses)\n",
    "    ax.plot(rhos, losses-min_val, label = f'ED_t={t}')\n",
    "    \n",
    "losses = []\n",
    "for rho in rhos:\n",
    "    gauss_mixture.set_parameter(rho)\n",
    "    val = gauss_mixture.negative_loglikelihood(data)\n",
    "    losses.append(val)\n",
    "\n",
    "losses = np.array(losses)\n",
    "\n",
    "min_val = min(losses)\n",
    "ax.plot(rhos, losses-min_val, label = f'MLE', linestyle = '--')\n",
    "    \n",
    "#ax.plot(0.2, 0.,  marker='*')\n",
    "plt.legend(fontsize=13, loc = 'upper right', frameon = True, edgecolor = 'white')\n",
    "plt.ylim(-0.01, 0.1)\n",
    "plt.xlabel(\"mixture weight\")\n",
    "plt.ylabel(\"loss (SM / ED/ CE)\")\n",
    "#ax.annotate('rho', xy=(0.2, 0.))\n",
    "ax.axvline(0.2, color = 'grey', linestyle = '--')\n",
    "plt.grid(linestyle='--', linewidth=0.5, alpha=0.5)\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.spines['right'].set_visible(False)\n",
    "\n",
    "#plt.savefig('Images/loss_Gauss_mixture')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "14e924c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "w = 1.\n",
    "\n",
    "ts = [0.8, 1.6, 3.2, 6.4, 12.8, 25.6]\n",
    "gauss_mixture.set_parameter(0.2)\n",
    "data = gauss_mixture.sample(4096)\n",
    "\n",
    "estimators_dict = {}\n",
    "estimators_dict64 = {}\n",
    "\n",
    "for i in range(50):\n",
    "    seed = np.random.randint(100000)\n",
    "    if 'mle' not in estimators_dict.keys():\n",
    "        estimators_dict['mle'] = []\n",
    "    else:\n",
    "        estimators_dict['mle'].append(minimize_scalar(mle_opt_fn, bounds = (0.01, 0.99), args = (data,), method='bounded').x) \n",
    "    m=8\n",
    "    for k in range(len(ts)):\n",
    "        t = ts[k]\n",
    "        if t not in estimators_dict.keys():\n",
    "            estimators_dict[t] = []\n",
    "        else:\n",
    "            estimators_dict[t].append(minimize_scalar(energy_discrepancy_opt_fn, bounds = (0.01, 0.99), args = (data, seed, m, t, w), method='bounded').x)  \n",
    "    m=64\n",
    "    for k in range(len(ts)):\n",
    "        t = ts[k]\n",
    "        if t not in estimators_dict64.keys():\n",
    "            estimators_dict64[t] = []\n",
    "        else:\n",
    "            estimators_dict64[t].append(minimize_scalar(energy_discrepancy_opt_fn, bounds = (0.01, 0.99), args = (data, seed, m, t, w), method='bounded').x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "46707281",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 12532 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "ts = [0.8, 1.6, 3.2, 6.4, 12.8, 25.6]\n",
    "set_pub()\n",
    "mse_dict = {}\n",
    "mse_dict64 = {}\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "for t in estimators_dict.keys():\n",
    "    results = torch.tensor(estimators_dict[t])\n",
    "    mean = results.mean()\n",
    "    std = results.std()\n",
    "    mse = (results - 0.2).pow(2).mean()\n",
    "    mse_dict[t] = mse\n",
    "\n",
    "for t in estimators_dict64.keys():\n",
    "    results = torch.tensor(estimators_dict64[t])\n",
    "    mean = results.mean()\n",
    "    std = results.std()\n",
    "    mse = (results - 0.2).pow(2).mean()\n",
    "    mse_dict64[t] = mse\n",
    "\n",
    "\n",
    "ax.plot(ts, [mse_dict[t] for t in ts], '+', markersize = 10, label = f'M = 8')\n",
    "ax.plot(ts, [mse_dict64[t] for t in ts], '*', markersize = 10, label = f'M = 64')\n",
    "ax.axhline(mse_dict['mle'], linestyle='--', label = f'MLE', color = 'green')\n",
    "ax.set_ylabel('MSE')\n",
    "ax.set_xlabel('t')\n",
    "ax.set_xscale('log')\n",
    "ax.set_yscale('log')\n",
    "ax.spines['top'].set_visible(False)\n",
    "ax.spines['right'].set_visible(False)\n",
    "plt.legend(fontsize=13, frameon=False)\n",
    "plt.grid(linestyle='--', linewidth=0.5, alpha=0.5)\n",
    "#plt.savefig('Images/Mixture_estimates_MSE_t')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09b56439",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
