pycontrails 0.56.0__cp310-cp310-macosx_10_9_x86_64.whl → 0.57.0__cp310-cp310-macosx_10_9_x86_64.whl

pycontrails/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.56.0'
-__version_tuple__ = version_tuple = (0, 56, 0)
+__version__ = version = '0.57.0'
+__version_tuple__ = version_tuple = (0, 57, 0)
 
-__commit_id__ = commit_id = 'gec7f244f4'
+__commit_id__ = commit_id = 'g7b8b60b87'

pycontrails/datalib/geo_utils.py

@@ -0,0 +1,261 @@
+"""Tooling and support for GEO satellites."""
+
+import numpy as np
+import numpy.typing as npt
+import xarray as xr
+
+from pycontrails.utils import dependencies
+
+
+def parallax_correct(
+    longitude: npt.NDArray[np.floating],
+    latitude: npt.NDArray[np.floating],
+    altitude: npt.NDArray[np.floating],
+    goes_da: xr.DataArray,
+) -> tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]:
+    r"""Apply parallax correction to WGS84 geodetic coordinates based on satellite perspective.
+
+    This function considers the ray from the satellite to the points of interest and finds
+    the intersection of this ray with the WGS84 ellipsoid. The intersection point is then
+    returned as the corrected longitude and latitude coordinates.
+
+    ::
+
+       @ satellite
+        \
+         \
+          \
+           \
+            \
+             * aircraft
+              \
+               \
+                x parallax corrected aircraft
+        -------------------------  surface
+
+    If the point of interest is not visible from the satellite (ie, on the opposite side of the
+    earth), the function returns nan for the corrected coordinates.
+
+    This function requires the :mod:`pyproj` package to be installed.
+
+    Parameters
+    ----------
+    longitude : npt.NDArray[np.floating]
+        A 1D array of longitudes in degrees.
+    latitude : npt.NDArray[np.floating]
+        A 1D array of latitudes in degrees.
+    altitude : npt.NDArray[np.floating]
+        A 1D array of altitudes in meters.
+    goes_da : xr.DataArray
+        DataArray containing the GOES projection information. Only the ``goes_imager_projection``
+        field of the :attr:`xr.DataArray.attrs` is used.
+
+    Returns
+    -------
+    tuple[npt.NDArray[np.floating], npt.NDArray[np.floating]]
+        A tuple containing the corrected longitude and latitude coordinates.
+
+    """
+    goes_imager_projection = goes_da.attrs["goes_imager_projection"]
+    sat_lon = goes_imager_projection["longitude_of_projection_origin"]
+    sat_lat = goes_imager_projection["latitude_of_projection_origin"]
+    sat_alt = goes_imager_projection["perspective_point_height"]
+
+    try:
+        import pyproj
+    except ModuleNotFoundError as exc:
+        dependencies.raise_module_not_found_error(
+            name="parallax_correct function",
+            package_name="pyproj",
+            module_not_found_error=exc,
+            pycontrails_optional_package="pyproj",
+        )
+
+    # Convert from WGS84 to ECEF coordinates
+    ecef_crs = pyproj.CRS("EPSG:4978")
+    wgs84_crs = pyproj.CRS("EPSG:4326")
+    transformer = pyproj.Transformer.from_crs(wgs84_crs, ecef_crs, always_xy=True)
+
+    p0 = np.array(transformer.transform([sat_lon], [sat_lat], [sat_alt]))
+    p1 = np.array(transformer.transform(longitude, latitude, altitude))
+
+    # Major and minor axes of the ellipsoid
+    a = ecef_crs.ellipsoid.semi_major_metre  # type: ignore[union-attr]
+    b = ecef_crs.ellipsoid.semi_minor_metre  # type: ignore[union-attr]
+    intersection = _intersection_with_ellipsoid(p0, p1, a, b)
+
+    # Convert back to WGS84 coordinates
+    inv_transformer = pyproj.Transformer.from_crs(ecef_crs, wgs84_crs, always_xy=True)
+    return inv_transformer.transform(*intersection)[:2]  # final coord is (close to) 0
+
+
+def _intersection_with_ellipsoid(
+    p0: npt.NDArray[np.floating],
+    p1: npt.NDArray[np.floating],
+    a: float,
+    b: float,
+) -> npt.NDArray[np.floating]:
+    """Find the intersection of a line with the surface of an ellipsoid."""
+    # Calculate the direction vector
+    px, py, pz = p0
+    v = p1 - p0
+    vx, vy, vz = v
+
+    # The line between p0 and p1 in parametric form is p(t) = p0 + t * v
+    # We need to find t such that p(t) lies on the ellipsoid
+    # x^2 / a^2 + y^2 / a^2 + z^2 / b^2 = 1
+    # (px + t * vx)^2 / a^2 + (py + t * vy)^2 / a^2 + (pz + t * vz)^2 / b^2 = 1
+    # Rearranging gives a quadratic in t
+
+    # Calculate the coefficients of this quadratic equation
+    A = vx**2 / a**2 + vy**2 / a**2 + vz**2 / b**2
+    B = 2 * (px * vx / a**2 + py * vy / a**2 + pz * vz / b**2)
+    C = px**2 / a**2 + py**2 / a**2 + pz**2 / b**2 - 1.0
+
+    # Calculate the discriminant
+    D = B**2 - 4 * A * C
+    sqrtD = np.sqrt(D, where=D >= 0, out=np.full_like(D, np.nan))
+
+    # Calculate the two possible solutions for t
+    t0 = (-B + sqrtD) / (2.0 * A)
+    t1 = (-B - sqrtD) / (2.0 * A)
+
+    # Calculate the intersection points
+    intersection0 = p0 + t0 * v
+    intersection1 = p0 + t1 * v
+
+    # Pick the intersection point that is closer to the aircraft (p1)
+    d0 = np.linalg.norm(intersection0 - p1, axis=0)
+    d1 = np.linalg.norm(intersection1 - p1, axis=0)
+    out = np.where(d0 < d1, intersection0, intersection1)
+
+    # Fill the points in which the aircraft is not visible by the satellite with nan
+    # This occurs when the earth is between the satellite and the aircraft
+    # In other words, we can check for t0 < 1 (or t1 < 1)
+    opposite_side = t0 < 1.0
+    out[:, opposite_side] = np.nan
+
+    return out
+
+
+def to_ash(da: xr.DataArray, convention: str = "SEVIRI") -> npt.NDArray[np.float32]:
+    """Compute 3d RGB array for the ASH color scheme.
+
+    Parameters
+    ----------
+    da : xr.DataArray
+        DataArray of GOES data with appropriate bands.
+    convention : str, optional
+        Convention for color space.
+
+        - SEVIRI convention requires bands C11, C14, C15.
+          Used in :cite:`kulikSatellitebasedDetectionContrails2019`.
+        - Standard convention requires bands C11, C13, C14, C15
+
+    Returns
+    -------
+    npt.NDArray[np.float32]
+        3d RGB array with ASH color scheme according to convention.
+
+    References
+    ----------
+    - `Ash RGB quick guide (the color space and color interpretations) <https://rammb.cira.colostate.edu/training/visit/quick_guides/GOES_Ash_RGB.pdf>`_
+    - :cite:`SEVIRIRGBCal`
+    - :cite:`kulikSatellitebasedDetectionContrails2019`
+
+    Examples
+    --------
+    >>> from pycontrails.datalib.goes import GOES
+    >>> goes = GOES(region="M2", bands=("C11", "C14", "C15"))
+    >>> da = goes.get("2022-10-03 04:34:00")
+    >>> rgb = to_ash(da)
+    >>> rgb.shape
+    (500, 500, 3)
+
+    >>> rgb[0, 0, :]
+    array([0.0127004 , 0.22793579, 0.3930847 ], dtype=float32)
+    """
+    if convention == "standard":
+        if not np.all(np.isin([11, 13, 14, 15], da["band_id"])):
+            msg = "DataArray must contain bands 11, 13, 14, and 15 for standard ash"
+            raise ValueError(msg)
+        c11 = da.sel(band_id=11).values  # 8.44
+        c13 = da.sel(band_id=13).values  # 10.33
+        c14 = da.sel(band_id=14).values  # 11.19
+        c15 = da.sel(band_id=15).values  # 12.27
+
+        red = c15 - c13
+        green = c14 - c11
+        blue = c13
+
+    elif convention in ("SEVIRI", "MIT"):  # retain MIT for backwards compatibility
+        if not np.all(np.isin([11, 14, 15], da["band_id"])):
+            msg = "DataArray must contain bands 11, 14, and 15 for SEVIRI ash"
+            raise ValueError(msg)
+        c11 = da.sel(band_id=11).values  # 8.44
+        c14 = da.sel(band_id=14).values  # 11.19
+        c15 = da.sel(band_id=15).values  # 12.27
+
+        red = c15 - c14
+        green = c14 - c11
+        blue = c14
+
+    else:
+        raise ValueError("Convention must be either 'SEVIRI' or 'standard'")
+
+    # See colostate pdf for slightly wider values
+    red = _clip_and_scale(red, -4.0, 2.0)
+    green = _clip_and_scale(green, -4.0, 5.0)
+    blue = _clip_and_scale(blue, 243.0, 303.0)
+    return np.dstack([red, green, blue])
+
+
+def _clip_and_scale(
+    arr: npt.NDArray[np.floating], low: float, high: float
+) -> npt.NDArray[np.floating]:
+    """Clip array and rescale to the interval [0, 1].
+
+    Array is first clipped to the interval [low, high] and then linearly rescaled
+    to the interval [0, 1] so that::
+
+        low -> 0
+        high -> 1
+
+    Parameters
+    ----------
+    arr : npt.NDArray[np.floating]
+        Array to clip and scale.
+    low : float
+        Lower clipping bound.
+    high : float
+        Upper clipping bound.
+
+    Returns
+    -------
+    npt.NDArray[np.floating]
+        Clipped and scaled array.
+    """
+    return (arr.clip(low, high) - low) / (high - low)
+
+
+def _coarsen_then_concat(da1: xr.DataArray, da2: xr.DataArray) -> xr.DataArray:
+    """Concatenate two DataArrays, averaging da2 to da1's resolution.
+
+    This function is hacky and should not be used publicly. It is used in goes.py
+    and himawari.py to combine data from different resolutions.
+
+    The function assumes that da2 has exactly twice the resolution of da1 in both
+    the x and y dimensions.
+    """
+    da2 = da2.coarsen(x=2, y=2, boundary="exact").mean()  # type: ignore[attr-defined]
+
+    # Gut check
+    np.testing.assert_allclose(da1["x"], da2["x"], atol=2e-5)
+    np.testing.assert_allclose(da1["y"], da2["y"], atol=2e-5)
+
+    # Assign the coarser coords to the coarsened coords to account for any small differences
+    da2["x"] = da1["x"]
+    da2["y"] = da1["y"]
+
+    # Finally, combine the datasets
+    return xr.concat([da1, da2], dim="band_id", coords="different", compat="equals")

pycontrails/datalib/gfs/gfs.py

@@ -15,7 +15,6 @@ import logging
 import pathlib
 import sys
 import warnings
-from collections.abc import Callable
 from datetime import datetime
 from typing import TYPE_CHECKING, Any
 
@@ -41,9 +40,8 @@ from pycontrails.datalib.gfs.variables import (
 from pycontrails.utils import dependencies, temp
 from pycontrails.utils.types import DatetimeLike
 
-# optional imports
 if TYPE_CHECKING:
-    import botocore
+    import s3fs
 
 logger = logging.getLogger(__name__)
 
@@ -125,10 +123,10 @@ class GFSForecast(metsource.MetDataSource):
     - `GFS Documentation <https://www.emc.ncep.noaa.gov/emc/pages/numerical_forecast_systems/gfs/documentation.php>`_
     """
 
-    __slots__ = ("cache_download", "cachestore", "client", "forecast_time", "grid", "show_progress")
+    __slots__ = ("cache_download", "cachestore", "forecast_time", "fs", "grid", "show_progress")
 
-    #: S3 client for accessing GFS bucket
-    client: botocore.client.S3
+    #: s3fs filesystem for anonymous access to GFS bucket
+    fs: s3fs.S3FileSystem | None
 
     #: Lat / Lon grid spacing. One of [0.25, 0.5, 1]
     grid: float
@@ -153,26 +151,6 @@ class GFSForecast(metsource.MetDataSource):
         show_progress: bool = False,
         cache_download: bool = False,
     ) -> None:
-        try:
-            import boto3
-        except ModuleNotFoundError as e:
-            dependencies.raise_module_not_found_error(
-                name="GFSForecast class",
-                package_name="boto3",
-                module_not_found_error=e,
-                pycontrails_optional_package="gfs",
-            )
-
-        try:
-            import botocore
-        except ModuleNotFoundError as e:
-            dependencies.raise_module_not_found_error(
-                name="GFSForecast class",
-                package_name="botocore",
-                module_not_found_error=e,
-                pycontrails_optional_package="gfs",
-            )
-
         # inputs
         self.paths = paths
         if cachestore is self.__marker:
@@ -196,13 +174,10 @@ class GFSForecast(metsource.MetDataSource):
         self.variables = metsource.parse_variables(variables, self.supported_variables)
         self.grid = metsource.parse_grid(grid, (0.25, 0.5, 1))
 
-        # note GFS allows unsigned requests (no credentials)
-        # https://stackoverflow.com/questions/34865927/can-i-use-boto3-anonymously/34866092#34866092
-        self.client = boto3.client(
-            "s3", config=botocore.client.Config(signature_version=botocore.UNSIGNED)
-        )
+        # s3 filesystem (created on first download)
+        self.fs = None
 
-        # set specific forecast time is requested
+        # set specific forecast time if requested, otherwise compute from timesteps
         if forecast_time is not None:
             forecast_time_pd = pd.to_datetime(forecast_time)
             if forecast_time_pd.hour % 6:
@@ -492,10 +467,35 @@ class GFSForecast(metsource.MetDataSource):
             ds.to_netcdf(cache_path)
 
     def _make_download(self, aws_key: str, target: str, filename: str) -> None:
+        """Download a single GRIB file using s3fs.
+
+        Parameters
+        ----------
+        aws_key : str
+            Key under GFS bucket forecast path.
+        target : str
+            Local filename to write.
+        filename : str
+            Original filename (used for progress label).
+        """
+        # Lazily import s3fs and create filesystem if needed
+        if self.fs is None:
+            try:
+                import s3fs
+            except ModuleNotFoundError as exc:
+                dependencies.raise_module_not_found_error(
+                    name="GFSForecast class",
+                    package_name="s3fs",
+                    module_not_found_error=exc,
+                    pycontrails_optional_package="gfs",
+                )
+            self.fs = s3fs.S3FileSystem(anon=True)
+
+        s3_path = f"s3://{GFS_FORECAST_BUCKET}/{aws_key}"
         if self.show_progress:
-            _download_with_progress(self.client, GFS_FORECAST_BUCKET, aws_key, target, filename)
+            _download_with_progress(self.fs, s3_path, target, filename)
         else:
-            self.client.download_file(Bucket=GFS_FORECAST_BUCKET, Key=aws_key, Filename=target)
+            self.fs.get(s3_path, target)
 
     def _open_gfs_dataset(self, filepath: str | pathlib.Path, t: datetime) -> xr.Dataset:
         """Open GFS grib file for one forecast timestep.
@@ -615,30 +615,20 @@ class GFSForecast(metsource.MetDataSource):
         return met.MetDataset(ds, **kwargs)
 
 
-def _download_with_progress(
-    client: botocore.client.S3, bucket: str, key: str, filename: str, label: str
-) -> None:
-    """Download with `tqdm` progress bar.
+def _download_with_progress(fs: s3fs.S3FileSystem, s3_path: str, target: str, label: str) -> None:
+    """Download with tqdm progress bar using s3fs.
 
     Parameters
     ----------
-    client : botocore.client.S3
-        S3 Client
-    bucket : str
-        AWS Bucket
-    key : str
-        Key within bucket to download
-    filename : str
-        Local filename to download to
+    fs : s3fs.S3FileSystem
+        Filesystem instance.
+    s3_path : str
+        Full s3 path (s3://bucket/key).
+    target : str
+        Local file path to write.
     label : str
-        Progress label
-
-    Raises
-    ------
-    ModuleNotFoundError
-        Raises if tqdm can't be found
+        Progress bar label.
     """
-
     try:
         from tqdm import tqdm
     except ModuleNotFoundError as e:
@@ -649,14 +639,18 @@ def _download_with_progress(
             pycontrails_optional_package="gfs",
         )
 
-    meta = client.head_object(Bucket=bucket, Key=key)
-    filesize = meta["ContentLength"]
-
-    def hook(t: Any) -> Callable:
-        def inner(bytes_amount: Any) -> None:
-            t.update(bytes_amount)
-
-        return inner
-
-    with tqdm(total=filesize, unit="B", unit_scale=True, desc=label) as t:
-        client.download_file(Bucket=bucket, Key=key, Filename=filename, Callback=hook(t))
+    # get object size via simple info call
+    info = fs.info(s3_path)
+    filesize = info.get("Size") or info.get("size")
+
+    with (
+        fs.open(s3_path, "rb") as fsrc,
+        open(target, "wb") as fdst,
+        tqdm(total=filesize, unit="B", unit_scale=True, desc=label) as t,
+    ):
+        # stream in chunks
+        chunk = fsrc.read(1024 * 1024)
+        while chunk:
+            fdst.write(chunk)
+            t.update(len(chunk))
+            chunk = fsrc.read(1024 * 1024)