{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from denoising_diffusion_pytorch import GaussianDiffusion1D\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "\n",
    "import sys\n",
    "sys.path.append('../')\n",
    "\n",
    "from data import load_dataset_from_str \n",
    "from data_generation import extract_signals, region_occupancy\n",
    "from model import TransformerDenoiser\n",
    "from evaluation import sample_heatmap, generate_heatmap, undo_rotation_and_scaling\n",
    "from dataset import Porto25, GeoLife25\n",
    "from evaluation import unscaled_occupancy, sample_synthetic_region\n",
    "from tqdm import tqdm\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_str = 'geolife-25'\n",
    "target_dataset_str = 'geolife-25'\n",
    "target_dataset_size = 10000\n",
    "\n",
    "if target_dataset_str == 'porto-25':\n",
    "    original_region = (41.1, 41.15, -8.65, -8.60)\n",
    "elif target_dataset_str == 'porto-100':\n",
    "    original_region = (41.1, 41.2, -8.65, -8.55)\n",
    "elif target_dataset_str == 'geolife-100':\n",
    "    original_region = (39.9, 40, 116.3, 116.4)\n",
    "elif target_dataset_str == 'geolife-25':\n",
    "    original_region = (39.9, 39.95, 116.3, 116.35)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 8\n",
    "num_workers = 10\n",
    "n_dim = 2\n",
    "beta_schedule = 'cosine'\n",
    "adam_betas = (0.9, 0.99)\n",
    "seq_len = 128  # For now the sequence length is baked into the data extraction\n",
    "\n",
    "hidden_dim = 64\n",
    "timesteps = 100\n",
    "objective = 'pred_noise'\n",
    "n_layers = 2\n",
    "n_heads = 4\n",
    "trial_number = 90\n",
    "PROTOTYPE_ID = 24\n",
    "\n",
    "checkpoint_path = f'model_weights/conditional-{dataset_str}.pth'\n",
    "\n",
    "model = TransformerDenoiser(n_dim, hidden_dim, seq_len=seq_len, n_heads=n_heads, n_layers=n_layers)\n",
    "\n",
    "diffusion = GaussianDiffusion1D(\n",
    "    model,\n",
    "    seq_length = seq_len,\n",
    "    timesteps = timesteps,\n",
    "    objective = objective,\n",
    "    beta_schedule = beta_schedule,\n",
    ")\n",
    "\n",
    "checkpoint = torch.load(checkpoint_path)\n",
    "\n",
    "diffusion.load_state_dict(checkpoint['diffusion_state_dict'])\n",
    "diffusion.cuda()\n",
    "diffusion.eval()\n",
    "\n",
    "dataset = load_dataset_from_str(target_dataset_str, prefix='../data/')\n",
    "\n",
    "if dataset_str == 'porto-25':\n",
    "    train_set = Porto25()\n",
    "elif dataset_str == 'geolife-25':\n",
    "    train_set = GeoLife25()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_heatmap(diffusion, query_heatmap, n_samples, sample_batch_size=256):\n",
    "    \"\"\"\n",
    "    Given a diffusion model and a heatmap, generate n_samples conditioned on the heatmap.\n",
    "    \"\"\"\n",
    "    sample_batch_size = min(sample_batch_size, n_samples)\n",
    "\n",
    "    query_heatmap = torch.tensor(query_heatmap).unsqueeze(0)\n",
    "    all_samples = []\n",
    "\n",
    "    n_batches = n_samples // sample_batch_size\n",
    "    remainder = n_samples % sample_batch_size\n",
    "\n",
    "    def sample(n_samples):\n",
    "        condition = query_heatmap.unsqueeze(0).expand(n_samples, 1, 64, 64).float().cuda()\n",
    "        condition = (condition - train_set.heatmaps_mean) / train_set.heatmaps_std\n",
    "        # Sample `n_samples` number of samples from the model\n",
    "        samples = diffusion.sample_conditional(condition, n_samples)\n",
    "        # Convert to numpy\n",
    "        samples = samples.cpu().numpy()\n",
    "\n",
    "        return samples\n",
    "\n",
    "    for i in range(n_batches):\n",
    "        all_samples.append(sample(sample_batch_size))\n",
    "\n",
    "    if remainder > 0:\n",
    "        all_samples.append(sample(remainder))\n",
    "\n",
    "    all_samples = np.concatenate(all_samples, axis=0)\n",
    "\n",
    "    return all_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "min_lat, max_lat, min_lon, max_lon = original_region\n",
    "map_height = max_lat - min_lat\n",
    "map_width = max_lon - min_lon\n",
    "\n",
    "width = 1/5 * map_width\n",
    "height = 1/5 * map_height\n",
    "rotation = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "regions = []\n",
    "\n",
    "latitudes = np.arange(min_lat, max_lat, height / 2)\n",
    "longitudes = np.arange(min_lon, max_lon, width / 2)\n",
    "indices = []\n",
    "\n",
    "for i, region_lon in enumerate(longitudes):\n",
    "    for j, region_lat in enumerate(latitudes):\n",
    "        regions.append((region_lat, region_lon, rotation, height, width))\n",
    "        indices.append((i, j))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2736552/2308002180.py:3: RuntimeWarning: divide by zero encountered in log\n",
      "  plt.imshow(np.log(heatmap_data), extent=[yedges[0], yedges[-1], xedges[0], xedges[-1]], origin='lower', interpolation='nearest')\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 134400 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "heatmap_data, xedges, yedges = unscaled_occupancy(dataset, 512)\n",
    "\n",
    "plt.imshow(np.log(heatmap_data), extent=[yedges[0], yedges[-1], xedges[0], xedges[-1]], origin='lower', interpolation='nearest')\n",
    "\n",
    "for region_lat, region_lon, rotation, height, width in regions:\n",
    "    rect = patches.Rectangle((region_lon, region_lat), width, height, angle=rotation, linewidth=1, edgecolor='r', facecolor='none')\n",
    "    plt.gca().add_patch(rect)\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 25700 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "heatmap_data, xedges, yedges = unscaled_occupancy(dataset, 10)\n",
    "\n",
    "plt.imshow(heatmap_data, extent=[yedges[0], yedges[-1], xedges[0], xedges[-1]], origin='lower', interpolation=None)\n",
    "\n",
    "for region_lat, region_lon, rotation, height, width in regions:\n",
    "    rect = patches.Rectangle((region_lon, region_lat), width, height, angle=rotation, linewidth=1, edgecolor='r', facecolor='none')\n",
    "    plt.gca().add_patch(rect)\n",
    "\n",
    "plt.colorbar()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def occupancy(trajectories, res):\n",
    "    \"\"\"\n",
    "    Unscaled variant of data_generation.region_occupancy, mainly useful for visualization.\n",
    "    \"\"\"\n",
    "    all_concatenated = np.concatenate(trajectories)\n",
    "    H, xedges, yedges = np.histogram2d(all_concatenated[:, 0], all_concatenated[:, 1], bins=(res, res), density=False)\n",
    "\n",
    "    return H, xedges, yedges\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "heatmap_data, xedges, yedges = occupancy(dataset, 10)\n",
    "\n",
    "counts = []\n",
    "\n",
    "for (i, j), (region_lat, region_lon, rotation, height, width) in zip(indices, regions):\n",
    "    counts.append(int(heatmap_data[j:(j+2), i:(i+2)].sum()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "counts = np.array(counts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "scaled_counts = (target_dataset_size * counts / counts.sum()).astype('int')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0, 305)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "min(scaled_counts), max(scaled_counts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([101, 144, 170, 203, 188, 180, 223, 225, 120,  26,  28,  46,  68,\n",
       "        77,  33,  49, 122,  93,  59,  44,  35,  51,  59,  55,  25,  87,\n",
       "       193, 178, 146,  81,  37,  47,  32,  22,  15,  84, 237, 221,  98,\n",
       "        37,  37,  50,  32,  36,  54,  77, 196, 156,   6,   0,  45,  55,\n",
       "        36,  35,  73,  82, 175, 163,  23,   2, 131, 118, 126, 169, 305,\n",
       "       269, 225, 175,  31,   7, 126, 117, 134, 190, 295, 261, 193, 128,\n",
       "        22,   5,  60,  91, 103,  96,  73,  90, 139, 166,  90,   8,  46,\n",
       "        73,  83,  70,  59,  75, 104, 135,  83,   7])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scaled_counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_synthetic_region(diffusion, dataset, region, n_samples, sample_batch_size=64):\n",
    "    \"\"\"\n",
    "    Extract subtrajectories from region, creates a heatmap of the real data, creates synthetic\n",
    "    trajectories conditioned on the heatmap and finally calculate a heatmap of the synthetic samples.\n",
    "    \"\"\"\n",
    "    assert n_samples >= sample_batch_size\n",
    "\n",
    "    trajectories_in_region = extract_signals(dataset, region)  # TODO: Set minimum length to same as the model was trained on\n",
    "    heatmap, _, _ = region_occupancy(trajectories_in_region, res=64)\n",
    "\n",
    "    samples = sample_heatmap(diffusion, heatmap, n_samples=n_samples, sample_batch_size=sample_batch_size)\n",
    "    synthetic_heatmap = generate_heatmap(samples)\n",
    "\n",
    "    return trajectories_in_region, heatmap, samples, synthetic_heatmap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/100 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|          | 1/100 [00:02<04:52,  2.96s/it]\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[17], line 9\u001b[0m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m n_sample \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[1;32m      7\u001b[0m     \u001b[38;5;66;03m# print(n_sample)\u001b[39;00m\n\u001b[1;32m      8\u001b[0m     batch_size \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmin\u001b[39m(\u001b[38;5;241m256\u001b[39m, n_sample)\n\u001b[0;32m----> 9\u001b[0m     results\u001b[38;5;241m.\u001b[39mappend(\u001b[43msample_synthetic_region\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdiffusion\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdataset\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mregion\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mn_sample\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m     10\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m     11\u001b[0m     results\u001b[38;5;241m.\u001b[39mappend([])\n",
      "Cell \u001b[0;32mIn[16], line 11\u001b[0m, in \u001b[0;36msample_synthetic_region\u001b[0;34m(diffusion, dataset, region, n_samples, sample_batch_size)\u001b[0m\n\u001b[1;32m      8\u001b[0m trajectories_in_region \u001b[38;5;241m=\u001b[39m extract_signals(dataset, region)  \u001b[38;5;66;03m# TODO: Set minimum length to same as the model was trained on\u001b[39;00m\n\u001b[1;32m      9\u001b[0m heatmap, _, _ \u001b[38;5;241m=\u001b[39m region_occupancy(trajectories_in_region, res\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m64\u001b[39m)\n\u001b[0;32m---> 11\u001b[0m samples \u001b[38;5;241m=\u001b[39m \u001b[43msample_heatmap\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdiffusion\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mheatmap\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mn_samples\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mn_samples\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msample_batch_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msample_batch_size\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     12\u001b[0m synthetic_heatmap \u001b[38;5;241m=\u001b[39m generate_heatmap(samples)\n\u001b[1;32m     14\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m trajectories_in_region, heatmap, samples, synthetic_heatmap\n",
      "Cell \u001b[0;32mIn[5], line 24\u001b[0m, in \u001b[0;36msample_heatmap\u001b[0;34m(diffusion, query_heatmap, n_samples, sample_batch_size)\u001b[0m\n\u001b[1;32m     21\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m samples\n\u001b[1;32m     23\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(n_batches):\n\u001b[0;32m---> 24\u001b[0m     all_samples\u001b[38;5;241m.\u001b[39mappend(\u001b[43msample\u001b[49m\u001b[43m(\u001b[49m\u001b[43msample_batch_size\u001b[49m\u001b[43m)\u001b[49m)\n\u001b[1;32m     26\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m remainder \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[1;32m     27\u001b[0m     all_samples\u001b[38;5;241m.\u001b[39mappend(sample(remainder))\n",
      "Cell \u001b[0;32mIn[5], line 17\u001b[0m, in \u001b[0;36msample_heatmap.<locals>.sample\u001b[0;34m(n_samples)\u001b[0m\n\u001b[1;32m     15\u001b[0m condition \u001b[38;5;241m=\u001b[39m (condition \u001b[38;5;241m-\u001b[39m train_set\u001b[38;5;241m.\u001b[39mheatmaps_mean) \u001b[38;5;241m/\u001b[39m train_set\u001b[38;5;241m.\u001b[39mheatmaps_std\n\u001b[1;32m     16\u001b[0m \u001b[38;5;66;03m# Sample `n_samples` number of samples from the model\u001b[39;00m\n\u001b[0;32m---> 17\u001b[0m samples \u001b[38;5;241m=\u001b[39m \u001b[43mdiffusion\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msample_conditional\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcondition\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mn_samples\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;66;03m# Convert to numpy\u001b[39;00m\n\u001b[1;32m     19\u001b[0m samples \u001b[38;5;241m=\u001b[39m samples\u001b[38;5;241m.\u001b[39mcpu()\u001b[38;5;241m.\u001b[39mnumpy()\n",
      "File \u001b[0;32m~/.conda/envs/tddpm/lib/python3.10/site-packages/torch/autograd/grad_mode.py:27\u001b[0m, in \u001b[0;36m_DecoratorContextManager.__call__.<locals>.decorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     24\u001b[0m \u001b[38;5;129m@functools\u001b[39m\u001b[38;5;241m.\u001b[39mwraps(func)\n\u001b[1;32m     25\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdecorate_context\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m     26\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mclone():\n\u001b[0;32m---> 27\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/repos/tddpm/tddpm/denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py:661\u001b[0m, in \u001b[0;36mGaussianDiffusion1D.sample_conditional\u001b[0;34m(self, condition, batch_size)\u001b[0m\n\u001b[1;32m    658\u001b[0m \u001b[38;5;129m@torch\u001b[39m\u001b[38;5;241m.\u001b[39mno_grad()\n\u001b[1;32m    659\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21msample_conditional\u001b[39m(\u001b[38;5;28mself\u001b[39m, condition, batch_size \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m16\u001b[39m):\n\u001b[1;32m    660\u001b[0m     seq_length, channels \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mseq_length, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mchannels\n\u001b[0;32m--> 661\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mp_sample_loop_conditional\u001b[49m\u001b[43m(\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbatch_size\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mchannels\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mseq_length\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcondition\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/.conda/envs/tddpm/lib/python3.10/site-packages/torch/autograd/grad_mode.py:27\u001b[0m, in \u001b[0;36m_DecoratorContextManager.__call__.<locals>.decorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     24\u001b[0m \u001b[38;5;129m@functools\u001b[39m\u001b[38;5;241m.\u001b[39mwraps(func)\n\u001b[1;32m     25\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdecorate_context\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m     26\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mclone():\n\u001b[0;32m---> 27\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/repos/tddpm/tddpm/denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py:611\u001b[0m, in \u001b[0;36mGaussianDiffusion1D.p_sample_loop_conditional\u001b[0;34m(self, shape, condition)\u001b[0m\n\u001b[1;32m    609\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m t \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mreversed\u001b[39m(\u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m0\u001b[39m, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnum_timesteps)):\n\u001b[1;32m    610\u001b[0m     self_cond \u001b[38;5;241m=\u001b[39m x_start \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mself_condition \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 611\u001b[0m     img, x_start \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mp_sample\u001b[49m\u001b[43m(\u001b[49m\u001b[43mimg\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mt\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcondition\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mself_cond\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    613\u001b[0m img \u001b[38;5;241m=\u001b[39m unnormalize_to_zero_to_one(img)\n\u001b[1;32m    614\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m img\n",
      "File \u001b[0;32m~/.conda/envs/tddpm/lib/python3.10/site-packages/torch/autograd/grad_mode.py:27\u001b[0m, in \u001b[0;36m_DecoratorContextManager.__call__.<locals>.decorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     24\u001b[0m \u001b[38;5;129m@functools\u001b[39m\u001b[38;5;241m.\u001b[39mwraps(func)\n\u001b[1;32m     25\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdecorate_context\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m     26\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mclone():\n\u001b[0;32m---> 27\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfunc\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/repos/tddpm/tddpm/denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py:581\u001b[0m, in \u001b[0;36mGaussianDiffusion1D.p_sample\u001b[0;34m(self, x, t, condition, x_self_cond, clip_denoised)\u001b[0m\n\u001b[1;32m    578\u001b[0m batched_times \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mfull((x\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m0\u001b[39m],), t, device \u001b[38;5;241m=\u001b[39m x\u001b[38;5;241m.\u001b[39mdevice, dtype \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mlong)\n\u001b[1;32m    579\u001b[0m \u001b[38;5;66;03m# batched_condition = condition.expand(x.shape[0], -1).to(x.device)\u001b[39;00m\n\u001b[0;32m--> 581\u001b[0m model_mean, _, model_log_variance, x_start \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mp_mean_variance_condition\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mt\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mbatched_times\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcondition\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mcondition\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx_self_cond\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mx_self_cond\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mclip_denoised\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mclip_denoised\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    582\u001b[0m noise \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mrandn_like(x) \u001b[38;5;28;01mif\u001b[39;00m t \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;241m0.\u001b[39m \u001b[38;5;66;03m# no noise if t == 0\u001b[39;00m\n\u001b[1;32m    583\u001b[0m pred_img \u001b[38;5;241m=\u001b[39m model_mean \u001b[38;5;241m+\u001b[39m (\u001b[38;5;241m0.5\u001b[39m \u001b[38;5;241m*\u001b[39m model_log_variance)\u001b[38;5;241m.\u001b[39mexp() \u001b[38;5;241m*\u001b[39m noise\n",
      "File \u001b[0;32m~/repos/tddpm/tddpm/denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py:565\u001b[0m, in \u001b[0;36mGaussianDiffusion1D.p_mean_variance_condition\u001b[0;34m(self, x, t, condition, x_self_cond, clip_denoised)\u001b[0m\n\u001b[1;32m    564\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mp_mean_variance_condition\u001b[39m(\u001b[38;5;28mself\u001b[39m, x, t, condition, x_self_cond \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m, clip_denoised \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m):\n\u001b[0;32m--> 565\u001b[0m     preds \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmodel_predictions\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mt\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx_self_cond\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mconditional\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mcondition\u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# TODO: Consistent naming\u001b[39;00m\n\u001b[1;32m    566\u001b[0m     x_start \u001b[38;5;241m=\u001b[39m preds\u001b[38;5;241m.\u001b[39mpred_x_start\n\u001b[1;32m    568\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m clip_denoised:\n",
      "File \u001b[0;32m~/repos/tddpm/tddpm/denoising_diffusion_pytorch/denoising_diffusion_pytorch_1d.py:527\u001b[0m, in \u001b[0;36mGaussianDiffusion1D.model_predictions\u001b[0;34m(self, x, t, x_self_cond, clip_x_start, conditional)\u001b[0m\n\u001b[1;32m    526\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mmodel_predictions\u001b[39m(\u001b[38;5;28mself\u001b[39m, x, t, x_self_cond \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m, clip_x_start \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m, conditional \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[0;32m--> 527\u001b[0m     model_output \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mt\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mx_self_cond\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mconditional\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    528\u001b[0m     maybe_clip \u001b[38;5;241m=\u001b[39m partial(torch\u001b[38;5;241m.\u001b[39mclamp, \u001b[38;5;28mmin\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1.\u001b[39m, \u001b[38;5;28mmax\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1.\u001b[39m) \u001b[38;5;28;01mif\u001b[39;00m clip_x_start \u001b[38;5;28;01melse\u001b[39;00m identity\n\u001b[1;32m    530\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mobjective \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mpred_noise\u001b[39m\u001b[38;5;124m'\u001b[39m:\n",
      "File \u001b[0;32m~/.conda/envs/tddpm/lib/python3.10/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_forward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43minput\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1131\u001b[0m \u001b[38;5;66;03m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[38;5;241m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/repos/tddpm/tddpm/model.py:110\u001b[0m, in \u001b[0;36mTransformerDenoiser.forward\u001b[0;34m(self, noisy_traj, time, x_self_cond, conditional, predict_conditional)\u001b[0m\n\u001b[1;32m    102\u001b[0m     embs\u001b[38;5;241m.\u001b[39mappend(i_emb)\n\u001b[1;32m    104\u001b[0m \u001b[38;5;66;03m# ic(x_emb.shape, y_emb.shape, self.time_step_embedding.shape)\u001b[39;00m\n\u001b[1;32m    105\u001b[0m \n\u001b[1;32m    106\u001b[0m \u001b[38;5;66;03m# ic| x_emb.shape: torch.Size([8, 128, 32])\u001b[39;00m\n\u001b[1;32m    107\u001b[0m \u001b[38;5;66;03m#     y_emb.shape: torch.Size([8, 128, 32])\u001b[39;00m\n\u001b[1;32m    108\u001b[0m \u001b[38;5;66;03m#     self.time_step_embedding.shape: torch.Size([128, 32])\u001b[39;00m\n\u001b[0;32m--> 110\u001b[0m time_step_embedding \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtime_step_embedding\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43munsqueeze\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mexpand\u001b[49m\u001b[43m(\u001b[49m\u001b[43membs\u001b[49m\u001b[43m[\u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mshape\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mto\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnoisy_traj\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    111\u001b[0m noisy_trajectory_token \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mnoisy_trajectory_token\u001b[38;5;241m.\u001b[39mreshape(\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m1\u001b[39m, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdim_embedding)\u001b[38;5;241m.\u001b[39mexpand(embs[\u001b[38;5;241m0\u001b[39m]\u001b[38;5;241m.\u001b[39mshape)\n\u001b[1;32m    113\u001b[0m \u001b[38;5;66;03m# ic(x_emb.shape, y_emb.shape, time_step_embedding.shape, noisy_trajectory_token.shape,\u001b[39;00m\n\u001b[1;32m    114\u001b[0m \u001b[38;5;66;03m#    x_emb.device, y_emb.device, time_step_embedding.device, noisy_trajectory_token.device)\u001b[39;00m\n\u001b[1;32m    115\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    122\u001b[0m \u001b[38;5;66;03m#     time_step_embedding.device: device(type='cuda', index=0)\u001b[39;00m\n\u001b[1;32m    123\u001b[0m \u001b[38;5;66;03m#     noisy_trajectory_token.device: device(type='cuda', index=0)\u001b[39;00m\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "results = []\n",
    "\n",
    "for scaled_count, region in tqdm(zip(scaled_counts, regions), total=len(regions)):\n",
    "    n_sample = scaled_count\n",
    "\n",
    "    if n_sample > 0:\n",
    "        # print(n_sample)\n",
    "        batch_size = min(256, n_sample)\n",
    "        results.append(sample_synthetic_region(diffusion, dataset, region, n_sample, batch_size))\n",
    "    else:\n",
    "        results.append([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = f'{target_dataset_str}.pickle'\n",
    "# with open(filename, 'wb') as f:\n",
    "    # pickle.dump(results, f, protocol=pickle.HIGHEST_PROTOCOL)\n",
    "\n",
    "# TODO: Read the file\n",
    "with open(filename, 'rb') as f:\n",
    "    results = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "nonzero_results = []\n",
    "nonzero_regions = []\n",
    "\n",
    "for result, region in zip(results, regions):\n",
    "    if len(result) != 0:\n",
    "        nonzero_results.append(result)\n",
    "        nonzero_regions.append(region)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "restored_originals = []\n",
    "restored_samples = []\n",
    "\n",
    "for (trajectories_in_region, heatmap, samples, synthetic_heatmap), region in zip(nonzero_results, nonzero_regions):\n",
    "    for traj in trajectories_in_region:\n",
    "        restored_originals.append(undo_rotation_and_scaling(traj, region))\n",
    "    \n",
    "    for traj in samples:\n",
    "        restored_samples.append(undo_rotation_and_scaling(traj.T, region))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "restored_samples = np.stack(restored_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(9950, 128, 2)"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "restored_samples.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[39.9, 39.949997, 116.3, 116.35]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "target_region_actual = [dataset[:, :, 0].min(), dataset[:, :, 0].max(), dataset[:, :, 1].min(), dataset[:, :, 1].max()]\n",
    "target_region_actual"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(39.9, 39.95, 116.3, 116.35)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "original_region"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "restored_scaled_samples = restored_samples - [target_region_actual[0], target_region_actual[2]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "restored_scaled_samples = restored_scaled_samples / [target_region_actual[1] - target_region_actual[0],\n",
    "                                                     target_region_actual[3] - target_region_actual[2]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "len(restored_scaled_samples_filtered)=9819\n",
      "restored_scaled_samples_filtered.max(axis=(0, 1))=array([0.99988831, 0.99989631])\n",
      "restored_scaled_samples_filtered.min(axis=(0, 1))=array([0.00013333, 0.00052036])\n"
     ]
    }
   ],
   "source": [
    "is_outside = np.any(restored_scaled_samples > 1, axis=(1, 2))\n",
    "restored_scaled_samples_filtered = restored_scaled_samples[~is_outside]\n",
    "print(f'{len(restored_scaled_samples_filtered)=}')\n",
    "print(f'{restored_scaled_samples_filtered.max(axis=(0, 1))=}\\n{restored_scaled_samples_filtered.min(axis=(0, 1))=}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(f'tddpm-conditional-{target_dataset_str}.npy', restored_scaled_samples_filtered)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2721899/4042001916.py:8: RuntimeWarning: divide by zero encountered in log\n",
      "  axs[1].imshow(np.log(global_heatmap), extent=[yedges[0], yedges[-1], xedges[0], xedges[-1]], origin='lower', interpolation='nearest')\n",
      "/tmp/ipykernel_2721899/4042001916.py:21: RuntimeWarning: divide by zero encountered in log\n",
      "  axs[3].imshow(np.log(global_synth_heatmap), extent=[yedges[0], yedges[-1], xedges[0], xedges[-1]], origin='lower', interpolation='nearest')\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 5610384 characters]",
      "text/plain": [
       "<Figure size 4000x4000 with 4 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<Figure size 2000x2000 with 0 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "global_heatmap, xedges, yedges = unscaled_occupancy(restored_originals, res=1024)\n",
    "global_synth_heatmap, xedges, yedges = unscaled_occupancy(restored_samples, res=1024)\n",
    "\n",
    "f, axs = plt.subplots(ncols=2, nrows=2, figsize=(40, 40))\n",
    "axs = [ax for row in axs for ax in row]\n",
    "\n",
    "plt.figure(figsize=(20, 20))\n",
    "axs[1].imshow(np.log(global_heatmap), extent=[yedges[0], yedges[-1], xedges[0], xedges[-1]], origin='lower', interpolation='nearest')\n",
    "# axs[1].imshow(np.log(global_heatmap), extent=[yedges[0], yedges[-1], xedges[0], xedges[-1]], origin='lower', interpolation=None)\n",
    "\n",
    "for region_lat, region_lon, rotation, height, width in regions:\n",
    "    rect = patches.Rectangle((region_lon, region_lat), width, height, angle=rotation, linewidth=1, edgecolor='r', facecolor='none')\n",
    "    axs[1].add_patch(rect)\n",
    "\n",
    "for traj in restored_originals:\n",
    "    axs[0].plot(traj[:, 1], traj[:, 0], alpha=1)\n",
    "\n",
    "for traj in restored_samples:\n",
    "    axs[2].plot(traj[:, 1], traj[:, 0], alpha=1)\n",
    "\n",
    "axs[3].imshow(np.log(global_synth_heatmap), extent=[yedges[0], yedges[-1], xedges[0], xedges[-1]], origin='lower', interpolation='nearest')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "denoising",
   "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
