cuspatial_synthetic_people.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "cuspatial_synthetic_people",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Census block synthetic people using GPU\n",
        "\n",
        "In this notebook, we want to create a dataset of synthetic people based on Census block population data. Census blocks are provided as non-overlapping polygons. We'll represent each synthetic person as a point of lat-lon coordinates. Points will be first generated randomly using normal distribution and then will be checked if they belong to some polygons. To run this at scale, we'll utilize the GPU by using cupy for randomizing the point data, and cuspatial for checking points in polygons."
      ],
      "metadata": {
        "id": "FQxNi0gV_MoM"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import geopandas as gpd\n",
        "import cudf, cupy, cuspatial\n",
        "\n",
        "# disbale warning related to quadtree scale setting\n",
        "import warnings\n",
        "warnings.filterwarnings(\"ignore\")\n"
      ],
      "metadata": {
        "id": "cFxie1V1nELY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Load data\n",
        "\n",
        "We'll load Census population data from parquet, and create an unique integer geoid for each polygon."
      ],
      "metadata": {
        "id": "2cSUlUc3pe94"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "gdf = gpd.read_parquet('census-parquet/part.6.parquet')\n",
        "\n",
        "geoids = cupy.array(range(len(gdf)))\n"
      ],
      "metadata": {
        "id": "6hchE_JFphIS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "As our polygons do not overlap, a point can only belong to at most 1 polygon. In addition, points data can have different density over different locations. We'll use quadtree data structure to store our points data. Let's read the data in a proper format so that cuspatial can understand and help use build quadtrees easily."
      ],
      "metadata": {
        "id": "zZX6mcRspkPB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "poly_offsets, ring_offsets, coords = cuspatial.read_polygon_shapefile('part6.shp')\n"
      ],
      "metadata": {
        "id": "bp5z2dlbptGz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Extract xcoords and ycoords from input data."
      ],
      "metadata": {
        "id": "siCfk4EEqtqO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "xcoords = coords.x\n",
        "ycoords = coords.y\n"
      ],
      "metadata": {
        "id": "fDBUCxFApk7U"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Polygons to points\n",
        "\n",
        "Below is an util function to create points in polygons. We'll do it in a vectorize maner so that we can run multiple polygons at once."
      ],
      "metadata": {
        "id": "rwwvpVnHrKop"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def polygons_to_points(\n",
        "    gdf: gpd.GeoDataFrame,            # table of all polygons\n",
        "    val_col: str,                     # name of population column in the table\n",
        "    poly_offsets: cudf.Series,        # polygon offsets\n",
        "    ring_offsets: cudf.Series,        # polygon offsets\n",
        "    xcoords: cudf.Series,             # flatten xcoords of polygons\n",
        "    ycoords: cudf.Series,             # flatten ycoords of polygons\n",
        "    geoids: cupy.ndarray,             # indexes of polygons\n",
        "    max_depth: float,                 # max depth of a quadtree\n",
        "    min_size: float,                  # minimum number of points for a non-leaf quadtree node\n",
        "    max_len: int,                     # max number of points generated in an iteration\n",
        "):\n",
        "    # total number of points to generate\n",
        "    total_points = gdf[val_col].sum()\n",
        "    # number of points in each polygon\n",
        "    num_points = cupy.asarray(gdf[val_col])    \n",
        "    # polygons with missing points\n",
        "    keep_polygons = cupy.arange(len(gdf))\n",
        "    num_polygons = gdf.shape[0]\n",
        "    # number of points generated in each polygon\n",
        "    gen_points = cupy.zeros(num_polygons)\n",
        "    # cudf DataFrame to store all points that is within an arbitrary polygon \n",
        "    keep_points = cudf.DataFrame({'x': [], 'y': [], 'geoid': []})\n",
        "    # bounding box of each polygon\n",
        "    polygon_bboxes = cuspatial.polygon_bounding_boxes(\n",
        "        poly_offsets, ring_offsets, xcoords, ycoords\n",
        "    )\n",
        "    # num iterations\n",
        "    it = 0\n",
        "    while keep_polygons.shape[0] > 0:\n",
        "        it += 1\n",
        "        print('it', it, len(keep_polygons))                \n",
        "        # bounds of polygons\n",
        "        x_min, y_min, x_max, y_max = gdf.iloc[keep_polygons.get()].total_bounds\n",
        "        # random points\n",
        "        xs = cupy.random.uniform(x_min, x_max, max_len)\n",
        "        ys = cupy.random.uniform(y_min, y_max, max_len)\n",
        "        # bounds of all the points we have just generated\n",
        "        min_px, max_px, min_py, max_py = xs.min(), xs.max(), ys.min(), ys.max()\n",
        "        # calculate scale of the quadtree\n",
        "        scale = max(max_px - min_px, max_py - min_py) // (1 << max_depth)\n",
        "        # build a quadtree\n",
        "        key_to_point, quadtree = cuspatial.quadtree_on_points(\n",
        "            xs, ys, min_px, max_px, min_py, max_py, scale, max_depth, min_size\n",
        "        )\n",
        "        poly_quad_pairs = cuspatial.join_quadtree_and_bounding_boxes(\n",
        "            quadtree, polygon_bboxes, x_min, x_max, y_min, y_max, scale, max_depth\n",
        "        )\n",
        "        result = cuspatial.quadtree_point_in_polygon(\n",
        "            poly_quad_pairs, quadtree, key_to_point, xs, ys,\n",
        "            poly_offsets, ring_offsets, xcoords, ycoords\n",
        "        )\n",
        "        # filtering result to only keep points of polygons that not completely full\n",
        "        result_array = cupy.asarray(result.polygon_index)\n",
        "        keep_ids = cupy.isin(result_array, keep_polygons)\n",
        "        keep_ids = cupy.where(keep_ids > 0)[0]\n",
        "        # keep_ids = cupy.isin(result_array, keep_polygons)\n",
        "        # keep_ids = result_array[cupy.where(keep_ids==1)[0]]\n",
        "        result = result.iloc[keep_ids]\n",
        "        keep_points = cudf.concat([\n",
        "            keep_points,\n",
        "            cudf.DataFrame({\n",
        "                'x': xs[result['point_index']],\n",
        "                'y': ys[result['point_index']],\n",
        "                'geoid': geoids[result['polygon_index']]})\n",
        "            ], ignore_index=True)\n",
        "        count_df = result.groupby('polygon_index').count().reset_index()\n",
        "        gen_points[count_df['polygon_index']] += count_df['point_index']\n",
        "        missing_points = num_points - gen_points\n",
        "        keep_polygons = cupy.where(missing_points > 0)[0]\n",
        "    # select exact number of points in each polygon\n",
        "    selected_points = []\n",
        "    count_points = keep_points.groupby('geoid').count()\n",
        "    sorted_points = keep_points.sort_values(by='geoid').reset_index(drop=True)\n",
        "    # start index of first point in this polygon\n",
        "    point_id = 0\n",
        "    for i in range(gdf.shape[0]):\n",
        "        if i % 1000 == 0:\n",
        "            print(f'Selecting exact number of points for the {i}-th polygon')\n",
        "        num_expected_points = gdf[val_col][i]\n",
        "        points = sorted_points.iloc[range(point_id, point_id + num_expected_points)]    \n",
        "        geoid = geoids[i]\n",
        "        num_generated_points = count_points['x'][geoid]\n",
        "        point_id += num_generated_points\n",
        "        selected_points.append(points)\n",
        "    result = cudf.concat(selected_points, ignore_index=True)\n",
        "    return result\n",
        "\n"
      ],
      "metadata": {
        "id": "YN80In-DQqnp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "As Census block data can be large, running them all at once can take long time we'll divide it into smaller batches to run it more efficiently."
      ],
      "metadata": {
        "id": "vE0OTMllsXbR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def get_subset(\n",
        "    begin: int,                  # index of first row in the subset\n",
        "    end: int,                    # index of last row in the subset\n",
        "    gdf: gpd.GeoDataFrame,       # table of all polygons\n",
        "    geoids: cupy.ndarray,        # indexes of all polygons\n",
        "    poly_offsets: cudf.Series,   # polygon offsets\n",
        "    ring_offsets: cudf.Series,   # ring offsets\n",
        "    xcoords: cudf.Series,        # xcoords of polygons\n",
        "    ycoords: cudf.Series,        # ycoords of polygons\n",
        "):\n",
        "    # select subset gdf\n",
        "    subset_gdf = gdf.iloc[range(begin, end)]\n",
        "    subset_geoids = geoids[begin:end]\n",
        "    # index of the first ring in the subset\n",
        "    begin_ring = poly_offsets[begin]\n",
        "    # index of the last ring in the subset\n",
        "    end_ring = poly_offsets[end]\n",
        "    # select subset ring offsets\n",
        "    subset_ring_offsets = ring_offsets.iloc[range(begin_ring, end_ring)].reset_index(drop=True)\n",
        "    subset_ring_offsets = subset_ring_offsets - subset_ring_offsets[0]\n",
        "    # subset poly offsets\n",
        "    subset_poly_offsets = poly_offsets.iloc[range(begin, end)].reset_index(drop=True)  - poly_offsets[begin]\n",
        "    # index of first coords\n",
        "    begin_coords = ring_offsets[begin_ring]\n",
        "    # index of last coords\n",
        "    if end_ring < len(ring_offsets) - 1:\n",
        "        end_coords = ring_offsets[end_ring]\n",
        "    else:\n",
        "        end_coords = len(coords)\n",
        "    # select subset coords\n",
        "    subset_xcoords = xcoords[begin_coords: end_coords].reset_index(drop=True)\n",
        "    subset_ycoords = ycoords[begin_coords: end_coords].reset_index(drop=True)\n",
        "    return (\n",
        "        subset_gdf,\n",
        "        subset_poly_offsets,\n",
        "        subset_ring_offsets,\n",
        "        subset_xcoords,\n",
        "        subset_ycoords,\n",
        "        subset_geoids\n",
        "    )\n",
        "\n"
      ],
      "metadata": {
        "id": "yuUwZMdYsXuC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Experiments\n",
        "\n",
        "We'll select run our dataset batch by batch with 2000 polygons each."
      ],
      "metadata": {
        "id": "2k56PdFgsGRk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "val_col = 'POP'\n",
        "\n",
        "batch_size = 2000\n",
        "all_batch_results = []\n",
        "for i in range(0, gdf.shape[0], batch_size):\n",
        "    print(f'Processing the batch {i // batch_size + 1}')\n",
        "    begin, end = i, i + batch_size\n",
        "    subset_gdf, subset_poly_offsets, subset_ring_offsets, subset_xcoords, subset_ycoords, subset_geoids = get_subset(\n",
        "        begin,\n",
        "        end,\n",
        "        gdf,\n",
        "        geoids,\n",
        "        poly_offsets,\n",
        "        ring_offsets,\n",
        "        xcoords,\n",
        "        ycoords\n",
        "    )\n",
        "    subset_result = polygons_to_points(\n",
        "        subset_gdf,\n",
        "        val_col,\n",
        "        subset_poly_offsets,\n",
        "        subset_ring_offsets,\n",
        "        subset_xcoords,\n",
        "        subset_ycoords,\n",
        "        subset_geoids,\n",
        "        max_depth=10,\n",
        "        min_size=500,\n",
        "        max_len=3_000_000\n",
        "    )\n",
        "    all_batch_results.append(subset_result)"
      ],
      "metadata": {
        "id": "xCF4wQZHBGjx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Combine all the results of all batches into a single dataframe."
      ],
      "metadata": {
        "id": "aOBpu2galvsD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "result = cudf.concat(all_batch_results, ignore_index=True)\n",
        "result"
      ],
      "metadata": {
        "id": "DiI3mU90luBu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**TODO**\n",
        "\n",
        "- Select subset/batches of polygons by extent of `(xmin:xmax, ymin:ymax)` to avoid a batch where polygons are far away from some other, which leads to a big bounding box so this vectorization method becomes less efficient. (See batch 6: polygon 10000 to polygon 12000)\n",
        "\n",
        "- Improve performance for selecting exact number of points in each polygon.\n",
        "\n",
        "- Try cudf.read_parquet\n",
        "\n",
        "- Move input to a public readable S3 bucket\n",
        "\n",
        "- Save output to a S3 bucket\n",
        "\n",
        "- Fine tune settings for quadtree: `max_depth`, `min_size`.\n",
        "\n",
        "- Fine tune the number of random points `max_len` generated at each iteration."
      ],
      "metadata": {
        "id": "ptQK6C9Ao_74"
      }
    }
  ]
}