src/1_feature_generation/0.1_download_satellite_colab.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimized Downloader for Google Colab\n",
    "Basically the same as `0_download_satellite.ipynb`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "from collections.abc import Mapping\n",
    "from typing import Any, Optional\n",
    "\n",
    "import ee\n",
    "import pandas as pd\n",
    "import time\n",
    "from tqdm.auto import tqdm\n",
    "\n",
    "# taken from https://github.com/sustainlab-group/africa_poverty (slightly modified)\n",
    "\n",
    "def df_to_fc(df: pd.DataFrame, lat_colname: str = 'lat',\n",
    "             lon_colname: str = 'lon') -> ee.FeatureCollection:\n",
    "    '''Create a ee.FeatureCollection from a pd.DataFrame.\n",
    "\n",
    "    Args\n",
    "    - csv_path: str, path to CSV file that includes at least two columns for\n",
    "        latitude and longitude coordinates\n",
    "    - lat_colname: str, name of latitude column\n",
    "    - lon_colname: str, name of longitude column\n",
    "\n",
    "    Returns: ee.FeatureCollection, contains one feature per row in the CSV file\n",
    "    '''\n",
    "    # convert values to Python native types\n",
    "    # see https://stackoverflow.com/a/47424340\n",
    "    df = df.astype('object')\n",
    "\n",
    "    ee_features = []\n",
    "    for i in range(len(df)):\n",
    "        props = df.iloc[i].to_dict()\n",
    "\n",
    "        # oddly EE wants (lon, lat) instead of (lat, lon)\n",
    "        _geometry = ee.Geometry.Point([\n",
    "            props[lon_colname],\n",
    "            props[lat_colname],\n",
    "        ])\n",
    "        ee_feat = ee.Feature(_geometry, props)\n",
    "        ee_features.append(ee_feat)\n",
    "\n",
    "    return ee.FeatureCollection(ee_features)\n",
    "\n",
    "\n",
    "def decode_qamask(img: ee.Image) -> ee.Image:\n",
    "    '''\n",
    "    Args\n",
    "    - img: ee.Image, Landsat 5/7/8 image containing 'pixel_qa' band\n",
    "\n",
    "    Returns\n",
    "    - masks: ee.Image, contains 5 bands of masks\n",
    "\n",
    "    Pixel QA Bit Flags (universal across Landsat 5/7/8)\n",
    "    Bit  Attribute\n",
    "    0    Fill\n",
    "    1    Clear\n",
    "    2    Water\n",
    "    3    Cloud Shadow\n",
    "    4    Snow\n",
    "    5    Cloud\n",
    "    '''\n",
    "    qa = img.select('pixel_qa')\n",
    "    clear = qa.bitwiseAnd(2).neq(0)  # 0 = not clear, 1 = clear\n",
    "    clear = clear.updateMask(clear).rename(['pxqa_clear'])\n",
    "\n",
    "    water = qa.bitwiseAnd(4).neq(0)  # 0 = not water, 1 = water\n",
    "    water = water.updateMask(water).rename(['pxqa_water'])\n",
    "\n",
    "    cloud_shadow = qa.bitwiseAnd(8).eq(0)  # 0 = shadow, 1 = not shadow\n",
    "    cloud_shadow = cloud_shadow.updateMask(cloud_shadow).rename(['pxqa_cloudshadow'])\n",
    "\n",
    "    snow = qa.bitwiseAnd(16).eq(0)  # 0 = snow, 1 = not snow\n",
    "    snow = snow.updateMask(snow).rename(['pxqa_snow'])\n",
    "\n",
    "    cloud = qa.bitwiseAnd(32).eq(0)  # 0 = cloud, 1 = not cloud\n",
    "    cloud = cloud.updateMask(cloud).rename(['pxqa_cloud'])\n",
    "\n",
    "    masks = ee.Image.cat([clear, water, cloud_shadow, snow, cloud])\n",
    "    return masks\n",
    "\n",
    "\n",
    "def mask_qaclear(img: ee.Image) -> ee.Image:\n",
    "    '''\n",
    "    Args\n",
    "    - img: ee.Image, Landsat 5/7/8 image containing 'pixel_qa' band\n",
    "\n",
    "    Returns\n",
    "    - img: ee.Image, input image with cloud-shadow, snow, cloud, and unclear\n",
    "        pixels masked out\n",
    "    '''\n",
    "    qam = decode_qamask(img)\n",
    "    cloudshadow_mask = qam.select('pxqa_cloudshadow')\n",
    "    snow_mask = qam.select('pxqa_snow')\n",
    "    cloud_mask = qam.select('pxqa_cloud')\n",
    "    return img.updateMask(cloudshadow_mask).updateMask(snow_mask).updateMask(cloud_mask)\n",
    "\n",
    "\n",
    "def add_latlon(img: ee.Image) -> ee.Image:\n",
    "    '''Creates a new ee.Image with 2 added bands of longitude and latitude\n",
    "    coordinates named 'LON' and 'LAT', respectively\n",
    "    '''\n",
    "    latlon = ee.Image.pixelLonLat().select(\n",
    "        opt_selectors=['longitude', 'latitude'],\n",
    "        opt_names=['LON', 'LAT'])\n",
    "    return img.addBands(latlon)\n",
    "\n",
    "\n",
    "def composite_nl(year: int) -> ee.Image:\n",
    "    '''Creates a median-composite nightlights (NL) image.\n",
    "\n",
    "    Args\n",
    "    - year: int, start year of survey\n",
    "\n",
    "    Returns: ee.Image, contains a single band named 'NIGHTLIGHTS'\n",
    "    '''\n",
    "    if year <= 2013:\n",
    "        img_col = ee.ImageCollection('NOAA/DMSP-OLS/NIGHTTIME_LIGHTS')\n",
    "        \n",
    "    else:\n",
    "        img_col = ee.ImageCollection('NOAA/VIIRS/DNB/MONTHLY_V1/VCMSLCFG')\n",
    "    \n",
    "    start_date, end_date = f'{year}-01-01', f'{year}-12-31'\n",
    "    return img_col.filterDate(start_date, end_date).median().select([0], ['NIGHTLIGHTS'])\n",
    "\n",
    "\n",
    "def tfexporter(collection: ee.FeatureCollection, export: str, prefix: str,\n",
    "               fname: str, selectors: Optional[ee.List] = None,\n",
    "               dropselectors: Optional[ee.List] = None,\n",
    "               bucket: Optional[str] = None) -> ee.batch.Task:\n",
    "    '''Creates and starts a task to export a ee.FeatureCollection to a TFRecord\n",
    "    file in Google Drive or Google Cloud Storage (GCS).\n",
    "\n",
    "    GCS:   gs://bucket/prefix/fname.tfrecord\n",
    "    Drive: prefix/fname.tfrecord\n",
    "\n",
    "    Args\n",
    "    - collection: ee.FeatureCollection\n",
    "    - export: str, 'drive' for Drive, 'gcs' for GCS\n",
    "    - prefix: str, folder name in Drive or GCS to export to, no trailing '/'\n",
    "    - fname: str, filename\n",
    "    - selectors: None or ee.List of str, names of properties to include in\n",
    "        output, set to None to include all properties\n",
    "    - dropselectors: None or ee.List of str, names of properties to exclude\n",
    "    - bucket: None or str, name of GCS bucket, only used if export=='gcs'\n",
    "\n",
    "    Returns\n",
    "    - task: ee.batch.Task\n",
    "    '''\n",
    "    if dropselectors is not None:\n",
    "        if selectors is None:\n",
    "            selectors = collection.first().propertyNames()\n",
    "\n",
    "        selectors = selectors.removeAll(dropselectors)\n",
    "\n",
    "    if export == 'gcs':\n",
    "        task = ee.batch.Export.table.toCloudStorage(\n",
    "            collection=collection,\n",
    "            description=fname,\n",
    "            bucket=bucket,\n",
    "            fileNamePrefix=f'{prefix}/{fname}',\n",
    "            fileFormat='TFRecord',\n",
    "            selectors=selectors)\n",
    "\n",
    "    elif export == 'drive':\n",
    "        task = ee.batch.Export.table.toDrive(\n",
    "            collection=collection,\n",
    "            description=fname,\n",
    "            folder=prefix,\n",
    "            fileNamePrefix=fname,\n",
    "            fileFormat='TFRecord',\n",
    "            selectors=selectors)\n",
    "\n",
    "    else:\n",
    "        raise ValueError(f'export \"{export}\" is not one of [\"gcs\", \"drive\"]')\n",
    "\n",
    "    task.start()\n",
    "    return task\n",
    "\n",
    "\n",
    "def sample_patch(point: ee.Feature, patches_array: ee.Image,\n",
    "                 scale: float) -> ee.Feature:\n",
    "    '''Extracts an image patch at a specific point.\n",
    "\n",
    "    Args\n",
    "    - point: ee.Feature\n",
    "    - patches_array: ee.Image, Array Image\n",
    "    - scale: int or float, scale in meters of the projection to sample in\n",
    "\n",
    "    Returns: ee.Feature, 1 property per band from the input image\n",
    "    '''\n",
    "    arrays_samples = patches_array.sample(\n",
    "        region=point.geometry(),\n",
    "        scale=scale,\n",
    "        projection='EPSG:3857',\n",
    "        factor=None,\n",
    "        numPixels=None,\n",
    "        dropNulls=False,\n",
    "        tileScale=12)\n",
    "    return arrays_samples.first().copyProperties(point)\n",
    "\n",
    "\n",
    "def get_array_patches(img: ee.Image,\n",
    "                      scale: float,\n",
    "                      ksize: float,\n",
    "                      points: ee.FeatureCollection,\n",
    "                      export: str,\n",
    "                      prefix: str,\n",
    "                      fname: str,\n",
    "                      selectors: Optional[ee.List] = None,\n",
    "                      dropselectors: Optional[ee.List] = None,\n",
    "                      bucket: Optional[str] = None\n",
    "                      ) -> ee.batch.Task:\n",
    "    '''Creates and starts a task to export square image patches in TFRecord\n",
    "    format to Google Drive or Google Cloud Storage (GCS). The image patches are\n",
    "    sampled from the given ee.Image at specific coordinates.\n",
    "\n",
    "    Args\n",
    "    - img: ee.Image, image covering the entire region of interest\n",
    "    - scale: int or float, scale in meters of the projection to sample in\n",
    "    - ksize: int or float, radius of square image patch\n",
    "    - points: ee.FeatureCollection, coordinates from which to sample patches\n",
    "    - export: str, 'drive' for Google Drive, 'gcs' for GCS\n",
    "    - prefix: str, folder name in Drive or GCS to export to, no trailing '/'\n",
    "    - fname: str, filename for export\n",
    "    - selectors: None or ee.List, names of properties to include in output,\n",
    "        set to None to include all properties\n",
    "    - dropselectors: None or ee.List, names of properties to exclude\n",
    "    - bucket: None or str, name of GCS bucket, only used if export=='gcs'\n",
    "\n",
    "    Returns: ee.batch.Task\n",
    "    '''\n",
    "    kern = ee.Kernel.square(radius=ksize, units='pixels')\n",
    "    patches_array = img.neighborhoodToArray(kern)\n",
    "\n",
    "    # ee.Image.sampleRegions() does not cut it for larger collections,\n",
    "    # using mapped sample instead\n",
    "    samples = points.map(lambda pt: sample_patch(pt, patches_array, scale))\n",
    "\n",
    "    # export to a TFRecord file which can be loaded directly in TensorFlow\n",
    "    return tfexporter(collection=samples, export=export, prefix=prefix,\n",
    "                      fname=fname, selectors=selectors,\n",
    "                      dropselectors=dropselectors, bucket=bucket)\n",
    "\n",
    "\n",
    "def wait_on_tasks(tasks: Mapping[Any, ee.batch.Task],\n",
    "                  show_probar: bool = True,\n",
    "                  poll_interval: int = 20,\n",
    "                  ) -> None:\n",
    "    '''Displays a progress bar of task progress.\n",
    "\n",
    "    Args\n",
    "    - tasks: dict, maps task ID to a ee.batch.Task\n",
    "    - show_progbar: bool, whether to display progress bar\n",
    "    - poll_interval: int, # of seconds between each refresh\n",
    "    '''\n",
    "    remaining_tasks = list(tasks.keys())\n",
    "    done_states = {ee.batch.Task.State.COMPLETED,\n",
    "                   ee.batch.Task.State.FAILED,\n",
    "                   ee.batch.Task.State.CANCEL_REQUESTED,\n",
    "                   ee.batch.Task.State.CANCELLED}\n",
    "\n",
    "    progbar = tqdm(total=len(remaining_tasks))\n",
    "    while len(remaining_tasks) > 0:\n",
    "        new_remaining_tasks = []\n",
    "        for taskID in remaining_tasks:\n",
    "            status = tasks[taskID].status()\n",
    "            state = status['state']\n",
    "\n",
    "            if state in done_states:\n",
    "                progbar.update(1)\n",
    "\n",
    "                if state == ee.batch.Task.State.FAILED:\n",
    "                    state = (state, status['error_message'])\n",
    "                elapsed_ms = status['update_timestamp_ms'] - status['creation_timestamp_ms']\n",
    "                elapsed_min = int((elapsed_ms / 1000) / 60)\n",
    "                progbar.write(f'Task {taskID} finished in {elapsed_min} min with state: {state}')\n",
    "            else:\n",
    "                new_remaining_tasks.append(taskID)\n",
    "        remaining_tasks = new_remaining_tasks\n",
    "        time.sleep(poll_interval)\n",
    "    progbar.close()\n",
    "\n",
    "\n",
    "class LandsatSR:\n",
    "    def __init__(self, filterpoly: ee.Geometry, start_date: str,\n",
    "                 end_date: str) -> None:\n",
    "        '''\n",
    "        Args\n",
    "        - filterpoly: ee.Geometry\n",
    "        - start_date: str, string representation of start date\n",
    "        - end_date: str, string representation of end date\n",
    "        '''\n",
    "        self.filterpoly = filterpoly\n",
    "        self.start_date = start_date\n",
    "        self.end_date = end_date\n",
    "\n",
    "        self.l8 = self.init_coll('LANDSAT/LC08/C01/T1_SR').map(self.rename_l8).map(self.rescale_l8)\n",
    "        self.l7 = self.init_coll('LANDSAT/LE07/C01/T1_SR').map(self.rename_l57).map(self.rescale_l57)\n",
    "        self.l5 = self.init_coll('LANDSAT/LT05/C01/T1_SR').map(self.rename_l57).map(self.rescale_l57)\n",
    "\n",
    "        self.merged = self.l5.merge(self.l7).merge(self.l8).sort('system:time_start')\n",
    "\n",
    "    def init_coll(self, name: str) -> ee.ImageCollection:\n",
    "        '''\n",
    "        Creates a ee.ImageCollection containing images of desired points\n",
    "        between the desired start and end dates.\n",
    "        Args\n",
    "        - name: str, name of collection\n",
    "        Returns: ee.ImageCollection\n",
    "        '''\n",
    "        return (ee.ImageCollection(name)\n",
    "                .filterBounds(self.filterpoly)\n",
    "                .filterDate(self.start_date, self.end_date))\n",
    "\n",
    "    @staticmethod\n",
    "    def rename_l8(img: ee.Image) -> ee.Image:\n",
    "        '''\n",
    "        Args\n",
    "        - img: ee.Image, Landsat 8 image\n",
    "        Returns\n",
    "        - img: ee.Image, with bands renamed\n",
    "        See: https://developers.google.com/earth-engine/datasets/catalog/LANDSAT_LC08_C01_T1_SR\n",
    "        Name       Scale Factor Description\n",
    "        B1         0.0001       Band 1 (Ultra Blue) surface reflectance, 0.435-0.451 um\n",
    "        B2         0.0001       Band 2 (Blue) surface reflectance, 0.452-0.512 um\n",
    "        B3         0.0001       Band 3 (Green) surface reflectance, 0.533-0.590 um\n",
    "        B4         0.0001       Band 4 (Red) surface reflectance, 0.636-0.673 um\n",
    "        B5         0.0001       Band 5 (Near Infrared) surface reflectance, 0.851-0.879 um\n",
    "        B6         0.0001       Band 6 (Shortwave Infrared 1) surface reflectance, 1.566-1.651 um\n",
    "        B7         0.0001       Band 7 (Shortwave Infrared 2) surface reflectance, 2.107-2.294 um\n",
    "        B10        0.1          Band 10 brightness temperature (Kelvin), 10.60-11.19 um\n",
    "        B11        0.1          Band 11 brightness temperature (Kelvin), 11.50-12.51 um\n",
    "        sr_aerosol              Aerosol attributes, see Aerosol QA table\n",
    "        pixel_qa                Pixel quality attributes, see Pixel QA table\n",
    "        radsat_qa               Radiometric saturation QA, see Radsat QA table\n",
    "        '''\n",
    "        newnames = ['AEROS', 'BLUE', 'GREEN', 'RED', 'NIR', 'SWIR1', 'SWIR2',\n",
    "                    'TEMP1', 'TEMP2', 'sr_aerosol', 'pixel_qa', 'radsat_qa']\n",
    "        return img.rename(newnames)\n",
    "\n",
    "    @staticmethod\n",
    "    def rescale_l8(img: ee.Image) -> ee.Image:\n",
    "        '''\n",
    "        Args\n",
    "        - img: ee.Image, Landsat 8 image, with bands already renamed\n",
    "            by rename_l8()\n",
    "        Returns\n",
    "        - img: ee.Image, with bands rescaled\n",
    "        '''\n",
    "        opt = img.select(['AEROS', 'BLUE', 'GREEN', 'RED', 'NIR', 'SWIR1', 'SWIR2'])\n",
    "        therm = img.select(['TEMP1', 'TEMP2'])\n",
    "        masks = img.select(['sr_aerosol', 'pixel_qa', 'radsat_qa'])\n",
    "\n",
    "        opt = opt.multiply(0.0001)\n",
    "        therm = therm.multiply(0.1)\n",
    "\n",
    "        scaled = ee.Image.cat([opt, therm, masks]).copyProperties(img)\n",
    "        # system properties are not copied\n",
    "        scaled = scaled.set('system:time_start', img.get('system:time_start'))\n",
    "        return scaled\n",
    "\n",
    "    @staticmethod\n",
    "    def rename_l57(img: ee.Image) -> ee.Image:\n",
    "        '''\n",
    "        Args\n",
    "        - img: ee.Image, Landsat 5/7 image\n",
    "        Returns\n",
    "        - img: ee.Image, with bands renamed\n",
    "        See: https://developers.google.com/earth-engine/datasets/catalog/LANDSAT_LT05_C01_T1_SR\n",
    "             https://developers.google.com/earth-engine/datasets/catalog/LANDSAT_LE07_C01_T1_SR\n",
    "        Name             Scale Factor Description\n",
    "        B1               0.0001       Band 1 (blue) surface reflectance, 0.45-0.52 um\n",
    "        B2               0.0001       Band 2 (green) surface reflectance, 0.52-0.60 um\n",
    "        B3               0.0001       Band 3 (red) surface reflectance, 0.63-0.69 um\n",
    "        B4               0.0001       Band 4 (near infrared) surface reflectance, 0.77-0.90 um\n",
    "        B5               0.0001       Band 5 (shortwave infrared 1) surface reflectance, 1.55-1.75 um\n",
    "        B6               0.1          Band 6 brightness temperature (Kelvin), 10.40-12.50 um\n",
    "        B7               0.0001       Band 7 (shortwave infrared 2) surface reflectance, 2.08-2.35 um\n",
    "        sr_atmos_opacity 0.001        Atmospheric opacity; < 0.1 = clear; 0.1 - 0.3 = average; > 0.3 = hazy\n",
    "        sr_cloud_qa                   Cloud quality attributes, see SR Cloud QA table. Note:\n",
    "                                          pixel_qa is likely to present more accurate results\n",
    "                                          than sr_cloud_qa for cloud masking. See page 14 in\n",
    "                                          the LEDAPS product guide.\n",
    "        pixel_qa                      Pixel quality attributes generated from the CFMASK algorithm,\n",
    "                                          see Pixel QA table\n",
    "        radsat_qa                     Radiometric saturation QA, see Radiometric Saturation QA table\n",
    "        '''\n",
    "        newnames = ['BLUE', 'GREEN', 'RED', 'NIR', 'SWIR1', 'TEMP1', 'SWIR2',\n",
    "                    'sr_atmos_opacity', 'sr_cloud_qa', 'pixel_qa', 'radsat_qa']\n",
    "        return img.rename(newnames)\n",
    "\n",
    "    @staticmethod\n",
    "    def rescale_l57(img: ee.Image) -> ee.Image:\n",
    "        '''\n",
    "        Args\n",
    "        - img: ee.Image, Landsat 5/7 image, with bands already renamed\n",
    "            by rename_157()\n",
    "        Returns\n",
    "        - img: ee.Image, with bands rescaled\n",
    "        '''\n",
    "        opt = img.select(['BLUE', 'GREEN', 'RED', 'NIR', 'SWIR1', 'SWIR2'])\n",
    "        atmos = img.select(['sr_atmos_opacity'])\n",
    "        therm = img.select(['TEMP1'])\n",
    "        masks = img.select(['sr_cloud_qa', 'pixel_qa', 'radsat_qa'])\n",
    "\n",
    "        opt = opt.multiply(0.0001)\n",
    "        atmos = atmos.multiply(0.001)\n",
    "        therm = therm.multiply(0.1)\n",
    "\n",
    "        scaled = ee.Image.cat([opt, therm, masks, atmos]).copyProperties(img)\n",
    "        # system properties are not copied\n",
    "        scaled = scaled.set('system:time_start', img.get('system:time_start'))\n",
    "        return scaled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ee\n",
    "import math\n",
    "import pandas as pd\n",
    "\n",
    "from lib import satellite_utils\n",
    "from __future__ import annotations\n",
    "from typing import Any, Optional\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ee.Authenticate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ee.Initialize()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EXPORT = 'drive'\n",
    "BUCKET = None\n",
    "\n",
    "LSMS_EXPORT_FOLDER = 'tfrecords_raw' # defined where to store in your drive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LSMS_CSV_PATH = '' # path to the csv file in the drive\n",
    "\n",
    "# band names\n",
    "MS_BANDS = ['BLUE', 'GREEN', 'RED', 'NIR', 'SWIR1', 'SWIR2', 'TEMP1']\n",
    "\n",
    "# image parameters\n",
    "PROJECTION = 'EPSG:3857'  # see https://epsg.io/3857\n",
    "SCALE = 30                # export resolution: 30m/px\n",
    "EXPORT_TILE_RADIUS = 127  # image dimension = (2*EXPORT_TILE_RADIUS) + 1 = 255px\n",
    "\n",
    "CHUNK_SIZE = None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def export_images(df: pd.DataFrame,\n",
    "                  country: str,\n",
    "                  year: int,\n",
    "                  export_folder: str,\n",
    "                  chunk_size: Optional[int] = None\n",
    "                  ) -> dict[tuple[str, str, int, int], ee.batch.Task]:\n",
    "    '''\n",
    "    Args\n",
    "    - df: pd.DataFrame, contains columns ['lat', 'lon', 'country', 'year']\n",
    "    - country: str, together with `year` determines the survey to export\n",
    "    - year: int, together with `country` determines the survey to export\n",
    "    - export_folder: str, name of folder for export\n",
    "    - chunk_size: int, optionally set a limit to the # of images exported per TFRecord file\n",
    "        - set to a small number (<= 50) if Google Earth Engine reports memory errors\n",
    "\n",
    "    Returns: dict, maps task name tuple (export_folder, country, year, chunk) to ee.batch.Task\n",
    "    '''\n",
    "    subset_df = df[(df['country'] == country) & (df['year'] == year)].reset_index(drop=True)\n",
    "    if chunk_size is None:\n",
    "        chunk_size = len(subset_df)\n",
    "    num_chunks = int(math.ceil(len(subset_df) / chunk_size))\n",
    "    tasks = {}\n",
    "\n",
    "    for i in range(num_chunks):\n",
    "        chunk_slice = slice(i * chunk_size, (i+1) * chunk_size - 1)  # df.loc[] is inclusive\n",
    "        fc = satellite_utils.df_to_fc(subset_df.loc[chunk_slice, :])\n",
    "        start_date, end_date = f'{year}-01-01', f'{year}-12-31'\n",
    "        \n",
    "        roi = fc.geometry()\n",
    "        imgcol = satellite_utils.LandsatSR(roi, start_date=start_date, end_date=end_date).merged\n",
    "        imgcol = imgcol.map(satellite_utils.mask_qaclear).select(MS_BANDS)\n",
    "        img = imgcol.median()\n",
    "\n",
    "        # add nightlights, latitude, and longitude bands\n",
    "        img = satellite_utils.add_latlon(img)\n",
    "        img = img.addBands(satellite_utils.composite_nl(year))\n",
    "\n",
    "        fname = f'{country}_{year}_{i:02d}'\n",
    "        tasks[(export_folder, country, year, i)] = satellite_utils.get_array_patches(\n",
    "            img=img, scale=SCALE, ksize=EXPORT_TILE_RADIUS,\n",
    "            points=fc, export=EXPORT,\n",
    "            prefix=export_folder, fname=fname,\n",
    "            bucket=BUCKET)\n",
    "    return tasks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tasks: dict[tuple[str, str, int, int], ee.batch.Task] = {}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lsms_df = pd.read_csv(LSMS_CSV_PATH, float_precision='high', index_col=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lsms_surveys = list(lsms_df.groupby(['country', 'year']).groups.keys())\n",
    "for country, year in lsms_surveys:\n",
    "    new_tasks = export_images(\n",
    "        df=lsms_df, country=country, year=year,\n",
    "        export_folder=LSMS_EXPORT_FOLDER, chunk_size=CHUNK_SIZE)\n",
    "    tasks.update(new_tasks)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}