{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "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": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_str = 'geolife-25'\n",
    "target_dataset_str = 'geolife-100'\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": 3,
   "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": 4,
   "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": 5,
   "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/10 * map_width\n",
    "height = 1/10 * map_width\n",
    "rotation = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2736093/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, 180108 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": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 40040 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "heatmap_data, xedges, yedges = unscaled_occupancy(dataset, 20)\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": 9,
   "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": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "heatmap_data, xedges, yedges = occupancy(dataset, 20)\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": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "counts = np.array(counts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "scaled_counts = (target_dataset_size * counts / counts.sum()).astype('int')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0, 322)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "min(scaled_counts), max(scaled_counts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([  7,  11,  11,  13,  14,  14,  18,  21,  18,  14,  14,  19,  49,\n",
       "        67,  86,  80,  61,  55,  27,   9,   2,   3,   4,   6,   2,   4,\n",
       "         9,   6,   7,  10,   5,   9,  50,  88,  91,  86,  69,  49,  34,\n",
       "        13,   2,   3,   5,   5,   1,   7,  17,  20,  26,  33,  24,  16,\n",
       "        44,  81,  94,  82,  41,  23,  25,   9,   2,   3,   2,   2,   1,\n",
       "         7,  23,  25,  19,  26,  24,  14,  24,  40,  60, 123,  89,  27,\n",
       "        34,  13,   3,   4,   2,   2,   4,   6,  18,  14,   0,  10,  17,\n",
       "        10,  26,  56, 146, 310, 240, 148, 163,  58,   3,   4,   3,   2,\n",
       "         5,   6,  15,  13,   2,   8,  14,   9,  29, 145, 295, 322, 195,\n",
       "       198, 205,  57,  10,   8,   9,  12,  23,  21,  19,  14,   3,   1,\n",
       "         1,   3,  24, 121, 220, 152,  56,  84,  70,  11,   9,   9,  10,\n",
       "        14,  23,  22,  17,  11,   2,   3,   5,   5,  27,  52,  78,  67,\n",
       "        50,  40,  12,   5,   8,  11,   9,  10,   9,  15,  26,  30,  20,\n",
       "        14,  18,  15,  44,  70,  62,  46,  63,  55,   9,   2,  14,  20,\n",
       "        15,  16,  14,  15,  24,  36,  31,  15,  14,  16,  47,  52,  33,\n",
       "        67, 122,  82,   7,   1,  12,  15,   8,  10,   7,   2,   4,  15,\n",
       "        17,   5,   5,   7,  28,  27,  15,  51, 102,  68,   3,   0,   7,\n",
       "         8,   0,   2,   2,   1,   4,   8,  10,  14,  18,  13,  31,  29,\n",
       "        18,  21,  41,  38,   3,   0,  14,  16,   6,   8,   8,   7,  11,\n",
       "         8,  14,  23,  16,   9,  29,  27,  13,  14,  36,  37,   4,   1,\n",
       "        15,  16,   6,   8,   8,   8,   9,   7,  17,  21,  11,   9,  35,\n",
       "        35,  15,  14,  38,  39,   6,   1,   7,   8,   0,   1,   2,   4,\n",
       "         3,   4,  13,  14,   8,   9,  42,  44,  17,  14,  38,  36,   3,\n",
       "         1,   7,   6,   0,   1,   2,   6,   5,   4,   9,   6,   0,   1,\n",
       "        33,  34,   7,   6,  33,  34,   4,   1,   8,   7,   0,   1,   3,\n",
       "         8,  12,   7,   7,   6,   0,   1,  30,  31,   5,   6,  28,  33,\n",
       "        11,   3,   9,   7,   1,   1,   2,   6,  11,   7,   5,   4,   0,\n",
       "         1,  25,  26,   4,   4,  19,  24,  10,   3,   6,   5,   1,   1,\n",
       "         1,   5,   8,   3,   2,   2,   1,   3,  20,  19,   1,   0,  11,\n",
       "        13,   3,   1,   1,   2,   0,   0,   0,   3,   5,   1,   0,   0,\n",
       "         1,   2,   9,   9,   1,   0,   2,   2,   0,   0])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scaled_counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "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": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/400 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/400 [00:01<?, ?it/s]\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[16], 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[15], line 8\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      2\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;124;03mExtract subtrajectories from region, creates a heatmap of the real data, creates synthetic\u001b[39;00m\n\u001b[1;32m      4\u001b[0m \u001b[38;5;124;03mtrajectories conditioned on the heatmap and finally calculate a heatmap of the synthetic samples.\u001b[39;00m\n\u001b[1;32m      5\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m n_samples \u001b[38;5;241m>\u001b[39m\u001b[38;5;241m=\u001b[39m sample_batch_size\n\u001b[0;32m----> 8\u001b[0m trajectories_in_region \u001b[38;5;241m=\u001b[39m \u001b[43mextract_signals\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[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[1;32m     11\u001b[0m samples \u001b[38;5;241m=\u001b[39m sample_heatmap(diffusion, heatmap, n_samples\u001b[38;5;241m=\u001b[39mn_samples, sample_batch_size\u001b[38;5;241m=\u001b[39msample_batch_size)\n",
      "File \u001b[0;32m~/repos/tddpm/tddpm/data_generation.py:57\u001b[0m, in \u001b[0;36mextract_signals\u001b[0;34m(dataset, region, min_len)\u001b[0m\n\u001b[1;32m     53\u001b[0m \u001b[38;5;66;03m# origin = (region_lat, region_lon)\u001b[39;00m\n\u001b[1;32m     54\u001b[0m \u001b[38;5;66;03m# size = (height, width)\u001b[39;00m\n\u001b[1;32m     56\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m traj \u001b[38;5;129;01min\u001b[39;00m dataset:\n\u001b[0;32m---> 57\u001b[0m     parts \u001b[38;5;241m=\u001b[39m \u001b[43mcrop_signal_with_rotation\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtraj\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mregion\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     58\u001b[0m     all_parts\u001b[38;5;241m.\u001b[39mextend(parts)\n\u001b[1;32m     60\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m all_parts\n",
      "File \u001b[0;32m~/repos/tddpm/tddpm/data_generation.py:33\u001b[0m, in \u001b[0;36mcrop_signal_with_rotation\u001b[0;34m(signal, region, min_len)\u001b[0m\n\u001b[1;32m     30\u001b[0m lon \u001b[38;5;241m=\u001b[39m signal[i, \u001b[38;5;241m1\u001b[39m]\n\u001b[1;32m     31\u001b[0m \u001b[38;5;66;03m# print(lat, lon)\u001b[39;00m\n\u001b[1;32m     32\u001b[0m \u001b[38;5;66;03m# print(min_lat, max_lat, min_lon, max_lon)\u001b[39;00m\n\u001b[0;32m---> 33\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m lat \u001b[38;5;241m<\u001b[39m \u001b[38;5;241m0\u001b[39m \u001b[38;5;129;01mor\u001b[39;00m lat \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m1\u001b[39m \\\n\u001b[1;32m     34\u001b[0m     \u001b[38;5;129;01mor\u001b[39;00m lon \u001b[38;5;241m<\u001b[39m \u001b[38;5;241m0\u001b[39m \u001b[38;5;129;01mor\u001b[39;00m lon \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[1;32m     35\u001b[0m     \n\u001b[1;32m     36\u001b[0m     \u001b[38;5;66;03m# if i - last_outside_region + 1 >= min_len:\u001b[39;00m\n\u001b[1;32m     37\u001b[0m     part \u001b[38;5;241m=\u001b[39m signal[last_outside_region\u001b[38;5;241m+\u001b[39m\u001b[38;5;241m1\u001b[39m:i]\n\u001b[1;32m     38\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(part) \u001b[38;5;241m>\u001b[39m\u001b[38;5;241m=\u001b[39m min_len:\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": 17,
   "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": 18,
   "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": 19,
   "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": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "restored_samples = np.stack(restored_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[39.90001, 40.0, 116.3, 116.39997]"
      ]
     },
     "execution_count": 21,
     "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": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(39.9, 40, 116.3, 116.4)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "original_region"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "restored_scaled_samples = restored_samples - [target_region_actual[0], target_region_actual[2]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "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": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "len(restored_scaled_samples_filtered)=9782\n",
      "restored_scaled_samples_filtered.max(axis=(0, 1))=array([0.99993252, 0.99935958])\n",
      "restored_scaled_samples_filtered.min(axis=(0, 1))=array([ 1.65505363e-03, -3.05273602e-05])\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": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(f'tddpm-conditional-{target_dataset_str}.npy', restored_scaled_samples_filtered)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2719416/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_2719416/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, 5606572 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
}
