train-zurich.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Multimodal Neural Network\n",
    "\n",
    "```\n",
    "     __\n",
    "    (__)_\n",
    "    (____)_\n",
    "    (______)\n",
    "...//(  00 )\\......version 1.0\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-02T08:43:48.670745Z",
     "start_time": "2020-04-02T08:43:47.606925Z"
    },
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "# Python Standard Libraries\n",
    "from datetime import datetime\n",
    "import glob\n",
    "import itertools\n",
    "import math\n",
    "import os\n",
    "import random\n",
    "import re\n",
    "import time\n",
    "import warnings\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\" # Select your GPU(s)\n",
    "\n",
    "# Deep Learning libraries\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "# Other libraries\n",
    "# ~ Scientific\n",
    "import numpy as np\n",
    "np.bool = np.bool_ # Since dropped from NumPy >=1.24\n",
    "import scipy.stats as st\n",
    "# ~ Image manipulation / visualisation\n",
    "import imgaug\n",
    "from imgaug import augmenters as iaa\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "import skimage\n",
    "import skimage.io as skio\n",
    "import skimage.transform as sktr\n",
    "# ~ Other\n",
    "from tqdm.notebook import tqdm\n",
    "\n",
    "# Local libraries\n",
    "from utils.image import *\n",
    "from utils.plotting import *\n",
    "from utils.torch import *\n",
    "\n",
    "# IPython\n",
    "from IPython.display import clear_output, HTML\n",
    "\n",
    "# Path to Zurich dataset:\n",
    "ZURICH = \"Zurich_dataset_v1.0\"\n",
    "# Download at: https://doi.org/10.5281/zenodo.5914759"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "count = torch.cuda.device_count()\n",
    "print(f\"{count} GPU device(s) available.\")\n",
    "print()\n",
    "print(\"List of GPUs:\")\n",
    "for i in range(count):\n",
    "    print(f\"* {torch.cuda.get_device_name(i)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-02T08:43:49.460057Z",
     "start_time": "2020-04-02T08:43:49.394241Z"
    }
   },
   "outputs": [],
   "source": [
    "###### METHOD RELATED\n",
    "# The place where the models will be saved\n",
    "export_folder = \"results\" # Add this path to the .gitignore\n",
    "# The number of channels in the latent space (best is 1 or 3 for visualization)\n",
    "latent_channels = 3\n",
    "# Modality slicing\n",
    "# You can choose a set of channels per modality (RGB for instance)\n",
    "# Modality A\n",
    "modA = slice(0, 3)\n",
    "modA_name = \"RGB\"\n",
    "modA_len = modA.stop - modA.start\n",
    "# Modality B\n",
    "modB = slice(3, 4)\n",
    "modB_name = \"NIR\"\n",
    "modB_len = modB.stop - modB.start\n",
    "# Distance function\n",
    "simfunctions = {\n",
    "    \"euclidean\" : lambda x, y: -torch.norm(x - y, p=2, dim=1).mean(),\n",
    "    \"L1\"        : lambda x, y: -torch.norm(x - y, p=1, dim=1).mean(),\n",
    "    \"huber\"     : lambda x, y: -torch.sqrt(1 + (x - y).pow(2)).mean(),\n",
    "    \"MSE\"       : lambda x, y: -(x - y).pow(2).mean(),\n",
    "    \"L3\"        : lambda x, y: -torch.norm(x - y, p=3, dim=1).mean(),\n",
    "    \"Linf\"      : lambda x, y: -torch.norm(x - y, p=float(\"inf\"), dim=1).mean(),\n",
    "    \"soft_corr\" : lambda x, y: F.softplus(x*y).sum(axis=1),\n",
    "    \"corr\"      : lambda x, y: (x*y).sum(axis=1),\n",
    "    \"cosine\"    : lambda x, y: F.cosine_similarity(x, y, dim=1, eps=1e-8).mean(),\n",
    "    \"angular\"   : lambda x, y: F.cosine_similarity(x, y, dim=1, eps=1e-8).acos().mean() / math.pi,\n",
    "}\n",
    "sim_func = simfunctions[\"MSE\"]\n",
    "# Temperature (tau) of the loss\n",
    "tau =  .1\n",
    "# L1/L2 activation regularization\n",
    "act_l1 = 0\n",
    "act_l2 = 0\n",
    "# p4 Equivariance (should always be True, unless you want to see how everything breaks visually otherwise)\n",
    "equivariance = False\n",
    "# Create video (slightly slower training)\n",
    "video = False\n",
    "video_path = export_folder\n",
    "\n",
    "# DEEP LEARNING RELATED\n",
    "# Device to train on (inference is done on cpu)\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "# Use two GPUs?\n",
    "device1 = device2 = device # 1 gpu for 2 modalities\n",
    "#device1, device2 = \"cuda:0\", \"cuda:1\" # 1 gpu per modality\n",
    "# Arguments for the tiramisu neural network\n",
    "tiramisu_args = {\n",
    "    # Number of convolutional filters for the first convolution\n",
    "    \"init_conv_filters\": 16,\n",
    "    # Number and depth of down blocks\n",
    "    \"down_blocks\": (2, 2, 2),\n",
    "    # Number and depth of up blocks\n",
    "    \"up_blocks\": (2, 2, 2),\n",
    "    # Number of dense layers in the bottleneck\n",
    "    \"bottleneck_layers\": 3,\n",
    "    # Upsampling type of layer (upsample has no grid artefacts)\n",
    "    \"upsampling_type\": \"upsample\",\n",
    "    # Type of max pooling, blurpool has better shift-invariance\n",
    "    \"transition_pooling\": \"max\",\n",
    "    # Dropout rate for the convolution\n",
    "    \"dropout_rate\": 0.0,\n",
    "    # Early maxpooling to reduce the input size\n",
    "    \"early_transition\": False,\n",
    "    # Activation function at the last layer\n",
    "    \"activation_func\": None,\n",
    "    # How much the conv layers should be compressed? (Memory saving)\n",
    "    \"compression\": 1.0,\n",
    "    # Memory efficient version of the tiramisu network (trades memory for computes)\n",
    "    # Gains of memory are enormous compared to the speed decrease.\n",
    "    # See: https://arxiv.org/pdf/1707.06990.pdf\n",
    "    \"efficient\": False,\n",
    "}\n",
    "# Epochs\n",
    "epochs = 5\n",
    "# Batch size\n",
    "batch_size = 24\n",
    "# Steps per epoch\n",
    "steps_per_epoch = 32\n",
    "# Number of steps\n",
    "steps = steps_per_epoch * epochs\n",
    "# How many unique patches are fed duriing one epoch\n",
    "samples_per_epoch = steps_per_epoch * batch_size\n",
    "# Optimiser\n",
    "#from lars.lars import LARS\n",
    "#optimiser = LARS\n",
    "optimiser = optim.Adam\n",
    "# Optimizer arguments\n",
    "opt_args = {\n",
    "    \"lr\": 1e-3,\n",
    "    \"weight_decay\": 1e-4,\n",
    "    #\"momentum\": 0.9\n",
    "}\n",
    "# Gradient norm. (limit on how big the gradients can get)\n",
    "grad_norm = 1.0\n",
    "\n",
    "# DATASET RELATED\n",
    "def worker_init_fn(worker_id):\n",
    "    base_seed = int(torch.randint(2**32, (1,)).item())\n",
    "    lib_seed = (base_seed + worker_id) % (2**32)\n",
    "    imgaug.seed(lib_seed)\n",
    "    np.random.seed(lib_seed)\n",
    "\n",
    "dataloader_args = {\n",
    "    \"batch_size\": batch_size,\n",
    "    \"shuffle\": False,\n",
    "    \"num_workers\": 16,\n",
    "    \"pin_memory\": True,\n",
    "    \"worker_init_fn\": worker_init_fn,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(export_folder):\n",
    "    os.makedirs(export_folder)\n",
    "    print(\"Created export folder!\")\n",
    "if video:\n",
    "    paths = [\"videos/similarity\", \"videos/latentspaces\"]\n",
    "    for path in paths:\n",
    "        fpath = os.path.join(video_path, path)\n",
    "        if os.path.exists(fpath):\n",
    "            continue\n",
    "        os.makedirs(fpath)\n",
    "        print(f\"Created {fpath} folder!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-02T08:43:48.691474Z",
     "start_time": "2020-04-02T08:43:48.672706Z"
    }
   },
   "outputs": [],
   "source": [
    "class ZurichDataset(Dataset):\n",
    "    def __init__(self, folder_path, transform=None):\n",
    "        self.path = folder_path\n",
    "        self.image_paths = sorted(glob.glob(folder_path))\n",
    "        self.transform = transform\n",
    "        # Preloading the dataset in the RAM\n",
    "        self.images = []\n",
    "        for i in range(len(self.image_paths)):\n",
    "            img = skio.imread(self.image_paths[i])\n",
    "            img = (img / 2**12).astype(np.float32) # 12-bit data\n",
    "            self.images.append(img)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.image_paths)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.get(idx)\n",
    "\n",
    "    def get(self, idx, augment=True):\n",
    "        if augment and self.transform:\n",
    "            return self.transform(self.images[idx])\n",
    "        return self.images[idx]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hyper-parameters / Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-02T08:44:04.553916Z",
     "start_time": "2020-04-02T08:43:50.262869Z"
    }
   },
   "outputs": [],
   "source": [
    "class ImgAugTransform:\n",
    "    def __init__(self, testing=False):\n",
    "        if not testing:\n",
    "            self.aug = iaa.Sequential([\n",
    "                iaa.CropToFixedSize(128,128),\n",
    "                iaa.Fliplr(0.5),\n",
    "                iaa.Affine(rotate=(-180, 180), order=[0, 1, 3], mode=\"symmetric\"),\n",
    "                iaa.Sometimes(0.2, iaa.OneOf([\n",
    "                    iaa.AdditiveGaussianNoise(loc=0, scale=(0., 0.05)),\n",
    "                    iaa.GaussianBlur(sigma=(0, 1.0)),\n",
    "                    iaa.EdgeDetect(alpha=1.0),\n",
    "                    iaa.CoarseDropout(0.1, size_percent=0.05, per_channel=True),\n",
    "                ])),\n",
    "                iaa.Multiply((0.9, 1.1), per_channel=0.3),\n",
    "            ])\n",
    "        else:\n",
    "            self.aug = iaa.Sequential([\n",
    "                iaa.CropToFixedSize(128,128),\n",
    "            ])\n",
    "\n",
    "    def __call__(self, img):\n",
    "        img = np.array(img)\n",
    "        return self.aug.augment_image(img)\n",
    "\n",
    "dset = ZurichDataset(ZURICH + \"/images_tif/zh10.tif\", transform=ImgAugTransform())\n",
    "dset_test = ZurichDataset(ZURICH + \"/images_tif/zh12.tif\", transform=ImgAugTransform(testing=True))\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    dset,\n",
    "    sampler=OverSampler(dset, samples_per_epoch),\n",
    "    **dataloader_args\n",
    ")\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "    dset_test,\n",
    "    sampler=OverSampler(dset_test, samples_per_epoch),\n",
    "    **dataloader_args\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-02T08:44:05.028261Z",
     "start_time": "2020-04-02T08:44:04.555503Z"
    }
   },
   "outputs": [],
   "source": [
    "#imgaug.seed(0)\n",
    "idx = 0 # Img to display\n",
    "fig, ax = plt.subplots(2, 2, figsize=(7, 7), sharex=True, sharey=True)\n",
    "fig.suptitle(\"Sample from the dataset\")\n",
    "ax[0, 0].set_title(modA_name)\n",
    "ax[0, 0].set_ylabel(\"Training set\")\n",
    "sample = dset[idx]\n",
    "ax[0, 0].imshow(mode2img(sample, modA))\n",
    "ax[0, 1].set_title(modB_name)\n",
    "ax[0, 1].imshow(mode2img(sample, modB))\n",
    "ax[1, 0].set_ylabel(\"Testing set\")\n",
    "sample = dset_test[idx]\n",
    "ax[1, 0].imshow(mode2img(sample, modA))\n",
    "ax[1, 1].imshow(mode2img(sample, modB))\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-02T08:44:08.022693Z",
     "start_time": "2020-04-02T08:44:05.032437Z"
    }
   },
   "outputs": [],
   "source": [
    "from models.tiramisu import DenseUNet\n",
    "\n",
    "class ModNet(DenseUNet):\n",
    "    def __init__(self, **args):\n",
    "        super(ModNet, self).__init__(**args, include_top=False)\n",
    "        out_channels = self.get_channels_count()[-1]\n",
    "        self.final_conv = torch.nn.Conv2d(out_channels, latent_channels, 1, bias=False)\n",
    "\n",
    "    def forward(self, x):\n",
    "        # Penultimate layer\n",
    "        L_hat = super(ModNet, self).forward(x)\n",
    "        # Final convolution\n",
    "        return self.final_conv(L_hat)\n",
    "\n",
    "torch.manual_seed(0)\n",
    "modelA = ModNet(in_channels=modA_len, nb_classes=latent_channels, **tiramisu_args).to(device1)\n",
    "modelB = ModNet(in_channels=modB_len, nb_classes=latent_channels, **tiramisu_args).to(device2)\n",
    "\n",
    "optimizerA = optimiser(modelA.parameters(), **opt_args)\n",
    "optimizerB = optimiser(modelB.parameters(), **opt_args)\n",
    "\n",
    "print(\"*** MODEL A ***\")\n",
    "modelA.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_pairwise_loss(Ls, similarity_fn, tau=1.0, device=None):\n",
    "    \"\"\"Computation of the final loss.\n",
    "    \n",
    "    Args:\n",
    "        Ls (list): the latent spaces.\n",
    "        similarity_fn (func): the similarity function between two datapoints x and y.\n",
    "        tau (float): the temperature to apply to the similarities.\n",
    "        device (str): the torch device to store the data and perform the computations.\n",
    "    \n",
    "    Returns (list of float):\n",
    "        softmaxes: the loss for each positive sample (length=2N, with N=batch size).\n",
    "        similarities: the similarity matrix with all pairwise similarities (2N, 2N)\n",
    "\n",
    "    Note:\n",
    "        This implementation works in the case where only 2 modalities are of\n",
    "        interest (M=2). Please refer to the paper for the full algorithm.\n",
    "    \"\"\"\n",
    "    # Computation of the similarity matrix\n",
    "    # The matrix contains the pairwise similarities between each sample of the full batch\n",
    "    # and each modalities.\n",
    "    points = torch.cat([L.to(device) for L in Ls])\n",
    "    N = batch_size\n",
    "    similarities = torch.zeros(2*N, 2*N).to(device)\n",
    "    for i in range(2*N):\n",
    "        for j in range(i+1):\n",
    "            s = similarity_fn(points[i], points[j])/tau\n",
    "            similarities[i, j] = s\n",
    "            similarities[j, i] = s\n",
    "\n",
    "    # Computation of the loss, one row after the other.\n",
    "    irange = np.arange(2*N)\n",
    "    softmaxes = torch.empty(2*N).to(device)\n",
    "    for i in range(2*N):\n",
    "        j = (i + N) % (2 * N)\n",
    "        pos = similarities[i, j]\n",
    "        # The negative examples are all the remaining points\n",
    "        # excluding self-similarity\n",
    "        neg = similarities[i][irange != i]\n",
    "        softmaxes[i] = -pos + torch.logsumexp(neg, dim=0)\n",
    "    return softmaxes, similarities\n",
    "\n",
    "def test():\n",
    "    \"\"\"Runs the model on the test data.\"\"\"\n",
    "    modelA.eval()\n",
    "    modelB.eval()\n",
    "    test_loss = []\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, data in enumerate(test_loader):\n",
    "            data = data.permute(0, 3, 1, 2)\n",
    "            dataA = data[:, modA].float().to(device1)\n",
    "            dataB = data[:, modB].float().to(device2)\n",
    "            L1 = modelA(dataA)\n",
    "            L2 = modelB(dataB)\n",
    "\n",
    "            softmaxes, similarities = compute_pairwise_loss(\n",
    "                [L1, L2],\n",
    "                similarity_fn=sim_func,\n",
    "                tau=tau,\n",
    "                device=device1\n",
    "            )\n",
    "            loss_test = softmaxes.mean()\n",
    "\n",
    "            if act_l1 > 0.:\n",
    "                loss_test += act_l1 * activation_decay([L1, L2], p=1, device=device1)\n",
    "            if act_l2 > 0.:\n",
    "                loss_test += act_l2 * activation_decay([L1, L2], p=2, device=device1)\n",
    "            test_loss.append(loss_test.item())\n",
    "    losses[\"test\"].append(np.mean(test_loss))\n",
    "\n",
    "    # Video creation if necessary\n",
    "    if video:\n",
    "        with torch.no_grad():\n",
    "            imgaug.seed(0)\n",
    "            batch = np.stack([dset[i] for i in range(batch_size)]).transpose(0, 3, 1, 2)\n",
    "            tbatchA = torch.tensor(batch[:, modA]).float().to(device1)\n",
    "            tbatchB = torch.tensor(batch[:, modB]).float().to(device2)\n",
    "\n",
    "            L1 = modelA(tbatchA)\n",
    "            L2 = modelB(tbatchB)\n",
    "\n",
    "            # Saving the latent spaces\n",
    "            idx = 0\n",
    "            fig = plot_input_latent(tbatchA[idx], tbatchB[idx], L1[idx], L2[idx])\n",
    "            plt.savefig(os.path.join(export_folder, f\"videos/latentspaces/{epoch:03d}.png\"))\n",
    "            plt.close(fig)\n",
    "            # Saving the similarity matrix\n",
    "            _, similarities = compute_pairwise_loss(\n",
    "                [L1, L2],\n",
    "                similarity_fn=sim_func,\n",
    "                tau=tau,\n",
    "                device=device1\n",
    "            )\n",
    "            fig = plt.figure(figsize=(5, 5))\n",
    "            plot_similarity_matrix(similarities, \"Similarity matrix (test set)\")\n",
    "            plt.savefig(os.path.join(export_folder, f\"videos/similarity/{epoch:03d}.png\"))\n",
    "            plt.close(fig)\n",
    "\n",
    "    return loss_test, similarities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-02T08:44:18.988971Z",
     "start_time": "2020-04-02T08:44:08.024668Z"
    }
   },
   "outputs": [],
   "source": [
    "import time\n",
    "start = time.time()\n",
    "\n",
    "modelA = modelA.to(device1)\n",
    "modelB = modelB.to(device2)\n",
    "torch.manual_seed(0)\n",
    "\n",
    "losses = {\"train\": [], \"test\": []}\n",
    "epoch = 0\n",
    "test() # First round of testing\n",
    "for epoch in range(1, epochs+1):\n",
    "    modelA.train()\n",
    "    modelB.train()\n",
    "    for batch_idx, data in enumerate(train_loader):\n",
    "        # Preparing the batch\n",
    "        data = data.permute(0, 3, 1, 2)\n",
    "        dataA = data[:, modA].float().to(device1)\n",
    "        dataB = data[:, modB].float().to(device2)\n",
    "        # Reseting the optimizer (gradients set to zero)\n",
    "        optimizerA.zero_grad()\n",
    "        optimizerB.zero_grad()\n",
    "\n",
    "        if equivariance:\n",
    "            # Applies random 90 degrees rotations to the data (group p4)\n",
    "            # This step enforces the formula of equivariance: d(f(T(x)), T^{-1}(f(x)))\n",
    "            # With f(x) the neural network, T(x) a transformation, T^{-1}(x) the inverse transformation\n",
    "            random_rotA = np.random.randint(4, size=batch_size)\n",
    "            random_rotB = np.random.randint(4, size=batch_size)\n",
    "            dataA_p4 = batch_rotate_p4(dataA, random_rotA, device1)\n",
    "            dataB_p4 = batch_rotate_p4(dataB, random_rotB, device2)\n",
    "\n",
    "            # Compute the forward pass\n",
    "            L1 = modelA(dataA_p4)\n",
    "            L2 = modelB(dataB_p4)\n",
    "\n",
    "            # Applies the inverse of the 90 degree rotation to recover the right positions\n",
    "            L1_ungrouped = batch_rotate_p4(L1, -random_rotA, device1)\n",
    "            L2_ungrouped = batch_rotate_p4(L2, -random_rotB, device2)\n",
    "        else:\n",
    "            L1_ungrouped = modelA(dataA)\n",
    "            L2_ungrouped = modelB(dataB)\n",
    "\n",
    "        # Computes the loss\n",
    "        loss = compute_pairwise_loss(\n",
    "            [L1_ungrouped, L2_ungrouped],\n",
    "            similarity_fn=sim_func,\n",
    "            tau=tau,\n",
    "            device=device1\n",
    "        )[0].mean()\n",
    "\n",
    "        # Activation regularization\n",
    "        if act_l1 > 0.:\n",
    "            loss += act_l1 * activation_decay([L1, L2], p=1., device=device1)\n",
    "        if act_l2 > 0.:\n",
    "            loss += act_l2 * activation_decay([L1, L2], p=2., device=device1)\n",
    "\n",
    "        # Computing the gradients\n",
    "        loss.backward()\n",
    "\n",
    "        # Clipping the the gradients if they are too big\n",
    "        torch.nn.utils.clip_grad_norm_(modelA.parameters(), grad_norm)\n",
    "        torch.nn.utils.clip_grad_norm_(modelB.parameters(), grad_norm)\n",
    "        # Performing the gradient descent\n",
    "        optimizerA.step()\n",
    "        optimizerB.step()\n",
    "\n",
    "        losses[\"train\"].append(loss.item())\n",
    "        clear_output(wait=True)\n",
    "        plot_losses(losses, steps, steps_per_epoch)\n",
    "        plt.show()\n",
    "\n",
    "    # Testing after each epoch\n",
    "    _, similarities = test()\n",
    "    clear_output(wait=True)\n",
    "    plot_losses(losses, steps, steps_per_epoch)\n",
    "    plt.show()\n",
    "\n",
    "modelA.eval();\n",
    "modelB.eval();\n",
    "\n",
    "duration = time.time() - start\n",
    "print(duration)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-04-02T08:31:36.551278Z",
     "start_time": "2020-04-02T08:31:36.513039Z"
    }
   },
   "outputs": [],
   "source": [
    "date = datetime.now().strftime(\"%Y%d%m_%H%M%S\")\n",
    "model_path = os.path.join(export_folder, f\"model_{date}.pt\")\n",
    "with warnings.catch_warnings():\n",
    "    warnings.simplefilter(\"ignore\")\n",
    "    torch.save({\n",
    "        \"modelA\": modelA,\n",
    "        \"modelB\": modelB,\n",
    "    }, model_path)\n",
    "print(f\"model saved as: {model_path}\")\n",
    "HTML(f'<a href=\"{model_path}\" title=\"{model_path}\" download target=\"_blank\" style=\"text-decoration:none\">🤖 Download model 🤖</a>')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sanity check of the training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-30T10:37:58.158384Z",
     "start_time": "2020-03-30T10:37:58.051780Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(5, 5))\n",
    "plot_similarity_matrix(similarities, \"Similarity matrix (test set)\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-30T13:41:48.802506Z",
     "start_time": "2020-03-30T13:41:47.822696Z"
    }
   },
   "outputs": [],
   "source": [
    "L1 = modelA.to(\"cpu\")(dataA.to(\"cpu\"))\n",
    "L2 = modelB.to(\"cpu\")(dataB.to(\"cpu\"))\n",
    "\n",
    "# Example id (can range from 0 to batch_size)\n",
    "for b in range(2):\n",
    "    plot_input_latent(dataA[b], dataB[b], L1[b], L2[b])\n",
    "    plt.savefig(os.path.join(export_folder, f\"latentspace{b:0>2d}.png\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random_rotA = np.random.randint(4, size=batch_size)\n",
    "random_rotB = np.random.randint(4, size=batch_size)\n",
    "dataA_p4 = batch_rotate_p4(dataA, random_rotA, device1)\n",
    "dataB_p4 = batch_rotate_p4(dataB, random_rotB, device2)\n",
    "L1 = modelA.to(\"cpu\")(dataA_p4.to(\"cpu\"))\n",
    "L2 = modelB.to(\"cpu\")(dataB_p4.to(\"cpu\"))\n",
    "\n",
    "# Example id (can range from 0 to batch_size)\n",
    "for b in range(2):\n",
    "    plot_input_latent(dataA_p4[b], dataB_p4[b], L1[b], L2[b])\n",
    "    plt.savefig(os.path.join(export_folder, f\"latentspace{b:0>2d}.png\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-30T13:41:57.919478Z",
     "start_time": "2020-03-30T13:41:57.148054Z"
    }
   },
   "outputs": [],
   "source": [
    "N = 2 # must be between 1 and batch_size\n",
    "\n",
    "cmap = \"jet\"\n",
    "fig, ax = plt.subplots(2*N, latent_channels, sharex=True, sharey=True, figsize=(latent_channels*3, 2*N*3))\n",
    "for i in range(N):\n",
    "    for j in range(latent_channels):\n",
    "        if i == 0:\n",
    "            ax[i, j].set_title(f\"Channel #{j+1}\")\n",
    "        if j == 0:\n",
    "            ax[i*2+0, j].set_ylabel(\"Latent space A\")\n",
    "            ax[i*2+1, j].set_ylabel(\"Latent space B\")\n",
    "        ax[i*2+0, j].imshow(tensor2np(L1[i, j]), cmap=cmap)\n",
    "        ax[i*2+1, j].imshow(tensor2np(L2[i, j]), cmap=cmap)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-29T21:23:45.563179Z",
     "start_time": "2020-03-29T21:23:45.295949Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15, 5))\n",
    "plt.title(\"Distribution of the intensities of the latent spaces\")\n",
    "_, bins1, _ = plt.hist(tensor2np(L1).ravel(), density=True, bins=20, label=\"Latent space A\")\n",
    "_, bins2, _ = plt.hist(tensor2np(L2).ravel(), density=True, bins=20, alpha=0.5, label=\"Latent space B\")\n",
    "posmax = max(abs(bins1.min()), abs(bins2.min()), bins1.max(), bins2.max())\n",
    "plt.xlim(-1.05*posmax, 1.05*posmax)\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checking the registration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Models available in the release file\n",
    "model = torch.load(\"models/model_zurich_rot90_disabled.pt\")\n",
    "modelA = model[\"modelA\"].cpu()\n",
    "modelA.eval()\n",
    "modelB = model[\"modelB\"].cpu()\n",
    "modelB.eval();\n",
    "\n",
    "model90 = torch.load(\"models/model_zurich_rot90_enabled.pt\")\n",
    "modelA90 = model90[\"modelA\"].cpu()\n",
    "modelA90.eval()\n",
    "modelB90 = model90[\"modelB\"].cpu()\n",
    "modelB90.eval();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_patch(img, x, y, w, h):\n",
    "    return img[y-h//2:y+h//2, x-w//2:x+w//2]\n",
    "\n",
    "def get_patch(img, patch_x, patch_y, patch_rot, patch_w = 128, patch_h = 128, interpolation=3):\n",
    "    max_patch_w = int(np.ceil(np.sqrt(2) * patch_w))\n",
    "    max_patch_h = int(np.ceil(np.sqrt(2) * patch_h))\n",
    "\n",
    "    unrotated_patch = extract_patch(img, patch_x, patch_y, max_patch_w, max_patch_h)\n",
    "    rotated_patch = sktr.rotate(unrotated_patch, patch_rot, order=interpolation)\n",
    "    final_rotated_patch = extract_patch(rotated_patch, max_patch_w//2, max_patch_h//2, patch_w, patch_h)\n",
    "    return final_rotated_patch\n",
    "\n",
    "img = skio.imread(ZURICH + \"/images_tif/zh12.tif\")\n",
    "img = (img / 2**12).astype(np.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def weighted_mean(x, w):\n",
    "    return np.sum(w * x) / np.sum(w)\n",
    "\n",
    "def weighted_std(x, w):\n",
    "    return np.sqrt(np.sum(w * (x - weighted_mean(x, w))**2) / np.sum(w))\n",
    "\n",
    "def weighted_corrcoef(a, b, w):\n",
    "    return weighted_mean((a - weighted_mean(a, w))*(b - weighted_mean(b, w)), w)/(weighted_std(a, w)*weighted_std(b, w))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This part pregenerates the correlation for plotting in the video below.\n",
    "import matplotlib.lines as lines\n",
    "\n",
    "aa = 16 # Antialiasing for the mask\n",
    "mx, my = np.mgrid[-1:1:2/128/aa, -1:1:2/128/aa]\n",
    "mask = (mx**2 + my**2)<1\n",
    "mask = sktr.rescale(mask.astype(float), 1/aa, anti_aliasing=True, anti_aliasing_sigma=1)\n",
    "maskRGB = np.repeat(mask[..., np.newaxis], 3, axis=-1)\n",
    "\n",
    "patchA = get_patch(img[..., modA], 610, 150, 0)\n",
    "patchB = get_patch(img[..., modB], 610, 150, 0)\n",
    "dataA = torch.tensor(patchA[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "dataB = torch.tensor(patchB[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "L1_0ang = modelA(dataA).detach().numpy()[0].transpose(1, 2, 0)\n",
    "L2_0ang = modelB(dataB).detach().numpy()[0].transpose(1, 2, 0)\n",
    "\n",
    "L1_0ang90 = modelA90(dataA).detach().numpy()[0].transpose(1, 2, 0)\n",
    "L2_0ang90 = modelB90(dataB).detach().numpy()[0].transpose(1, 2, 0)\n",
    "\n",
    "norot_L1 = []\n",
    "norot_L2 = []\n",
    "rot90_L1 = []\n",
    "rot90_L2 = []\n",
    "for angle in tqdm(range(360)):\n",
    "    patchA = get_patch(img[..., modA], 610, 150, angle)\n",
    "    patchB = get_patch(img[..., modB], 610, 150, angle)\n",
    "    dataA = torch.tensor(patchA[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "    dataB = torch.tensor(patchB[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "\n",
    "    L1 = modelA(dataA).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L2 = modelB(dataB).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L1 = sktr.rotate(L1, -angle, order=3)\n",
    "    L2 = sktr.rotate(L2, -angle, order=3)\n",
    "    norot_L1.append(weighted_corrcoef(L1_0ang.ravel(), L1.ravel(), maskRGB.ravel()))\n",
    "    norot_L2.append(weighted_corrcoef(L2_0ang.ravel(), L2.ravel(), maskRGB.ravel()))\n",
    "\n",
    "    L1 = modelA90(dataA).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L2 = modelB90(dataB).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L1 = sktr.rotate(L1, -angle, order=3)\n",
    "    L2 = sktr.rotate(L2, -angle, order=3)\n",
    "    rot90_L1.append(weighted_corrcoef(L1_0ang90.ravel(), L1.ravel(), maskRGB.ravel()))\n",
    "    rot90_L2.append(weighted_corrcoef(L2_0ang90.ravel(), L2.ravel(), maskRGB.ravel()))\n",
    "    #break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "patchRGB = np.concatenate([patchA[..., i:i+1] for i in [2,1,0]], axis=-1)\n",
    "patchRGB -= np.percentile(patchRGB.reshape(-1, 3), 0.01, axis=0)\n",
    "patchRGB /= np.percentile(patchRGB.reshape(-1, 3), 99.9, axis=0)\n",
    "patchRGB = np.concatenate([patchRGB, mask[..., np.newaxis]], axis=-1)\n",
    "plt.imshow(patchRGB)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generates each frames for a figure of the paper\n",
    "import matplotlib.lines as lines\n",
    "import matplotlib as mpl\n",
    "import matplotlib.collections as coll\n",
    "plt.rcParams.update({'font.size': 13})\n",
    "\n",
    "aa = 8 # Antialiasing for the mask\n",
    "mx, my = np.mgrid[-1:1:2/128/aa, -1:1:2/128/aa]\n",
    "mask = (mx**2 + my**2)<1\n",
    "mask = sktr.rescale(mask.astype(float), 1/aa, anti_aliasing=True, anti_aliasing_sigma=1)\n",
    "\n",
    "def lineToSegments(x, y):\n",
    "    segments = np.array([x, y]).T.reshape(-1, 1, 2)\n",
    "    return np.concatenate([segments[:-1], segments[1:]], axis=1)\n",
    "\n",
    "for angle in tqdm(range(360)):\n",
    "    patchA = get_patch(img[..., modA], 610, 150, 0)\n",
    "    patchB = get_patch(img[..., modB], 610, 150, 0)\n",
    "    patchRGB = np.concatenate([patchA[..., i:i+1] for i in [2,1,0]], axis=-1)\n",
    "    patchRGB -= np.percentile(patchRGB.reshape(-1, 3), 0.01, axis=0)\n",
    "    patchRGB /= np.percentile(patchRGB.reshape(-1, 3), 99.9, axis=0)\n",
    "    patchRGB = np.concatenate([np.clip(patchRGB, 0, 1.), mask[..., np.newaxis]], axis=-1)\n",
    "\n",
    "    fig, ax = plt.subplots(2, 5, figsize=(12, 5), dpi=300)\n",
    "    # polar\n",
    "    gs = ax[0, 3].get_gridspec()\n",
    "    for a in ax[:, 3:].ravel():\n",
    "        a.remove()\n",
    "    axbig = fig.add_subplot(gs[:, 3:], polar=True)\n",
    "    # Get the color\n",
    "    N = 360\n",
    "    #cycler = mpl.rcParams['axes.prop_cycle'][0]\n",
    "    cm = np.array(mpl.colors.to_rgba('#1f77b4'))\n",
    "    c = np.repeat(cm[np.newaxis], N, axis=0)\n",
    "    transparency = np.minimum(np.linspace(0, 1, N)**2 + .1, 1.0)\n",
    "    c[:, 3] = np.roll(transparency, angle)\n",
    "    # Plot the results\n",
    "    x = np.linspace(0, 2*np.pi, N)\n",
    "    lc1 = coll.LineCollection(lineToSegments(x, np.abs(norot_L1)), colors=c)\n",
    "    axbig.add_artist(lc1)\n",
    "    cm = np.array(mpl.colors.to_rgba('#ff7f0e'))\n",
    "    c = np.repeat(cm[np.newaxis], N, axis=0)\n",
    "    c[:, 3] = np.roll(transparency, angle)\n",
    "    # Plot the results\n",
    "    x = np.linspace(0, 2*np.pi, N)\n",
    "    lc2 = coll.LineCollection(lineToSegments(x, np.abs(rot90_L1)), colors=c)\n",
    "    axbig.add_artist(lc2)\n",
    "    axbig.set_ylim(0, 1.1)\n",
    "    axbig.set_title(f\"Correlation $\\\\left[\\\\rho\\\\right]_+$ between\\nCoMIR$_{{\\\\theta=0\\\\degree}}$ and CoMIR$_{{\\\\theta={angle:3d}\\\\degree}}$\")\n",
    "    axbig.axes.spines['polar'].set_visible(False)\n",
    "\n",
    "    cm = plt.get_cmap(\"viridis\")\n",
    "    patchNIR = cm(patchB[...,0] / 0.2)\n",
    "    patchNIR[..., -1] = mask\n",
    "    ax[0, 0].imshow(patchRGB, vmin=0, vmax=1)\n",
    "    ax[0, 0].axis(\"off\")\n",
    "    ax[1, 0].imshow(patchNIR, vmin=0, vmax=1)\n",
    "    ax[1, 0].axis(\"off\")\n",
    "    #line = lines.Line2D([64, 64+np.cos(-angle*np.pi/180)*64], [64, 64+np.sin(-angle*np.pi/180)*64], lw=1, color='white', axes=ax[0, 0])\n",
    "    #ax[0, 0].add_line(line)\n",
    "    ax[0, 0].arrow(64, 64, np.cos(-angle*np.pi/180)*64, np.sin(-angle*np.pi/180)*64, color=\"white\", width=2, clip_on=False, length_includes_head=True)\n",
    "    ax[1, 0].arrow(64, 64, np.cos(-angle*np.pi/180)*64, np.sin(-angle*np.pi/180)*64, color=\"white\", width=2, clip_on=False, length_includes_head=True)\n",
    "\n",
    "    patchA = get_patch(img[..., modA], 610, 150, angle)\n",
    "    patchB = get_patch(img[..., modB], 610, 150, angle)\n",
    "    dataA = torch.tensor(patchA[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "    dataB = torch.tensor(patchB[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "\n",
    "    L1 = modelA(dataA).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L2 = modelB(dataB).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L1 = sktr.rotate(L1, -angle, order=3)\n",
    "    L2 = sktr.rotate(L2, -angle, order=3)\n",
    "\n",
    "    minint, maxint = min(L1.min(), L2.min()), max(L1.max(), L2.max())\n",
    "    L1n, L2n = normalize_common(L1, L2)\n",
    "    L1n = np.concatenate([L1n, mask[..., np.newaxis]], axis=-1)\n",
    "    L2n = np.concatenate([L2n, mask[..., np.newaxis]], axis=-1)\n",
    "    ax[0, 1].imshow(L1n, vmin=0, vmax=1)\n",
    "    ax[0, 1].axis(\"off\")\n",
    "    circle1 = plt.Circle((63.5, 63.5), 64, color=\"#1f77b4\", linewidth=4, fill=False, clip_on=False)\n",
    "    ax[0, 1].add_artist(circle1)\n",
    "    ax[1, 1].imshow(L2n, vmin=0, vmax=1)\n",
    "    ax[1, 1].axis(\"off\")\n",
    "\n",
    "    L1 = modelA90(dataA).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L2 = modelB90(dataB).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L1 = sktr.rotate(L1, -angle, order=3)\n",
    "    L2 = sktr.rotate(L2, -angle, order=3)\n",
    "\n",
    "    minint, maxint = min(L1.min(), L2.min()), max(L1.max(), L2.max())\n",
    "    L1n, L2n = normalize_common(L1, L2)\n",
    "    L1n = np.concatenate([L1n, mask[..., np.newaxis]], axis=-1)\n",
    "    L2n = np.concatenate([L2n, mask[..., np.newaxis]], axis=-1)\n",
    "    ax[0, 2].imshow(L1n, vmin=0, vmax=1)\n",
    "    ax[0, 2].axis(\"off\")\n",
    "    circle1 = plt.Circle((63.5, 63.5), 64, color=\"#ff7f0e\", linewidth=4, fill=False, clip_on=False)\n",
    "    ax[0, 2].add_artist(circle1)\n",
    "    ax[1, 2].imshow(L2n, vmin=0, vmax=1)\n",
    "    ax[1, 2].axis(\"off\")\n",
    "\n",
    "    ax[0, 0].set_title(\"Input images\")\n",
    "    ax[0, 1].set_title(\"CoMIR\\n$\\\\mathcal{C}_4$ disabled\")\n",
    "    ax[0, 2].set_title(\"CoMIR\\n$\\\\mathcal{C}_4$ enabled\")\n",
    "    \n",
    "    ax[0, 0].text(-15, 64, \"RGB\", fontsize=\"large\", rotation=90, va=\"center\")\n",
    "    ax[1, 0].text(-15, 64, \"NIR\", fontsize=\"large\", rotation=90, va=\"center\")\n",
    "    plt.tight_layout()\n",
    "\n",
    "    plt.subplots_adjust(0.0, 0.0, 0.9, 0.9, 0.05, 0.05)\n",
    "    plt.savefig(f\"results/videos/rot90/frame{angle:03d}.png\", bbox_inches='tight', pad_inches=0.10)\n",
    "    plt.close(fig)\n",
    "    #break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generating the video above is done with ffmpeg:\n",
    "#> ffmpeg -framerate 36 -i rot90/frame%03d.png -r 30 -pix_fmt yuv420p -filter:v \"crop=in_w-1:in_h-1\" rotation.mp4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generates the figure (associated to the video generated in the previous cell) appearing in the paper\n",
    "import matplotlib.lines as lines\n",
    "plt.rcParams.update({'font.size': 20})\n",
    "\n",
    "aa = 8 # Antialiasing for the mask\n",
    "mx, my = np.mgrid[-1:1:2/128/aa, -1:1:2/128/aa]\n",
    "mask = (mx**2 + my**2)<1\n",
    "mask = sktr.rescale(mask.astype(float), 1/aa, anti_aliasing=True, anti_aliasing_sigma=1)\n",
    "\n",
    "def lineToSegments(x, y):\n",
    "    segments = np.array([x, y]).T.reshape(-1, 1, 2)\n",
    "    return np.concatenate([segments[:-1], segments[1:]], axis=1)\n",
    "\n",
    "patchA = get_patch(img[..., modA], 610, 150, 0)\n",
    "patchB = get_patch(img[..., modB], 610, 150, 0)\n",
    "patchA90 = get_patch(img[..., modA], 610, 150, 90)\n",
    "patchB90 = get_patch(img[..., modB], 610, 150, 90)\n",
    "\n",
    "patchRGB = np.concatenate([patchA[..., i:i+1] for i in [2,1,0]], axis=-1)\n",
    "patchRGB -= np.percentile(patchRGB.reshape(-1, 3), 0.01, axis=0)\n",
    "patchRGB /= np.percentile(patchRGB.reshape(-1, 3), 99.9, axis=0)\n",
    "patchRGB = np.concatenate([np.clip(patchRGB, 0, 1.), mask[..., np.newaxis]], axis=-1)\n",
    "\n",
    "angle = 90\n",
    "fig, ax = plt.subplots(2, 7, figsize=(20, 6))\n",
    "# polar\n",
    "gs = ax[0, 5].get_gridspec()\n",
    "for a in ax[:, 5:].ravel():\n",
    "    a.remove()\n",
    "axbig = fig.add_subplot(gs[:, 5:], polar=True)\n",
    "# Get the color\n",
    "N = 360\n",
    "#cycler = mpl.rcParams['axes.prop_cycle'][0]\n",
    "cm = np.array(mpl.colors.to_rgba('#1f77b4'))\n",
    "c = np.repeat(cm[np.newaxis], N, axis=0)\n",
    "# Plot the results\n",
    "x = np.linspace(0, 2*np.pi, N)\n",
    "axbig.plot(x, np.abs(norot_L1), color=cm, label=\"$\\\\mathcal{C}_4$ disabled\")\n",
    "cm = np.array(mpl.colors.to_rgba('#ff7f0e'))\n",
    "c = np.repeat(cm[np.newaxis], N, axis=0)\n",
    "# Plot the results\n",
    "x = np.linspace(0, 2*np.pi, N)\n",
    "axbig.plot(x, np.abs(rot90_L1), color=cm, label=\"$\\\\mathcal{C}_4$ enabled\")\n",
    "axbig.set_ylim(0, 1.1)\n",
    "axbig.set_title(\"Correlation $\\\\left[\\\\rho\\\\right]_+$ between\\nCoMIR$_{{\\\\theta=0\\\\degree}}$ and CoMIR$_{\\\\theta \\\\in [0; 360)}$\")\n",
    "marker_size = 100\n",
    "axbig.scatter([np.pi/40], [.9], zorder=3, c=\"#1f77b4\", marker=\"s\", s=marker_size)\n",
    "axbig.scatter([np.pi/2], [np.abs(norot_L1)[90]], zorder=3, c=\"#1f77b4\", s=marker_size)\n",
    "axbig.scatter([-np.pi/40], [.9], zorder=3, c=\"#ff7f0e\", marker=\"s\", s=marker_size)\n",
    "axbig.scatter([np.pi/2], [np.abs(rot90_L1)[90]], zorder=3, c=\"#ff7f0e\", s=marker_size)\n",
    "axbig.plot([np.pi/40, 0, -np.pi/40], [.9, 1., .9], c=\"k\", alpha=0.4)\n",
    "axbig.axes.spines['polar'].set_visible(False)\n",
    "axbig.legend(loc='upper center', bbox_to_anchor=(0.5, -0.05), ncol=2)\n",
    "\n",
    "cm = plt.get_cmap(\"viridis\")\n",
    "patchNIR = cm(patchB[...,0] / 0.2)\n",
    "patchNIR[..., -1] = mask\n",
    "ax[0, 0].imshow(patchRGB, vmin=0, vmax=1, interpolation=\"nearest\")\n",
    "ax[0, 0].axis(\"off\")\n",
    "ax[1, 0].imshow(patchNIR, vmin=0, vmax=1, interpolation=\"nearest\")\n",
    "ax[1, 0].axis(\"off\")\n",
    "\n",
    "dataA = torch.tensor(patchA[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "dataB = torch.tensor(patchB[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "dataA90 = torch.tensor(patchA90[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "dataB90 = torch.tensor(patchB90[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "\n",
    "L1 = modelA(dataA).detach().numpy()[0].transpose(1, 2, 0)\n",
    "L2 = modelB(dataB).detach().numpy()[0].transpose(1, 2, 0)\n",
    "\n",
    "L190 = modelA(dataA90).detach().numpy()[0].transpose(1, 2, 0)\n",
    "L290 = modelB(dataB90).detach().numpy()[0].transpose(1, 2, 0)\n",
    "L190 = sktr.rotate(L190, -angle, order=3)\n",
    "L290 = sktr.rotate(L290, -angle, order=3)\n",
    "\n",
    "L1n, L2n = normalize_common(L1, L2)\n",
    "L1n = np.concatenate([L1n, mask[..., np.newaxis]], axis=-1)\n",
    "L2n = np.concatenate([L2n, mask[..., np.newaxis]], axis=-1)\n",
    "ax[0, 1].imshow(L1n, vmin=0, vmax=1)\n",
    "ax[0, 1].axis(\"off\")\n",
    "ax[0, 1].scatter([10], [10], c=\"#1f77b4\", marker=\"s\", s=marker_size)\n",
    "circle1 = plt.Circle((63.5, 63.5), 64, color=\"#1f77b4\", linewidth=4, fill=False, clip_on=False)\n",
    "ax[0, 1].add_artist(circle1)\n",
    "ax[1, 1].imshow(L2n, vmin=0, vmax=1)\n",
    "ax[1, 1].axis(\"off\")\n",
    "L1n, L2n = normalize_common(L190, L290)\n",
    "L1n = np.concatenate([L1n, mask[..., np.newaxis]], axis=-1)\n",
    "L2n = np.concatenate([L2n, mask[..., np.newaxis]], axis=-1)\n",
    "ax[0, 2].imshow(L1n, vmin=0, vmax=1, interpolation=\"nearest\")\n",
    "ax[0, 2].axis(\"off\")\n",
    "circle1 = plt.Circle((63.5, 63.5), 64, color=\"#1f77b4\", linewidth=4, fill=False, clip_on=False)\n",
    "ax[0, 2].add_artist(circle1)\n",
    "ax[0, 2].scatter([10], [10], c=\"#1f77b4\", s=marker_size)\n",
    "ax[1, 2].imshow(L2n, vmin=0, vmax=1)\n",
    "ax[1, 2].axis(\"off\")\n",
    "L1 = modelA90(dataA).detach().numpy()[0].transpose(1, 2, 0)\n",
    "L2 = modelB90(dataB).detach().numpy()[0].transpose(1, 2, 0)\n",
    "\n",
    "L190 = modelA90(dataA90).detach().numpy()[0].transpose(1, 2, 0)\n",
    "L290 = modelB90(dataB90).detach().numpy()[0].transpose(1, 2, 0)\n",
    "L190 = sktr.rotate(L190, -angle, order=3)\n",
    "L290 = sktr.rotate(L290, -angle, order=3)\n",
    "\n",
    "L1n, L2n = normalize_common(L1, L2)\n",
    "L1n = np.concatenate([L1n, mask[..., np.newaxis]], axis=-1)\n",
    "L2n = np.concatenate([L2n, mask[..., np.newaxis]], axis=-1)\n",
    "ax[0, 3].imshow(L1n, vmin=0, vmax=1)\n",
    "ax[0, 3].axis(\"off\")\n",
    "circle2 = plt.Circle((63.5, 63.5), 64, color=\"#ff7f0e\", linewidth=4, fill=False, clip_on=False)\n",
    "ax[0, 3].add_artist(circle2)\n",
    "ax[0, 3].scatter([10], [10], c=\"#ff7f0e\", marker=\"s\", s=marker_size)\n",
    "ax[1, 3].imshow(L2n, vmin=0, vmax=1)\n",
    "ax[1, 3].axis(\"off\")\n",
    "L1n, L2n = normalize_common(L190, L290)\n",
    "L1n = np.concatenate([L1n, mask[..., np.newaxis]], axis=-1)\n",
    "L2n = np.concatenate([L2n, mask[..., np.newaxis]], axis=-1)\n",
    "ax[0, 4].imshow(L1n, vmin=0, vmax=1)\n",
    "ax[0, 4].axis(\"off\")\n",
    "circle2 = plt.Circle((63.5, 63.5), 64, color=\"#ff7f0e\", linewidth=4, fill=False, clip_on=False)\n",
    "ax[0, 4].add_artist(circle2)\n",
    "ax[0, 4].scatter([10], [10], c=\"#ff7f0e\", s=marker_size)\n",
    "ax[1, 4].imshow(L2n, vmin=0, vmax=1)\n",
    "ax[1, 4].axis(\"off\")\n",
    "\n",
    "ax[0, 0].set_title(\"Input images\")\n",
    "ax[0, 1].set_title(\"CoMIR$_{\\\\theta=0\\\\degree}$\\n$\\\\mathcal{C}_4$ disabled\")\n",
    "ax[0, 2].set_title(\"CoMIR$_{\\\\theta=90\\\\degree}$\\n$\\\\mathcal{C}_4$ disabled\")\n",
    "ax[0, 3].set_title(\"CoMIR$_{\\\\theta=0\\\\degree}$\\n$\\\\mathcal{C}_4$ enabled\")\n",
    "ax[0, 4].set_title(\"CoMIR$_{\\\\theta=90\\\\degree}$\\n$\\\\mathcal{C}_4$ enabled\")\n",
    "\n",
    "ax[0, 0].text(-20, 64, \"RGB\", fontsize=\"large\", rotation=90, va=\"center\")\n",
    "ax[1, 0].text(-20, 64, \"NIR\", fontsize=\"large\", rotation=90, va=\"center\")\n",
    "#plt.tight_layout()\n",
    "\n",
    "#plt.subplots_adjust(0.125, 0.1, 0.9, 0.9, 0.2, 0.2)\n",
    "plt.subplots_adjust(0.0, 0.0, 0.9, 0.9, 0.05, 0.05)\n",
    "plt.savefig(f\"results/rot_equivariance.pdf\", bbox_inches='tight', pad_inches=0.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Testing the differrent tau configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !!! IMPORTANT !!!\n",
    "# This needs to be re-run every time to test the different losses:\n",
    "simfuncname = \"mse\"\n",
    "#simfuncname = \"cosine\"\n",
    "#simfuncname = \"corr\"\n",
    "\n",
    "###### METHOD RELATED\n",
    "# The place where the models will be saved\n",
    "export_folder = \"results\" # Add this path to the .gitignore\n",
    "# The number of channels in the latent space (best is 1 or 3 for visualization)\n",
    "latent_channels = 3\n",
    "# Modality slicing\n",
    "# You can choose a set of channels per modality (RGB for instance)\n",
    "# Modality A\n",
    "modA = slice(0, 3)\n",
    "modA_name = \"RGB\"\n",
    "modA_len = modA.stop - modA.start\n",
    "# Modality B\n",
    "modB = slice(3, 4)\n",
    "modB_name = \"NIF\"\n",
    "modB_len = modB.stop - modB.start\n",
    "# Distance function\n",
    "simfunctions = {\n",
    "    \"euclidean\" : lambda x, y: -torch.norm(x - y, p=2, dim=1).mean(),\n",
    "    \"L1\"        : lambda x, y: -torch.norm(x - y, p=1, dim=1).mean(),\n",
    "    \"huber\"     : lambda x, y: -torch.sqrt(1 + (x - y).pow(2)).mean(),\n",
    "    \"mse\"       : lambda x, y: -(x - y).pow(2).mean(),\n",
    "    \"L3\"        : lambda x, y: -torch.norm(x - y, p=3, dim=1).mean(),\n",
    "    \"Linf\"      : lambda x, y: -torch.norm(x - y, p=float(\"inf\"), dim=1).mean(),\n",
    "    \"soft_corr\" : lambda x, y: F.softplus(x*y).sum(axis=1).mean(),\n",
    "    \"corr\"      : lambda x, y: (x*y).sum(axis=1).mean(),\n",
    "    \"cosine\"    : lambda x, y: F.cosine_similarity(x, y, dim=1, eps=1e-8).mean(),\n",
    "    \"angular\"   : lambda x, y: F.cosine_similarity(x, y, dim=1, eps=1e-8).acos().mean() / math.pi,\n",
    "}\n",
    "sim_func = simfunctions[simfuncname]\n",
    "# Temperature (tau) of the loss\n",
    "tau =  .1\n",
    "# L1/L2 activation regularization\n",
    "act_l1 = 0\n",
    "act_l2 = 0\n",
    "# p4 Equivariance (should always be True, unless you want to see how everything breaks visually otherwise)\n",
    "equivariance = True\n",
    "# Create video (slightly slower training)\n",
    "video = False\n",
    "video_path = export_folder\n",
    "\n",
    "# DEEP LEARNING RELATED\n",
    "# Device to train on (inference is done on cpu)\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "# Use two GPUs?\n",
    "device1 = device2 = device # 1 gpu for 2 modalities\n",
    "#device1, device2 = \"cuda:0\", \"cuda:1\" # 1 gpu per modality\n",
    "# Arguments for the tiramisu neural network\n",
    "tiramisu_args = {\n",
    "    # Number of convolutional filters for the first convolution\n",
    "    \"init_conv_filters\": 16,\n",
    "    # Number and depth of down blocks\n",
    "    \"down_blocks\": (2, 2, 2),\n",
    "    # Number and depth of up blocks\n",
    "    \"up_blocks\": (2, 2, 2),\n",
    "    # Number of dense layers in the bottleneck\n",
    "    \"bottleneck_layers\": 3,\n",
    "    # Upsampling type of layer (upsample has no grid artefacts)\n",
    "    \"upsampling_type\": \"upsample\",\n",
    "    # Type of max pooling, blurpool has better shift-invariance\n",
    "    \"transition_pooling\": \"max\",\n",
    "    # Dropout rate for the convolution\n",
    "    \"dropout_rate\": 0.0,\n",
    "    # Early maxpooling to reduce the input size\n",
    "    \"early_transition\": False,\n",
    "    # Activation function at the last layer\n",
    "    \"activation_func\": None,\n",
    "    # How much the conv layers should be compressed? (Memory saving)\n",
    "    \"compression\": 1.0,\n",
    "    # Memory efficient version of the tiramisu network (trades memory for computes)\n",
    "    # Gains of memory are enormous compared to the speed decrease.\n",
    "    # See: https://arxiv.org/pdf/1707.06990.pdf\n",
    "    \"efficient\": False,\n",
    "}\n",
    "# Epochs\n",
    "epochs = 5\n",
    "# Batch size\n",
    "batch_size = 24\n",
    "# Steps per epoch\n",
    "steps_per_epoch = 32\n",
    "# Number of steps\n",
    "steps = steps_per_epoch * epochs\n",
    "# How many unique patches are fed duriing one epoch\n",
    "samples_per_epoch = steps_per_epoch * batch_size\n",
    "# Optimiser\n",
    "#from lars.lars import LARS\n",
    "#optimiser = LARS\n",
    "optimiser = optim.Adam\n",
    "# Optimizer arguments\n",
    "opt_args = {\n",
    "    \"lr\": 1e-3,\n",
    "    \"weight_decay\": 1e-4,\n",
    "    #\"momentum\": 0.9\n",
    "}\n",
    "# Gradient norm. (limit on how big the gradients can get)\n",
    "grad_norm = 1.0\n",
    "\n",
    "# DATASET RELATED\n",
    "def worker_init_fn(worker_id):\n",
    "    base_seed = int(torch.randint(2**32, (1,)).item())\n",
    "    lib_seed = (base_seed + worker_id) % (2**32)\n",
    "    imgaug.seed(lib_seed)\n",
    "    np.random.seed(lib_seed)\n",
    "\n",
    "dataloader_args = {\n",
    "    \"batch_size\": batch_size,\n",
    "    \"shuffle\": False,\n",
    "    \"num_workers\": 16,\n",
    "    \"pin_memory\": True,\n",
    "    \"worker_init_fn\": worker_init_fn,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-29T21:23:55.735252Z",
     "start_time": "2020-03-29T21:23:54.921632Z"
    }
   },
   "outputs": [],
   "source": [
    "from models.tiramisu import DenseUNet\n",
    "# Taus being tested\n",
    "taus = [1e-3, 2e-3, 5e-3, 1e-2, 2e-2, 5e-2, 1e-1, 2e-1, 5e-1, 1, 2, 5, 10, 20, 50]\n",
    "for tau in taus:\n",
    "    print(\"TAU:\", tau)\n",
    "    class ModNet(DenseUNet):\n",
    "        def __init__(self, **args):\n",
    "            super(ModNet, self).__init__(**args, include_top=False)\n",
    "            out_channels = self.get_channels_count()[-1]\n",
    "            self.final_conv = torch.nn.Conv2d(out_channels, latent_channels, 1, bias=False)\n",
    "\n",
    "        def forward(self, x):\n",
    "            # Penultimate layer\n",
    "            L_hat = super(ModNet, self).forward(x)\n",
    "            # Final convolution\n",
    "            return self.final_conv(L_hat)\n",
    "\n",
    "    torch.manual_seed(0)\n",
    "    modelA = ModNet(in_channels=modA_len, nb_classes=latent_channels, **tiramisu_args).to(device1)\n",
    "    modelB = ModNet(in_channels=modB_len, nb_classes=latent_channels, **tiramisu_args).to(device2)\n",
    "\n",
    "    optimizerA = optimiser(modelA.parameters(), **opt_args)\n",
    "    optimizerB = optimiser(modelB.parameters(), **opt_args)\n",
    "\n",
    "    modelA = modelA.to(device1)\n",
    "    modelB = modelB.to(device2)\n",
    "    torch.manual_seed(0)\n",
    "\n",
    "    losses = {\"train\": [], \"test\": []}\n",
    "    epoch = 0\n",
    "    #test() # First round of testing\n",
    "    for epoch in range(1, epochs+1):\n",
    "        modelA.train()\n",
    "        modelB.train()\n",
    "        for batch_idx, data in enumerate(train_loader):\n",
    "            # Preparing the batch\n",
    "            data = data.permute(0, 3, 1, 2)\n",
    "            dataA = data[:, modA].float().to(device1)\n",
    "            dataB = data[:, modB].float().to(device2)\n",
    "            # Reseting the optimizer (gradients set to zero)\n",
    "            optimizerA.zero_grad()\n",
    "            optimizerB.zero_grad()\n",
    "\n",
    "            if equivariance:\n",
    "                # Applies random 90 degrees rotations to the data (group p4)\n",
    "                # This step enforces the formula of equivariance: d(f(T(x)), T^{-1}(f(x)))\n",
    "                # With f(x) the neural network, T(x) a transformation, T^{-1}(x) the inverse transformation\n",
    "                random_rotA = np.random.randint(4, size=batch_size)\n",
    "                random_rotB = np.random.randint(4, size=batch_size)\n",
    "                dataA_p4 = batch_rotate_p4(dataA, random_rotA, device1)\n",
    "                dataB_p4 = batch_rotate_p4(dataB, random_rotB, device2)\n",
    "\n",
    "                # Compute the forward pass\n",
    "                L1 = modelA(dataA_p4)\n",
    "                L2 = modelB(dataB_p4)\n",
    "\n",
    "                # Applies the inverse of the 90 degree rotation to recover the right positions\n",
    "                L1_ungrouped = batch_rotate_p4(L1, -random_rotA, device1)\n",
    "                L2_ungrouped = batch_rotate_p4(L2, -random_rotB, device2)\n",
    "            else:\n",
    "                L1_ungrouped = modelA(dataA)\n",
    "                L2_ungrouped = modelB(dataB)\n",
    "\n",
    "            # Computes the loss\n",
    "            loss = compute_pairwise_loss(\n",
    "                [L1_ungrouped, L2_ungrouped],\n",
    "                similarity_fn=sim_func,\n",
    "                tau=tau,\n",
    "                device=device1\n",
    "            )[0].mean()\n",
    "\n",
    "            # Activation regularization\n",
    "            if act_l1 > 0.:\n",
    "                loss += act_l1 * activation_decay([L1, L2], p=1., device=device1)\n",
    "            if act_l2 > 0.:\n",
    "                loss += act_l2 * activation_decay([L1, L2], p=2., device=device1)\n",
    "\n",
    "            # Computing the gradients\n",
    "            loss.backward()\n",
    "\n",
    "            # Clipping the the gradients if they are too big\n",
    "            torch.nn.utils.clip_grad_norm_(modelA.parameters(), grad_norm)\n",
    "            torch.nn.utils.clip_grad_norm_(modelB.parameters(), grad_norm)\n",
    "            # Performing the gradient descent\n",
    "            optimizerA.step()\n",
    "            optimizerB.step()\n",
    "\n",
    "            losses[\"train\"].append(loss.item())\n",
    "            clear_output(wait=True)\n",
    "            plot_losses(losses, steps, steps_per_epoch)\n",
    "            plt.show()\n",
    "\n",
    "        # Testing after each epoch\n",
    "        #_, similarities = test()\n",
    "        #clear_output(wait=True)\n",
    "        #plot_losses(losses, steps, steps_per_epoch)\n",
    "        #plt.show()\n",
    "\n",
    "    modelA.eval();\n",
    "    modelB.eval();\n",
    "\n",
    "    date = datetime.now().strftime(\"%Y%d%m_%H%M%S\")\n",
    "    model_path = os.path.join(export_folder, f\"model_zurich_{simfuncname}_tau_{tau:06.3f}.pt\")\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.simplefilter(\"ignore\")\n",
    "        torch.save({\n",
    "            \"modelA\": modelA,\n",
    "            \"modelB\": modelB,\n",
    "        }, model_path)\n",
    "    print(f\"model saved as: {model_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-29T21:24:07.219627Z",
     "start_time": "2020-03-29T21:24:01.836923Z"
    }
   },
   "outputs": [],
   "source": [
    "simfuncname = \"mse\"\n",
    "import matplotlib as mpl\n",
    "mpl.rc(\"font\", size=40)\n",
    "taus = [1e-3, 2e-3, 5e-3, 1e-2, 2e-2, 5e-2, 1e-1, 2e-1, 5e-1, 1, 2, 5, 10, 20, 50]\n",
    "fig, ax = plt.subplots(2, len(taus), figsize=(len(taus)*4, 8), sharex=True, sharey=True)\n",
    "L1s, L2s = [], []\n",
    "for i, tau in tqdm(enumerate(taus)):\n",
    "    model = torch.load(f\"models/model_zurich_{simfuncname}_tau_{tau:06.3f}.pt\")\n",
    "    modelA = model[\"modelA\"].cpu()\n",
    "    modelA.eval()\n",
    "    modelB = model[\"modelB\"].cpu()\n",
    "    modelB.eval();\n",
    "\n",
    "    \n",
    "    patchA = get_patch(img[..., modA], 610, 150, 0)\n",
    "    patchB = get_patch(img[..., modB], 610, 150, 0)\n",
    "    dataA = torch.tensor(patchA[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "    dataB = torch.tensor(patchB[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "\n",
    "    L1 = modelA(dataA).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L2 = modelB(dataB).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L1s.append(L1)\n",
    "    L2s.append(L2)\n",
    "\n",
    "    L1n, L2n = normalize_common(L1, L2)\n",
    "    ax[0, i].set_title(f\"$\\\\tau={tau}$\")\n",
    "    ax[0, i].imshow(L1n)\n",
    "    ax[1, i].imshow(L2n)\n",
    "L1s = np.array(L1s)\n",
    "L2s = np.array(L2s)\n",
    "\n",
    "plt.subplots_adjust(0.0, 0.0, 0.9, 0.9, 0.05, 0.05)\n",
    "plt.savefig(f\"results/temperature_{simfuncname}.pdf\", bbox_inches=\"tight\", pad_inches=0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generating the correlation matrix\n",
    "C = []\n",
    "for i in range(15):\n",
    "    C.append(np.corrcoef(L1s.reshape(15, -1)[i], L2s.reshape(15, -1)[i])[0, 1])\n",
    "\n",
    "# Bootstrapping\n",
    "mean = np.mean(C)\n",
    "means = []\n",
    "for i in range(100000): # Adjust to your needs\n",
    "    Cstar = np.random.choice(C, size=15)\n",
    "    means.append(np.mean(Cstar))\n",
    "\n",
    "lo, hi = np.percentile(means, [2.5, 97.5])\n",
    "lo, hi = 2 * mean - np.array([hi, lo])\n",
    "\n",
    "print(f\"Bootstrapped mean pairwise correlation: {mean:.1%} [{lo:.1%}; {hi:.1%}]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Testing the persistence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !!! IMPORTANT !!!\n",
    "# This needs to be re-run every time to test the two differrent settings:\n",
    "same_seed = True # All models have the same seed if True\n",
    "same_train_image = True # All moodels have the same training image if True\n",
    "# Also the name of the saved models needs to be changed.\n",
    "\n",
    "###### METHOD RELATED\n",
    "# The place where the models will be saved\n",
    "export_folder = \"results\" # Add this path to the .gitignore\n",
    "# The number of channels in the latent space (best is 1 or 3 for visualization)\n",
    "latent_channels = 3\n",
    "# Modality slicing\n",
    "# You can choose a set of channels per modality (RGB for instance)\n",
    "# Modality A\n",
    "modA = slice(0, 3)\n",
    "modA_name = \"RGB\"\n",
    "modA_len = modA.stop - modA.start\n",
    "# Modality B\n",
    "modB = slice(3, 4)\n",
    "modB_name = \"NIF\"\n",
    "modB_len = modB.stop - modB.start\n",
    "# Distance function\n",
    "simfunctions = {\n",
    "    \"euclidean\" : lambda x, y: -torch.norm(x - y, p=2, dim=1).mean(),\n",
    "    \"L1\"        : lambda x, y: -torch.norm(x - y, p=1, dim=1).mean(),\n",
    "    \"huber\"     : lambda x, y: -torch.sqrt(1 + (x - y).pow(2)).mean(),\n",
    "    \"MSE\"       : lambda x, y: -(x - y).pow(2).mean(),\n",
    "    \"L3\"        : lambda x, y: -torch.norm(x - y, p=3, dim=1).mean(),\n",
    "    \"Linf\"      : lambda x, y: -torch.norm(x - y, p=float(\"inf\"), dim=1).mean(),\n",
    "    \"soft_corr\" : lambda x, y: F.softplus(x*y).sum(axis=1),\n",
    "    \"corr\"      : lambda x, y: (x*y).sum(axis=1),\n",
    "    \"cosine\"    : lambda x, y: F.cosine_similarity(x, y, dim=1, eps=1e-8).mean(),\n",
    "    \"angular\"   : lambda x, y: F.cosine_similarity(x, y, dim=1, eps=1e-8).acos().mean() / math.pi,\n",
    "}\n",
    "sim_func = simfunctions[\"MSE\"]\n",
    "# Temperature (tau) of the loss\n",
    "tau =  .1\n",
    "# L1/L2 activation regularization\n",
    "act_l1 = 0\n",
    "act_l2 = 0\n",
    "# p4 Equivariance (should always be True, unless you want to see how everything breaks visually otherwise)\n",
    "equivariance = True\n",
    "# Create video (slightly slower training)\n",
    "video = False\n",
    "video_path = export_folder\n",
    "\n",
    "# DEEP LEARNING RELATED\n",
    "# Device to train on (inference is done on cpu)\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "# Use two GPUs?\n",
    "device1 = device2 = device # 1 gpu for 2 modalities\n",
    "#device1, device2 = \"cuda:0\", \"cuda:1\" # 1 gpu per modality\n",
    "# Arguments for the tiramisu neural network\n",
    "tiramisu_args = {\n",
    "    # Number of convolutional filters for the first convolution\n",
    "    \"init_conv_filters\": 16,\n",
    "    # Number and depth of down blocks\n",
    "    \"down_blocks\": (2, 2, 2),\n",
    "    # Number and depth of up blocks\n",
    "    \"up_blocks\": (2, 2, 2),\n",
    "    # Number of dense layers in the bottleneck\n",
    "    \"bottleneck_layers\": 3,\n",
    "    # Upsampling type of layer (upsample has no grid artefacts)\n",
    "    \"upsampling_type\": \"upsample\",\n",
    "    # Type of max pooling, blurpool has better shift-invariance\n",
    "    \"transition_pooling\": \"max\",\n",
    "    # Dropout rate for the convolution\n",
    "    \"dropout_rate\": 0.0,\n",
    "    # Early maxpooling to reduce the input size\n",
    "    \"early_transition\": False,\n",
    "    # Activation function at the last layer\n",
    "    \"activation_func\": None,\n",
    "    # How much the conv layers should be compressed? (Memory saving)\n",
    "    \"compression\": 1.0,\n",
    "    # Memory efficient version of the tiramisu network (trades memory for computes)\n",
    "    # Gains of memory are enormous compared to the speed decrease.\n",
    "    # See: https://arxiv.org/pdf/1707.06990.pdf\n",
    "    \"efficient\": False,\n",
    "}\n",
    "# Epochs\n",
    "epochs = 5\n",
    "# Batch size\n",
    "batch_size = 24\n",
    "# Steps per epoch\n",
    "steps_per_epoch = 32\n",
    "# Number of steps\n",
    "steps = steps_per_epoch * epochs\n",
    "# How many unique patches are fed duriing one epoch\n",
    "samples_per_epoch = steps_per_epoch * batch_size\n",
    "# Optimiser\n",
    "#from lars.lars import LARS\n",
    "#optimiser = LARS\n",
    "optimiser = optim.Adam\n",
    "# Optimizer arguments\n",
    "opt_args = {\n",
    "    \"lr\": 1e-3,\n",
    "    \"weight_decay\": 1e-4,\n",
    "    #\"momentum\": 0.9\n",
    "}\n",
    "# Gradient norm. (limit on how big the gradients can get)\n",
    "grad_norm = 1.0\n",
    "\n",
    "# DATASET RELATED\n",
    "def worker_init_fn(worker_id):\n",
    "    base_seed = int(torch.randint(2**32, (1,)).item())\n",
    "    lib_seed = (base_seed + worker_id) % (2**32)\n",
    "    imgaug.seed(lib_seed)\n",
    "    np.random.seed(lib_seed)\n",
    "\n",
    "dataloader_args = {\n",
    "    \"batch_size\": batch_size,\n",
    "    \"shuffle\": False,\n",
    "    \"num_workers\": 16,\n",
    "    \"pin_memory\": True,\n",
    "    \"worker_init_fn\": worker_init_fn,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-03-29T21:24:12.377425Z",
     "start_time": "2020-03-29T21:24:11.918571Z"
    }
   },
   "outputs": [],
   "source": [
    "from models.tiramisu import DenseUNet\n",
    "for mdl in range(20):\n",
    "    print(\"iteration:\", mdl)\n",
    "    if same_train_image:\n",
    "        dset = ZurichDataset(f\"../zurich/images/zh10.tif\", transform=ImgAugTransform())\n",
    "    else:\n",
    "        dset = ZurichDataset(f\"../zurich/images/zh{mdl}.tif\", transform=ImgAugTransform())\n",
    "    train_loader = torch.utils.data.DataLoader(\n",
    "        dset,\n",
    "        sampler=OverSampler(dset, samples_per_epoch),\n",
    "        **dataloader_args\n",
    "    )\n",
    "\n",
    "    class ModNet(DenseUNet):\n",
    "        def __init__(self, **args):\n",
    "            super(ModNet, self).__init__(**args, include_top=False)\n",
    "            out_channels = self.get_channels_count()[-1]\n",
    "            self.final_conv = torch.nn.Conv2d(out_channels, latent_channels, 1, bias=False)\n",
    "\n",
    "        def forward(self, x):\n",
    "            # Penultimate layer\n",
    "            L_hat = super(ModNet, self).forward(x)\n",
    "            # Final convolution\n",
    "            return self.final_conv(L_hat)\n",
    "\n",
    "    if same_seed:\n",
    "        torch.manual_seed(0)\n",
    "    else:\n",
    "        torch.manual_seed(mdl)\n",
    "    modelA = ModNet(in_channels=modA_len, nb_classes=latent_channels, **tiramisu_args).to(device1)\n",
    "    modelB = ModNet(in_channels=modB_len, nb_classes=latent_channels, **tiramisu_args).to(device2)\n",
    "\n",
    "    optimizerA = optimiser(modelA.parameters(), **opt_args)\n",
    "    optimizerB = optimiser(modelB.parameters(), **opt_args)\n",
    "\n",
    "    modelA = modelA.to(device1)\n",
    "    modelB = modelB.to(device2)\n",
    "    torch.manual_seed(0)\n",
    "\n",
    "    losses = {\"train\": [], \"test\": []}\n",
    "    epoch = 0\n",
    "    #test() # First round of testing\n",
    "    for epoch in range(1, epochs+1):\n",
    "        modelA.train()\n",
    "        modelB.train()\n",
    "        for batch_idx, data in enumerate(train_loader):\n",
    "            # Preparing the batch\n",
    "            data = data.permute(0, 3, 1, 2)\n",
    "            dataA = data[:, modA].float().to(device1)\n",
    "            dataB = data[:, modB].float().to(device2)\n",
    "            # Reseting the optimizer (gradients set to zero)\n",
    "            optimizerA.zero_grad()\n",
    "            optimizerB.zero_grad()\n",
    "\n",
    "            if equivariance:\n",
    "                # Applies random 90 degrees rotations to the data (group p4)\n",
    "                # This step enforces the formula of equivariance: d(f(T(x)), T^{-1}(f(x)))\n",
    "                # With f(x) the neural network, T(x) a transformation, T^{-1}(x) the inverse transformation\n",
    "                random_rotA = np.random.randint(4, size=batch_size)\n",
    "                random_rotB = np.random.randint(4, size=batch_size)\n",
    "                dataA_p4 = batch_rotate_p4(dataA, random_rotA, device1)\n",
    "                dataB_p4 = batch_rotate_p4(dataB, random_rotB, device2)\n",
    "\n",
    "                # Compute the forward pass\n",
    "                L1 = modelA(dataA_p4)\n",
    "                L2 = modelB(dataB_p4)\n",
    "\n",
    "                # Applies the inverse of the 90 degree rotation to recover the right positions\n",
    "                L1_ungrouped = batch_rotate_p4(L1, -random_rotA, device1)\n",
    "                L2_ungrouped = batch_rotate_p4(L2, -random_rotB, device2)\n",
    "            else:\n",
    "                L1_ungrouped = modelA(dataA)\n",
    "                L2_ungrouped = modelB(dataB)\n",
    "\n",
    "            # Computes the loss\n",
    "            loss = compute_pairwise_loss(\n",
    "                [L1_ungrouped, L2_ungrouped],\n",
    "                similarity_fn=sim_func,\n",
    "                tau=tau,\n",
    "                device=device1\n",
    "            )[0].mean()\n",
    "\n",
    "            # Activation regularization\n",
    "            if act_l1 > 0.:\n",
    "                loss += act_l1 * activation_decay([L1, L2], p=1., device=device1)\n",
    "            if act_l2 > 0.:\n",
    "                loss += act_l2 * activation_decay([L1, L2], p=2., device=device1)\n",
    "\n",
    "            # Computing the gradients\n",
    "            loss.backward()\n",
    "\n",
    "            # Clipping the the gradients if they are too big\n",
    "            torch.nn.utils.clip_grad_norm_(modelA.parameters(), grad_norm)\n",
    "            torch.nn.utils.clip_grad_norm_(modelB.parameters(), grad_norm)\n",
    "            # Performing the gradient descent\n",
    "            optimizerA.step()\n",
    "            optimizerB.step()\n",
    "\n",
    "            #losses[\"train\"].append(loss.item())\n",
    "            #clear_output(wait=True)\n",
    "            #plot_losses(losses, steps, steps_per_epoch)\n",
    "            #plt.show()\n",
    "\n",
    "        # Testing after each epoch\n",
    "        #_, similarities = test()\n",
    "        #clear_output(wait=True)\n",
    "        #plot_losses(losses, steps, steps_per_epoch)\n",
    "        #plt.show()\n",
    "\n",
    "    modelA.eval();\n",
    "    modelB.eval();\n",
    "\n",
    "    date = datetime.now().strftime(\"%Y%d%m_%H%M%S\")\n",
    "    model_path = os.path.join(export_folder, f\"model_zurich_persist_{mdl:02d}.pt\")\n",
    "    with warnings.catch_warnings():\n",
    "        warnings.simplefilter(\"ignore\")\n",
    "        torch.save({\n",
    "            \"modelA\": modelA,\n",
    "            \"modelB\": modelB,\n",
    "        }, model_path)\n",
    "    print(f\"model saved as: {model_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L1s_fw_fi, L2s_fw_fi = [], [] # Fixed Weight (fw), Fixed Image (fi)\n",
    "L1s_dw_fi, L2s_dw_fi = [], [] # Different Weights (dw), Fixed Image (fi)\n",
    "L1s_fw_di, L2s_fw_di = [], [] # Fixed Weight (fw), Different train Image (di)\n",
    "L1s_dw_di, L2s_dw_di = [], [] # Different Weights (dw), Different train Image (di)\n",
    "\n",
    "for i in tqdm(range(50)):\n",
    "    model = torch.load(f\"models/model_zurich_persist_fw_fi_{i:02d}.pt\") # Fixed weights, 1 train image\n",
    "    modelA = model[\"modelA\"].cpu()\n",
    "    modelA.eval()\n",
    "    modelB = model[\"modelB\"].cpu()\n",
    "    modelB.eval();\n",
    "    \n",
    "    patchA = get_patch(img[..., modA], 610, 150, 0)\n",
    "    patchB = get_patch(img[..., modB], 610, 150, 0)\n",
    "    dataA = torch.tensor(patchA[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "    dataB = torch.tensor(patchB[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "\n",
    "    L1 = modelA(dataA).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L2 = modelB(dataB).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L1s_fw_fi.append(L1)\n",
    "    L2s_fw_fi.append(L2)\n",
    "L1s_fw_fi = np.array(L1s_fw_fi)\n",
    "L2s_fw_fi = np.array(L2s_fw_fi)\n",
    "\n",
    "for i in tqdm(range(20)):\n",
    "    model = torch.load(f\"models/model_zurich_persist_dw_fi_{i:02d}.pt\")  # Rand weights, 1 train image\n",
    "    modelA = model[\"modelA\"].cpu()\n",
    "    modelA.eval()\n",
    "    modelB = model[\"modelB\"].cpu()\n",
    "    modelB.eval();\n",
    "    \n",
    "    patchA = get_patch(img[..., modA], 610, 150, 0)\n",
    "    patchB = get_patch(img[..., modB], 610, 150, 0)\n",
    "    dataA = torch.tensor(patchA[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "    dataB = torch.tensor(patchB[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "\n",
    "    L1 = modelA(dataA).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L2 = modelB(dataB).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L1s_dw_fi.append(L1)\n",
    "    L2s_dw_fi.append(L2)\n",
    "L1s_dw_fi = np.array(L1s_dw_fi)\n",
    "L2s_dw_fi = np.array(L2s_dw_fi)\n",
    "\n",
    "for i in tqdm(range(20)):\n",
    "    model = torch.load(f\"models/model_zurich_persist_fw_di_{i:02d}.pt\") # Fixed weights, 1 image per nn\n",
    "    modelA = model[\"modelA\"].cpu()\n",
    "    modelA.eval()\n",
    "    modelB = model[\"modelB\"].cpu()\n",
    "    modelB.eval();\n",
    "    \n",
    "    patchA = get_patch(img[..., modA], 610, 150, 0)\n",
    "    patchB = get_patch(img[..., modB], 610, 150, 0)\n",
    "    dataA = torch.tensor(patchA[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "    dataB = torch.tensor(patchB[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "\n",
    "    L1 = modelA(dataA).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L2 = modelB(dataB).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L1s_fw_di.append(L1)\n",
    "    L2s_fw_di.append(L2)\n",
    "L1s_fw_di = np.array(L1s_fw_di)\n",
    "L2s_fw_di = np.array(L2s_fw_di)\n",
    "\n",
    "for i in tqdm(range(20)):\n",
    "    model = torch.load(f\"models/model_zurich_persist_dw_di_{i:02d}.pt\") # Rand weights, 1 image per nn\n",
    "    modelA = model[\"modelA\"].cpu()\n",
    "    modelA.eval()\n",
    "    modelB = model[\"modelB\"].cpu()\n",
    "    modelB.eval();\n",
    "    \n",
    "    patchA = get_patch(img[..., modA], 610, 150, 0)\n",
    "    patchB = get_patch(img[..., modB], 610, 150, 0)\n",
    "    dataA = torch.tensor(patchA[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "    dataB = torch.tensor(patchB[np.newaxis].transpose(0, 3, 1, 2), device=\"cpu\")\n",
    "\n",
    "    L1 = modelA(dataA).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L2 = modelB(dataB).detach().numpy()[0].transpose(1, 2, 0)\n",
    "    L1s_dw_di.append(L1)\n",
    "    L2s_dw_di.append(L2)\n",
    "L1s_dw_di = np.array(L1s_dw_di)\n",
    "L2s_dw_di = np.array(L2s_dw_di)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 0\n",
    "np.random.seed(0)\n",
    "for Rs in [\n",
    "    (L1s_fw_fi, L2s_fw_fi, \"fw_fi\"),\n",
    "    (L1s_dw_fi, L2s_dw_fi, \"dw_fi\"),\n",
    "    (L1s_fw_di, L2s_fw_di, \"fw_di\"),\n",
    "    (L1s_dw_di, L2s_dw_di, \"dw_di\")]:\n",
    "    name = Rs[2]\n",
    "    rhosm1 = []\n",
    "    rhosm2 = []\n",
    "    for j in range(1000):\n",
    "        R1 = Rs[0].reshape(Rs[0].shape[0], -1)\n",
    "        R2 = Rs[1].reshape(Rs[1].shape[0], -1)\n",
    "        if i > 1:\n",
    "            R1 = np.delete(R1, 11, 0)\n",
    "            R2 = np.delete(R2, 11, 0)\n",
    "        N = R1.shape[0]\n",
    "        idx = np.random.choice(R1.shape[0], size=N)\n",
    "        R1 = R1[idx]\n",
    "        idx = np.random.choice(R1.shape[0], size=N)\n",
    "        R2 = R2[idx]\n",
    "        rhos1 = np.tril(np.corrcoef(R1), -1).ravel()\n",
    "        rhos1 = rhos1[rhos1 != 0.0]\n",
    "        rhos2 = np.tril(np.corrcoef(R2), -1).ravel()\n",
    "        rhos2 = rhos2[rhos2 != 0.0]\n",
    "        rhosm1.append(np.mean(rhos1))\n",
    "        rhosm1.append(np.mean(rhos2))\n",
    "    \n",
    "    R1 = Rs[0].reshape(Rs[0].shape[0], -1)\n",
    "    R2 = Rs[1].reshape(Rs[1].shape[0], -1)\n",
    "    if i > 1:\n",
    "        R1 = np.delete(R1, 11, 0)\n",
    "        R2 = np.delete(R2, 11, 0)\n",
    "    N = R1.shape[0]\n",
    "    rhos1 = np.tril(np.corrcoef(R1), -1).ravel()\n",
    "    rhos1 = rhos1[rhos1 != 0.0]\n",
    "    rhos2 = np.tril(np.corrcoef(R2), -1).ravel()\n",
    "    rhos2 = rhos2[rhos2 != 0.0]\n",
    "    rho = .5 * (np.mean(rhos1) + np.mean(rhos2))\n",
    "    rhoci = 2 * rho - np.array(np.percentile(rhosm1, [97.5, 2.5]))\n",
    "    print(f\"{name}: {rho: 4.1%} (95% C.I.: [{rhoci[0]: 4.1%}; {rhoci[1]: 4.1%}], N={2*N})\")\n",
    "    i += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mpl.rc(\"font\", size=20)\n",
    "j = 0\n",
    "for Rs in [\n",
    "    (L1s_fw_fi, L2s_fw_fi, \"fw_fi\"),\n",
    "    (L1s_dw_fi, L2s_dw_fi, \"dw_fi\"),\n",
    "    (L1s_fw_di, L2s_fw_di, \"fw_di\"),\n",
    "    (L1s_dw_di, L2s_dw_di, \"dw_di\")]:\n",
    "    R1, R2, tp = Rs\n",
    "    N = R1.shape[0]\n",
    "    rows = int(np.ceil(N / 10))\n",
    "    fig, axs = plt.subplots(rows, 10, figsize=(20, 2*rows), sharex=True, sharey=True)\n",
    "    axs = axs.ravel()\n",
    "    for i in range(len(axs)):\n",
    "        if j > 1 and i == 11:\n",
    "            fig.delaxes(axs[i])\n",
    "            continue\n",
    "        \n",
    "        L1n, L2n = normalize_common(R1[i], R2[i])\n",
    "        axs[i].imshow(L1n)\n",
    "    plt.savefig(f\"results/embeddings_{tp}.pdf\", bbox_inches='tight', pad_inches=0.10)\n",
    "    j += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  },
  "notify_time": "10"
 },
 "nbformat": 4,
 "nbformat_minor": 4
}