{
 "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 = 50000\n",
    "k_anon = 50\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_620937/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": [],
   "source": [
    "def histogram(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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 39648 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "heatmap_data, xedges, yedges = histogram(dataset, 20)\n",
    "\n",
    "# print((heatmap_data < k_anon).sum() - (heatmap_data == 0).sum())\n",
    "\n",
    "heatmap_data[heatmap_data < k_anon] = 0\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": [
    "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, 1612)"
      ]
     },
     "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([  37,   55,   57,   69,   70,   71,   92,  108,   92,   73,   70,\n",
       "         95,  247,  337,  432,  401,  308,  278,  138,   47,   11,   15,\n",
       "         23,   30,   14,   20,   46,   34,   39,   53,   26,   46,  251,\n",
       "        444,  457,  430,  346,  247,  171,   65,   10,   16,   29,   28,\n",
       "          9,   36,   89,  100,  133,  169,  124,   83,  223,  409,  470,\n",
       "        413,  207,  117,  125,   46,   14,   16,   14,   10,    5,   35,\n",
       "        115,  126,   99,  133,  121,   70,  123,  203,  304,  615,  446,\n",
       "        135,  173,   67,   16,   20,   13,   14,   21,   32,   90,   72,\n",
       "          2,   51,   85,   50,  133,  284,  731, 1552, 1203,  740,  817,\n",
       "        292,   19,   23,   15,   14,   26,   32,   77,   68,   11,   43,\n",
       "         70,   48,  147,  725, 1479, 1612,  977,  993, 1029,  286,   50,\n",
       "         43,   47,   62,  119,  109,   97,   74,   15,    8,    5,   16,\n",
       "        120,  609, 1102,  761,  284,  420,  352,   56,   49,   46,   52,\n",
       "         71,  116,  111,   88,   59,   10,   15,   25,   26,  138,  262,\n",
       "        392,  336,  251,  204,   60,   26,   43,   56,   49,   50,   48,\n",
       "         77,  130,  153,  104,   72,   90,   79,  224,  353,  312,  234,\n",
       "        319,  275,   49,   10,   74,  100,   76,   84,   73,   77,  122,\n",
       "        183,  155,   76,   71,   81,  235,  261,  169,  335,  614,  411,\n",
       "         36,    6,   61,   76,   40,   52,   38,   12,   20,   76,   85,\n",
       "         25,   25,   35,  144,  137,   75,  255,  513,  341,   19,    4,\n",
       "         39,   40,    4,   14,   14,    8,   24,   41,   54,   73,   90,\n",
       "         67,  156,  147,   92,  108,  207,  193,   16,    2,   71,   80,\n",
       "         31,   43,   40,   35,   57,   40,   73,  117,   82,   49,  145,\n",
       "        137,   69,   71,  183,  189,   22,    6,   75,   84,   32,   43,\n",
       "         41,   40,   49,   35,   88,  107,   57,   46,  176,  177,   78,\n",
       "         74,  194,  195,   30,    9,   39,   40,    1,    7,   12,   22,\n",
       "         18,   20,   66,   71,   44,   47,  211,  220,   87,   72,  192,\n",
       "        183,   18,    6,   35,   33,    1,    7,   13,   30,   28,   20,\n",
       "         45,   30,    0,    6,  166,  172,   35,   32,  166,  171,   21,\n",
       "          6,   44,   36,    2,    7,   15,   40,   60,   37,   38,   30,\n",
       "          2,    9,  153,  155,   27,   34,  142,  166,   57,   15,   49,\n",
       "         39,    5,    7,   12,   30,   56,   37,   28,   24,    2,    9,\n",
       "        128,  131,   21,   19,   95,  122,   54,   16,   31,   28,    8,\n",
       "          7,    7,   25,   41,   18,   12,   12,    8,   17,  100,   97,\n",
       "          8,    2,   57,   66,   15,    6,    8,   11,    4,    3,    2,\n",
       "         16,   26,    9,    4,    4,    8,   14,   47,   45,    5,    0,\n",
       "         12,   12,    0,    0])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "scaled_counts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def heatmap_k_anonymity(trajectories, k_anon, res=64):\n",
    "    all_concatenated = np.concatenate(trajectories)\n",
    "    H, xedges, yedges = np.histogram2d(all_concatenated[:, 1], all_concatenated[:, 0], bins=(res, res), range=((0, 1), (0, 1)), density=False)\n",
    "    H[H<k_anon] = 0\n",
    "    H /= H.sum()\n",
    "    H *= res**2\n",
    "    H = H.T\n",
    "\n",
    "    return H, xedges, yedges"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_synthetic_region(diffusion, dataset, region, n_samples, k_anon, 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)\n",
    "    heatmap, _, _ = heatmap_k_anonymity(trajectories_in_region, k_anon, 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": null,
   "metadata": {},
   "outputs": [],
   "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, k_anon, batch_size))\n",
    "    else:\n",
    "        results.append([])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "trajectories_in_region, heatmap, samples, synthetic_heatmap = results[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4096"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array(heatmap.shape).prod()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.0, 203.2488900496213, 4096.0)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "heatmap.min(), heatmap.max(), heatmap.sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.image.AxesImage at 0x7f7405e71b10>"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 44560 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "f, axs = plt.subplots(ncols=2)\n",
    "for train_t in trajectories_in_region:\n",
    "    axs[0].plot(train_t[:, 1], train_t[:, 0], alpha=.1)\n",
    "axs[0].imshow(heatmap, extent=[0, 1, 0, 1], origin='lower')\n",
    "\n",
    "for synth_t in samples:\n",
    "    axs[1].plot(synth_t[1, :], synth_t[0, :], alpha=.1)\n",
    "axs[1].imshow(synthetic_heatmap, extent=[0, 1, 0, 1], origin='lower')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "filename = f'{target_dataset_str}-{k_anon}.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": 23,
   "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": 24,
   "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": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "restored_samples = np.stack(restored_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[39.90001, 40.0, 116.3, 116.39997]"
      ]
     },
     "execution_count": 26,
     "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": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(39.9, 40, 116.3, 116.4)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "original_region"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "restored_scaled_samples = restored_samples - [target_region_actual[0], target_region_actual[2]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "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": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "len(restored_scaled_samples_filtered)=49647\n",
      "restored_scaled_samples_filtered.max(axis=(0, 1))=array([nan, nan])\n",
      "restored_scaled_samples_filtered.min(axis=(0, 1))=array([nan, nan])\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": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save(f'tddpm-conditional-{target_dataset_str}-k-anon-{k_anon}.npy', restored_scaled_samples_filtered)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "ename": "ValueError",
     "evalue": "autodetected range of [nan, nan] is not finite",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[32], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m global_heatmap, xedges, yedges \u001b[38;5;241m=\u001b[39m unscaled_occupancy(restored_originals, res\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1024\u001b[39m)\n\u001b[0;32m----> 2\u001b[0m global_synth_heatmap, xedges, yedges \u001b[38;5;241m=\u001b[39m \u001b[43munscaled_occupancy\u001b[49m\u001b[43m(\u001b[49m\u001b[43mrestored_samples\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mres\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1024\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m      4\u001b[0m f, axs \u001b[38;5;241m=\u001b[39m plt\u001b[38;5;241m.\u001b[39msubplots(ncols\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m2\u001b[39m, nrows\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m2\u001b[39m, figsize\u001b[38;5;241m=\u001b[39m(\u001b[38;5;241m40\u001b[39m, \u001b[38;5;241m40\u001b[39m))\n\u001b[1;32m      5\u001b[0m axs \u001b[38;5;241m=\u001b[39m [ax \u001b[38;5;28;01mfor\u001b[39;00m row \u001b[38;5;129;01min\u001b[39;00m axs \u001b[38;5;28;01mfor\u001b[39;00m ax \u001b[38;5;129;01min\u001b[39;00m row]\n",
      "File \u001b[0;32m~/repos/tddpm/tddpm/evaluation.py:33\u001b[0m, in \u001b[0;36munscaled_occupancy\u001b[0;34m(trajectories, res)\u001b[0m\n\u001b[1;32m     29\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     30\u001b[0m \u001b[38;5;124;03mUnscaled variant of data_generation.region_occupancy, mainly useful for visualization.\u001b[39;00m\n\u001b[1;32m     31\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m     32\u001b[0m all_concatenated \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mconcatenate(trajectories)\n\u001b[0;32m---> 33\u001b[0m H, xedges, yedges \u001b[38;5;241m=\u001b[39m \u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mhistogram2d\u001b[49m\u001b[43m(\u001b[49m\u001b[43mall_concatenated\u001b[49m\u001b[43m[\u001b[49m\u001b[43m:\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mall_concatenated\u001b[49m\u001b[43m[\u001b[49m\u001b[43m:\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbins\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43mres\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mres\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdensity\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43;01mTrue\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m     35\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m H, xedges, yedges\n",
      "File \u001b[0;32m<__array_function__ internals>:200\u001b[0m, in \u001b[0;36mhistogram2d\u001b[0;34m(*args, **kwargs)\u001b[0m\n",
      "File \u001b[0;32m~/.conda/envs/tddpm/lib/python3.10/site-packages/numpy/lib/twodim_base.py:820\u001b[0m, in \u001b[0;36mhistogram2d\u001b[0;34m(x, y, bins, range, density, weights)\u001b[0m\n\u001b[1;32m    818\u001b[0m     xedges \u001b[38;5;241m=\u001b[39m yedges \u001b[38;5;241m=\u001b[39m asarray(bins)\n\u001b[1;32m    819\u001b[0m     bins \u001b[38;5;241m=\u001b[39m [xedges, yedges]\n\u001b[0;32m--> 820\u001b[0m hist, edges \u001b[38;5;241m=\u001b[39m \u001b[43mhistogramdd\u001b[49m\u001b[43m(\u001b[49m\u001b[43m[\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43my\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbins\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mrange\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdensity\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mweights\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    821\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m hist, edges[\u001b[38;5;241m0\u001b[39m], edges[\u001b[38;5;241m1\u001b[39m]\n",
      "File \u001b[0;32m<__array_function__ internals>:200\u001b[0m, in \u001b[0;36mhistogramdd\u001b[0;34m(*args, **kwargs)\u001b[0m\n",
      "File \u001b[0;32m~/.conda/envs/tddpm/lib/python3.10/site-packages/numpy/lib/histograms.py:1001\u001b[0m, in \u001b[0;36mhistogramdd\u001b[0;34m(sample, bins, range, density, weights)\u001b[0m\n\u001b[1;32m    998\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m bins[i] \u001b[38;5;241m<\u001b[39m \u001b[38;5;241m1\u001b[39m:\n\u001b[1;32m    999\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m   1000\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;124m`bins[\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m]` must be positive, when an integer\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;241m.\u001b[39mformat(i))\n\u001b[0;32m-> 1001\u001b[0m smin, smax \u001b[38;5;241m=\u001b[39m \u001b[43m_get_outer_edges\u001b[49m\u001b[43m(\u001b[49m\u001b[43msample\u001b[49m\u001b[43m[\u001b[49m\u001b[43m:\u001b[49m\u001b[43m,\u001b[49m\u001b[43mi\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43mrange\u001b[39;49m\u001b[43m[\u001b[49m\u001b[43mi\u001b[49m\u001b[43m]\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1002\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m   1003\u001b[0m     n \u001b[38;5;241m=\u001b[39m operator\u001b[38;5;241m.\u001b[39mindex(bins[i])\n",
      "File \u001b[0;32m~/.conda/envs/tddpm/lib/python3.10/site-packages/numpy/lib/histograms.py:323\u001b[0m, in \u001b[0;36m_get_outer_edges\u001b[0;34m(a, range)\u001b[0m\n\u001b[1;32m    321\u001b[0m     first_edge, last_edge \u001b[38;5;241m=\u001b[39m a\u001b[38;5;241m.\u001b[39mmin(), a\u001b[38;5;241m.\u001b[39mmax()\n\u001b[1;32m    322\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (np\u001b[38;5;241m.\u001b[39misfinite(first_edge) \u001b[38;5;129;01mand\u001b[39;00m np\u001b[38;5;241m.\u001b[39misfinite(last_edge)):\n\u001b[0;32m--> 323\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    324\u001b[0m             \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mautodetected range of [\u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m, \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m] is not finite\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(first_edge, last_edge))\n\u001b[1;32m    326\u001b[0m \u001b[38;5;66;03m# expand empty range to avoid divide by zero\u001b[39;00m\n\u001b[1;32m    327\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m first_edge \u001b[38;5;241m==\u001b[39m last_edge:\n",
      "\u001b[0;31mValueError\u001b[0m: autodetected range of [nan, nan] is not finite"
     ]
    }
   ],
   "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
}
