find_lakes.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LFvW_IZZHEKf"
   },
   "outputs": [],
   "source": [
    "import osmium\n",
    "from tqdm import tqdm\n",
    "from shapely.geometry import Point, LineString, Polygon\n",
    "from shapely.ops import transform\n",
    "from shapely import intersects\n",
    "import pyproj\n",
    "from rtree import index\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import utils.projections as projections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bCO8EEDwHEKg"
   },
   "outputs": [],
   "source": [
    "# Define a class to handle OSM data parsing\n",
    "class OSMHandler(osmium.SimpleHandler):\n",
    "    def __init__(self):\n",
    "        super(OSMHandler, self).__init__()\n",
    "        self.lakes = []\n",
    "        self.bodies = ['lake', 'reservoir', 'pond']\n",
    "\n",
    "    def way(self, w):\n",
    "        if 'natural' in w.tags and w.tags['natural'] == 'water' and 'water' in w.tags and w.tags['water'] in self.bodies:\n",
    "            lake_coords = []\n",
    "            for n in w.nodes:\n",
    "                if n.location.valid():\n",
    "                    lake_coords.append(Point(n.location.lat, n.location.lon))\n",
    "            if len(lake_coords) > 1:\n",
    "                lake_meta = {\n",
    "                    'id': w.id,\n",
    "                    'name': w.tags['name'] if 'name' in w.tags else '',\n",
    "                }\n",
    "                line = LineString(lake_coords)\n",
    "                if line.is_closed:\n",
    "                    polygon = Polygon(lake_coords)\n",
    "                    self.lakes.append((polygon, lake_meta))\n",
    "\n",
    "# Load OSM data using OSMHandler\n",
    "osm_handler = OSMHandler()\n",
    "osm_handler.apply_file(\"data/extract.osm.pbf\", locations=True, idx='flex_mem') # use idx=dense_mmap_array for large files on linux"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SVhyEKDvHEKh"
   },
   "outputs": [],
   "source": [
    "lakes = osm_handler.lakes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7I55EFPEHEKi"
   },
   "outputs": [],
   "source": [
    "BUFFER = 25\n",
    "buffered_lakes = []\n",
    "for (poly, meta) in tqdm(lakes):\n",
    "    buffered_poly = projections.calculate_buffer(poly, BUFFER)\n",
    "    buffered_lakes.append((buffered_poly, meta))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "44AFLk5sHEKi"
   },
   "outputs": [],
   "source": [
    "# Create an R-tree index\n",
    "lake_idx = index.Index()\n",
    "\n",
    "# Insert polygons into the R-tree index with their bounding boxes\n",
    "for i, (poly, _) in enumerate(buffered_lakes):\n",
    "    bbox = poly.bounds\n",
    "    lake_idx.insert(i, bbox)\n",
    "\n",
    "# Function to find the polygon a line is in (O(log(N)) time complexity)\n",
    "def find_lake_for_trail(line):\n",
    "    for i in lake_idx.intersection(line.bounds):\n",
    "        if intersects(buffered_lakes[i][0], line):\n",
    "            return i  # Index of the polygon"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-v_BhP5yHEKj"
   },
   "outputs": [],
   "source": [
    "def map_trail_to_lake(trail_subset):\n",
    "    first_encountered_lake_idx = []\n",
    "    for (line, _) in tqdm(trail_subset):\n",
    "        polygon_index = find_lake_for_trail(line)\n",
    "        first_encountered_lake_idx.append(polygon_index)\n",
    "    return first_encountered_lake_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RsdliHpjHEKk"
   },
   "outputs": [],
   "source": [
    "def print_trail_and_lake(trail, lake, buffered_lake):\n",
    "    fig, ax = plt.subplots()\n",
    "\n",
    "    rendered_trail = projections.project_line_to_meters(trail[0])\n",
    "    rendered_lake = projections.project_poly_to_meters(lake[0])\n",
    "    rendered_buffered_lake = projections.project_poly_to_meters(buffered_lake[0])\n",
    "\n",
    "    # Plot polygons\n",
    "    x, y = rendered_lake.exterior.xy\n",
    "    ax.plot(x, y, label='Lake', color='r')\n",
    "    x, y = rendered_buffered_lake.exterior.xy\n",
    "    ax.plot(x, y, label='Buffered Lake', color='y')\n",
    "\n",
    "    # Plot line\n",
    "    x, y = rendered_trail.xy\n",
    "    plt.plot(x, y, label='Trail', color='b')\n",
    "\n",
    "    # Set equal aspect ratio\n",
    "    plt.axis('equal')\n",
    "\n",
    "    # Set labels and legend\n",
    "    ax.set_xlabel('X')\n",
    "    ax.set_ylabel('Y')\n",
    "    ax.set_title(\"{} and {}\".format(trail[1]['name'], lake[1]['name']))\n",
    "    ax.legend()\n",
    "\n",
    "    # Show the plot\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "ZpbotSgXHEKk"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "\n",
    "filename = os.path.join(\"data\", 'all_trail_segments.pickle')\n",
    "\n",
    "with open(filename, 'rb') as f:\n",
    "    hiking_trails = pickle.load(f)\n",
    "\n",
    "trail_subset = hiking_trails"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "M53G4NyIHEKl"
   },
   "outputs": [],
   "source": [
    "lake_map = map_trail_to_lake(trail_subset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "LDhbGNMnHEKm"
   },
   "outputs": [],
   "source": [
    "LIMIT = 15\n",
    "count = 0\n",
    "for i in range(len(trail_subset)):\n",
    "    trail = trail_subset[i]\n",
    "\n",
    "    lake_idx = lake_map[i]\n",
    "    if lake_idx == None:\n",
    "        continue\n",
    "\n",
    "    lake = lakes[lake_idx]\n",
    "    buffered_lake = buffered_lakes[lake_idx]\n",
    "    print_trail_and_lake(trail, lake, buffered_lake)\n",
    "\n",
    "    count += 1\n",
    "    if count == LIMIT:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7nRBMi0XH1N9"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "\n",
    "filename = os.path.join(\"data\", 'all_lakes.pickle')\n",
    "with open(filename, 'rb') as f:\n",
    "    pickle.dump(lakes)\n",
    "\n",
    "filename = os.path.join(\"data\", 'trail_to_lake_map.pickle')\n",
    "\n",
    "with open(filename, 'rb') as f:\n",
    "    pickle.dump(lake_map)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}