GAN_GRF.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d034d268",
   "metadata": {},
   "source": [
    "# GAN trained on GRF data on 2D domain\n",
    "\n",
    "The inputs to the generative model are samples of a GRF\n",
    "The data is also samples of another GRF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "separated-allowance",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import pylab as plt\n",
    "import torch.nn.functional as F\n",
    "import torch.nn as nn\n",
    "from random_fields import *\n",
    "from scipy.stats import binned_statistic"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddbd6d0e",
   "metadata": {},
   "source": [
    "#### Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "invisible-silly",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = 60-8 #resolution\n",
    "in_value_dim = 1 # dimension of the co-domain of input functions\n",
    "ntrain = 8192 # number of training samples\n",
    "width = 32 # width of the initial CNN layer.\n",
    "lr = 1e-3  # learning rate of the optimizer\n",
    "device = 'cuda:0'\n",
    "epochs = 400\n",
    "λ_grad = 10.0 # Lagrange coefficinet for gradient penalty\n",
    "n_critic = 10 # every n_critic iteration the generator is updated\n",
    "batch_size = 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial-charge",
   "metadata": {},
   "outputs": [],
   "source": [
    "# normalization, pointwise gaussian\n",
    "class InputNormalizer(object):\n",
    "    def __init__(self, x, eps=0.00001):\n",
    "        super(InputNormalizer, self).__init__()\n",
    "\n",
    "        # x could be in shape of ntrain*n or ntrain*T*n or ntrain*n*T\n",
    "        self.mean = torch.mean(x, dim=(1,2,3)).mean(dim=0)\n",
    "        self.std = torch.std(x, dim=(1,2,3)).mean(dim=0)\n",
    "        self.eps = eps\n",
    "\n",
    "    def encode(self, x):\n",
    "        x = (x - self.mean) / (self.std + self.eps)\n",
    "        return x\n",
    "\n",
    "    def decode(self, x, sample_idx=None):\n",
    "        if sample_idx is None:\n",
    "            std = self.std + self.eps # n\n",
    "            mean = self.mean\n",
    "        else:\n",
    "            if len(self.mean.shape) == len(sample_idx[0].shape):\n",
    "                std = self.std[sample_idx] + self.eps  # batch*n\n",
    "                mean = self.mean[sample_idx]\n",
    "            if len(self.mean.shape) > len(sample_idx[0].shape):\n",
    "                std = self.std[:,sample_idx]+ self.eps # T*batch*n\n",
    "                mean = self.mean[:,sample_idx]\n",
    "\n",
    "        # x is in shape of batch*n or T*batch*n\n",
    "        x = (x * std) + mean\n",
    "        return x\n",
    "\n",
    "    def cuda(self):\n",
    "        self.mean = self.mean.to(device)\n",
    "        self.std = self.std.to(device)\n",
    "\n",
    "    def cpu(self):\n",
    "        self.mean = self.mean.cpu()\n",
    "        self.std = self.std.cpu()\n",
    "        \n",
    "# auto correlation\n",
    "def compute_acovf(z):\n",
    "    z_hat = torch.fft.rfft2(z)\n",
    "    acf = torch.fft.irfft2(torch.conj(z_hat) * z_hat)\n",
    "    acf = torch.fft.fftshift(acf).mean(dim=0) / z[0].numel()\n",
    "    acf_r = acf.view(-1).cpu().detach().numpy()\n",
    "    lags_x, lags_y = torch.meshgrid(torch.arange(res) - res//2, torch.arange(res) - res//2)\n",
    "    lags_r = torch.sqrt(lags_x**2 + lags_y**2).view(-1).cpu().detach().numpy()\n",
    "\n",
    "    idx = np.argsort(lags_r)\n",
    "    lags_r = lags_r[idx]\n",
    "    acf_r = acf_r[idx]\n",
    "\n",
    "    bin_means, bin_edges, binnumber = binned_statistic(lags_r, acf_r, 'mean', bins=np.linspace(0.0, res, 50))\n",
    "    return bin_edges[:-1], bin_means"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3f3af1a",
   "metadata": {},
   "source": [
    "### Samples of input GRF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "romance-elevation",
   "metadata": {},
   "outputs": [],
   "source": [
    "grf = GaussianRF_idct(2, res, alpha=1.5, tau=1.0, device=device)\n",
    "z = grf.sample(10000)\n",
    "print(z.mean(), z.std())\n",
    "z = z.detach().cpu()\n",
    "fig, ax = plt.subplots(1,4, figsize=(16,4))\n",
    "for i in range(4):\n",
    "    bar = ax[i].imshow(z[i])\n",
    "plt.colorbar(bar)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6e8b7c5",
   "metadata": {},
   "source": [
    "### Samples of output GRF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "brown-label",
   "metadata": {},
   "outputs": [],
   "source": [
    "grf_x = GaussianRF_idct(2, res, alpha=3.0, tau=5.0, device=device)\n",
    "x_train = grf_x.sample(ntrain)\n",
    "x_train[:ntrain//2] += 1.0\n",
    "x_train[ntrain//2:] += 1.0\n",
    "idx = torch.randperm(x_train.shape[0])\n",
    "x_train = x_train[idx]\n",
    "print(x_train.mean(), x_train.std())\n",
    "x_train -= x_train.mean()\n",
    "x_train /= x_train.std()\n",
    "x_train = x_train.detach().cpu().unsqueeze(-1)\n",
    "fig, ax = plt.subplots(1,4, figsize=(16,4))\n",
    "for i in range(4):\n",
    "    bar = ax[i].imshow(x_train[i])\n",
    "plt.colorbar(bar)\n",
    "plt.show()\n",
    "\n",
    "x_hist = x_train.view(-1).cpu().detach().numpy()\n",
    "plt.hist(x_hist, bins=100, histtype='step', density=True)\n",
    "plt.show()\n",
    "\n",
    "lags_ref, acf_ref = compute_acovf(x_train.squeeze())\n",
    "plt.plot(lags_ref, acf_ref)\n",
    "plt.show()\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(x_train), batch_size=batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e16bd030",
   "metadata": {},
   "outputs": [],
   "source": [
    "def weights_init(m):\n",
    "    classname = m.__class__.__name__\n",
    "    if classname.find('Conv') != -1:\n",
    "        nn.init.normal_(m.weight.data, 0.0, 0.02)\n",
    "    elif classname.find('BatchNorm') != -1:\n",
    "        nn.init.normal_(m.weight.data, 1.0, 0.02)\n",
    "        nn.init.constant_(m.bias.data, 0)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce65d47e",
   "metadata": {},
   "source": [
    "### Generator and Discriminator functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fafc1f66",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class Generator_net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Generator_net, self).__init__()\n",
    "        \n",
    "        self.Conv1 = nn.Conv2d(in_value_dim, width, kernel_size=7, stride = 1, padding = 3)\n",
    "\n",
    "        self.Conv2 = nn.Conv2d(width, 2*width, kernel_size=7, stride = 1, padding = 3)\n",
    "\n",
    "        self.Conv3 = nn.Conv2d(2*width, 2*width, kernel_size=7, stride = 1, padding = 3)\n",
    "\n",
    "        self.Conv4 = nn.Conv2d(2* width, 1* width, kernel_size=7, stride = 1, padding = 3)\n",
    "        \n",
    "#         self.Conv5 = nn.Conv2d(2* width, 2* width, kernel_size=7, stride = 1, padding = 3)\n",
    "        \n",
    "#         self.Conv6 = nn.Conv2d(1* width, width, kernel_size=3, stride = 1, padding = 1)\n",
    "        \n",
    "        self.Conv7 = nn.Conv2d(width, in_value_dim, kernel_size=3, stride = 1, padding = 1)\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = torch.permute(x,(0,3,1,2))\n",
    "        x = F.relu(self.Conv1(x))\n",
    "        x = F.relu(self.Conv2(x))\n",
    "        x = F.relu(self.Conv3(x))\n",
    "        x = F.relu(self.Conv4(x))\n",
    "#         x = F.relu(self.Conv5(x))\n",
    "#         x = F.relu(self.Conv6(x))\n",
    "        x = self.Conv7(x)\n",
    "        x = torch.permute(x,(0,2,3,1))\n",
    "        return x\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "class Discriminator_net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Discriminator_net, self).__init__()\n",
    "\n",
    "        self.Conv1 = nn.Conv2d(in_value_dim, width, kernel_size=5, stride = 2, padding = 2)        \n",
    "        \n",
    "        self.Conv2 = nn.Conv2d(width, 1* width, kernel_size=5, stride = 2, padding = 2)\n",
    "        \n",
    "        self.Conv3 = nn.Conv2d(1* width, 2* width, kernel_size=7, stride = 2, padding = 3)\n",
    "\n",
    "#         self.Conv4 = nn.Conv2d(4* width, 2*width, kernel_size=6, stride = 4, padding = 2)\n",
    "\n",
    "        self.Conv5 = nn.Conv2d(2*width, 1*width, kernel_size=7, stride = 3, padding = 3)\n",
    "        self.fc = nn.Linear(288, 1)\n",
    "\n",
    "#         self.Conv4 = nn.Conv2d(4* width, width, kernel_size=4, stride = 2, padding = 2)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x = torch.permute(x,(0,3,1,2))\n",
    "        x = F.leaky_relu(self.Conv1(x))\n",
    "        x = F.leaky_relu(self.Conv2(x))\n",
    "        x = F.leaky_relu(self.Conv3(x))\n",
    "#         x = F.leaky_relu(self.Conv4(x))\n",
    "        x = F.leaky_relu(self.Conv5(x))\n",
    "        x = torch.flatten(x, start_dim=1)\n",
    "        x = self.fc(x)\n",
    "\n",
    "        return x\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3581da3",
   "metadata": {},
   "outputs": [],
   "source": [
    "D_net = Discriminator_net().to(device)\n",
    "G_net = Generator_net().to(device)\n",
    "D_net.apply(weights_init)\n",
    "G_net.apply(weights_init)\n",
    "nn_params = sum(p.numel() for p in D_net.parameters() if p.requires_grad)\n",
    "print(\"Number discriminator parameters: \", nn_params)\n",
    "nn_params = sum(p.numel() for p in G_net.parameters() if p.requires_grad)\n",
    "print(\"Number generator parameters: \", nn_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85dbfd60",
   "metadata": {},
   "source": [
    "### Samples generated with initilized generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15d29ca2",
   "metadata": {},
   "outputs": [],
   "source": [
    "z = grf.sample(4).unsqueeze(-1)\n",
    "print(z.shape)\n",
    "x = G_net(z)\n",
    "print(x.shape)\n",
    "fig, ax = plt.subplots(1,4, figsize=(16,4))\n",
    "for i in range(4):\n",
    "    ax[i].imshow(x[i].cpu().detach().numpy())\n",
    "plt.show()\n",
    "print(D_net(x))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a620882f",
   "metadata": {},
   "outputs": [],
   "source": [
    "G_optimizer_net = torch.optim.Adam(G_net.parameters(), lr=lr)\n",
    "D_optimizer_net = torch.optim.Adam(D_net.parameters(), lr=lr)\n",
    "\n",
    "D_net.train()\n",
    "G_net.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "soviet-distinction",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_gradient_penalty(model, real_images, fake_images, device):\n",
    "    \"\"\"Calculates the gradient penalty loss for GAN\"\"\"\n",
    "    # Random weight term for interpolation between real and fake data\n",
    "    alpha = torch.randn((real_images.size(0), 1, 1, 1), device=device)\n",
    "    # Get random interpolation between real and fake data\n",
    "    interpolates = (alpha * real_images + ((1 - alpha) * fake_images)).requires_grad_(True)\n",
    "\n",
    "    model_interpolates = model(interpolates)\n",
    "    grad_outputs = torch.ones(model_interpolates.size(), device=device, requires_grad=False)\n",
    "\n",
    "    # Get gradient w.r.t. interpolates\n",
    "    gradients = torch.autograd.grad(\n",
    "        outputs=model_interpolates,\n",
    "        inputs=interpolates,\n",
    "        grad_outputs=grad_outputs,\n",
    "        create_graph=True,\n",
    "        retain_graph=True,\n",
    "        only_inputs=True,\n",
    "    )[0]\n",
    "    gradients = gradients.view(gradients.size(0), -1)\n",
    "    gradient_penalty = torch.mean((gradients.norm(2, dim=1) - 1/np.sqrt(res * res)) ** 2)\n",
    "    return gradient_penalty"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acute-collect",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_GAN(D, G, train_data, epochs, D_optim, G_optim, scheduler=None):\n",
    "    losses_D = np.zeros(epochs)\n",
    "    losses_G = np.zeros(epochs)\n",
    "    for i in range(epochs):\n",
    "        loss_D = 0.0\n",
    "        loss_G = 0.0\n",
    "        for j, data in enumerate(train_data):\n",
    "            # Train D\n",
    "            x = data[0].to(device)\n",
    "            D_optimizer.zero_grad()\n",
    "\n",
    "            x_syn = G(grf.sample(x.shape[0]).unsqueeze(-1))\n",
    "\n",
    "            W_loss = -torch.mean(D(x)) + torch.mean(D(x_syn.detach()))\n",
    "\n",
    "            gradient_penalty = calculate_gradient_penalty(D, x.data, x_syn.data, device)\n",
    "\n",
    "            loss = W_loss + λ_grad * gradient_penalty\n",
    "            loss.backward()\n",
    "\n",
    "            loss_D += loss.item()\n",
    "\n",
    "            D_optim.step()\n",
    "            \n",
    "            # Train G\n",
    "            if (j + 1) % n_critic == 0:\n",
    "                G_optimizer.zero_grad()\n",
    "\n",
    "                x_syn = G(grf.sample(x.shape[0]).unsqueeze(-1))\n",
    "\n",
    "                loss = -torch.mean(D(x_syn))\n",
    "                loss.backward()\n",
    "                loss_G += loss.item()\n",
    "\n",
    "                G_optim.step()\n",
    "            \n",
    "        losses_D[i] = loss_D / (ntrain / batch_size)\n",
    "        losses_G[i] = loss_G / (ntrain / batch_size * n_critic)\n",
    "            \n",
    "        with torch.no_grad():\n",
    "            z = grf.sample(1000).unsqueeze(-1)\n",
    "            x = G(z)\n",
    "            fig, ax = plt.subplots(1,5, figsize=(20,4))\n",
    "            for j in range(3):\n",
    "                ax[j].imshow(x[j].cpu().detach().numpy())\n",
    "            lags, acf = compute_acovf(x.squeeze())\n",
    "            ax[3].hist(x.view(-1).cpu().detach().numpy(), bins=100, histtype='step', density=True)\n",
    "            ax[3].hist(x_hist, bins=100, histtype='step', density=True)\n",
    "            ax[4].plot(lags_ref, acf_ref, c='r')\n",
    "            ax[4].plot(lags, acf, c='k')\n",
    "            plt.show()\n",
    "        print(i, \"D: \", losses_D[i], \"G: \", losses_G[i], \"mean: \", x.mean().item(), \"std: \", x.std().item())\n",
    "        storedata(i,G)\n",
    "\n",
    "    return losses_D, losses_G"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e0911bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def storedata(ite,G):\n",
    "    high = 2\n",
    "    low = 2\n",
    "    numb_fig = 5\n",
    "    fig, ax = plt.subplots(1,numb_fig, figsize=(16,4))\n",
    "    for i in range(numb_fig):\n",
    "        bar = ax[i].imshow(x_train[i],extent=[0,1,0,1])\n",
    "#         bar.set_clim(low, high)\n",
    "    cax = fig.add_axes([ax[numb_fig-1].get_position().x1+0.01,ax[numb_fig-1].get_position().y0,0.02,ax[numb_fig-1].get_position().height])\n",
    "    plt.colorbar(bar, cax=cax) # Similar to fig.colorbar(im, cax = cax)\n",
    "\n",
    "    plt.savefig('~/GANO/Figures/GAN_GRF/{}GP.pdf'.format(ite))  \n",
    "#     plt.show()\n",
    "\n",
    "\n",
    "############################################\n",
    "\n",
    "    with torch.no_grad():\n",
    "        z = grf.sample(numb_fig).unsqueeze(-1)\n",
    "        x = G(z)\n",
    "        fig, ax = plt.subplots(1,numb_fig, figsize=(16,4))\n",
    "    #     ax.set_xlabel([0,1])\n",
    "        for j in range(numb_fig):\n",
    "            bar = ax[j].imshow(x[j].cpu().detach().numpy(),extent=[0,1,0,1])\n",
    "#             bar.set_clim(low, high)\n",
    "        cax = fig.add_axes([ax[numb_fig-1].get_position().x1+0.01,ax[numb_fig-1].get_position().y0,0.02,ax[numb_fig-1].get_position().height])\n",
    "        plt.colorbar(bar, cax=cax) # Similar to fig.colorbar(im, cax = cax)\n",
    "        plt.savefig('~/GANO/Figures/GAN_GRF/{}GRF_GAN.pdf'.format(ite))  \n",
    "#         plt.show()\n",
    "        \n",
    "        \n",
    "############################################\n",
    "        \n",
    "    with torch.no_grad():\n",
    "        z = grf.sample(1000).unsqueeze(-1)\n",
    "        x = G(z)\n",
    "        lags, acf = compute_acovf(x.squeeze())\n",
    "        \n",
    "\n",
    "\n",
    "    fig, ax = plt.subplots(1,1, figsize=(4,4),tight_layout=True)\n",
    "    ax.hist(x_hist, bins=1000, histtype='step', density=True, color='#ff7f0e', label='Ground truth')\n",
    "    ax.hist(x.view(-1).cpu().detach().numpy(), bins=1000, color='#1f77b4' , histtype='step', density=True, label='GAN')\n",
    "\n",
    "    # ax[1].plot(lags_ref, acf_ref, c='r')\n",
    "    # ax[1].plot(lags, acf, c='k')\n",
    "    plt.xlabel('value')\n",
    "    # plt.title('Histogram')\n",
    "    plt.ylabel('Histogram')\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "    plt.legend()\n",
    "\n",
    "    plt.savefig('~/GANO/Figures/GAN_GRF/{}GRF_HistogramGAN.pdf'.format(ite))  \n",
    "#     plt.show()\n",
    "    \n",
    "############################################\n",
    "\n",
    "    fig, ax = plt.subplots(1,1, figsize=(4,4),tight_layout=True)\n",
    "    ax.plot(lags_ref/max(lags_ref), acf_ref, c='#ff7f0e',label='Ground truth'.format(ite))\n",
    "    ax.plot(lags/max(lags_ref), acf,c='#1f77b4' , label='GAN')\n",
    "    plt.xlabel('Position')\n",
    "    plt.ylabel('Auto correlation')\n",
    "    # plt.title('Auto correlation')\n",
    "    plt.legend()\n",
    "    plt.savefig('~/GANO/Figures/GAN_GRF/{}GRF_AutoCorrelationGAN.pdf'.format(ite))  \n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aging-merchandise",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "losses_D, losses_G = train_GAN(D_net, G_net, train_loader, epochs, D_optimizer_net, G_optimizer_net)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "limiting-colon",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(np.arange(epochs), losses_D, c='k', label='D')\n",
    "plt.plot(np.arange(epochs), losses_G, c='b', label='G')\n",
    "plt.legend()\n",
    "plt.ylabel(\"Loss\")\n",
    "plt.xlabel(\"Epochs\")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "“my_environment”",
   "language": "python",
   "name": "myenv"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}