people-and-planet-ai/weather-forecasting/serving/weather-data/weather/data.py

# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Data utilities to grab data from Earth Engine.
Meant to be used for both training and prediction so the model is
trained on exactly the same data that will be used for predictions.
"""

from __future__ import annotations

from datetime import datetime, timedelta
import io

import ee
from google.api_core import exceptions, retry
import google.auth
import numpy as np
from numpy.lib.recfunctions import structured_to_unstructured
import requests

# Constants.
SCALE = 10000  # meters per pixel
INPUT_HOUR_DELTAS = [-4, -2, 0]
OUTPUT_HOUR_DELTAS = [2, 6]
WINDOW = timedelta(days=1)

# Authenticate and initialize Earth Engine with the default credentials.
credentials, project = google.auth.default(
    scopes=[
        "https://www.googleapis.com/auth/cloud-platform",
        "https://www.googleapis.com/auth/earthengine",
    ]
)

# Use the Earth Engine High Volume endpoint.
#   https://developers.google.com/earth-engine/cloud/highvolume
ee.Initialize(
    credentials.with_quota_project(None),
    project=project,
    opt_url="https://earthengine-highvolume.googleapis.com",
)


def get_gpm(date: datetime) -> ee.Image:
    """Gets a Global Precipitation Measurement image for the selected date.

    For more information:
        https://developers.google.com/earth-engine/datasets/catalog/NASA_GPM_L3_IMERG_V06

    Args:
        date: Date to take a snapshot from.

    Returns: An Earth Engine image.
    """
    window_start = (date - WINDOW).isoformat()
    window_end = date.isoformat()
    return (
        ee.ImageCollection("NASA/GPM_L3/IMERG_V06")
        .filterDate(window_start, window_end)
        .select("precipitationCal")
        .sort("system:time_start", False)
        .mosaic()
        .unmask(0)
        .float()
    )


def get_gpm_sequence(dates: list[datetime]) -> ee.Image:
    """Gets a Global Precipitation Measurement sequence for the selected dates.

    Args:
        dates: List of dates to get images from.

    Returns: An Earth Engine image.
    """
    images = [get_gpm(date) for date in dates]
    return ee.ImageCollection(images).toBands()


def get_goes16(date: datetime) -> ee.Image:
    """Gets a GOES 16 image for the selected date.

    For more information:
        https://developers.google.com/earth-engine/datasets/catalog/NOAA_GOES_16_MCMIPF

    Args:
        date: Date to take a snapshot from.

    Returns: An Earth Engine image.
    """
    window_start = (date - WINDOW).isoformat()
    window_end = date.isoformat()
    return (
        ee.ImageCollection("NOAA/GOES/16/MCMIPF")
        .filterDate(window_start, window_end)
        .select("CMI_C.*")
        .sort("system:time_start", False)
        .mosaic()
        .unmask(0)
        .float()
    )


def get_goes16_sequence(dates: list[datetime]) -> ee.Image:
    """Gets a GOES 16 sequence for the selected dates.

    Args:
        dates: List of dates to get images from.

    Returns: An Earth Engine image.
    """
    images = [get_goes16(date) for date in dates]
    return ee.ImageCollection(images).toBands()


def get_elevation() -> ee.Image:
    """Gets a digital elevation map.

    For more information:
        https://developers.google.com/earth-engine/datasets/catalog/MERIT_DEM_v1_0_3

    Returns: An Earth Engine image.
    """
    return ee.Image("MERIT/DEM/v1_0_3").rename("elevation").unmask(0).float()


def get_inputs_image(date: datetime) -> ee.Image:
    """Gets an Earth Engine image with all the inputs for the model.

    Args:
        date: Date to take a snapshot from.

    Returns: An Earth Engine image.
    """
    dates = [date + timedelta(hours=h) for h in INPUT_HOUR_DELTAS]
    precipitation = get_gpm_sequence(dates)
    cloud_and_moisture = get_goes16_sequence(dates)
    elevation = get_elevation()
    return ee.Image([precipitation, cloud_and_moisture, elevation])


def get_labels_image(date: datetime) -> ee.Image:
    """Gets an Earth Engine image with the labels to train the model.

    Args:
        date: Date to take a snapshot from.

    Returns: An Earth Engine image.
    """
    dates = [date + timedelta(hours=h) for h in OUTPUT_HOUR_DELTAS]
    return get_gpm_sequence(dates)


def get_inputs_patch(date: datetime, point: tuple, patch_size: int) -> np.ndarray:
    """Gets the patch of pixels for the inputs.

    Args:
        date: The date of interest.
        point: A (longitude, latitude) coordinate.
        patch_size: Size in pixels of the surrounding square patch.

    Returns: The pixel values of a patch as a NumPy array.
    """
    image = get_inputs_image(date)
    patch = get_patch(image, point, patch_size, SCALE)
    return structured_to_unstructured(patch)


def get_labels_patch(date: datetime, point: tuple, patch_size: int) -> np.ndarray:
    """Gets the patch of pixels for the labels.

    Args:
        date: The date of interest.
        point: A (longitude, latitude) coordinate.
        patch_size: Size in pixels of the surrounding square patch.

    Returns: The pixel values of a patch as a NumPy array.
    """
    image = get_labels_image(date)
    patch = get_patch(image, point, patch_size, SCALE)
    return structured_to_unstructured(patch)


@retry.Retry()
def get_patch(image: ee.Image, point: tuple, patch_size: int, scale: int) -> np.ndarray:
    """Fetches a patch of pixels from Earth Engine.

    It retries if we get error "429: Too Many Requests".

    Args:
        image: Image to get the patch from.
        point: A (longitude, latitude) pair for the point of interest.
        patch_size: Size in pixels of the surrounding square patch.
        scale: Number of meters per pixel.

    Raises:
        requests.exceptions.RequestException

    Returns:
        The requested patch of pixels as a structured
        NumPy array with shape (width, height).
    """
    geometry = ee.Geometry.Point(point)
    url = image.getDownloadURL(
        {
            "region": geometry.buffer(scale * patch_size / 2, 1).bounds(1),
            "dimensions": [patch_size, patch_size],
            "format": "NPY",
        }
    )

    # If we get "429: Too Many Requests" errors, it's safe to retry the request.
    # The Retry library only works with `google.api_core` exceptions.
    response = requests.get(url)
    if response.status_code == 429:
        raise exceptions.TooManyRequests(response.text)

    # Still raise any other exceptions to make sure we got valid data.
    response.raise_for_status()
    return np.load(io.BytesIO(response.content), allow_pickle=True)