pycontrails 0.58.0__cp314-cp314-macosx_10_13_x86_64.whl

pycontrails/datalib/ecmwf/variables.py

@@ -0,0 +1,268 @@
+"""ECMWF Parameter Support.
+
+Sourced from the ECMWF Parameter DB:
+    https://apps.ecmwf.int/codes/grib/param-db
+"""
+
+from __future__ import annotations
+
+from pycontrails.core import met_var
+from pycontrails.core.met import MetVariable
+
+RelativeVorticity = MetVariable(
+    short_name="vo",
+    standard_name="atmosphere_upward_relative_vorticity",
+    long_name="Vorticity (relative)",
+    units="s**-1",
+    level_type="isobaricInhPa",
+    grib1_id=43,
+    ecmwf_id=138,
+    grib2_id=(0, 2, 12),
+    description=(
+        "Atmosphere upward relative vorticity is the vertical component of the 3D air vorticity"
+        " vector. The vertical component arises from horizontal velocity only. 'Relative' in this"
+        " context means the vorticity of the air relative to the rotating solid earth reference"
+        " frame, i.e. excluding the Earth's own rotation. In contrast, the quantity with standard"
+        " name atmosphere_upward_absolute_vorticity includes the Earth's rotation. 'Upward'"
+        " indicates a vector component which is positive when directed upward (negative downward)."
+        " A positive value of atmosphere_upward_relative_vorticity indicates anticlockwise rotation"
+        " when viewed from above."
+    ),
+)
+
+PotentialVorticity = MetVariable(
+    short_name="pv",
+    standard_name="potential_vorticity",
+    long_name="Potential vorticity (K m^2 / kg s)",
+    units="K m**2 kg**-1 s**-1",
+    level_type="isobaricInhPa",
+    grib1_id=128,
+    ecmwf_id=60,
+    grib2_id=(0, 2, 14),
+    amip="pvu",
+    description=(
+        "Potential vorticity is a measure of the capacity for air to rotate in the atmosphere.If we"
+        " ignore the effects of heating and friction, potential vorticity is conserved following an"
+        " air parcel.It is used to look for places where large wind storms are likely to originate"
+        " and develop.Potential vorticity increases strongly above the tropopause and therefore, it"
+        " can also be used in studiesrelated to the stratosphere and stratosphere-troposphere"
+        " exchanges. Large wind storms develop when a columnof air in the atmosphere starts to"
+        " rotate. Potential vorticity is calculated from the wind, temperature andpressure across a"
+        " column of air in the atmosphere."
+    ),
+)
+
+CloudAreaFractionInLayer = MetVariable(
+    short_name="cc",
+    standard_name="fraction_of_cloud_cover",
+    long_name="Cloud area fraction in atmosphere layer",
+    ecmwf_id=248,
+    level_type="isobaricInhPa",
+    grib2_id=(0, 6, 32),
+    units="[0 - 1]",
+    amip="cl",
+    description=(
+        "This parameter is the proportion of a grid box covered by cloud (liquid or ice) at a"
+        " specific pressure level."
+    ),
+)
+
+SpecificCloudLiquidWaterContent = MetVariable(
+    short_name="clwc",
+    standard_name="specific_cloud_liquid_water_content",
+    long_name="Specific cloud liquid water content",
+    units="kg kg**-1",
+    level_type="isobaricInhPa",
+    ecmwf_id=246,
+    grib2_id=(0, 1, 83),
+    description=(
+        "This parameter is the mass of cloud liquid water droplets per kilogram of the total mass"
+        " of moist air. The 'total mass of moist air' is the sum of the dry air, water vapour,"
+        " cloud liquid, cloud ice, rain and falling snow. This parameter represents the average"
+        " value for a grid box."
+    ),
+)
+
+
+SpecificCloudIceWaterContent = MetVariable(
+    short_name="ciwc",
+    standard_name="specific_cloud_ice_water_content",
+    long_name="Specific cloud ice water content",
+    units="kg kg**-1",
+    level_type="isobaricInhPa",
+    ecmwf_id=247,
+    grib2_id=(0, 1, 84),
+    description=(
+        "This parameter is the mass of cloud ice particles per kilogram of the total mass of moist"
+        " air. The 'total mass of moist air' is the sum of the dry air, water vapour, cloud liquid,"
+        " cloud ice, rain and falling snow. This parameter represents the average value for a grid"
+        " box."
+    ),
+)
+
+# Override units and description on Relative humidity
+RelativeHumidity = MetVariable(
+    short_name=met_var.RelativeHumidity.short_name,
+    standard_name=met_var.RelativeHumidity.standard_name,
+    long_name=met_var.RelativeHumidity.long_name,
+    units="%",
+    level_type=met_var.RelativeHumidity.level_type,
+    grib1_id=met_var.RelativeHumidity.grib1_id,
+    ecmwf_id=met_var.RelativeHumidity.ecmwf_id,
+    grib2_id=met_var.RelativeHumidity.grib2_id,
+    description=(
+        "This parameter is the water vapour pressure as a percentage of the value at which the air"
+        " becomes saturated "
+        "(the point at which water vapour begins to condense into liquid water or deposition into"
+        " ice)."
+        "For temperatures over 0°C (273.15 K) it is calculated for saturation over water. "
+        "At temperatures below -23°C it is calculated for saturation over ice. "
+        "Between -23°C and 0°C this parameter is calculated by interpolating between the ice and"
+        " water values using a quadratic function."
+        "See https://www.ecmwf.int/sites/default/files/elibrary/2016/17117-part-iv-physical-processes.pdf#subsection.7.4.2"
+    ),
+)
+
+
+Divergence = MetVariable(
+    short_name="d",
+    standard_name="divergence_of_wind",
+    long_name="Divergence of wind",
+    units="s**-1",
+    level_type="isobaricInhPa",
+    ecmwf_id=155,
+    grib2_id=(0, 2, 13),
+    description=(
+        "This parameter is the horizontal divergence of velocity. It is the rate "
+        "at which air is spreading out horizontally from a point, per square metre. "
+        "This parameter is positive for air that is spreading out, or diverging, and "
+        "negative for air that is moving inward toward a point, or converging."
+    ),
+)
+
+
+TOAIncidentSolarRadiation = MetVariable(
+    short_name="tisr",
+    standard_name="toa_incident_solar_radiation",
+    long_name="Top of atmosphere incident shortwave radiation",
+    units="J m**-2",
+    level_type="nominalTop",
+    ecmwf_id=212,
+    grib2_id=(192, 128, 212),  # reference ECMWF
+    description="Top of atmosphere incident solar radiation. Accumulated field.",
+)
+
+TopNetSolarRadiation = MetVariable(
+    short_name="tsr",
+    standard_name="top_net_solar_radiation",
+    long_name="Top of atmosphere net solar (shortwave) radiation",
+    units="J m**-2",
+    level_type="nominalTop",
+    ecmwf_id=178,
+    grib2_id=(0, 4, 1),
+    description=(
+        "This parameter is the incoming solar radiation (also known as shortwave radiation) "
+        "minus the outgoing solar radiation at the top of the atmosphere. "
+        "It is the amount of radiation passing through a horizontal plane. "
+        "The incoming solar radiation is the amount received from the Sun. "
+        "The outgoing solar radiation is the amount reflected and scattered by the Earth's"
+        " atmosphere and surface"
+        "See https://www.ecmwf.int/sites/default/files/elibrary/2015/18490-radiation-quantities-ecmwf-model-and-mars.pdf"
+    ),
+)
+
+TopNetThermalRadiation = MetVariable(
+    short_name="ttr",
+    standard_name="top_net_thermal_radiation",
+    long_name="Top of atmosphere net thermal (longwave) radiation",
+    units="J m**-2",
+    level_type="nominalTop",
+    ecmwf_id=179,
+    grib2_id=(0, 5, 5),
+    description=(
+        "The thermal (also known as terrestrial or longwave) "
+        "radiation emitted to space at the top of the atmosphere is commonly known as the Outgoing"
+        " Longwave Radiation (OLR). "
+        "The top net thermal radiation (this parameter) is equal to the negative of OLR."
+        "See https://www.ecmwf.int/sites/default/files/elibrary/2015/18490-radiation-quantities-ecmwf-model-and-mars.pdf"
+    ),
+)
+
+SurfaceSolarDownwardRadiation = MetVariable(
+    short_name="ssrd",
+    standard_name="surface_solar_downward_radiation",
+    long_name="Surface Solar Downward Radiation",
+    units="J m**-2",
+    level_type="surface",
+    ecmwf_id=169,
+    grib2_id=(0, 4, 7),
+    description=(
+        "This parameter is the amount of solar radiation (also known as shortwave radiation) that"
+        " reaches a horizontal plane at the surface of the Earth. This parameter comprises both"
+        " direct and diffuse solar radiation."
+    ),
+)
+
+CloudAreaFraction = MetVariable(
+    short_name="tcc",
+    standard_name="total_cloud_cover",
+    long_name="Cloud area fraction (total)",
+    level_type="surface",
+    grib1_id=71,
+    ecmwf_id=164,
+    grib2_id=(192, 128, 164),  # reference ECMWF
+    units="[0 - 1]",
+    amip="clt",
+    description=(
+        "This parameter is the proportion of a grid box covered by cloud (liquid or ice) for a"
+        " whole atmosphere column."
+    ),
+)
+
+OzoneMassMixingRatio = MetVariable(
+    short_name="o3",
+    standard_name="mass_fraction_of_ozone_in_air",
+    long_name="Ozone mass mixing ratio",
+    level_type="isobaricInhPa",
+    units="kg kg**-1",
+    ecmwf_id=203,
+    description="This parameter is the mass of ozone per kilogram of air.",
+)
+
+
+PRESSURE_LEVEL_VARIABLES = [
+    met_var.AirTemperature,
+    met_var.SpecificHumidity,
+    met_var.Geopotential,
+    met_var.EastwardWind,
+    met_var.NorthwardWind,
+    met_var.VerticalVelocity,
+    RelativeHumidity,
+    RelativeVorticity,
+    CloudAreaFractionInLayer,
+    SpecificCloudIceWaterContent,
+    SpecificCloudLiquidWaterContent,
+    PotentialVorticity,
+]
+SURFACE_VARIABLES = [
+    met_var.SurfacePressure,
+    TOAIncidentSolarRadiation,
+    TopNetSolarRadiation,
+    TopNetThermalRadiation,
+    CloudAreaFraction,
+    SurfaceSolarDownwardRadiation,
+]
+MODEL_LEVEL_VARIABLES = [
+    met_var.AirTemperature,
+    met_var.SpecificHumidity,
+    met_var.VerticalVelocity,
+    met_var.EastwardWind,
+    met_var.NorthwardWind,
+    RelativeVorticity,
+    Divergence,
+    CloudAreaFractionInLayer,
+    SpecificCloudIceWaterContent,
+    SpecificCloudLiquidWaterContent,
+]
+
+ECMWF_VARIABLES = PRESSURE_LEVEL_VARIABLES + SURFACE_VARIABLES

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/__init__.py

@@ -0,0 +1,28 @@
+"""GFS Data Access."""
+
+from __future__ import annotations
+
+from pycontrails.datalib.gfs.gfs import GFS_FORECAST_BUCKET, GFSForecast
+from pycontrails.datalib.gfs.variables import (
+    GFS_VARIABLES,
+    PRESSURE_LEVEL_VARIABLES,
+    SURFACE_VARIABLES,
+    CloudIceWaterMixingRatio,
+    TOAUpwardLongwaveRadiation,
+    TOAUpwardShortwaveRadiation,
+    TotalCloudCoverIsobaric,
+    Visibility,
+)
+
+__all__ = [
+    "GFS_FORECAST_BUCKET",
+    "GFS_VARIABLES",
+    "PRESSURE_LEVEL_VARIABLES",
+    "SURFACE_VARIABLES",
+    "CloudIceWaterMixingRatio",
+    "GFSForecast",
+    "TOAUpwardLongwaveRadiation",
+    "TOAUpwardShortwaveRadiation",
+    "TotalCloudCoverIsobaric",
+    "Visibility",
+]