GAN_volcano.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "08fb2831",
   "metadata": {},
   "source": [
    "# GAN trained on volcano data on 2D domain\n",
    "\n",
    "The inputs to the generative model are samples of a GRF\n"
   ]
  },
  {
   "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 *"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e38b763a",
   "metadata": {},
   "source": [
    "#### Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "invisible-silly",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = 128-8 #resolution\n",
    "ntrain = 4096 # number of training samples\n",
    "width= 96 # the dimension of the co-domain of the initial U-NO layer.\n",
    "lr = 1e-4 # learning rate of the optimizer\n",
    "device = 'cuda:0'\n",
    "epochs = 400\n",
    "λ_grad = 10.0 # Lagrange coefficinet for gradient penalty\n",
    "n_critic = 20 # every n_critic iteration the generator is updated\n",
    "in_value_dim = 1 # dimension of the co-domain of input functions\n",
    "out_value_dim = 2 # dimension of the co-domain of output functions\n",
    "batch_size = 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "tropical-upgrade",
   "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",
    "def compute_acovf(z):\n",
    "    from scipy.stats import binned_statistic\n",
    "    z_hat = torch.fft.rfft2(z)\n",
    "    acf = torch.fft.irfft2(torch.conj(z_hat) * z_hat) / (z[0].numel())\n",
    "    acf = torch.fft.fftshift(acf).mean(dim=0)\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": "0a79f110",
   "metadata": {},
   "source": [
    "### Load the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "noble-association",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from IPython.core.display import display, HTML\n",
    "display(HTML(\"<style>.container { width:100% !important; }</style>\"))\n",
    "import glob\n",
    "files = glob.glob('~/GANO/volcano_data/**/*.int', recursive=True)[:ntrain]\n",
    "x_train = torch.zeros(ntrain, res, res, 2).float()\n",
    "for i, f in enumerate(files):\n",
    "    dtype = np.float32\n",
    "    nline = 128\n",
    "    nsamp = 128\n",
    "\n",
    "    with open(f, 'rb') as fn:\n",
    "        load_arr = np.frombuffer(fn.read(), dtype=dtype)\n",
    "        img = np.array(load_arr.reshape((nline, nsamp, -1)))\n",
    "\n",
    "    phi = np.angle(img[:,:,0] + img[:,:,1]*1j)\n",
    "    x_train[i,:,:,0] = torch.cos(torch.tensor(phi[:res, :res]))\n",
    "    x_train[i,:,:,1] = torch.sin(torch.tensor(phi[:res, :res]))    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "environmental-retirement",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Directional statistics\n",
    "def circular_var(x, dim=None):\n",
    "    #R = torch.sqrt((x.mean(dim=(1,2))**2).sum(dim=1))\n",
    "    phase = torch.atan2(x[:,:,:,1], x[:,:,:,0])\n",
    "    phase = (phase + np.pi) % (2 * np.pi) - np.pi\n",
    "    \n",
    "    C1 = torch.cos(phase).sum(dim=(1,2))\n",
    "    S1 = torch.sin(phase).sum(dim=(1,2))\n",
    "    R1 = torch.sqrt(C1**2 + S1**2) / (phase.shape[1]*phase.shape[2])\n",
    "    return 1 - R1\n",
    "\n",
    "def circular_skew(x):\n",
    "    phase = torch.atan2(x[:,:,:,1], x[:,:,:,0])\n",
    "    phase = (phase + np.pi) % (2 * np.pi) - np.pi\n",
    "    \n",
    "    C1 = torch.cos(phase).sum(dim=(1,2))\n",
    "    S1 = torch.sin(phase).sum(dim=(1,2))\n",
    "    R1 = torch.sqrt(C1**2 + S1**2) / (phase.shape[1]*phase.shape[2])\n",
    "    \n",
    "    C2 = torch.cos(2*phase).sum(dim=(1,2))\n",
    "    S2 = torch.sin(2*phase).sum(dim=(1,2))\n",
    "    R2 = torch.sqrt(C2**2 + S2**2) / (phase.shape[1]*phase.shape[2])\n",
    "    \n",
    "    T1 = torch.atan2(S1, C1)\n",
    "    T2 = torch.atan2(S2, C2)\n",
    "\n",
    "    return R2 * torch.sin(T2 - 2*T1) / (1 - R1)**(3/2)\n",
    "\n",
    "\n",
    "var = circular_var(x_train)\n",
    "plt.hist(var.numpy(), bins=np.linspace(0.0, 1.0, 25), histtype='step', linewidth=2.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "recreational-narrative",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "print(x_train.shape)\n",
    "for i in range(10):\n",
    "    fig, ax = plt.subplots(1,1, figsize=(20,4))\n",
    "    j = torch.randint(x_train.shape[0], size=(1,))[0]\n",
    "    phase = torch.atan2(x_train[j,:,:,1], x_train[j,:,:,0]).cpu().detach().numpy()\n",
    "    phase = (phase + np.pi) % (2 * np.pi) - np.pi\n",
    "    bar = ax.imshow(phase,  cmap='RdYlBu', vmin = -np.pi, vmax=np.pi, interpolation=None)\n",
    "    plt.colorbar(bar)\n",
    "plt.show()\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(x_train), batch_size=batch_size, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "moral-rough",
   "metadata": {},
   "outputs": [],
   "source": [
    "numb_fig = 5\n",
    "ite = 200\n",
    "fig, ax = plt.subplots(1,numb_fig, figsize=(16,4))\n",
    "for j in range(numb_fig):\n",
    "    phase = torch.atan2(x_train[j,:,:,1], x_train[j,:,:,0]).cpu().detach().numpy()\n",
    "    phase = (phase + np.pi) % (2 * np.pi) - np.pi\n",
    "    bar = ax[j].imshow(phase,  cmap='RdYlBu', vmin = -np.pi, vmax=np.pi,extent=[0,1,0,1])\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/volcano/{}real.pdf'.format(ite))  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfddbb59",
   "metadata": {},
   "source": [
    "### Data visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "compact-limitation",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "datapoint = 1 #index of a data point\n",
    "numb_fig = 4\n",
    "ite = 200\n",
    "fig, ax = plt.subplots(5,numb_fig, figsize=(14,16))\n",
    "for i in range(5):\n",
    "    for j in range(numb_fig):\n",
    "        phase = torch.atan2(x_train[datapoint+i*numb_fig+j,:,:,1], x_train[datapoint+i*numb_fig+j,:,:,0]).cpu().detach().numpy()\n",
    "        phase = (phase + np.pi) % (2 * np.pi) - np.pi\n",
    "        bar = ax[i,j].imshow(phase,  cmap='RdYlBu', vmin = -np.pi, vmax=np.pi,extent=[0,1,0,1])\n",
    "#cax = fig.add_axes([ax[i,numb_fig-1].get_position().x1+0.01,ax[i,numb_fig-1].get_position().y0,0.02,ax[i,numb_fig-1].get_position().height])\n",
    "#plt.colorbar(bar, cax=cax) # Similar to fig.colorbar(im, cax = cax)\n",
    "plt.savefig('~/GANO/Figures/volcano/{}real.pdf'.format(ite)) \n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06779ba8",
   "metadata": {},
   "source": [
    "### Samples of input GRF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "romance-elevation",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "grf = GaussianRF_idct(2, res, alpha=1.5, tau=1.0, device=device)\n",
    "numb_fig = 5\n",
    "z = grf.sample(5)\n",
    "print(z.mean(), z.std())\n",
    "z = z.detach().cpu()\n",
    "fig, ax = plt.subplots(1,numb_fig, figsize=(16,4))\n",
    "for i in range(5):\n",
    "    bar = ax[i].imshow(z[i], extent=[0,1,0,1])\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)\n",
    "plt.savefig('~/GANO/Figures/volcano/{}GRF.pdf'.format(ite))  \n",
    "#plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8338709f",
   "metadata": {},
   "source": [
    "### Generator and Discriminator functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b83d359f",
   "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",
    "\n",
    "\n",
    "        \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, 2* width, kernel_size=7, stride = 1, padding = 3)\n",
    "        self.Conv4_ = nn.Conv2d(2* width, 4* width, kernel_size=7, stride = 1, padding = 3)\n",
    "        self.Conv4__ = nn.Conv2d(4* width, 4* width, kernel_size=7, stride = 1, padding = 3)\n",
    "        self.Conv4___ = nn.Conv2d(4* width, 2* 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(2* width, width, kernel_size=3, stride = 1, padding = 1)\n",
    "        \n",
    "        self.Conv7 = nn.Conv2d(width, out_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.Conv4_(x))\n",
    "        x = F.relu(self.Conv4__(x))\n",
    "        x = F.relu(self.Conv4___(x))\n",
    "        x = F.relu(self.Conv5(x))\n",
    "        x = F.relu(self.Conv6(x))\n",
    "        x = 1.1*F.tanh(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(out_value_dim, width, kernel_size=5, stride = 2, padding = 2)        \n",
    "        \n",
    "        self.Conv2 = nn.Conv2d(width, 2* width, kernel_size=5, stride = 2, padding = 2)\n",
    "        \n",
    "        self.Conv2_ = nn.Conv2d(2*width, 2* width, kernel_size=5, stride = 1, padding = 2)\n",
    "        \n",
    "        self.Conv2__ = nn.Conv2d(2*width, 2* width, kernel_size=5, stride = 1, padding = 2)\n",
    "        \n",
    "        self.Conv3 = nn.Conv2d(2* width, 4* 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(384, 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.Conv2_(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",
    "#         print(x.shape)\n",
    "        x = self.fc(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bd7efc4",
   "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)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68827f12",
   "metadata": {},
   "source": [
    "### Samples generated with initilized generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "interesting-rogers",
   "metadata": {},
   "outputs": [],
   "source": [
    "z = grf.sample(4).unsqueeze(-1)\n",
    "x = G_net(z)\n",
    "D_net(x)\n",
    "fig, ax = plt.subplots(1,4, figsize=(16,4))\n",
    "for i in range(4):\n",
    "    phase = torch.atan2(x[i,:,:,1], x[i,:,:,0]).cpu().detach().numpy()\n",
    "    phase = (phase + np.pi) % (2 * np.pi) - np.pi\n",
    "    bar = ax[i].imshow(phase,  cmap='RdYlBu', vmin = -np.pi, vmax=np.pi, interpolation=None)\n",
    "plt.colorbar(bar)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "published-warner",
   "metadata": {},
   "outputs": [],
   "source": [
    "G_optimizer_net = torch.optim.Adam(G_net.parameters(), lr=lr) #, weight_decay=1e-4)\n",
    "D_optimizer_net = torch.optim.Adam(D_net.parameters(), lr=lr) #, weight_decay=1e-4)\n",
    "\n",
    "D_net.train()\n",
    "G_net.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "italic-paper",
   "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.reshape(gradients.size(0), -1)\n",
    "    gradient_penalty = torch.mean((gradients.norm(2, dim=1) - 1.0/res) ** 2)\n",
    "    return gradient_penalty"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "geographic-florence",
   "metadata": {},
   "outputs": [],
   "source": [
    "def storedata(ite,G):\n",
    "    high = 2\n",
    "    low = 0.3\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/volcano/{}GRF.pdf'.format(ite))  \n",
    "    plt.show()\n",
    "\n",
    "\n",
    "############################################\n",
    "\n",
    "    with torch.no_grad():\n",
    "        z = grf.sample(5*4).unsqueeze(-1)\n",
    "        x = G(z)\n",
    "        numb_fig = 4\n",
    "        fig, ax = plt.subplots(5,numb_fig, figsize=(14,16))\n",
    "        for i in range(5):\n",
    "            for j in range(numb_fig):\n",
    "                phase = torch.atan2(x[i*numb_fig+j,:,:,1], x[i*numb_fig+j,:,:,0]).cpu().detach().numpy()\n",
    "                phase = (phase + np.pi) % (2 * np.pi) - np.pi\n",
    "                bar = ax[i,j].imshow(phase,  cmap='RdYlBu', vmin = -np.pi, vmax=np.pi,extent=[0,1,0,1])\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//volcano/{}GRF_GAN.pdf'.format(ite))  \n",
    "        plt.show()\n",
    "        \n",
    "############################################\n",
    "    numb_fig = 5\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        var = torch.zeros(10000, device=device).float()\n",
    "        skew = torch.zeros(10000, device=device).float()\n",
    "        for j in range(10000//100):\n",
    "            z = grf.sample(100).unsqueeze(-1)\n",
    "            x = G(z)\n",
    "            var[j*100:(j+1)*100] = circular_var(x)\n",
    "            skew[j*100:(j+1)*100] = circular_skew(x)\n",
    "        \n",
    "    var_train = circular_var(x_train)\n",
    "    skew_train = circular_skew(x_train)\n",
    "\n",
    "    fig, ax = plt.subplots(1,1, figsize=(4,4),tight_layout=True)\n",
    "    ax.hist(var_train.cpu().detach().numpy(), bins=np.linspace(0.0, 1.0, 50), histtype='step', density=True, color='#ff7f0e', label='Ground truth')\n",
    "    ax.hist(var.view(-1).cpu().detach().numpy(), bins=np.linspace(0.0, 1.0, 50), color='#1f77b4' , histtype='step', density=True, label='GAN')\n",
    "    ax.set_xlim(0.0, 1.0)\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/volcano/{}GRF_varianceGAN.pdf'.format(ite))  \n",
    "    plt.show()\n",
    "    \n",
    "############################################\n",
    "\n",
    "    fig, ax = plt.subplots(1,1, figsize=(4,4),tight_layout=True)\n",
    "    ax.hist(skew_train.cpu().detach().numpy(), bins=np.linspace(-4.0, 4.0, 50), histtype='step', density=True, color='#ff7f0e', label='Ground truth')\n",
    "    ax.hist(skew.view(-1).cpu().detach().numpy(), bins=np.linspace(-4.0, 4.0, 50), color='#1f77b4' , histtype='step', density=True, label='GAN')\n",
    "    plt.xlabel('value')\n",
    "    plt.ylabel('Histogram')\n",
    "    ax.set_xlim(-4.0, 4.0)\n",
    "    # plt.title('Auto correlation')\n",
    "    plt.legend()\n",
    "    plt.savefig('~/GANO/Figures/volcano/{}GRF_skewnessGAN.pdf'.format(ite))\n",
    "    plt.show()\n",
    "    \n",
    "#     plt.close('all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "touched-brunswick",
   "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",
    "    losses_W = np.zeros(epochs)\n",
    "    for i in range(epochs):\n",
    "        loss_D = 0.0\n",
    "        loss_G = 0.0\n",
    "        loss_W = 0.0\n",
    "        for j, data in enumerate(train_data):\n",
    "            # Train D\n",
    "            x = data[0].to(device)\n",
    "            D_optim.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",
    "            loss_W += W_loss.item()\n",
    "\n",
    "            D_optim.step()\n",
    "            \n",
    "            # Train G\n",
    "            if (j + 1) % n_critic == 0:\n",
    "                G_optim.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 / batch_size\n",
    "        losses_G[i] = loss_G / batch_size\n",
    "        losses_W[i] = loss_W / batch_size\n",
    "        print(i, \"D: \", losses_D[i], \"G: \", losses_G[i], \"W: \", losses_W[i])\n",
    "\n",
    "        storedata(i, G)\n",
    "        \n",
    "    return losses_D, losses_G, losses_W"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "defined-binary",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "losses_D, losses_G, losses_W = train_GAN(D_net, G_net, train_loader, epochs, D_optimizer_net, G_optimizer_net)\n",
    "G.load_state_dict(torch.load(\"~/GANO/volcano/G.pt\".format()))"
   ]
  }
 ],
 "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
}