GANO_mixGRF.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "10df7e68",
   "metadata": {},
   "source": [
    "# GANO trained on mixture of GRFs data on 2D domain\n",
    "\n",
    "The inputs to the generative model are samples of a GRF\n",
    "The data is sampled from mixture of GRFs"
   ]
  },
  {
   "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": "9ecdcd3f",
   "metadata": {},
   "source": [
    "#### Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "invisible-silly",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = 64-8 #resolution\n",
    "npad = 8  # padding of latio 8/128 in U-NO\n",
    "ntrain = 8192 # number of training samples\n",
    "modes = 20 # number of Fourier modes in the initial FNO layer\n",
    "d_co_domain = 20 #the dimension of the co-domain of the initial U-NO layer.\n",
    "lr = 1e-4 #learning rate of the optimizer\n",
    "device = 'cuda:5'\n",
    "epochs = 200\n",
    "λ_grad = 10.0 # Lagrange coefficinet for gradient penalty\n",
    "n_critic = 5 # every n_critic iteration the generator is updated\n",
    "batch_size = 64\n"
   ]
  },
  {
   "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",
    "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)\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": "d8d2939b",
   "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": "c42dcf0e",
   "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": "dedicated-writing",
   "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\n",
    "\n",
    "class SpecDConv2d(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels, modes1, modes2):\n",
    "        super(SpecDConv2d, self).__init__()\n",
    "        \"\"\"\n",
    "        2D Fourier layer. It does FFT, then outputs truncated nodess    \n",
    "        \"\"\"\n",
    "        self.in_channels = in_channels\n",
    "        self.out_channels = out_channels\n",
    "        self.modes1 = modes1\n",
    "        self.modes2 = modes2\n",
    "\n",
    "        self.scale = (1 / (in_channels * out_channels))\n",
    "        self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))\n",
    "        self.weights2 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))\n",
    "\n",
    "    # Complex multiplication\n",
    "    def compl_mul2d(self, input, weights):\n",
    "        return torch.einsum(\"bixy,ioxy->boxy\", input, weights)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x_ft = torch.fft.rfft2(x)\n",
    "\n",
    "        # Multiply relevant Fourier modes\n",
    "        c1_mul = self.compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1)\n",
    "        c2_mul = self.compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2)\n",
    "        c_out = torch.cat((c1_mul.abs(), c2_mul.abs()), dim=1)\n",
    "\n",
    "        return c_out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02a7abe4",
   "metadata": {},
   "source": [
    "### Generator operator and Discriminator functional"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "honest-rating",
   "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, 1)\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",
    "        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, in_d_co_domain-2)\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",
    "        kx = kx.view(batch_size,-1, 1)\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": "worth-exhibit",
   "metadata": {},
   "outputs": [],
   "source": [
    "# D = Discriminator(modes, modes, d_co_domain, padding=npad).to(device)\n",
    "# G = Generator(modes, modes, d_co_domain, padding=npad).to(device)\n",
    "D = Discriminator(1+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": "7df2e00b",
   "metadata": {},
   "source": [
    "### Samples generated with initilized generator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "plastic-poker",
   "metadata": {},
   "outputs": [],
   "source": [
    "z = grf.sample(4).unsqueeze(-1)\n",
    "x = G(z)\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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abandoned-crowd",
   "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": "numeric-instrumentation",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_JSGANO(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 x, in train_data:\n",
    "            x = x.to(device)\n",
    "            D_optimizer.zero_grad()\n",
    "            D_optimizer.zero_grad()\n",
    "            \n",
    "            y_real = D(x)\n",
    "            loss_D_real = fn_loss(y_real, torch.ones_like(y_real, device=device))\n",
    "\n",
    "            x_syn = G(grf.sample(x.shape[0]).unsqueeze(-1))\n",
    "            y_syn = D(x_syn)\n",
    "            loss_D_syn = fn_loss(y_syn, torch.zeros_like(y_syn, device=device))\n",
    "            \n",
    "            loss = loss_D_real + loss_D_syn\n",
    "            loss.backward()\n",
    "            loss_D += loss.item()\n",
    "            \n",
    "            D_optim.step()\n",
    "            \n",
    "            x_syn = G(grf.sample(x.shape[0]).unsqueeze(-1))\n",
    "            y_syn = D(x_syn)\n",
    "            \n",
    "            loss = fn_loss(y_syn, torch.ones_like(y_syn, device=device))\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",
    "\n",
    "        print(i, \"D: \", losses_D[i], \"G: \", losses_G[i])\n",
    "            \n",
    "        with torch.no_grad():\n",
    "            if i % 1 == 0:\n",
    "                z = grf.sample(4).unsqueeze(-1)\n",
    "                x = G(z)\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",
    "\n",
    "    return losses_D, losses_G"
   ]
  },
  {
   "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 WGAN GRF\"\"\"\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_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",
    "    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 / batch_size\n",
    "        losses_G[i] = loss_G / batch_size\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": "d2c7092c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def storedata(ite,G):\n",
    "    high = 2\n",
    "    low = 0.3\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/GANO_Mix_GRF/{}GRF_mix.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/GANO_Mix_GRF/{}GRF_GANO_mix.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='GANO')\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/GANO_Mix_GRF/{}GRF_HistogramGANO_mix.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='GANO')\n",
    "    plt.xlabel('Position')\n",
    "    plt.ylabel('Auto correlation')\n",
    "    # plt.title('Auto correlation')\n",
    "    plt.legend()\n",
    "    plt.savefig('~/GANO/Figures/GANO_Mix_GRF/{}GRF_AutoCorrelationGANO_mix.pdf'.format(ite))  \n",
    "\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aging-merchandise",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "losses_D, losses_G = train_GANO(D, G, train_loader, epochs, D_optimizer, G_optimizer)"
   ]
  },
  {
   "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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "handled-football",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(D, \"models/gano_mixture/D_lr{}_d_co_domain{}_ntrain{}_ncritic{}.pt\".format(lr, d_co_domain, ntrain, n_critic))\n",
    "torch.save(G, \"models/gano_mixture/G_lr{}_d_co_domain{}_ntrain{}_ncritic{}.pt\".format(lr, d_co_domain, ntrain, n_critic))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "collected-flooring",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}