GANO_volcano.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "af34cad4",
   "metadata": {},
   "source": [
    "# GANO trained on volcano data on 2D domain\n",
    "\n",
    "The inputs to the generative model are samples of a GRF\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "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 neuralop.models import UNO"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dbfe216",
   "metadata": {},
   "source": [
    "#### Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "invisible-silly",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = 128-8 #resolution\n",
    "npad = 8 # padding of latio 8/128 in U-NO\n",
    "ntrain = 4096 # number of training samples\n",
    "modes = 20 # number of Fourier modes in the initial FNO layer\n",
    "d_co_domain= 32 #the dimension of the co-domain of the initial U-NO layer.\n",
    "lr = 1e-4 #learning rate of the optimizer\n",
    "device = 'cuda:7'\n",
    "epochs = 300\n",
    "λ_grad = 10.0 # Lagrange coefficinet for gradient penalty\n",
    "n_critic = 10 # every n_critic iteration the generator is updated\n",
    "batch_size = 32"
   ]
  },
  {
   "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": "79431b06",
   "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 { d_co_domain: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": "markdown",
   "id": "d7f8a396",
   "metadata": {},
   "source": [
    "### Directional statistics\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "environmental-retirement",
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "var = circular_var(x_train)\n",
    "plt.hist(var.numpy(), bins=np.linspace(0.0, 1.0, 25), histtype='step', lined_co_domain=2.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c18e221f",
   "metadata": {},
   "source": [
    "### Data visualization"
   ]
  },
  {
   "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": "markdown",
   "id": "4b548f74",
   "metadata": {},
   "source": [
    "### Data visualization"
   ]
  },
  {
   "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": "1e3b6bfb",
   "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": "d4fa72f2",
   "metadata": {},
   "source": [
    "### Generator operator and Discriminator functional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "french-watch",
   "metadata": {},
   "outputs": [],
   "source": [
    "class SpectralConv2d(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, dim1, dim2,modes1 = None, modes2 = None):\n",
    "        super(SpectralConv2d, self).__init__()\n",
    "\n",
    "        \"\"\"\n",
    "        2D Fourier layer. It does FFT, linear transform, and Inverse FFT.    \n",
    "        \"\"\"\n",
    "        in_channels = int(in_channels)\n",
    "        out_channels = int(out_channels)\n",
    "        self.in_channels = in_channels\n",
    "        self.out_channels = out_channels\n",
    "        self.dim1 = dim1 #output dimensions\n",
    "        self.dim2 = dim2\n",
    "        if modes1 is not None:\n",
    "            self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1\n",
    "            self.modes2 = modes2\n",
    "        else:\n",
    "            self.modes1 = dim1//2-1 #if not given take the highest number of modes can be taken\n",
    "            self.modes2 = dim2//2 \n",
    "        self.scale = (1 / (2*in_channels))**(1.0/2.0)\n",
    "        self.weights1 = nn.Parameter(self.scale * (torch.randn(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat)))\n",
    "        self.weights2 = nn.Parameter(self.scale * (torch.randn(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat)))\n",
    "\n",
    "    # Complex multiplication\n",
    "    def compl_mul2d(self, input, weights):\n",
    "        # (batch, in_channel, x,y ), (in_channel, out_channel, x,y) -> (batch, out_channel, x,y)\n",
    "        return torch.einsum(\"bixy,ioxy->boxy\", input, weights)\n",
    "\n",
    "    def forward(self, x, dim1 = None,dim2 = None):\n",
    "        if dim1 is not None:\n",
    "            self.dim1 = dim1\n",
    "            self.dim2 = dim2\n",
    "        batchsize = x.shape[0]\n",
    "        #Compute Fourier coeffcients up to factor of e^(- something constant)\n",
    "        x_ft = torch.fft.rfft2(x)\n",
    "\n",
    "        # Multiply relevant Fourier modes\n",
    "        out_ft = torch.zeros(batchsize, self.out_channels,  self.dim1, self.dim2//2 + 1 , dtype=torch.cfloat, device=x.device)\n",
    "        out_ft[:, :, :self.modes1, :self.modes2] = \\\n",
    "            self.compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1)\n",
    "        out_ft[:, :, -self.modes1:, :self.modes2] = \\\n",
    "            self.compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2)\n",
    "\n",
    "        #Return to physical space\n",
    "        x = torch.fft.irfft2(out_ft, s=(self.dim1, self.dim2))\n",
    "        return x\n",
    "\n",
    "\n",
    "class pointwise_op(nn.Module):\n",
    "    def __init__(self, in_channel, out_channel,dim1, dim2):\n",
    "        super(pointwise_op,self).__init__()\n",
    "        self.conv = nn.Conv2d(int(in_channel), int(out_channel), 1)\n",
    "        self.dim1 = int(dim1)\n",
    "        self.dim2 = int(dim2)\n",
    "\n",
    "    def forward(self,x, dim1 = None, dim2 = None):\n",
    "        if dim1 is None:\n",
    "            dim1 = self.dim1\n",
    "            dim2 = self.dim2\n",
    "        x_out = self.conv(x)\n",
    "        x_out = torch.nn.functional.interpolate(x_out, size = (dim1, dim2),mode = 'bicubic',align_corners=True)\n",
    "        return x_out\n",
    "\n",
    "class Generator(nn.Module):\n",
    "    def __init__(self, in_d_co_domain, d_co_domain, pad = 0, factor = 3/4):\n",
    "        super(Generator, self).__init__()\n",
    "\n",
    "        \"\"\"\n",
    "        The overall network. It contains 4 layers of the Fourier layer.\n",
    "        1. Lift the input to the desire channel dimension by self.fc0 .\n",
    "        2. 4 layers of the integral operators u' = (W + K)(u).\n",
    "            W defined by self.w; K defined by self.conv .\n",
    "        3. Project from the channel space to the output space by self.fc1 and self.fc2 .\n",
    "        \n",
    "        input: the solution of the coefficient function and locations (a(x, y), x, y)\n",
    "        input shape: (batchsize, x=s, y=s, c=3)\n",
    "        output: the solution \n",
    "        output shape: (batchsize, x=s, y=s, c=1)\n",
    "        \"\"\"\n",
    "        self.in_d_co_domain = in_d_co_domain # input channel\n",
    "        self.d_co_domain = d_co_domain \n",
    "        self.factor = factor\n",
    "        self.padding = pad  # pad the domain if input is non-periodic\n",
    "\n",
    "        self.fc0 = nn.Linear(self.in_d_co_domain, self.d_co_domain) # input channel is 3: (a(x, y), x, y)\n",
    "\n",
    "        self.conv0 = SpectralConv2d(self.d_co_domain, 2*factor*self.d_co_domain, 48, 48, 24, 24)\n",
    "\n",
    "        self.conv1 = SpectralConv2d(2*factor*self.d_co_domain, 4*factor*self.d_co_domain, 32, 32, 16,16)\n",
    "\n",
    "        self.conv2 = SpectralConv2d(4*factor*self.d_co_domain, 8*factor*self.d_co_domain, 16, 16,8,8)\n",
    "        \n",
    "        self.conv2_1 = SpectralConv2d(8*factor*self.d_co_domain, 16*factor*self.d_co_domain, 8, 8,4,4)\n",
    "        \n",
    "        self.conv2_9 = SpectralConv2d(16*factor*self.d_co_domain, 8*factor*self.d_co_domain, 16, 16,4,4)\n",
    "        \n",
    "\n",
    "        self.conv3 = SpectralConv2d(16*factor*self.d_co_domain, 4*factor*self.d_co_domain, 32, 32,8,8)\n",
    "\n",
    "        self.conv4 = SpectralConv2d(8*factor*self.d_co_domain, 2*factor*self.d_co_domain, 48, 48,16,16)\n",
    "\n",
    "        self.conv5 = SpectralConv2d(4*factor*self.d_co_domain, self.d_co_domain, 64, 64,24,24) # will be reshaped\n",
    "\n",
    "        self.w0 = pointwise_op(self.d_co_domain,2*factor*self.d_co_domain,48, 48) #\n",
    "        \n",
    "        self.w1 = pointwise_op(2*factor*self.d_co_domain, 4*factor*self.d_co_domain, 32, 32) #\n",
    "        \n",
    "        self.w2 = pointwise_op(4*factor*self.d_co_domain, 8*factor*self.d_co_domain, 16, 16) #\n",
    "        \n",
    "        self.w2_1 = pointwise_op(8*factor*self.d_co_domain, 16*factor*self.d_co_domain, 8, 8)\n",
    "        \n",
    "        self.w2_9 = pointwise_op(16*factor*self.d_co_domain, 8*factor*self.d_co_domain, 16, 16)\n",
    "        \n",
    "        self.w3 = pointwise_op(16*factor*self.d_co_domain, 4*factor*self.d_co_domain, 32, 32) #\n",
    "        \n",
    "        self.w4 = pointwise_op(8*factor*self.d_co_domain, 2*factor*self.d_co_domain, 48, 48)\n",
    "        \n",
    "        self.w5 = pointwise_op(4*factor*self.d_co_domain, self.d_co_domain, 64, 64) # will be reshaped\n",
    "\n",
    "        self.fc1 = nn.Linear(2*self.d_co_domain, 4*self.d_co_domain)\n",
    "        self.fc2 = nn.Linear(4*self.d_co_domain, 2)\n",
    "\n",
    "    def forward(self, x):\n",
    "\n",
    "        grid = self.get_grid(x.shape, x.device)\n",
    "        x = torch.cat((x, grid), dim=-1)\n",
    "\n",
    "        x_fc0 = self.fc0(x)\n",
    "        x_fc0 = F.gelu(x_fc0)\n",
    "        \n",
    "        x_fc0 = x_fc0.permute(0, 3, 1, 2)\n",
    "        \n",
    "        \n",
    "        x_fc0 = F.pad(x_fc0, [0,self.padding, 0,self.padding])\n",
    "        \n",
    "        D1,D2 = x_fc0.shape[-2],x_fc0.shape[-1]\n",
    "        \n",
    "\n",
    "        x1_c0 = self.conv0(x_fc0,int(D1*self.factor),int(D2*self.factor))\n",
    "        x2_c0 = self.w0(x_fc0,int(D1*self.factor),int(D2*self.factor))\n",
    "        x_c0 = x1_c0 + x2_c0\n",
    "        x_c0 = F.gelu(x_c0)\n",
    "        #print(x.shape)\n",
    "\n",
    "        x1_c1 = self.conv1(x_c0 ,D1//2,D2//2)\n",
    "        x2_c1 = self.w1(x_c0 ,D1//2,D2//2)\n",
    "        x_c1 = x1_c1 + x2_c1\n",
    "        x_c1 = F.gelu(x_c1)\n",
    "        #print(x.shape)\n",
    "\n",
    "        x1_c2 = self.conv2(x_c1 ,D1//4,D2//4)\n",
    "        x2_c2 = self.w2(x_c1 ,D1//4,D2//4)\n",
    "        x_c2 = x1_c2 + x2_c2\n",
    "        x_c2 = F.gelu(x_c2 )\n",
    "        #print(x.shape)\n",
    "        \n",
    "        x1_c2_1 = self.conv2_1(x_c2,D1//8,D2//8)\n",
    "        x2_c2_1 = self.w2_1(x_c2,D1//8,D2//8)\n",
    "        x_c2_1 = x1_c2_1 + x2_c2_1\n",
    "        x_c2_1 = F.gelu(x_c2_1)\n",
    "        \n",
    "        x1_c2_9 = self.conv2_9(x_c2_1,D1//4,D2//4)\n",
    "        x2_c2_9 = self.w2_9(x_c2_1,D1//4,D2//4)\n",
    "        x_c2_9 = x1_c2_9 + x2_c2_9\n",
    "        x_c2_9 = F.gelu(x_c2_9)\n",
    "        x_c2_9 = torch.cat([x_c2_9, x_c2], dim=1) \n",
    "\n",
    "        x1_c3 = self.conv3(x_c2_9,D1//2,D2//2)\n",
    "        x2_c3 = self.w3(x_c2_9,D1//2,D2//2)\n",
    "        x_c3 = x1_c3 + x2_c3\n",
    "        x_c3 = F.gelu(x_c3)\n",
    "        x_c3 = torch.cat([x_c3, x_c1], dim=1)\n",
    "\n",
    "        x1_c4 = self.conv4(x_c3,int(D1*self.factor),int(D2*self.factor))\n",
    "        x2_c4 = self.w4(x_c3,int(D1*self.factor),int(D2*self.factor))\n",
    "        x_c4 = x1_c4 + x2_c4\n",
    "        x_c4 = F.gelu(x_c4)\n",
    "        x_c4 = torch.cat([x_c4, x_c0], dim=1)\n",
    "\n",
    "        x1_c5 = self.conv5(x_c4,D1,D2)\n",
    "        x2_c5 = self.w5(x_c4,D1,D2)\n",
    "        x_c5 = x1_c5 + x2_c5\n",
    "        x_c5 = F.gelu(x_c5)\n",
    "        \n",
    "\n",
    "        x_c5 = torch.cat([x_c5, x_fc0], dim=1)\n",
    "        if self.padding!=0:\n",
    "            x_c5 = x_c5[..., :-self.padding, :-self.padding]\n",
    "\n",
    "        x_c5 = x_c5.permute(0, 2, 3, 1)\n",
    "        \n",
    "        x_fc1 = self.fc1(x_c5)\n",
    "        x_fc1 = F.gelu(x_fc1)\n",
    "        \n",
    "        x_out = self.fc2(x_fc1)\n",
    "        \n",
    "        x_out = torch.tanh(x_out)\n",
    "        \n",
    "        return x_out\n",
    "    \n",
    "    def get_grid(self, shape, device):\n",
    "        batchsize, size_x, size_y = shape[0], shape[1], shape[2]\n",
    "        gridx = torch.tensor(np.linspace(0, 1, size_x), dtype=torch.float)\n",
    "        gridx = gridx.reshape(1, size_x, 1, 1).repeat([batchsize, 1, size_y, 1])\n",
    "        gridy = torch.tensor(np.linspace(0, 1, size_y), dtype=torch.float)\n",
    "        gridy = gridy.reshape(1, 1, size_y, 1).repeat([batchsize, size_x, 1, 1])\n",
    "        return torch.cat((gridx, gridy), dim=-1).to(device)\n",
    "\n",
    "\n",
    "\n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, in_d_co_domain, d_co_domain, kernel_dim=16, pad = 0, factor = 3/4):\n",
    "        super(Discriminator, self).__init__()\n",
    "\n",
    "        \"\"\"\n",
    "        The overall network. It contains 4 layers of the Fourier layer.\n",
    "        1. Lift the input to the desire channel dimension by self.fc0 .\n",
    "        2. 4 layers of the integral operators u' = (W + K)(u).\n",
    "            W defined by self.w; K defined by self.conv .\n",
    "        3. Project from the channel space to the output space by self.fc1 and self.fc2 .\n",
    "        \n",
    "        input: the solution of the coefficient function and locations (a(x, y), x, y)\n",
    "        input shape: (batchsize, x=s, y=s, c=3)\n",
    "        output: the solution \n",
    "        output shape: (batchsize, x=s, y=s, c=1)\n",
    "        \"\"\"\n",
    "        self.in_d_co_domain = in_d_co_domain # input channel\n",
    "        self.d_co_domain = d_co_domain \n",
    "        self.factor = factor\n",
    "        self.padding = pad  # pad the domain if input is non-periodic\n",
    "        self.kernel_dim = kernel_dim\n",
    "        self.fc0 = nn.Linear(self.in_d_co_domain, self.d_co_domain) # input channel is 3: (a(x, y), x, y)\n",
    "\n",
    "        self.conv0 = SpectralConv2d(self.d_co_domain, 2*factor*self.d_co_domain, 48, 48, 24, 24)\n",
    "\n",
    "        self.conv1 = SpectralConv2d(2*factor*self.d_co_domain, 4*factor*self.d_co_domain, 32, 32, 16,16)\n",
    "\n",
    "        self.conv2 = SpectralConv2d(4*factor*self.d_co_domain, 8*factor*self.d_co_domain, 16, 16,8,8)\n",
    "        \n",
    "        self.conv2_1 = SpectralConv2d(8*factor*self.d_co_domain, 16*factor*self.d_co_domain, 8, 8,4,4)\n",
    "        \n",
    "        self.conv2_9 = SpectralConv2d(16*factor*self.d_co_domain, 8*factor*self.d_co_domain, 16, 16,4,4)\n",
    "        \n",
    "\n",
    "        self.conv3 = SpectralConv2d(16*factor*self.d_co_domain, 4*factor*self.d_co_domain, 32, 32,8,8)\n",
    "\n",
    "        self.conv4 = SpectralConv2d(8*factor*self.d_co_domain, 2*factor*self.d_co_domain, 48, 48,16,16)\n",
    "\n",
    "        self.conv5 = SpectralConv2d(4*factor*self.d_co_domain, self.d_co_domain, 64, 64,24,24) # will be reshaped\n",
    "\n",
    "        self.w0 = pointwise_op(self.d_co_domain,2*factor*self.d_co_domain,48, 48) #\n",
    "        \n",
    "        self.w1 = pointwise_op(2*factor*self.d_co_domain, 4*factor*self.d_co_domain, 32, 32) #\n",
    "        \n",
    "        self.w2 = pointwise_op(4*factor*self.d_co_domain, 8*factor*self.d_co_domain, 16, 16) #\n",
    "        \n",
    "        self.w2_1 = pointwise_op(8*factor*self.d_co_domain, 16*factor*self.d_co_domain, 8, 8)\n",
    "        \n",
    "        self.w2_9 = pointwise_op(16*factor*self.d_co_domain, 8*factor*self.d_co_domain, 16, 16)\n",
    "        \n",
    "        self.w3 = pointwise_op(16*factor*self.d_co_domain, 4*factor*self.d_co_domain, 32, 32) #\n",
    "        \n",
    "        self.w4 = pointwise_op(8*factor*self.d_co_domain, 2*factor*self.d_co_domain, 48, 48)\n",
    "        \n",
    "        self.w5 = pointwise_op(4*factor*self.d_co_domain, self.d_co_domain, 64, 64) # will be reshaped\n",
    "\n",
    "        self.fc1 = nn.Linear(2*self.d_co_domain, 4*self.d_co_domain)\n",
    "        self.fc2 = nn.Linear(4*self.d_co_domain, 1)\n",
    "        \n",
    "        # kernel for last functional operation\n",
    "\n",
    "        self.knet = kernel(2, self.kernel_dim)\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        res1 = x.shape[1]\n",
    "        res2 = x.shape[2]\n",
    "        grid = self.get_grid(x.shape, x.device)\n",
    "        x = torch.cat((x, grid), dim=-1)\n",
    "         \n",
    "        x_fc0 = self.fc0(x)\n",
    "        x_fc0 = F.gelu(x_fc0)\n",
    "        \n",
    "        x_fc0 = x_fc0.permute(0, 3, 1, 2)\n",
    "        \n",
    "        \n",
    "        x_fc0 = F.pad(x_fc0, [0,self.padding, 0,self.padding])\n",
    "        \n",
    "        D1,D2 = x_fc0.shape[-2],x_fc0.shape[-1]\n",
    "        \n",
    "\n",
    "        x1_c0 = self.conv0(x_fc0,int(D1*self.factor),int(D2*self.factor))\n",
    "        x2_c0 = self.w0(x_fc0,int(D1*self.factor),int(D2*self.factor))\n",
    "        x_c0 = x1_c0 + x2_c0\n",
    "        x_c0 = F.gelu(x_c0)\n",
    "        #print(x.shape)\n",
    "\n",
    "        x1_c1 = self.conv1(x_c0 ,D1//2,D2//2)\n",
    "        x2_c1 = self.w1(x_c0 ,D1//2,D2//2)\n",
    "        x_c1 = x1_c1 + x2_c1\n",
    "        x_c1 = F.gelu(x_c1)\n",
    "        #print(x.shape)\n",
    "\n",
    "        x1_c2 = self.conv2(x_c1 ,D1//4,D2//4)\n",
    "        x2_c2 = self.w2(x_c1 ,D1//4,D2//4)\n",
    "        x_c2 = x1_c2 + x2_c2\n",
    "        x_c2 = F.gelu(x_c2 )\n",
    "        #print(x.shape)\n",
    "        \n",
    "        x1_c2_1 = self.conv2_1(x_c2,D1//8,D2//8)\n",
    "        x2_c2_1 = self.w2_1(x_c2,D1//8,D2//8)\n",
    "        x_c2_1 = x1_c2_1 + x2_c2_1\n",
    "        x_c2_1 = F.gelu(x_c2_1)\n",
    "        \n",
    "        x1_c2_9 = self.conv2_9(x_c2_1,D1//4,D2//4)\n",
    "        x2_c2_9 = self.w2_9(x_c2_1,D1//4,D2//4)\n",
    "        x_c2_9 = x1_c2_9 + x2_c2_9\n",
    "        x_c2_9 = F.gelu(x_c2_9)\n",
    "        x_c2_9 = torch.cat([x_c2_9, x_c2], dim=1) \n",
    "\n",
    "        x1_c3 = self.conv3(x_c2_9,D1//2,D2//2)\n",
    "        x2_c3 = self.w3(x_c2_9,D1//2,D2//2)\n",
    "        x_c3 = x1_c3 + x2_c3\n",
    "        x_c3 = F.gelu(x_c3)\n",
    "        x_c3 = torch.cat([x_c3, x_c1], dim=1)\n",
    "\n",
    "        x1_c4 = self.conv4(x_c3,int(D1*self.factor),int(D2*self.factor))\n",
    "        x2_c4 = self.w4(x_c3,int(D1*self.factor),int(D2*self.factor))\n",
    "        x_c4 = x1_c4 + x2_c4\n",
    "        x_c4 = F.gelu(x_c4)\n",
    "        x_c4 = torch.cat([x_c4, x_c0], dim=1)\n",
    "\n",
    "        x1_c5 = self.conv5(x_c4,D1,D2)\n",
    "        x2_c5 = self.w5(x_c4,D1,D2)\n",
    "        x_c5 = x1_c5 + x2_c5\n",
    "        x_c5 = F.gelu(x_c5)\n",
    "        \n",
    "\n",
    "        x_c5 = torch.cat([x_c5, x_fc0], dim=1)\n",
    "        if self.padding!=0:\n",
    "            x_c5 = x_c5[..., :-self.padding, :-self.padding]\n",
    "\n",
    "        x = x_c5\n",
    "        x = x.permute(0, 2, 3, 1)\n",
    "\n",
    "        x = self.fc1(x)\n",
    "        x = F.gelu(x)\n",
    "\n",
    "        x = self.fc2(x)\n",
    "        \n",
    "        kx = self.knet(grid)\n",
    "\n",
    "        kx = kx.view(batch_size, -1, 1)\n",
    "\n",
    "        x = x.view(batch_size, -1, 1)\n",
    "        x = torch.einsum('bik,bik->bk', kx, x)/(res1*res2)\n",
    "\n",
    "        return x\n",
    "\n",
    "    def get_grid(self, shape, device):\n",
    "        batchsize, size_x, size_y = shape[0], shape[1], shape[2]\n",
    "        gridx = torch.tensor(np.linspace(0, 1, size_x), dtype=torch.float)\n",
    "        gridx = gridx.reshape(1, size_x, 1, 1).repeat([batchsize, 1, size_y, 1])\n",
    "        gridy = torch.tensor(np.linspace(0, 1, size_y), dtype=torch.float)\n",
    "        gridy = gridy.reshape(1, 1, size_y, 1).repeat([batchsize, size_x, 1, 1])\n",
    "        return torch.cat((gridx, gridy), dim=-1).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "endangered-graphics",
   "metadata": {},
   "outputs": [],
   "source": [
    "def kernel(in_chan=2, up_dim=32):\n",
    "    \"\"\"\n",
    "        Kernel network apply on grid\n",
    "    \"\"\"\n",
    "    layers = nn.Sequential(\n",
    "                nn.Linear(in_chan, up_dim, bias=True), torch.nn.GELU(),\n",
    "                nn.Linear(up_dim, up_dim, bias=True), torch.nn.GELU(),\n",
    "                nn.Linear(up_dim, 1, bias=False)\n",
    "            )\n",
    "    return layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "innocent-franklin",
   "metadata": {},
   "outputs": [],
   "source": [
    "D = Discriminator(2+2, d_co_domain, pad=npad).to(device)\n",
    "G = Generator(1+2, d_co_domain, pad=npad).to(device)\n",
    "nn_params = sum(p.numel() for p in D.parameters() if p.requires_grad)\n",
    "print(\"Number discriminator parameters: \", nn_params)\n",
    "nn_params = sum(p.numel() for p in G.parameters() if p.requires_grad)\n",
    "print(\"Number generator parameters: \", nn_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f17747ec",
   "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(z)\n",
    "D(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 = torch.optim.Adam(G.parameters(), lr=lr) #, weight_decay=1e-4)\n",
    "D_optimizer = torch.optim.Adam(D.parameters(), lr=lr) #, weight_decay=1e-4)\n",
    "fn_loss = nn.BCEWithLogitsLoss()\n",
    "\n",
    "D.train()\n",
    "G.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": "small-norman",
   "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_GANO.pdf'.format(ite))  \n",
    "\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='GANO')\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_varianceGANO.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='GANO')\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_skewnessGANO.pdf'.format(ite))\n",
    "    \n",
    "    plt.close('all')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "touched-brunswick",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_GANO(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",
    "            \n",
    "        torch.save(D.state_dict(), \"~/GANO/models/volcano/D.pt\".format())\n",
    "        torch.save(G.state_dict(), \"~/GANO/models/volcano/G.pt\".format())\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_GANO(D, G, train_loader, epochs, D_optimizer, G_optimizer)\n",
    "G.load_state_dict(torch.load(\"~/GANO/models/volcano/G.pt\".format()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "greatest-trading",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1,1, figsize=(6,4),tight_layout=True)\n",
    "ax.plot(np.arange(epochs), losses_D, c='k', label='D')\n",
    "ax.plot(np.arange(epochs), losses_G, c='b', label='G')\n",
    "plt.legend(loc = 'lower right')\n",
    "plt.ylabel(\"Loss\")\n",
    "plt.xlabel(\"Epoch\")\n",
    "plt.xlim(0.0, 300.0)\n",
    "plt.savefig(\"~/GANO/Figures/volcano/\")"
   ]
  }
 ],
 "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
}