pycontrails 0.55.0__cp313-cp313-macosx_11_0_arm64.whl → 0.57.0__cp313-cp313-macosx_11_0_arm64.whl

pycontrails/datalib/landsat.py

@@ -1,16 +1,17 @@
-"""Support for LANDSAT 8-9 imagery retrieval through Google Cloud Platform."""
+"""Support for Landsat 8 Collection 1 imagery retrieval through Google Cloud Platform."""
 
 from __future__ import annotations
 
 from collections.abc import Iterable
 
 import numpy as np
+import numpy.typing as npt
 import pandas as pd
 import xarray as xr
 
 from pycontrails.core import Flight, cache
-from pycontrails.datalib._leo_utils import search
-from pycontrails.datalib._leo_utils.vis import equalize, normalize
+from pycontrails.datalib.leo_utils import search
+from pycontrails.datalib.leo_utils.vis import equalize, normalize
 from pycontrails.utils import dependencies
 
 try:
@@ -60,24 +61,28 @@ def query(
     extent: str | None = None,
     columns: list[str] | None = None,
 ) -> pd.DataFrame:
-    """Find Landsat 8 and 9 imagery within spatiotemporal region of interest.
+    """Find Landsat 8 Collection 1 imagery within spatiotemporal region of interest.
 
     This function requires access to the
     `Google BigQuery API <https://cloud.google.com/bigquery?hl=en>`__
     and uses the `BigQuery python library <https://cloud.google.com/python/docs/reference/bigquery/latest/index.html>`__.
 
+    See :func:`pycontrails.datalib.leo_utils.landsat_metadata.open_landsat_metadata`
+    to download and parse the daily bulk Landsat metadata CSV file from USGS. This CSV holds
+    Collection 2 metadata, so includes the most recent scenes from Landsat 8 and 9.
+
     Parameters
     ----------
     start_time : np.datetime64
         Start of time period for search
     end_time : np.datetime64
         End of time period for search
-    extent : str, optional
+    extent : str | None, optional
         Spatial region of interest as a GeoJSON string. If not provided, defaults
         to a global extent.
-    columns : list[str], optional.
+    columns : list[str] | None, optional
         Columns to return from Google
-        `BigQuery table <https://console.cloud.google.com/bigquery?p=bigquery-public-data&d=cloud_storage_geo_index&t=landsat_index&page=table&_ga=2.90807450.1051800793.1716904050-255800408.1705955196>`__.
+        `BigQuery table <https://console.cloud.google.com/bigquery?p=bigquery-public-data&d=cloud_storage_geo_index&t=landsat_index&page=table>`__.
         By default, returns imagery base URL and sensing time.
 
     Returns
@@ -99,7 +104,7 @@ def intersect(
     flight: Flight,
     columns: list[str] | None = None,
 ) -> pd.DataFrame:
-    """Find Landsat 8 and 9 imagery intersecting with flight track.
+    """Find Landsat 8 Collection 1 imagery intersecting with flight track.
 
     This function will return all scenes with a bounding box that includes flight waypoints
     both before and after the sensing time.
@@ -108,13 +113,17 @@ def intersect(
     `Google BigQuery API <https://cloud.google.com/bigquery?hl=en>`__
     and uses the `BigQuery python library <https://cloud.google.com/python/docs/reference/bigquery/latest/index.html>`__.
 
+    See :func:`pycontrails.datalib.leo_utils.landsat_metadata.open_landsat_metadata`
+    to download and parse the daily bulk Landsat metadata CSV file from USGS. This CSV holds
+    Collection 2 metadata, so includes the most recent scenes from Landsat 8 and 9.
+
     Parameters
     ----------
     flight : Flight
         Flight for intersection
-    columns : list[str], optional.
+    columns : list[str] | None, optional
         Columns to return from Google
-        `BigQuery table <https://console.cloud.google.com/bigquery?p=bigquery-public-data&d=cloud_storage_geo_index&t=landsat_index&page=table&_ga=2.90807450.1051800793.1716904050-255800408.1705955196>`__.
+        `BigQuery table <https://console.cloud.google.com/bigquery?p=bigquery-public-data&d=cloud_storage_geo_index&t=landsat_index&page=table>`__.
         By default, returns imagery base URL and sensing time.
 
     Returns
@@ -131,20 +140,31 @@ def intersect(
 
 
 class Landsat:
-    """Support for Landsat 8 and 9 data handling.
+    """Support for Landsat 8 Collection 1 data handling.
+
+    This interface does not support Landsat Collection 2, which includes all new
+    scenes from Landsat 8 and 9 and benefits from improved calibration and processing
+    algorithms. The USGS stopped updating Collection 1 in 2021, so this interface
+    works only with legacy data. In addition, Collection 1 does not include viewing angle
+    data, preventing scan angle or sensing time corrections.
+
+    To access Landsat Collection 2 data, use one of the following tools:
+
+    - `USGS M2M API <https://m2m.cr.usgs.gov/>`__ (requires registration)
+    - `USGS Earth Explorer <https://earthexplorer.usgs.gov/>`__ (requires registration;
+       includes a web interface)
+    - `Amazon Web Services (AWS) <https://registry.opendata.aws/usgs-landsat/>`__ (requester pays)
+    - `Google Earth Engine <https://developers.google.com/earth-engine/datasets/catalog/landsat>`__
+      (requires registration; stricter usage limits than the other options)
 
-    This class uses the `PROJ <https://proj.org/en/9.4/index.html>`__ coordinate
-    transformation software through the
-    `pyproj <https://pyproj4.github.io/pyproj/stable/index.html>`__ python interface.
-    pyproj is installed as part of the ``sat`` set of optional dependencies
-    (``pip install pycontrails[sat]``), but PROJ must be installed manually.
+    These services are not yet integrated with pycontrails.
 
     Parameters
     ----------
     base_url : str
         Base URL of Landsat scene. To find URLs for Landsat scenes at
         specific locations and times, see :func:`query` and :func:`intersect`.
-    bands : str | set[str] | None
+    bands : str | Iterable[str] | None
         Set of bands to retrieve. The 11 possible bands are represented by
         the string "B1" to "B11". For the Google Landsat contrails color scheme,
         set ``bands=("B9", "B10", "B11")``. For the true color scheme, set
@@ -155,7 +175,7 @@ class Landsat:
         - B8: 15 m
         - B10, B11: 30 m (upsampled from true resolution of 100 m)
 
-    cachestore : cache.CacheStore, optional
+    cachestore : cache.CacheStore | None, optional
         Cache store for Landsat data. If None, a :class:`DiskCacheStore` is used.
 
     See Also
@@ -202,17 +222,19 @@ class Landsat:
 
         Parameters
         ----------
-        reflective : str = {"raw", "radiance", "reflectance"}, optional
+        reflective : str, optional
+            One of {"raw", "radiance", "reflectance"}.
             Whether to return raw values or rescaled radiances or reflectances for reflective bands.
             By default, return reflectances.
-        thermal : str = {"raw", "radiance", "brightness_temperature"}, optional
+        thermal : str, optional
+            One of {"raw", "radiance", "brightness_temperature"}.
             Whether to return raw values or rescaled radiances or brightness temperatures
             for thermal bands. By default, return brightness temperatures.
 
         Returns
         -------
-        xr.DataArray
-            DataArray of Landsat data.
+        xr.Dataset
+            Dataset of Landsat data.
         """
         if reflective not in ["raw", "radiance", "reflectance"]:
             msg = "reflective band processing must be one of ['raw', 'radiance', 'reflectance']"
@@ -423,14 +445,15 @@ def _read_image_coordinates(meta: str, band: str) -> tuple[np.ndarray, np.ndarra
 
 def extract_landsat_visualization(
     ds: xr.Dataset, color_scheme: str = "true"
-) -> tuple[np.ndarray, pyproj.CRS, tuple[float, float, float, float]]:
+) -> tuple[npt.NDArray[np.float32], pyproj.CRS, tuple[float, float, float, float]]:
     """Extract artifacts for visualizing Landsat data with the given color scheme.
 
     Parameters
     ----------
     ds : xr.Dataset
         Dataset of Landsat data as returned by :meth:`Landsat.get`.
-    color_scheme : str = {"true", "google_contrails"}
+    color_scheme : str, optional
+        One of {"true", "google_contrails"}.
         Color scheme to use for visualization. The true color scheme
         requires reflectances for bands B2, B3, and B4; and the
         `Google contrails color scheme <https://research.google/pubs/a-human-labeled-landsat-contrails-dataset>`__
@@ -442,7 +465,7 @@ def extract_landsat_visualization(
         3D RGB array of shape ``(height, width, 3)``.
     src_crs : pyproj.CRS
         Imagery projection
-    src_extent : tuple[float,float,float,float]
+    src_extent : tuple[float, float, float, float]
         Imagery extent in projected coordinates
 
     References
@@ -551,7 +574,7 @@ def to_google_contrails(ds: xr.Dataset) -> tuple[np.ndarray, pyproj.CRS]:
     red = ((signal - lower) / (upper - lower)).clip(0.0, 1.0)
 
     # green: cirrus band transmittance
-    signal = 1 - rc.values
+    signal = 1.0 - rc.values
     lower = 0.8
     upper = 1.0
     green = adapt(((signal - lower) / (upper - lower)).clip(0.0, 1.0))

pycontrails/datalib/leo_utils/__init__.py

@@ -0,0 +1,5 @@
+"""Tools for working with Sentinel-2 and Landsat data."""
+
+from pycontrails.datalib.leo_utils.correction import estimate_scan_time, scan_angle_correction
+
+__all__ = ["estimate_scan_time", "scan_angle_correction"]

pycontrails/datalib/leo_utils/correction.py

@@ -0,0 +1,266 @@
+"""Support for overlaying flight and contrail data on Landsat & Sentinel images."""
+
+from typing import Literal, overload
+
+import numpy as np
+import numpy.typing as npt
+import pandas as pd
+import pyproj
+import shapely
+import xarray as xr
+
+
+def ephemeris_ecef_to_utm(ephemeris_df: pd.DataFrame, utm_crs: pyproj.CRS) -> pd.DataFrame:
+    """Convert ephemeris data from ECEF to UTM coordinates.
+
+    Parameters
+    ----------
+    ephemeris_df : pd.DataFrame
+        DataFrame containing the ephemeris data with columns:
+        - 'EPHEMERIS_ECEF_X': ECEF X coordinates (meters)
+        - 'EPHEMERIS_ECEF_Y': ECEF Y coordinates (meters)
+        - 'EPHEMERIS_ECEF_Z': ECEF Z coordinates (meters)
+        - 'EPHEMERIS_TIME': Timestamps (as datetime64[ns])
+    utm_crs : pyproj.CRS
+        The UTM coordinate reference system to convert to.
+
+    Returns
+    -------
+    pd.DataFrame
+        A DataFrame with columns:
+        - 'x': UTM easting (meters)
+        - 'y': UTM northing (meters)
+        - 'z': Altitude (meters)
+        - 't': Timestamps (as datetime64[ns])
+    """
+    # Define the source CRS: ECEF (Earth-Centered, Earth-Fixed) with WGS84 datum
+    source_crs = pyproj.CRS(proj="geocent", datum="WGS84")
+
+    # Create a transformer object to convert from source CRS to target CRS
+    # The default order for ECEF is (X, Y, Z) and for UTM is (Easting, Northing, Height)
+    transformer = pyproj.Transformer.from_crs(source_crs, utm_crs)
+
+    ecef_x = ephemeris_df["EPHEMERIS_ECEF_X"].to_numpy()
+    ecef_y = ephemeris_df["EPHEMERIS_ECEF_Y"].to_numpy()
+    ecef_z = ephemeris_df["EPHEMERIS_ECEF_Z"].to_numpy()
+    ecef_t = ephemeris_df["EPHEMERIS_TIME"].to_numpy()
+
+    x, y, h = transformer.transform(ecef_x, ecef_y, ecef_z)
+    return pd.DataFrame({"x": x, "y": y, "z": h, "t": ecef_t})
+
+
+@overload
+def scan_angle_correction(
+    ds: xr.Dataset,
+    x: npt.NDArray[np.floating],
+    y: npt.NDArray[np.floating],
+    z: npt.NDArray[np.floating],
+    *,
+    maxiter: int = ...,
+    tol: float = ...,
+    full_output: Literal[False] = ...,
+) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]: ...
+
+
+@overload
+def scan_angle_correction(
+    ds: xr.Dataset,
+    x: npt.NDArray[np.floating],
+    y: npt.NDArray[np.floating],
+    z: npt.NDArray[np.floating],
+    *,
+    maxiter: int = ...,
+    tol: float = ...,
+    full_output: Literal[True],
+) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.bool_]]: ...
+
+
+def scan_angle_correction(
+    ds: xr.Dataset,
+    x: npt.NDArray[np.floating],
+    y: npt.NDArray[np.floating],
+    z: npt.NDArray[np.floating],
+    *,
+    maxiter: int = 5,
+    tol: float = 10.0,
+    full_output: bool = False,
+) -> (
+    tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]
+    | tuple[npt.NDArray[np.floating], npt.NDArray[np.floating], npt.NDArray[np.bool_]]
+):
+    """Apply the scan angle correction to the given x, y, z coordinates.
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The dataset containing the viewing azimuth angle (VAA)
+        and viewing zenith angle (VZA) arrays. The units for both are degrees.
+    x : npt.NDArray[np.floating]
+        The x coordinates of the points to correct. Should be in the
+        correct UTM coordinate system
+    y : npt.NDArray[np.floating]
+        The y coordinates of the points to correct. Should be in the
+        correct UTM coordinate system.
+    z : npt.NDArray[np.floating]
+        The z coordinates (altitude in meters) of the points to correct.
+    maxiter : int, optional
+        Maximum number of iterations to perform. Default is 5.
+    tol : float, optional
+        Tolerance for convergence in meters. Default is 10.0.
+    full_output : bool, optional
+        If True, return an additional boolean array indicating which points
+        successfully converged. Default is False.
+
+    Returns
+    -------
+    tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]
+        The corrected x and y coordinates as numpy arrays in the UTM
+        coordinate system. Points that are not contained in the non-nan
+        region of the image will contain nan values in the output arrays.
+    """
+    # Confirm that x is monotonically increasing and y is decreasing
+    # (This is assumed in the filtering logic below)
+    if not np.all(np.diff(ds["x"]) > 0.0):
+        msg = "ds['x'] must be monotonically increasing"
+        raise ValueError(msg)
+    if not np.all(np.diff(ds["y"]) < 0.0):
+        msg = "ds['y'] must be monotonically decreasing"
+        raise ValueError(msg)
+
+    try:
+        ds = ds[["VZA", "VAA"]].load()  # nice to load these once here instead of repeatedly below
+    except KeyError as e:
+        raise KeyError("ds must contain the variables 'VZA' and 'VAA'") from e
+
+    x = np.atleast_1d(x).astype(np.float64, copy=False)
+    y = np.atleast_1d(y).astype(np.float64, copy=False)
+    z = np.atleast_1d(z).astype(np.float64, copy=False)
+
+    x_proj = xr.DataArray(x.copy(), dims="points")  # need to copy because we modify below
+    y_proj = xr.DataArray(y.copy(), dims="points")  # need to copy because we modify below
+
+    offset0 = np.zeros_like(x)
+
+    for _ in range(maxiter):
+        # Note that we often get nan values back after interpolation
+        # It's arguably better to propagate nans than to keep the original values
+        # because the original values may be in the nan region of the image
+        # (or outside the image entirely)
+        vza, vaa = _interpolate_angles(ds, x_proj, y_proj)
+
+        # Convert to radians
+        vza_rad = np.deg2rad(vza)
+        vaa_rad = np.deg2rad(vaa)
+
+        # Apply spherical projection offset
+        offset = z * np.tan(vza_rad)
+        dx_offset = offset * np.sin(vaa_rad)
+        dy_offset = offset * np.cos(vaa_rad)
+
+        # Update the newly predicted x and y locations
+        x_proj[:] = x - dx_offset
+        y_proj[:] = y - dy_offset
+
+        error = np.abs(offset - offset0)
+        converged = error < tol
+        if np.all(converged | np.isnan(error)):
+            break
+
+        offset0 = offset
+
+    if full_output:
+        return x_proj.values, y_proj.values, converged
+    return x_proj.values, y_proj.values
+
+
+def _interpolate_angles(
+    ds: xr.Dataset,
+    xi: xr.DataArray,
+    yi: xr.DataArray,
+) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
+    """
+    Interpolate view zenith angle (VZA) and view azimuth angle (VAA).
+
+    Parameters
+    ----------
+    ds : xr.Dataset
+        Dataset containing at least the variables "VZA" and "VAA",
+        with coordinates ``x`` and ``y`` that define the spatial grid.
+    xi : xr.DataArray
+        X-coordinates of the target points for interpolation.
+        Must be the same length as ``yi``.
+    yi : xr.DataArray
+        Y-coordinates of the target points for interpolation.
+        Must be the same length as ``xi``.
+
+    Returns
+    -------
+    vza : npt.NDArray[np.floating]
+        Interpolated view zenith angles at the given (xi, yi) points.
+    vaa : npt.NDArray[np.floating]
+        Interpolated view azimuth angles at the given (xi, yi) points.
+    """
+    interped = ds[["VZA", "VAA"]].interp(x=xi, y=yi)
+    return interped["VZA"].values, interped["VAA"].values
+
+
+def estimate_scan_time(
+    ephemeris_df: pd.DataFrame,
+    utm_crs: pyproj.CRS,
+    x: npt.NDArray[np.floating],
+    y: npt.NDArray[np.floating],
+) -> npt.NDArray[np.datetime64]:
+    """Estimate the scan time for the given x, y pixels.
+
+    Project the x, y coordinates (in UTM coordinate system) onto the
+    ephemeris track and interpolate the time.
+
+    Parameters
+    ----------
+    ephemeris_df : pd.DataFrame
+        DataFrame containing the ephemeris data with columns:
+        - 'EPHEMERIS_ECEF_X': ECEF X coordinates (meters)
+        - 'EPHEMERIS_ECEF_Y': ECEF Y coordinates (meters)
+        - 'EPHEMERIS_ECEF_Z': ECEF Z coordinates (meters)
+        - 'EPHEMERIS_TIME': Timestamps (as datetime64[ns])
+    utm_crs : pyproj.CRS
+        The UTM coordinate reference system used for projection.
+    x : npt.NDArray[np.floating]
+        The x coordinates of the points to estimate the scan time for. Should be in the
+        correct UTM coordinate system.
+    y : npt.NDArray[np.floating]
+        The y coordinates of the points to estimate the scan time for. Should be in the
+        correct UTM coordinate system.
+
+    Returns
+    -------
+    npt.NDArray[np.datetime64]
+        The estimated scan times as numpy datetime64[ns] array. Points for which
+        ``x`` or ``y`` are nan will have ``NaT`` as the corresponding output value.
+    """
+    ephemeris_utm = ephemeris_ecef_to_utm(ephemeris_df, utm_crs)
+
+    valid = np.isfinite(x) & np.isfinite(y)
+    points = shapely.points(x[valid], y[valid])
+
+    line = shapely.LineString(ephemeris_utm[["x", "y"]])
+
+    distance = line.project(points)
+    projected = line.interpolate(distance)
+    projected_x = shapely.get_coordinates(projected)[:, 0]
+
+    if ephemeris_utm["t"].dtype != "datetime64[ns]":
+        # This could be relaxed if needed, but datetime64[ns] is what we expect
+        raise ValueError("ephemeris_utm['t'] must have dtype 'datetime64[ns]'")
+    if not ephemeris_utm["x"].diff().iloc[1:].lt(0).all():
+        # This should always be the case for sun-synchronous satellites
+        raise ValueError("ephemeris_utm['x'] must be strictly decreasing for np.interp")
+
+    out = np.full(x.shape, np.datetime64("NaT", "ns"))
+    out[valid] = np.interp(
+        projected_x,
+        ephemeris_utm["x"].iloc[::-1],
+        ephemeris_utm["t"].iloc[::-1].astype(int),
+    ).astype("datetime64[ns]")
+
+    return out