swpc-kit 0.3.0__py3-none-any.whl

swpc_kit/export/geojson.py

@@ -0,0 +1,111 @@
+"""GeoJSON export of the gridded products (standard library only).
+
+:func:`to_geojson` turns an OVATION or D-RAP product into a GeoJSON
+``FeatureCollection`` (RFC 7946): one feature per grid cell, as the cell's
+polygon or its centre point, optionally only for cells at or above a
+threshold. Longitudes are moved into -180 to 180; a cell that straddles the
+antimeridian becomes a ``MultiPolygon`` of its two halves, and cells at the
+poles stop at ±90.
+"""
+
+import math
+from datetime import UTC, datetime
+from typing import Any, Literal
+
+from swpc_kit.models.grids import GridProduct
+
+Geometry = Literal["polygon", "point"]
+
+
+def _iso(value: datetime) -> str:
+    return value.astimezone(UTC).isoformat().replace("+00:00", "Z")
+
+
+def _ring(west: float, south: float, east: float, north: float) -> list[list[float]]:
+    # Counter-clockwise, as RFC 7946 asks for exterior rings.
+    return [[west, south], [east, south], [east, north], [west, north], [west, south]]
+
+
+def _cell(
+    latitude: float, longitude: float, half_height: float, half_width: float
+) -> dict[str, Any]:
+    south = max(latitude - half_height, -90.0)
+    north = min(latitude + half_height, 90.0)
+    west, east = longitude - half_width, longitude + half_width
+    if west < -180:
+        parts = [(west + 360, 180.0), (-180.0, east)]
+    elif east > 180:
+        parts = [(west, 180.0), (-180.0, east - 360)]
+    else:
+        return {"type": "Polygon", "coordinates": [_ring(west, south, east, north)]}
+    return {
+        "type": "MultiPolygon",
+        "coordinates": [[_ring(w, south, e, north)] for w, e in parts],
+    }
+
+
+def to_geojson(
+    product: GridProduct,
+    *,
+    threshold: float | None = None,
+    geometry: Geometry = "polygon",
+) -> dict[str, Any]:
+    """Return the cells of a gridded product as a GeoJSON FeatureCollection.
+
+    Each feature has the properties ``latitude`` and ``longitude`` (the cell
+    centre, longitude in [-180, 180)) and ``value``. Features come row by
+    row, from south to north and from west to east. The collection carries a
+    ``metadata`` member (a foreign member in RFC 7946 terms) with
+    ``quantity``, ``units``, ``source_url``, ``data_time``, each time of
+    ``product.times`` as ISO 8601 text, and ``threshold``.
+
+    Args:
+        product: An OVATION or D-RAP model, or any :class:`GridProduct`.
+        threshold: Keep only cells whose value is at least this; ``None``
+            keeps every cell.
+        geometry: ``"polygon"`` for the cell outline, ``"point"`` for its
+            centre.
+
+    Returns:
+        A dict ready for :func:`json.dumps`.
+
+    Raises:
+        ValueError: If ``geometry`` is unknown or ``threshold`` is not a
+            finite number.
+    """
+    if geometry not in ("polygon", "point"):
+        raise ValueError(f"geometry must be 'polygon' or 'point', got {geometry!r}")
+    if threshold is not None and not math.isfinite(threshold):
+        raise ValueError(f"threshold must be a finite number, got {threshold!r}")
+    grid = product.grid.wrapped()
+    half_height = grid.latitude_step / 2
+    half_width = grid.longitude_step / 2
+    features = []
+    for latitude, longitude, value in grid.cells():
+        if threshold is not None and value < threshold:
+            continue
+        shape = (
+            _cell(latitude, longitude, half_height, half_width)
+            if geometry == "polygon"
+            else {"type": "Point", "coordinates": [longitude, latitude]}
+        )
+        features.append(
+            {
+                "type": "Feature",
+                "geometry": shape,
+                "properties": {
+                    "latitude": latitude,
+                    "longitude": longitude,
+                    "value": value,
+                },
+            }
+        )
+    metadata: dict[str, Any] = {
+        "quantity": grid.quantity,
+        "units": grid.units,
+        "source_url": product.source_url,
+        "data_time": _iso(product.data_time),
+    }
+    metadata.update({name: _iso(time) for name, time in product.times.items()})
+    metadata["threshold"] = threshold
+    return {"type": "FeatureCollection", "metadata": metadata, "features": features}

swpc_kit/export/geotiff.py

@@ -0,0 +1,134 @@
+"""GeoTIFF export of the gridded products, with the standard library only.
+
+:func:`write_geotiff` writes an OVATION or D-RAP product as a one-band
+``float32`` GeoTIFF in EPSG:4326 (WGS 84 latitude/longitude), compressed
+with deflate. The raster is north-up, with longitudes moved into -180 to
+180, and each pixel is a grid cell: the transform puts pixel edges half a
+step around the cell centres, so OVATION, whose centres reach the poles and
+-180, spans -180.5 to 179.5 and -90.5 to 90.5.
+
+The times the content declares go into the dataset tags as ISO 8601 text
+(``DATA_TIME``, then one tag per time, such as ``OBSERVATION_TIME``,
+``FORECAST_TIME`` or ``VALID_AT``), with ``SOURCE_URL``, ``QUANTITY`` and
+``UNITS``; ``TIFFTAG_DATETIME`` holds the data time in the TIFF format.
+They are written as GDAL writes them, so GDAL, rasterio and QGIS read them
+back as dataset metadata.
+"""
+
+import math
+import os
+import struct
+import zlib
+from datetime import UTC, datetime
+from pathlib import Path
+from xml.sax.saxutils import escape
+
+from swpc_kit.export._tiff import ASCII, DOUBLE, LONG, SHORT, Entry, encode
+from swpc_kit.models.grids import GridProduct
+
+__all__ = ["write_geotiff"]
+
+_STRIP_BYTES = 8192
+"""About how much raw data goes in one strip, as GDAL chooses."""
+
+
+def _iso(value: datetime) -> str:
+    return value.astimezone(UTC).isoformat().replace("+00:00", "Z")
+
+
+def _float32(value: float) -> float:
+    return float(struct.unpack("<f", struct.pack("<f", value))[0])
+
+
+def write_geotiff(
+    product: GridProduct,
+    path: str | os.PathLike[str],
+    *,
+    nodata: float = math.nan,
+) -> None:
+    """Write a gridded product to a GeoTIFF file.
+
+    Cells the grid does not have (the missing longitudes of a reduced
+    OVATION file) are set to ``nodata``.
+
+    Args:
+        product: An OVATION or D-RAP model, or any :class:`GridProduct`.
+        path: Where to write the file; an existing file is replaced.
+        nodata: Value of the cells without data, declared as the nodata
+            value of the band. It must not be one of the grid's values.
+
+    Raises:
+        ValueError: If ``nodata`` is one of the grid's values, or does not
+            fit in a ``float32``.
+    """
+    grid = product.grid.wrapped()
+    try:
+        nodata_32 = _float32(nodata)
+    except OverflowError:
+        raise ValueError(f"nodata {nodata} does not fit in float32") from None
+    if not math.isnan(nodata) and any(
+        _float32(value) == nodata_32 for row in grid.values for value in row
+    ):
+        raise ValueError(f"nodata {nodata} is also a value of the grid")
+    dy, dx = grid.latitude_step, grid.longitude_step
+    south, north = grid.latitudes[0], grid.latitudes[-1]
+    west, east = grid.longitudes[0], grid.longitudes[-1]
+    height = round((north - south) / dy) + 1
+    width = round((east - west) / dx) + 1
+    raster = [[nodata] * width for _ in range(height)]
+    cols = [round((lon - west) / dx) for lon in grid.longitudes]
+    for lat, values in zip(grid.latitudes, grid.values, strict=True):
+        # Row 0 is the northernmost latitude; the grid is stored south first.
+        row = raster[height - 1 - round((lat - south) / dy)]
+        for col, value in zip(cols, values, strict=True):
+            row[col] = value
+    rows_per_strip = max(1, _STRIP_BYTES // (4 * width))
+    strips = [
+        zlib.compress(
+            b"".join(
+                struct.pack(f"<{width}f", *row)
+                for row in raster[top : top + rows_per_strip]
+            )
+        )
+        for top in range(0, height, rows_per_strip)
+    ]
+
+    data_time = product.data_time
+    tags = {
+        "SOURCE_URL": product.source_url,
+        "QUANTITY": grid.quantity,
+        "UNITS": grid.units,
+        "DATA_TIME": _iso(data_time),
+    }
+    tags.update({name.upper(): _iso(time) for name, time in product.times.items()})
+    items = [f'<Item name="{key}">{escape(text)}</Item>' for key, text in tags.items()]
+    items.append(
+        f'<Item name="DESCRIPTION" sample="0" role="description">'
+        f"{escape(grid.quantity)}</Item>"
+    )
+    items.append(
+        f'<Item name="UNITTYPE" sample="0" role="unittype">{escape(grid.units)}</Item>'
+    )
+
+    entries: list[Entry] = [
+        (256, LONG, [width]),  # ImageWidth
+        (257, LONG, [height]),  # ImageLength
+        (258, SHORT, [32]),  # BitsPerSample
+        (259, SHORT, [8]),  # Compression: deflate
+        (262, SHORT, [1]),  # PhotometricInterpretation: min is black
+        (277, SHORT, [1]),  # SamplesPerPixel
+        (278, LONG, [rows_per_strip]),  # RowsPerStrip
+        (279, LONG, [len(strip) for strip in strips]),  # StripByteCounts
+        (284, SHORT, [1]),  # PlanarConfiguration: contiguous
+        (306, ASCII, data_time.strftime("%Y:%m:%d %H:%M:%S")),  # DateTime
+        (339, SHORT, [3]),  # SampleFormat: IEEE floating point
+        (33550, DOUBLE, [dx, dy, 0.0]),  # ModelPixelScale
+        # ModelTiepoint: pixel (0, 0) is the north-west corner of the grid.
+        (33922, DOUBLE, [0.0, 0.0, 0.0, west - dx / 2, north + dy / 2, 0.0]),
+        # GeoKeyDirectory, version 1.1.0 with three keys: a geographic model
+        # (1024 = 2), pixels that are areas (1025 = 1), and EPSG:4326 (2048).
+        (34735, SHORT, [1, 1, 0, 3, 1024, 0, 1, 2, 1025, 0, 1, 1, 2048, 0, 1, 4326]),
+        (42112, ASCII, f"<GDALMetadata>{''.join(items)}</GDALMetadata>"),
+        (42113, ASCII, repr(nodata)),  # GDAL_NODATA
+    ]
+    Path(path).write_bytes(encode(entries, strips))

swpc_kit/models/__init__.py

@@ -0,0 +1,116 @@
+"""Models: frozen, typed dataclasses for SWPC products.
+
+Every time is a ``datetime`` in UTC. Each top-level model exposes
+``data_time`` (the most recent time declared by the content), ``source_url``,
+and ``issued``, ``observed``, ``valid_from`` or ``valid_to`` where the source
+provides them. The models are re-exported here from their modules:
+``indices``, ``solar_wind``, ``goes``, ``alerts``, ``bulletins``, ``grids``.
+"""
+
+from swpc_kit.models.alerts import (
+    AlertMessage,
+    Alerts,
+    MessageField,
+    MessageKind,
+    NoaaScale,
+    ScaleCategory,
+    field_key,
+)
+from swpc_kit.models.bulletins import (
+    DailyAp,
+    GeomagneticForecast,
+    KpPeriod,
+    OutlookDay,
+    ThreeDayForecast,
+    TwentySevenDayOutlook,
+)
+from swpc_kit.models.goes import (
+    FlareClass,
+    GoesIntegralProtons,
+    GoesXrays,
+    ProtonSample,
+    XrayBand,
+    XraySample,
+)
+from swpc_kit.models.grids import (
+    DrapFrequencies,
+    Grid,
+    GridProduct,
+    OvationAurora,
+    wrap_longitude,
+)
+from swpc_kit.models.indices import (
+    DstValue,
+    F107Flux,
+    F107Value,
+    GeomagneticStormScale,
+    Kp1mValue,
+    KpForecastStatus,
+    KpForecastValue,
+    KpValue,
+    KyotoDst,
+    NoaaScales,
+    NoaaScalesDay,
+    PlanetaryKIndex,
+    PlanetaryKIndex1m,
+    PlanetaryKIndexForecast,
+    RadioBlackoutScale,
+    SolarRadiationScale,
+)
+from swpc_kit.models.records import SkippedRecord
+from swpc_kit.models.solar_wind import (
+    RtswMag1m,
+    RtswMagSample,
+    RtswSource,
+    RtswWind1m,
+    RtswWindSample,
+)
+
+__all__ = [
+    "AlertMessage",
+    "Alerts",
+    "DailyAp",
+    "DrapFrequencies",
+    "DstValue",
+    "F107Flux",
+    "F107Value",
+    "FlareClass",
+    "GeomagneticForecast",
+    "GeomagneticStormScale",
+    "GoesIntegralProtons",
+    "GoesXrays",
+    "Grid",
+    "GridProduct",
+    "Kp1mValue",
+    "KpForecastStatus",
+    "KpForecastValue",
+    "KpPeriod",
+    "KpValue",
+    "KyotoDst",
+    "MessageField",
+    "MessageKind",
+    "NoaaScale",
+    "NoaaScales",
+    "NoaaScalesDay",
+    "OutlookDay",
+    "OvationAurora",
+    "PlanetaryKIndex",
+    "PlanetaryKIndex1m",
+    "PlanetaryKIndexForecast",
+    "ProtonSample",
+    "RadioBlackoutScale",
+    "RtswMag1m",
+    "RtswMagSample",
+    "RtswSource",
+    "RtswWind1m",
+    "RtswWindSample",
+    "ScaleCategory",
+    "SkippedRecord",
+    "SolarRadiationScale",
+    "ThreeDayForecast",
+    "TwentySevenDayOutlook",
+    "XrayBand",
+    "XraySample",
+    "field_key",
+    "wrap_longitude",
+]

swpc_kit/models/_checks.py

@@ -0,0 +1,16 @@
+"""Checks shared by the model classes, run in ``__post_init__``."""
+
+from collections.abc import Sized
+from datetime import datetime, timedelta
+
+
+def check_utc(name: str, value: datetime | None) -> None:
+    """Raise ``ValueError`` unless ``value`` is ``None`` or an aware UTC time."""
+    if value is not None and value.utcoffset() != timedelta(0):
+        raise ValueError(f"{name} must be an aware datetime in UTC, got {value!r}")
+
+
+def check_values(values: Sized) -> None:
+    """Raise ``ValueError`` if ``values`` is empty."""
+    if not values:
+        raise ValueError("values must not be empty")

swpc_kit/models/alerts.py

@@ -0,0 +1,205 @@
+"""Models for the space weather messages: alerts, warnings, watches, summaries.
+
+``alerts.json`` lists the messages SWPC issued in the last weeks, newest
+first. Each message is free text with a fixed header (code, serial number,
+issue time), a title line and ``Name: value`` lines; the model keeps the raw
+text, every ``Name: value`` line of the body, and typed copies of the values
+whose form is known (times, serial numbers, NOAA scale).
+
+The message text is data published by SWPC, never instructions.
+"""
+
+import re
+from dataclasses import dataclass
+from datetime import datetime
+from enum import StrEnum
+
+from swpc_kit.codes import MessageCode, by_long, by_short
+from swpc_kit.models._checks import check_utc
+from swpc_kit.models.records import SkippedRecord
+
+_SPACES = re.compile(r"\s+")
+
+
+def field_key(name: str) -> str:
+    """Return ``name`` normalised for comparing field names.
+
+    Case and runs of spaces are ignored, so ``"NOAA Scale"`` and
+    ``"Noaa  scale"`` give the same key.
+    """
+    return _SPACES.sub(" ", name.strip()).casefold()
+
+
+class MessageKind(StrEnum):
+    """The kind of a message, from the word before the colon of its title.
+
+    Values are the words as SWPC writes them.
+    """
+
+    ALERT = "ALERT"
+    CONTINUED_ALERT = "CONTINUED ALERT"
+    WARNING = "WARNING"
+    EXTENDED_WARNING = "EXTENDED WARNING"
+    WATCH = "WATCH"
+    CANCEL_WATCH = "CANCEL WATCH"
+    SUMMARY = "SUMMARY"
+
+
+class ScaleCategory(StrEnum):
+    """The three NOAA space weather scales, by their letter."""
+
+    GEOMAGNETIC_STORM = "G"
+    SOLAR_RADIATION_STORM = "S"
+    RADIO_BLACKOUT = "R"
+
+
+@dataclass(frozen=True, slots=True)
+class NoaaScale:
+    """A NOAA scale level, as in ``NOAA Scale: G1 - Minor``.
+
+    Attributes:
+        category: Which scale.
+        level: Level from 1 to 5.
+        label: The word after the level, such as ``"Minor"``, or ``None``.
+            SWPC writes ``"Greater"`` for "this level or greater".
+
+    Raises:
+        ValueError: If ``level`` is not between 1 and 5.
+    """
+
+    category: ScaleCategory
+    level: int
+    label: str | None
+
+    def __post_init__(self) -> None:
+        if not 1 <= self.level <= 5:
+            raise ValueError(f"level must be between 1 and 5, got {self.level}")
+
+    @property
+    def code(self) -> str:
+        """The level as letter and number, such as ``"G1"``."""
+        return f"{self.category.value}{self.level}"
+
+
+@dataclass(frozen=True, slots=True)
+class MessageField:
+    """One ``Name: value`` line of the body of a message.
+
+    Attributes:
+        name: The name as written, such as ``"Valid From"``.
+        value: The rest of the line, without surrounding spaces; may be empty.
+    """
+
+    name: str
+    value: str
+
+
+@dataclass(frozen=True, slots=True)
+class AlertMessage:
+    """One space weather message (alert, warning, watch or summary).
+
+    Attributes:
+        product_id: Short code of the message, such as ``"K04W"``.
+        code: Long code from the ``Space Weather Message Code:`` line, such
+            as ``"WARK04"``.
+        serial_number: From the ``Serial Number:`` line. Not unique: a
+            continued alert can repeat the serial number of an earlier one.
+        issued: ``issue_datetime`` of the record, to the millisecond.
+        issue_time: From the ``Issue Time:`` line, to the minute.
+        kind: Kind of message from its title, or ``None`` when unknown.
+        title: The title line (``"ALERT: Geomagnetic K-index of 5"``), or an
+            empty string when the message has none.
+        valid_from: ``Valid From:``, or ``None``.
+        valid_to: ``Now Valid Until:`` (extended warnings) or ``Valid To:``,
+            or ``None``.
+        threshold_reached: ``Threshold Reached:``, or ``None``.
+        begin_time: ``Begin Time:``, or ``None``.
+        maximum_time: ``Maximum Time:`` or ``Max time:``, or ``None``.
+        end_time: ``End Time:``, or ``None``.
+        observed: ``Observed:`` (sudden impulse summaries), or ``None``.
+        extension_to_serial: ``Extension to Serial Number:``, or ``None``.
+        continuation_of_serial: ``Continuation of Serial Number:``, or
+            ``None``.
+        cancel_serial: ``Cancel Serial Number:``, or ``None``.
+        noaa_scale: First valid ``NOAA Scale:`` or ``Noaa Scale:``, or
+            ``None``. Watches give their level in the title, not here.
+        fields: Every ``Name: value`` line after the title, in order,
+            repeats included. A typed attribute is ``None`` when its value
+            has an unknown form; the raw value stays here.
+        text: The message as served.
+
+    Raises:
+        ValueError: If a time is not an aware datetime in UTC.
+    """
+
+    product_id: str
+    code: str
+    serial_number: int
+    issued: datetime
+    issue_time: datetime
+    kind: MessageKind | None = None
+    title: str = ""
+    valid_from: datetime | None = None
+    valid_to: datetime | None = None
+    threshold_reached: datetime | None = None
+    begin_time: datetime | None = None
+    maximum_time: datetime | None = None
+    end_time: datetime | None = None
+    observed: datetime | None = None
+    extension_to_serial: int | None = None
+    continuation_of_serial: int | None = None
+    cancel_serial: int | None = None
+    noaa_scale: NoaaScale | None = None
+    fields: tuple[MessageField, ...] = ()
+    text: str = ""
+
+    def __post_init__(self) -> None:
+        for name in (
+            "issued",
+            "issue_time",
+            "valid_from",
+            "valid_to",
+            "threshold_reached",
+            "begin_time",
+            "maximum_time",
+            "end_time",
+            "observed",
+        ):
+            check_utc(name, getattr(self, name))
+
+    def field(self, name: str) -> str | None:
+        """Return the value of the first field called ``name``, or ``None``.
+
+        Case and runs of spaces in ``name`` are ignored.
+        """
+        key = field_key(name)
+        return next((f.value for f in self.fields if field_key(f.name) == key), None)
+
+    @property
+    def info(self) -> MessageCode | None:
+        """Catalogue entry of the message, by short code then long code.
+
+        ``None`` when neither code is in :mod:`swpc_kit.codes`.
+        """
+        return by_short(self.product_id) or by_long(self.code)
+
+
+@dataclass(frozen=True, slots=True)
+class Alerts:
+    """The space weather messages of the last weeks (``alerts.json``).
+
+    Attributes:
+        messages: The messages, in source order (newest first). May be
+            empty.
+        source_url: Where the content came from.
+        skipped: The records that could not be read, in source order.
+    """
+
+    messages: tuple[AlertMessage, ...]
+    source_url: str
+    skipped: tuple[SkippedRecord, ...] = ()
+
+    @property
+    def data_time(self) -> datetime | None:
+        """Issue time of the most recent message, or ``None`` when empty."""
+        return max((m.issued for m in self.messages), default=None)