rsplot 0.1.0__py3-none-any.whl

rsplot/plotting/styles.py

@@ -0,0 +1,35 @@
+"""Font registration and matplotlib style configuration."""
+
+from __future__ import annotations
+
+import os
+
+import matplotlib as mpl
+from matplotlib import font_manager
+
+from rsplot.plotting.colormaps import register_rsplot_colormaps
+
+
+def setup_fonts(font_dir: str) -> None:
+    """Register Chinese fonts and set matplotlib defaults.
+
+    Looks for Microsoft YaHei in font_dir. Falls back gracefully
+    if fonts are not found (labels will use default sans-serif).
+    """
+    candidates = [
+        os.path.join(font_dir, "MicrosoftYaHei", "微软雅黑.ttf"),
+        os.path.join(font_dir, "MicrosoftYaHei", "微软雅黑粗体.ttf"),
+    ]
+
+    registered = False
+    for fp in candidates:
+        if os.path.exists(fp):
+            font_manager.fontManager.addfont(fp)
+            registered = True
+
+    if registered:
+        mpl.rcParams["font.family"] = "sans-serif"
+        mpl.rcParams["font.sans-serif"] = ["Microsoft YaHei"]
+
+    mpl.rcParams["axes.unicode_minus"] = False
+    register_rsplot_colormaps()

rsplot/readers/__init__.py

@@ -0,0 +1,96 @@
+"""Satellite data readers and product registry."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import numpy as np
+
+from rsplot.readers.base import BaseReader
+from rsplot.readers.tropomi_hcho import TropomiHCHOReader
+from rsplot.readers.tropomi_no2 import TropomiNO2Reader
+from rsplot.readers.tropomi_o3 import TropomiO3Reader
+
+HCHO_CMAP = "rsplot_hcho"
+
+
+@dataclass
+class ProductInfo:
+    """Per-product metadata: reader class, display settings, auto-range config."""
+
+    reader_cls: type[BaseReader]
+    colorbar_label: str
+    cmap: str
+    vmin: float | None  # None = auto from data percentiles
+    vmax: float | None
+    min_range: float  # minimum vmax-vmin to avoid exaggerating small differences
+
+
+PRODUCT_REGISTRY: dict[str, ProductInfo] = {
+    "no2": ProductInfo(
+        reader_cls=TropomiNO2Reader,
+        colorbar_label="NO2 VCD (x10^15 molec/cm2)",
+        cmap="YlOrRd",
+        vmin=0.0,
+        vmax=20.0,
+        min_range=5.0,
+    ),
+    "o3": ProductInfo(
+        reader_cls=TropomiO3Reader,
+        colorbar_label="O3 Total Column (DU)",
+        cmap="YlOrRd",
+        vmin=None,
+        vmax=None,
+        min_range=30.0,
+    ),
+    "hcho": ProductInfo(
+        reader_cls=TropomiHCHOReader,
+        colorbar_label="HCHO VCD (x10^15 molec/cm2)",
+        cmap=HCHO_CMAP,
+        vmin=None,  # unit not yet confirmed — let auto-range pick
+        vmax=None,
+        min_range=5.0,
+    ),
+}
+
+
+def get_product_info(product: str) -> ProductInfo:
+    """Return ProductInfo for a product name."""
+    key = product.lower()
+    if key not in PRODUCT_REGISTRY:
+        available = ", ".join(PRODUCT_REGISTRY)
+        raise ValueError(
+            f"Unsupported product '{product}'. Available: {available}"
+        )
+    return PRODUCT_REGISTRY[key]
+
+
+def auto_vrange(
+    grid: np.ndarray, prod_info: ProductInfo
+) -> tuple[float, float]:
+    """Compute vmin/vmax from data, respecting min_range and rounding nicely."""
+    p2 = float(np.nanpercentile(grid, 2))
+    p98 = float(np.nanpercentile(grid, 98))
+    span = p98 - p2
+
+    # Enforce minimum range: expand symmetrically around the center
+    if span < prod_info.min_range:
+        center = (p2 + p98) / 2
+        p2 = center - prod_info.min_range / 2
+        p98 = center + prod_info.min_range / 2
+        span = prod_info.min_range
+
+    # Round to nice numbers: pick step from span magnitude
+    if span < 10:
+        step = 1.0
+    elif span < 50:
+        step = 5.0
+    elif span < 100:
+        step = 10.0
+    else:
+        step = 20.0
+
+    vmin = math.floor(p2 / step) * step
+    vmax = math.ceil(p98 / step) * step
+    return vmin, vmax

rsplot/readers/base.py

@@ -0,0 +1,96 @@
+"""Abstract base class for satellite data readers."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+import numpy as np
+import xarray as xr
+
+
+@dataclass
+class SwathData:
+    """Raw swath data after QA filtering and spatial subsetting."""
+
+    lon: np.ndarray  # 1-D
+    lat: np.ndarray  # 1-D
+    values: np.ndarray  # 1-D, in display units
+    n_pixels: int
+    n_files: int
+    n_broken: int
+
+
+class BaseReader(ABC):
+    """Interface for reading a single-day satellite product."""
+
+    @abstractmethod
+    def read(
+        self,
+        data_dir: str,
+        date: str,
+        extent: tuple[float, float, float, float],
+        qa_threshold: float = 0.5,
+    ) -> SwathData:
+        """Read and filter all orbits for a given date and spatial extent.
+
+        Parameters
+        ----------
+        data_dir : root directory containing product files.
+        date : date string, e.g. "20260208".
+        extent : (lon_min, lon_max, lat_min, lat_max).
+        qa_threshold : minimum QA value to keep.
+
+        Returns
+        -------
+        SwathData with concatenated 1-D arrays.
+        """
+        ...
+
+
+def read_tropomi_product_arrays(
+    path: str, value_name: str
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Read lon/lat/value/qa from either flattened or grouped TROPOMI files.
+
+    Some local archives flatten the ``/PRODUCT`` group into variables like
+    ``PRODUCT_longitude``; official OFFL files usually keep a real PRODUCT
+    NetCDF group with plain variable names. This helper accepts both layouts.
+    """
+    flat_names = {
+        "lon": "PRODUCT_longitude",
+        "lat": "PRODUCT_latitude",
+        "value": f"PRODUCT_{value_name}",
+        "qa": "PRODUCT_qa_value",
+    }
+    group_names = {
+        "lon": "longitude",
+        "lat": "latitude",
+        "value": value_name,
+        "qa": "qa_value",
+    }
+
+    errors: list[str] = []
+    for group, names in ((None, flat_names), ("PRODUCT", group_names)):
+        try:
+            kwargs = {} if group is None else {"group": group}
+            ds = xr.open_dataset(path, **kwargs)
+        except (OSError, ValueError) as e:
+            label = "root" if group is None else group
+            errors.append(f"{label}: {e}")
+            continue
+
+        try:
+            lon = ds[names["lon"]].squeeze().values
+            lat = ds[names["lat"]].squeeze().values
+            value = ds[names["value"]].squeeze().values
+            qa = ds[names["qa"]].squeeze().values
+            return lon, lat, value, qa
+        except KeyError as e:
+            label = "root" if group is None else group
+            errors.append(f"{label}: missing {e}")
+        finally:
+            ds.close()
+
+    msg = "; ".join(str(e).splitlines()[0] for e in errors)
+    raise ValueError(f"无法读取 TROPOMI PRODUCT 变量: {msg}")

rsplot/readers/guokong.py

@@ -0,0 +1,143 @@
+"""National ground monitoring station (国控站点) reader.
+
+Reads hourly combined NC files from /home/zzgsg/data/guokong_combine/.
+File naming: {YYYYMMDDHH}.nc, e.g. 2026041415.nc.
+
+Accepts both YYYYMMDDHH (exact hour) and YYYYMMDD (auto-picks latest hour).
+"""
+
+from __future__ import annotations
+
+import glob
+import os
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import numpy as np
+import xarray as xr
+
+if TYPE_CHECKING:
+    from shapely.geometry.base import BaseGeometry
+
+
+@dataclass
+class StationData:
+    """Filtered station observations for one variable."""
+
+    id: np.ndarray       # 1-D string array
+    lon: np.ndarray      # 1-D float
+    lat: np.ndarray      # 1-D float
+    values: np.ndarray   # 1-D float (NaN where invalid)
+    var_name: str
+    n_stations: int
+    n_valid: int
+
+
+def read_stations(
+    data_dir: str,
+    datetime_str: str,
+    var_name: str,
+    extent: tuple[float, float, float, float] | None = None,
+) -> tuple[StationData, str]:
+    """Read station data for a given hour and variable.
+
+    Parameters
+    ----------
+    data_dir : directory containing {YYYYMMDDHH}.nc files.
+    datetime_str : YYYYMMDDHH (exact hour) or YYYYMMDD (latest hour).
+    var_name : variable name in the NC file (e.g. "AQI", "PM2.5", "O3").
+    extent : optional (lon_min, lon_max, lat_min, lat_max) for spatial filter.
+
+    Returns
+    -------
+    tuple of (StationData, resolved_datetime_str).
+    """
+    if not datetime_str.isdigit() or len(datetime_str) not in (8, 10):
+        raise ValueError(
+            f"datetime 参数 '{datetime_str}' 格式不正确，"
+            f"请使用 YYYYMMDD 或 YYYYMMDDHH，如 20260414 或 2026041415。"
+        )
+
+    # If 8-digit date, find the latest available hourly file for that day
+    if len(datetime_str) == 8:
+        pattern = os.path.join(data_dir, f"{datetime_str}??.nc")
+        matches = sorted(glob.glob(pattern))
+        if not matches:
+            raise FileNotFoundError(
+                f"未找到 {datetime_str} 当天的站点数据文件: {pattern}\n"
+                f"请检查日期和 --data-dir 路径。"
+            )
+        path = matches[-1]  # latest hour
+        datetime_str = os.path.basename(path).removesuffix(".nc")
+    else:
+        path = os.path.join(data_dir, f"{datetime_str}.nc")
+
+    try:
+        ds = xr.open_dataset(path)
+    except FileNotFoundError:
+        raise FileNotFoundError(
+            f"未找到站点数据文件: {path}\n"
+            f"请检查 datetime 参数和 --data-dir 路径。"
+        )
+
+    try:
+        lon = ds["lon"].values.astype(float)
+        lat = ds["lat"].values.astype(float)
+        station_id = ds["id"].values
+
+        # Read variable — xarray auto-applies scale_factor/add_offset for CO
+        raw = ds[var_name].values.astype(float)
+    finally:
+        ds.close()
+
+    # Fill values (65535 for ushort) → NaN
+    raw[raw >= 65534] = np.nan
+
+    # Spatial filter
+    if extent is not None:
+        lon_min, lon_max, lat_min, lat_max = extent
+        mask = (
+            (lon >= lon_min)
+            & (lon <= lon_max)
+            & (lat >= lat_min)
+            & (lat <= lat_max)
+        )
+        lon = lon[mask]
+        lat = lat[mask]
+        raw = raw[mask]
+        station_id = station_id[mask]
+
+    n_valid = int(np.isfinite(raw).sum())
+
+    return StationData(
+        id=station_id,
+        lon=lon,
+        lat=lat,
+        values=raw,
+        var_name=var_name,
+        n_stations=len(lon),
+        n_valid=n_valid,
+    ), datetime_str
+
+
+def mask_stations_to_region(
+    data: StationData, geometry: BaseGeometry
+) -> StationData:
+    """Keep only stations inside the focus region geometry."""
+    from shapely.vectorized import contains
+
+    inside = contains(geometry.buffer(1e-9), data.lon, data.lat)
+    return _subset_stations(data, inside)
+
+
+def _subset_stations(data: StationData, mask: np.ndarray) -> StationData:
+    values = data.values[mask]
+    return StationData(
+        id=np.asarray(data.id)[mask],
+        lon=data.lon[mask],
+        lat=data.lat[mask],
+        values=values,
+        var_name=data.var_name,
+        n_stations=int(mask.sum()),
+        n_valid=int(np.isfinite(values).sum()),
+    )

rsplot/readers/tropomi_hcho.py

@@ -0,0 +1,126 @@
+"""TROPOMI OFFL L2 HCHO reader.
+
+Reads ``S5P_OFFL_L2__HCHO___`` per-orbit swath files, applies QA
+filtering, spatial subsetting, and unit conversion to ``10^15 molec/cm²``.
+
+Filename pattern (one file per orbit, ~16 files per day)::
+
+    S5P_OFFL_L2__HCHO___{YYYYMMDD}T{HHMMSS}_{end}_{orbit}_{coll}_{proc}_{tstamp}.nc.zip
+
+Two structural differences from the NO2 / O3 NRTI readers:
+
+1. **Real NetCDF groups, not flat names.** OFFL files keep the
+   ``/PRODUCT`` group hierarchy intact, so xarray's root-level dataset
+   is empty. We open the ``PRODUCT`` group explicitly and address
+   variables by their plain name (``longitude``, ``qa_value`` …) rather
+   than the underscore-flattened ``PRODUCT_longitude`` form the NRTI
+   readers use.
+2. **Misleading ``.nc.zip`` extension.** ``file(1)`` reports each as
+   native HDF5 / NetCDF4, so xarray reads them directly without any
+   unzip step. We also accept plain ``.nc`` in case files are
+   re-distributed without the suffix.
+
+Unit: the source attribute on the variable is ``mol m-2``. We convert
+to ``10^15 molec/cm²`` so HCHO matches the NO2 colorbar units and so
+FNR = HCHO / NO2 is a dimensionless ratio directly comparable with the
+literature.
+"""
+
+from __future__ import annotations
+
+import glob
+import os
+
+import numpy as np
+from rich.progress import track
+
+from rsplot.readers.base import (
+    BaseReader,
+    SwathData,
+    read_tropomi_product_arrays,
+)
+
+# mol/m² -> 10^15 molec/cm² (same convention as the NO2 reader)
+UNIT_CONVERSION = 6.02214076e23 * 1e-4 * 1e-15
+
+VALUE_NAME = "formaldehyde_tropospheric_vertical_column"
+
+
+class TropomiHCHOReader(BaseReader):
+    """Reader for TROPOMI OFFL L2 HCHO tropospheric column files."""
+
+    def read(
+        self,
+        data_dir: str,
+        date: str,
+        extent: tuple[float, float, float, float],
+        qa_threshold: float = 0.5,
+    ) -> SwathData:
+        # Match both `.nc` and `.nc.zip` (the latter is HDF5 underneath
+        # despite the misleading extension in this archive).
+        patterns = [
+            os.path.join(data_dir, f"S5P_*_L2__HCHO___{date}T*.nc"),
+            os.path.join(data_dir, f"S5P_*_L2__HCHO___{date}T*.nc.zip"),
+        ]
+        files: list[str] = []
+        for p in patterns:
+            files.extend(glob.glob(p))
+        files = sorted(set(files))
+        if not files:
+            raise FileNotFoundError(
+                f"未找到匹配文件: \n  "
+                + "\n  ".join(patterns)
+                + f"\n请检查 --hcho-dir / --data-dir 和 --date={date}。"
+            )
+
+        lon_min, lon_max, lat_min, lat_max = extent
+        all_lon, all_lat, all_hcho = [], [], []
+        broken = 0
+        broken_examples: list[str] = []
+
+        for f in track(files, description="[cyan]读取 HCHO 轨道..."):
+            try:
+                lon, lat, hcho, qa = read_tropomi_product_arrays(f, VALUE_NAME)
+            except (OSError, ValueError, KeyError) as e:
+                broken += 1
+                if len(broken_examples) < 3:
+                    broken_examples.append(
+                        f"{os.path.basename(f)}: {str(e).splitlines()[0]}"
+                    )
+                continue
+
+            mask = (
+                (qa > qa_threshold)
+                & (lon >= lon_min)
+                & (lon <= lon_max)
+                & (lat >= lat_min)
+                & (lat <= lat_max)
+                & np.isfinite(hcho)
+                & (hcho < 1e30)
+            )
+            if mask.sum() > 0:
+                all_lon.append(lon[mask])
+                all_lat.append(lat[mask])
+                all_hcho.append(hcho[mask] * UNIT_CONVERSION)
+
+        if not all_lon:
+            broken_msg = (
+                "\n损坏/不可读示例:\n  " + "\n  ".join(broken_examples)
+                if broken_examples
+                else ""
+            )
+            raise ValueError(
+                f"{date} 当天 HCHO 无有效数据覆盖目标区域 {extent}。\n"
+                f"共扫描 {len(files)} 个文件，损坏 {broken} 个；"
+                f"可尝试放宽 --qa 阈值。"
+                f"{broken_msg}"
+            )
+
+        return SwathData(
+            lon=np.concatenate(all_lon),
+            lat=np.concatenate(all_lat),
+            values=np.concatenate(all_hcho),
+            n_pixels=sum(len(a) for a in all_lon),
+            n_files=len(files),
+            n_broken=broken,
+        )

rsplot/readers/tropomi_no2.py

@@ -0,0 +1,102 @@
+"""TROPOMI NO2 reader.
+
+Reads S5P_*_L2__NO2____ files, applies QA filtering, spatial subsetting,
+and unit conversion to 10^15 molec/cm2. Supports both flattened NRT-style
+variables and official PRODUCT-group files.
+"""
+
+from __future__ import annotations
+
+import glob
+import os
+
+import numpy as np
+from rich.progress import track
+
+from rsplot.readers.base import (
+    BaseReader,
+    SwathData,
+    read_tropomi_product_arrays,
+)
+
+# mol/m2 -> 10^15 molec/cm2
+UNIT_CONVERSION = 6.02214076e23 * 1e-4 * 1e-15
+
+VALUE_NAME = "nitrogendioxide_tropospheric_column"
+
+
+class TropomiNO2Reader(BaseReader):
+    """Reader for TROPOMI L2 NO2 tropospheric column files."""
+
+    def read(
+        self,
+        data_dir: str,
+        date: str,
+        extent: tuple[float, float, float, float],
+        qa_threshold: float = 0.5,
+    ) -> SwathData:
+        patterns = [
+            os.path.join(data_dir, f"S5P_*_L2__NO2____{date}*.nc"),
+            os.path.join(data_dir, f"S5P_*_L2__NO2____{date}*.nc.zip"),
+        ]
+        files: list[str] = []
+        for p in patterns:
+            files.extend(glob.glob(p))
+        files = sorted(set(files))
+        if not files:
+            raise FileNotFoundError(
+                f"未找到匹配文件:\n  "
+                + "\n  ".join(patterns)
+                + "\n请检查 --data-dir 和 --date 参数。"
+            )
+
+        lon_min, lon_max, lat_min, lat_max = extent
+        all_lon, all_lat, all_no2 = [], [], []
+        broken = 0
+        broken_examples: list[str] = []
+
+        for f in track(files, description="[cyan]读取轨道数据..."):
+            try:
+                lon, lat, no2, qa = read_tropomi_product_arrays(f, VALUE_NAME)
+            except (OSError, ValueError, KeyError) as e:
+                broken += 1
+                if len(broken_examples) < 3:
+                    broken_examples.append(
+                        f"{os.path.basename(f)}: {str(e).splitlines()[0]}"
+                    )
+                continue
+
+            mask = (
+                (qa > qa_threshold)
+                & (lon >= lon_min)
+                & (lon <= lon_max)
+                & (lat >= lat_min)
+                & (lat <= lat_max)
+                & np.isfinite(no2)
+                & (no2 < 1e30)
+            )
+            if mask.sum() > 0:
+                all_lon.append(lon[mask])
+                all_lat.append(lat[mask])
+                all_no2.append(no2[mask] * UNIT_CONVERSION)
+
+        if not all_lon:
+            broken_msg = (
+                "\n损坏/不可读示例:\n  " + "\n  ".join(broken_examples)
+                if broken_examples
+                else ""
+            )
+            raise ValueError(
+                f"{date} 当天无有效数据覆盖目标区域 {extent}。\n"
+                f"共扫描 {len(files)} 个文件，损坏 {broken} 个。"
+                f"{broken_msg}"
+            )
+
+        return SwathData(
+            lon=np.concatenate(all_lon),
+            lat=np.concatenate(all_lat),
+            values=np.concatenate(all_no2),
+            n_pixels=sum(len(a) for a in all_lon),
+            n_files=len(files),
+            n_broken=broken,
+        )

rsplot/readers/tropomi_o3.py

@@ -0,0 +1,104 @@
+"""TROPOMI O3 reader.
+
+Reads S5P_*_L2__O3_____ files, applies QA filtering, spatial subsetting,
+and unit conversion to DU (Dobson Units). Supports both flattened NRT-style
+variables and official PRODUCT-group files.
+"""
+
+from __future__ import annotations
+
+import glob
+import os
+
+import numpy as np
+from rich.progress import track
+
+from rsplot.readers.base import (
+    BaseReader,
+    SwathData,
+    read_tropomi_product_arrays,
+)
+
+# mol/m2 -> DU
+# 1 DU = 2.6867e16 molec/cm2
+# mol/m2 * Avogadro * 1e-4 (m2->cm2) / 2.6867e16 = DU
+UNIT_CONVERSION = 6.02214076e23 * 1e-4 / 2.6867e16
+
+VALUE_NAME = "ozone_total_vertical_column"
+
+
+class TropomiO3Reader(BaseReader):
+    """Reader for TROPOMI L2 O3 total column files."""
+
+    def read(
+        self,
+        data_dir: str,
+        date: str,
+        extent: tuple[float, float, float, float],
+        qa_threshold: float = 0.5,
+    ) -> SwathData:
+        patterns = [
+            os.path.join(data_dir, f"S5P_*_L2__O3_____{date}*.nc"),
+            os.path.join(data_dir, f"S5P_*_L2__O3_____{date}*.nc.zip"),
+        ]
+        files: list[str] = []
+        for p in patterns:
+            files.extend(glob.glob(p))
+        files = sorted(set(files))
+        if not files:
+            raise FileNotFoundError(
+                f"未找到匹配文件:\n  "
+                + "\n  ".join(patterns)
+                + "\n请检查 --data-dir 和 --date 参数。"
+            )
+
+        lon_min, lon_max, lat_min, lat_max = extent
+        all_lon, all_lat, all_o3 = [], [], []
+        broken = 0
+        broken_examples: list[str] = []
+
+        for f in track(files, description="[cyan]读取轨道数据..."):
+            try:
+                lon, lat, o3, qa = read_tropomi_product_arrays(f, VALUE_NAME)
+            except (OSError, ValueError, KeyError) as e:
+                broken += 1
+                if len(broken_examples) < 3:
+                    broken_examples.append(
+                        f"{os.path.basename(f)}: {str(e).splitlines()[0]}"
+                    )
+                continue
+
+            mask = (
+                (qa > qa_threshold)
+                & (lon >= lon_min)
+                & (lon <= lon_max)
+                & (lat >= lat_min)
+                & (lat <= lat_max)
+                & np.isfinite(o3)
+                & (o3 < 1e30)
+            )
+            if mask.sum() > 0:
+                all_lon.append(lon[mask])
+                all_lat.append(lat[mask])
+                all_o3.append(o3[mask] * UNIT_CONVERSION)
+
+        if not all_lon:
+            broken_msg = (
+                "\n损坏/不可读示例:\n  " + "\n  ".join(broken_examples)
+                if broken_examples
+                else ""
+            )
+            raise ValueError(
+                f"{date} 当天无有效数据覆盖目标区域 {extent}。\n"
+                f"共扫描 {len(files)} 个文件，损坏 {broken} 个。"
+                f"{broken_msg}"
+            )
+
+        return SwathData(
+            lon=np.concatenate(all_lon),
+            lat=np.concatenate(all_lat),
+            values=np.concatenate(all_o3),
+            n_pixels=sum(len(a) for a in all_lon),
+            n_files=len(files),
+            n_broken=broken,
+        )