rsplot 0.1.0__py3-none-any.whl

rsplot/geo/boundaries.py

@@ -0,0 +1,281 @@
+"""Administrative boundary loading and region name resolution.
+
+Resolves user input into a RegionInfo with level, geometry,
+sub-boundaries for overlay, etc.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from functools import lru_cache
+from typing import TYPE_CHECKING
+
+import geopandas as gpd
+import pandas as pd
+from shapely.ops import unary_union
+
+from rsplot.config import LEVEL_DEFAULTS, AppConfig, LevelParams
+
+if TYPE_CHECKING:
+    from shapely.geometry.base import BaseGeometry
+
+
+# Suffixes to strip for fuzzy matching
+_SUFFIXES = re.compile(
+    r"(省|市|区|县|自治区|自治州|自治县|地区|盟|林区|特别行政区)$"
+)
+
+NAME_FIELD = "name"
+
+# 直辖市 — resolved as province but displayed with county-level sub-boundaries
+MUNICIPALITIES = {"北京市", "天津市", "上海市", "重庆市"}
+
+
+@dataclass
+class RegionInfo:
+    """Everything the plotting pipeline needs about a resolved region."""
+
+    name: str  # display name (e.g. "芜湖市")
+    level: str  # "key_region" | "province" | "city" | "county"
+    geometry: BaseGeometry  # union geometry for masking
+    extent: tuple[
+        float, float, float, float
+    ]  # (lon_min, lon_max, lat_min, lat_max)
+    params: LevelParams  # plotting parameters for this level
+
+    # sub-boundary GeoDataFrames for overlay
+    main_boundary_gdf: gpd.GeoDataFrame | None = None  # red border
+    sub_boundary_gdf: gpd.GeoDataFrame | None = None  # grey sub-divisions
+    detail_boundary_gdf: gpd.GeoDataFrame | None = None  # finer sub-divisions (counties in province view)
+    sub_name_field: str = NAME_FIELD
+
+
+def _strip_suffix(name: str) -> str:
+    return _SUFFIXES.sub("", name)
+
+
+def _fuzzy_match(
+    query: str, candidates: gpd.GeoDataFrame
+) -> gpd.GeoDataFrame | None:
+    """Match query against NAME_FIELD: exact first, then suffix-stripped."""
+    # exact
+    exact = candidates[candidates[NAME_FIELD] == query]
+    if len(exact) > 0:
+        return exact
+
+    # with common suffixes appended
+    for suffix in ["省", "市", "区", "县", "自治区"]:
+        padded = candidates[candidates[NAME_FIELD] == query + suffix]
+        if len(padded) > 0:
+            return padded
+
+    # strip suffix from both sides
+    q_stripped = _strip_suffix(query)
+    for _, row in candidates.iterrows():
+        if _strip_suffix(row[NAME_FIELD]) == q_stripped:
+            return candidates[candidates[NAME_FIELD] == row[NAME_FIELD]]
+
+    return None
+
+
+# ---------------------------------------------------------------------------
+# GeoJSON loaders (cached)
+# ---------------------------------------------------------------------------
+@lru_cache(maxsize=1)
+def _load_province(geojson_dir: str) -> gpd.GeoDataFrame:
+    return gpd.read_file(f"{geojson_dir}/中国_省.geojson")
+
+
+@lru_cache(maxsize=1)
+def _load_city(geojson_dir: str) -> gpd.GeoDataFrame:
+    return gpd.read_file(f"{geojson_dir}/中国_市.geojson")
+
+
+@lru_cache(maxsize=1)
+def _load_county(geojson_dir: str) -> gpd.GeoDataFrame:
+    return gpd.read_file(f"{geojson_dir}/中国_县.geojson")
+
+
+# ---------------------------------------------------------------------------
+# Core resolver
+# ---------------------------------------------------------------------------
+def resolve_region(
+    name: str, cfg: AppConfig, level_override: str | None = None
+) -> RegionInfo:
+    """Resolve a user-supplied region name into a RegionInfo.
+
+    Cascade order: key_region → province → city → county.
+    If level_override is given, skip cascade and match only that level.
+    """
+    gj = cfg.geojson_dir
+    gdf_prov = _load_province(gj)
+    gdf_city = _load_city(gj)
+    gdf_county = _load_county(gj)
+
+    # --- 0. Country level (中国 / China) ---
+    if (level_override is None or level_override == "country") and name in (
+        "中国",
+        "China",
+        "china",
+        "CHINA",
+    ):
+        union_geom = unary_union(gdf_prov.geometry)
+        info = _build_info(
+            name="中国",
+            level="country",
+            geometry=union_geom,
+            main_gdf=gdf_prov,
+            sub_gdf=gdf_prov,
+        )
+        # Override extent: focus on mainland (exclude deep South China Sea)
+        info.extent = (73.0, 136.0, 17.0, 54.0)
+        return info
+
+    # --- Handle "city/county" format for county disambiguation ---
+    if "/" in name:
+        city_name, county_name = name.split("/", 1)
+        city_match = _fuzzy_match(city_name, gdf_city)
+        if city_match is None:
+            raise ValueError(f"无法识别城市 '{city_name}'。")
+        city_geom = unary_union(city_match.geometry)
+        sub_counties = gdf_county[
+            gdf_county.geometry.within(city_geom.buffer(0.01))
+        ]
+        county_match = _fuzzy_match(county_name, sub_counties)
+        if county_match is None:
+            available = ", ".join(sub_counties[NAME_FIELD].tolist())
+            raise ValueError(
+                f"在 {city_match.iloc[0][NAME_FIELD]} 下未找到 '{county_name}'。"
+                f"\n可用县区: {available}"
+            )
+        union_geom = unary_union(county_match.geometry)
+        return _build_info(
+            name=county_match.iloc[0][NAME_FIELD],
+            level="county",
+            geometry=union_geom,
+            main_gdf=county_match,
+            sub_gdf=None,
+        )
+
+    # --- 1. Key regions (config presets) ---
+    if level_override is None or level_override == "key_region":
+        upper = name.upper()
+        for kr_name, name_list in cfg.key_regions.items():
+            if upper == kr_name.upper() or name == kr_name:
+                # Match against both provinces and cities
+                matched_provs = gdf_prov[gdf_prov[NAME_FIELD].isin(name_list)]
+                matched_cities = gdf_city[gdf_city[NAME_FIELD].isin(name_list)]
+                if len(matched_provs) == 0 and len(matched_cities) == 0:
+                    raise ValueError(
+                        f"重点区域 '{kr_name}' 包含的名称在 GeoJSON 中找不到: {name_list}"
+                    )
+                main_gdf = pd.concat([matched_provs, matched_cities])
+                union_geom = unary_union(main_gdf.geometry)
+                # Key regions show cities only (no counties) to avoid label clutter.
+                # Pure city list (e.g. PRD): use matched cities directly.
+                # Mixed / province list (e.g. BTH, YRD): enumerate ALL cities
+                # within the union so each province also contributes its cities.
+                if len(matched_provs) == 0:
+                    sub = matched_cities
+                else:
+                    sub = gdf_city[gdf_city.geometry.intersects(union_geom)]
+                return _build_info(
+                    name=kr_name,
+                    level="key_region",
+                    geometry=union_geom,
+                    main_gdf=main_gdf,
+                    sub_gdf=sub,
+                )
+
+    # --- 2. Province ---
+    if level_override is None or level_override == "province":
+        match = _fuzzy_match(name, gdf_prov)
+        if match is not None:
+            prov_name = match.iloc[0][NAME_FIELD]
+            union_geom = unary_union(match.geometry)
+
+            if prov_name in MUNICIPALITIES:
+                # 直辖市: show county/district boundaries like a city view
+                sub_counties = gdf_county[
+                    gdf_county.geometry.within(union_geom.buffer(0.01))
+                ]
+                return _build_info(
+                    name=prov_name,
+                    level="province",
+                    geometry=union_geom,
+                    main_gdf=match,
+                    sub_gdf=sub_counties,
+                )
+            else:
+                # Regular province: cities + counties within each city
+                sub_cities = gdf_city[
+                    gdf_city.geometry.intersects(union_geom)
+                ]
+                detail_counties = gdf_county[
+                    gdf_county.geometry.intersects(union_geom)
+                ]
+                return _build_info(
+                    name=prov_name,
+                    level="province",
+                    geometry=union_geom,
+                    main_gdf=match,
+                    sub_gdf=sub_cities,
+                    detail_gdf=detail_counties,
+                )
+
+    # --- 3. City ---
+    if level_override is None or level_override == "city":
+        match = _fuzzy_match(name, gdf_city)
+        if match is not None:
+            union_geom = unary_union(match.geometry)
+            sub_counties = gdf_county[
+                gdf_county.geometry.within(union_geom.buffer(0.01))
+            ]
+            return _build_info(
+                name=match.iloc[0][NAME_FIELD],
+                level="city",
+                geometry=union_geom,
+                main_gdf=match,
+                sub_gdf=sub_counties,
+            )
+
+    # --- 4. County ---
+    if level_override == "county":
+        # "/" case already handled above; bare name with --level county is ambiguous
+        raise ValueError(
+            f"县/区级别需要指定所属城市，格式: 城市/县区，如 合肥/蜀山、成都/高新"
+        )
+
+    raise ValueError(
+        f"无法识别区域 '{name}'。请检查名称是否正确，"
+        "或使用 --level 指定级别 (key_region/province/city/county)。"
+    )
+
+
+def _build_info(
+    name: str,
+    level: str,
+    geometry: BaseGeometry,
+    main_gdf: gpd.GeoDataFrame | None,
+    sub_gdf: gpd.GeoDataFrame | None,
+    detail_gdf: gpd.GeoDataFrame | None = None,
+) -> RegionInfo:
+    bounds = geometry.bounds  # (minx, miny, maxx, maxy)
+    params = LEVEL_DEFAULTS[level]
+    extent = (
+        bounds[0],
+        bounds[2],
+        bounds[1],
+        bounds[3],
+    )  # (lon_min, lon_max, lat_min, lat_max)
+    return RegionInfo(
+        name=name,
+        level=level,
+        geometry=geometry,
+        extent=extent,
+        params=params,
+        main_boundary_gdf=main_gdf,
+        sub_boundary_gdf=sub_gdf,
+        detail_boundary_gdf=detail_gdf,
+    )

rsplot/geo/gridding.py

@@ -0,0 +1,120 @@
+"""Gridding and smoothing for satellite swath data.
+
+Bins irregular satellite pixels onto a regular lat/lon grid,
+applies region masking and optional NaN-aware Gaussian smoothing.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+from scipy.stats import binned_statistic_2d
+
+if TYPE_CHECKING:
+    from shapely.geometry.base import BaseGeometry
+
+
+def grid_data(
+    lon: np.ndarray,
+    lat: np.ndarray,
+    values: np.ndarray,
+    extent: tuple[float, float, float, float],
+    res: float,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Bin irregular data onto a regular grid.
+
+    Parameters
+    ----------
+    lon, lat, values : 1-D arrays of the same length.
+    extent : (lon_min, lon_max, lat_min, lat_max).
+    res : grid resolution in degrees.
+
+    Returns
+    -------
+    LON, LAT : 2-D meshgrid arrays.
+    grid : 2-D array of mean values (NaN where no data).
+    """
+    lon_min, lon_max, lat_min, lat_max = extent
+    lon_bins = np.arange(lon_min, lon_max + res, res)
+    lat_bins = np.arange(lat_min, lat_max + res, res)
+
+    grid = (
+        binned_statistic_2d(
+            lon, lat, values, statistic="mean", bins=[lon_bins, lat_bins]
+        ).statistic.T
+    )  # transpose: binned_statistic_2d returns (x, y), we want (lat, lon)
+
+    lon_c = 0.5 * (lon_bins[:-1] + lon_bins[1:])
+    lat_c = 0.5 * (lat_bins[:-1] + lat_bins[1:])
+    LON, LAT = np.meshgrid(lon_c, lat_c)
+
+    return LON, LAT, grid
+
+
+def fill_nan_gaps(
+    LON: np.ndarray, LAT: np.ndarray, grid: np.ndarray, max_gap: int = 3
+) -> np.ndarray:
+    """Fill small NaN gaps using nearest-neighbor interpolation.
+
+    Only fills NaN cells within *max_gap* cells of valid data, preventing
+    long-range extrapolation across large empty swath gaps.
+    """
+    from scipy.interpolate import NearestNDInterpolator
+    from scipy.ndimage import binary_dilation
+
+    valid = np.isfinite(grid)
+    if valid.all() or not valid.any():
+        return grid
+
+    # Only fill NaN cells near existing data
+    fill_zone = binary_dilation(valid, iterations=max_gap) & ~valid
+    if not fill_zone.any():
+        return grid
+
+    interp = NearestNDInterpolator(
+        np.column_stack((LON[valid], LAT[valid])), grid[valid]
+    )
+    filled = grid.copy()
+    filled[fill_zone] = interp(LON[fill_zone], LAT[fill_zone])
+    return filled
+
+
+def mask_to_region(
+    LON: np.ndarray,
+    LAT: np.ndarray,
+    grid: np.ndarray,
+    geometry: BaseGeometry,
+) -> np.ndarray:
+    """Mask grid to region boundary using shapely vectorized contains."""
+    from shapely.vectorized import contains
+
+    region_mask = contains(geometry, LON, LAT)
+    return np.where(region_mask, grid, np.nan)
+
+
+def smooth_grid(
+    grid: np.ndarray,
+    geometry: BaseGeometry,
+    LON: np.ndarray,
+    LAT: np.ndarray,
+    sigma: float = 2.0,
+) -> np.ndarray:
+    """NaN-aware Gaussian smoothing using astropy, re-masked to region boundary.
+
+    Uses astropy's convolve with nan_treatment='interpolate' to avoid
+    boundary dilution that scipy.ndimage.gaussian_filter causes.
+    """
+    from astropy.convolution import Gaussian2DKernel, convolve
+    from shapely.vectorized import contains
+
+    kernel = Gaussian2DKernel(x_stddev=sigma)
+    smoothed = convolve(
+        grid, kernel, nan_treatment="interpolate", preserve_nan=False
+    )
+
+    # Re-mask to region boundary (smoothing bleeds outside)
+    region_mask = contains(geometry, LON, LAT)
+    smoothed[~region_mask] = np.nan
+
+    return smoothed

rsplot/plotting/__init__.py

@@ -0,0 +1 @@
+"""Plotting module: raster maps, styles, annotations, basemaps."""

rsplot/plotting/colormaps.py

@@ -0,0 +1,67 @@
+"""Project colormap definitions."""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import matplotlib as mpl
+import matplotlib.colors as mcolors
+import numpy as np
+
+HCHO_CMAP = "rsplot_hcho"
+STATION_CONCENTRATION_CMAP = "rsplot_station_concentration"
+
+DEFAULT_HCHO_CMAP_FILE = (
+    "/exports/XCZ/python/tmp/pycharm_project_24/colormap_22.txt"
+)
+HCHO_CMAP_FILE_ENV = "RSPLOT_HCHO_CMAP_FILE"
+
+
+def register_rsplot_colormaps() -> None:
+    """Register custom colormaps used by rsplot plots."""
+    _register(_build_hcho_cmap())
+    _register(
+        mcolors.LinearSegmentedColormap.from_list(
+            STATION_CONCENTRATION_CMAP,
+            [
+                "#fff7ec",
+                "#fee8c8",
+                "#fdbb84",
+                "#fc8d59",
+                "#e34a33",
+                "#b30000",
+            ],
+        )
+    )
+
+
+def _build_hcho_cmap() -> mcolors.Colormap:
+    cmap_file = Path(os.environ.get(HCHO_CMAP_FILE_ENV, DEFAULT_HCHO_CMAP_FILE))
+    if cmap_file.exists():
+        data = np.loadtxt(cmap_file)
+        if data.shape != (256, 3):
+            raise ValueError(
+                f"{cmap_file} must contain a 256x3 RGB colormap table, "
+                f"got shape {data.shape}"
+            )
+        cmap = mcolors.ListedColormap(data, name=HCHO_CMAP)
+    else:
+        cmap = mcolors.LinearSegmentedColormap.from_list(
+            HCHO_CMAP,
+            [
+                "#ffffff",
+                "#fff7bc",
+                "#fec44f",
+                "#fd8d3c",
+                "#e31a1c",
+                "#800026",
+            ],
+        )
+    cmap.set_under("w")
+    return cmap
+
+
+def _register(cmap: mcolors.Colormap) -> None:
+    if cmap.name not in mpl.colormaps:
+        mpl.colormaps.register(cmap)

rsplot/plotting/fnr.py

@@ -0,0 +1,200 @@
+"""FNR (HCHO/NO2) map plotting with discrete regime colormap.
+
+Uses the same boundary / basemap / gridline conventions as plot_raster,
+but swaps the continuous colorbar for a 3-segment discrete one aligned
+with the FNR_VOC_LIMIT / FNR_NOX_LIMIT thresholds so the regime readout
+is directly visible on the map.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import cartopy.crs as ccrs
+import cartopy.feature as cfeature
+import matplotlib.pyplot as plt
+import matplotlib.ticker as mticker
+import numpy as np
+from cartopy.feature import ShapelyFeature
+from cartopy.mpl.gridliner import LATITUDE_FORMATTER, LONGITUDE_FORMATTER
+from matplotlib.colors import BoundaryNorm, ListedColormap
+
+from rsplot.fnr import FNR_DISPLAY_MAX, FNR_NOX_LIMIT, FNR_VOC_LIMIT
+from rsplot.plotting.raster import _add_sub_labels, _figsize_for_level
+from rsplot.tiles.tianditu import TianDiTuTiles
+
+if TYPE_CHECKING:
+    from rsplot.geo.boundaries import RegionInfo
+
+
+# Fixed regime palette — deliberately kept consistent across levels so
+# readers recognise the same semantics at first glance.
+_REGIME_COLORS = [
+    "#d7301f",  # VOC-limited  (FNR < 1) — red
+    "#fdae61",  # transition   (1-2)     — orange/yellow
+    "#2c7fb8",  # NOx-limited  (FNR > 2) — blue
+]
+
+
+def _build_fnr_cmap() -> tuple[ListedColormap, BoundaryNorm]:
+    cmap = ListedColormap(_REGIME_COLORS, name="fnr_regime")
+    cmap.set_bad(color="#dddddd", alpha=0.0)  # NaN → transparent
+    bounds = [0.0, FNR_VOC_LIMIT, FNR_NOX_LIMIT, FNR_DISPLAY_MAX]
+    norm = BoundaryNorm(bounds, ncolors=cmap.N)
+    return cmap, norm
+
+
+def plot_fnr(
+    LON: np.ndarray,
+    LAT: np.ndarray,
+    fnr: np.ndarray,
+    region: RegionInfo,
+    *,
+    dpi: int = 300,
+    basemap: bool | None = None,
+    tianditu_key: str | None = None,
+    title: str | None = None,
+    output: str | None = None,
+    subtitle: str | None = None,
+) -> None:
+    """Draw an FNR regime map."""
+    params = region.params
+    use_basemap = basemap if basemap is not None else params.basemap
+
+    if use_basemap and tianditu_key:
+        tiles = TianDiTuTiles(tk=tianditu_key, layer="img")
+        projection = tiles.crs
+    else:
+        tiles = None
+        projection = ccrs.PlateCarree()
+
+    fig, ax = plt.subplots(
+        figsize=_figsize_for_level(region.level),
+        dpi=dpi,
+        subplot_kw={"projection": projection},
+    )
+
+    cmap, norm = _build_fnr_cmap()
+    alpha = 0.65 if use_basemap else 0.9
+    im = ax.pcolormesh(
+        LON,
+        LAT,
+        fnr,
+        transform=ccrs.PlateCarree(),
+        cmap=cmap,
+        norm=norm,
+        shading="auto",
+        alpha=alpha,
+    )
+
+    if tiles is not None:
+        ax.add_image(tiles, params.zoom)
+
+    # --- Boundaries (same layering as plot_raster) ---
+    if (
+        region.detail_boundary_gdf is not None
+        and len(region.detail_boundary_gdf) > 0
+    ):
+        ax.add_feature(
+            ShapelyFeature(
+                region.detail_boundary_gdf.geometry,
+                ccrs.PlateCarree(),
+                edgecolor="#aaa",
+                facecolor="none",
+                linewidth=0.25,
+                linestyle=":",
+            )
+        )
+    if (
+        region.sub_boundary_gdf is not None
+        and len(region.sub_boundary_gdf) > 0
+    ):
+        is_county_sub = region.level == "city"
+        ax.add_feature(
+            ShapelyFeature(
+                region.sub_boundary_gdf.geometry,
+                ccrs.PlateCarree(),
+                edgecolor="#888" if not is_county_sub else "#666",
+                facecolor="none",
+                linewidth=0.4 if not is_county_sub else 0.6,
+                linestyle="-" if not is_county_sub else "--",
+            )
+        )
+        _add_sub_labels(ax, region)
+    if region.main_boundary_gdf is not None:
+        ax.add_feature(
+            ShapelyFeature(
+                region.main_boundary_gdf.geometry,
+                ccrs.PlateCarree(),
+                edgecolor="red",
+                facecolor="none",
+                linewidth=params.boundary_linewidth,
+            )
+        )
+
+    if not use_basemap:
+        ax.coastlines(linewidth=0.5, color="#333")
+        ax.add_feature(cfeature.OCEAN, facecolor="#e6f2ff", zorder=0)
+
+    lon_min, lon_max, lat_min, lat_max = region.extent
+    pad = params.pad
+    ax.set_extent(
+        [lon_min - pad, lon_max + pad, lat_min - pad, lat_max + pad],
+        crs=ccrs.PlateCarree(),
+    )
+
+    gl = ax.gridlines(
+        draw_labels=True, linewidth=0.2, linestyle="--", alpha=0.5
+    )
+    gl.top_labels = gl.right_labels = False
+    gl.xformatter = LONGITUDE_FORMATTER
+    gl.yformatter = LATITUDE_FORMATTER
+    gl.xlocator = mticker.MultipleLocator(params.grid_step)
+    gl.ylocator = mticker.MultipleLocator(params.grid_step)
+
+    # --- Discrete colorbar with regime labels centered in each band ---
+    shrink = 0.6 if region.level in ("key_region", "province") else 0.7
+    cb = plt.colorbar(
+        im,
+        ax=ax,
+        shrink=shrink,
+        pad=0.05,
+        ticks=[
+            FNR_VOC_LIMIT / 2,
+            (FNR_VOC_LIMIT + FNR_NOX_LIMIT) / 2,
+            (FNR_NOX_LIMIT + FNR_DISPLAY_MAX) / 2,
+        ],
+    )
+    cb.ax.set_yticklabels(
+        [
+            f"VOC-limited\n(<{FNR_VOC_LIMIT:.0f})",
+            f"Transition\n({FNR_VOC_LIMIT:.0f}-{FNR_NOX_LIMIT:.0f})",
+            f"NOx-limited\n(>{FNR_NOX_LIMIT:.0f})",
+        ],
+        fontsize=8,
+    )
+    cb.set_label("FNR = HCHO / NO2", fontsize=9)
+
+    # --- Title + subtitle ---
+    if title is None:
+        title = f"FNR regime — {region.name}"
+    ax.set_title(title, fontsize=14)
+    if subtitle:
+        ax.text(
+            0.5,
+            -0.08,
+            subtitle,
+            transform=ax.transAxes,
+            ha="center",
+            va="top",
+            fontsize=9,
+            color="#444",
+        )
+
+    plt.tight_layout()
+
+    if output:
+        fig.savefig(output, dpi=dpi, bbox_inches="tight")
+        plt.close(fig)
+    else:
+        plt.show()