klimapy 1.1.0__py3-none-win_amd64.whl

klimapy/BCSD/BiasCorrection.py

@@ -0,0 +1,209 @@
+"""Calendar-aligned precipitation quantile correction using Rust FITPACK."""
+
+import logging
+from pathlib import Path
+from contextlib import ExitStack
+
+import numpy as np
+
+from ..native import Field, TimeAxis
+from ..native.correction import Transfer
+from .._core.common import WINDOW_SIZE, N_QUANTILES
+from .._core.native_adapter import open_precip_dataset
+from .._core.engine import (
+    calendar_info,
+    align_reference,
+    mapped_days,
+    GridWeights,
+    allocate,
+)
+
+logger = logging.getLogger(__name__)
+DEBUG_CHECK = False
+__all__ = ["BC", "create_cdf_spline", "close_result"]
+
+
+def create_cdf_spline(gcm_data, obs_data, n_quantiles=50):
+    try:
+        return Transfer(gcm_data, obs_data, n_quantiles)
+    except ValueError as exc:
+        if "20 finite observations" in str(exc):
+            return None
+        raise
+
+
+def BC(
+    obs_path,
+    gcm_hist_path,
+    gcm_fut_path,
+    model_name,
+    ref_start_year,
+    ref_end_year,
+    *,
+    strict=False,
+    time_chunk_size=128,
+    work_dir=None,
+    window_size=None,
+    n_quantiles=None,
+):
+    buffers = []
+    try:
+        window = WINDOW_SIZE if window_size is None else window_size
+        quantiles = N_QUANTILES if n_quantiles is None else n_quantiles
+        if (
+            isinstance(time_chunk_size, bool)
+            or not isinstance(time_chunk_size, int)
+            or time_chunk_size < 1
+            or isinstance(window, bool)
+            or not isinstance(window, int)
+            or window < 0
+            or isinstance(quantiles, bool)
+            or not isinstance(quantiles, int)
+            or quantiles < 4
+        ):
+            raise ValueError(
+                "Positive time_chunk_size, nonnegative window_size and n_quantiles>=4 required"
+            )
+        with ExitStack() as stack:
+            obs, historical, future = [
+                stack.enter_context(open_precip_dataset(p))
+                for p in (obs_path, gcm_hist_path, gcm_fut_path)
+            ]
+            for d in ("lat", "lon"):
+                if not np.array_equal(historical.coords[d], future.coords[d]):
+                    raise ValueError(
+                        "Historical and future model grids must match exactly"
+                    )
+            cal, ndays = calendar_info(historical.time)
+            if calendar_info(future.time)[0] != cal:
+                raise ValueError("Historical and future calendars must match")
+            oi, hi = align_reference(
+                obs.time, historical.time, ref_start_year, ref_end_year
+            )
+            weights = stack.enter_context(
+                GridWeights(
+                    obs.coords["lat"],
+                    obs.coords["lon"],
+                    historical.coords["lat"],
+                    historical.coords["lon"],
+                )
+            )
+            reference = allocate(
+                (len(oi), *historical.shape[1:]), np.float64, work_dir, "reference"
+            )
+            buffers.append(reference)
+            for start in range(0, len(oi), time_chunk_size):
+                indices = oi[start : start + time_chunk_size]
+                reference[start : start + len(indices)] = weights.apply(
+                    obs.read(time=indices)
+                )
+            ref_time = TimeAxis(
+                historical.coords["time"][hi],
+                historical.coordinate_attrs.get("time", {}),
+            )
+            reference_days = mapped_days(ref_time, cal)
+            history_days, future_days = (
+                mapped_days(historical.time, cal),
+                mapped_days(future.time, cal),
+            )
+            windows = {}
+            for day in range(1, ndays + 1):
+                distance = np.abs(reference_days - day)
+                windows[day] = np.flatnonzero(
+                    (reference_days > 0)
+                    & (np.minimum(distance, ndays - distance) <= window)
+                )
+            outputs = [
+                allocate(a.shape, np.float64, work_dir, "corrected")
+                for a in (historical, future)
+            ]
+            buffers.extend(outputs)
+            for out in outputs:
+                out[:] = np.nan
+            stats = {
+                "matching_reference_days": len(oi),
+                "fitted_functions": 0,
+                "unfitted_functions": 0,
+                "approximate_functions": 0,
+                "reference_period": [ref_start_year, ref_end_year],
+                "calendar": cal,
+            }
+            for row in range(historical.shape[1]):
+                histrow, futrow = historical.read(lat=row), future.read(lat=row)
+                for col in range(historical.shape[2]):
+                    if not np.isfinite(reference[:, row, col]).any():
+                        stats["unfitted_functions"] += ndays
+                        continue
+                    train = histrow[hi, col]
+                    for day, indices in windows.items():
+                        function = create_cdf_spline(
+                            train[indices], reference[indices, row, col], quantiles
+                        )
+                        stats[
+                            "fitted_functions"
+                            if function is not None
+                            else "unfitted_functions"
+                        ] += 1
+                        if (
+                            function is not None
+                            and getattr(function, "fit_status", 0) > 0
+                        ):
+                            stats["approximate_functions"] += 1
+                        try:
+                            for series, days, out in [
+                                (histrow[:, col], history_days, outputs[0]),
+                                (futrow[:, col], future_days, outputs[1]),
+                            ]:
+                                mask = days == day
+                                out[mask, row, col] = (
+                                    function(series[mask])
+                                    if function is not None
+                                    else series[mask]
+                                )
+                        finally:
+                            if function is not None:
+                                function.close()
+            answer = {
+                "corrected_hist_pr": historical.result(outputs[0]),
+                "corrected_fut_pr": future.result(outputs[1]),
+                "calendar_type": cal,
+                "max_doy": ndays,
+                "gcm_time_coord": "time",
+                "gcm_lat_coord": "lat",
+                "gcm_lon_coord": "lon",
+                "original_hist_path": str(gcm_hist_path),
+                "original_fut_path": str(gcm_fut_path),
+                "model_name": model_name,
+                "diagnostics": stats,
+            }
+            for array in (answer["corrected_hist_pr"], answer["corrected_fut_pr"]):
+                array.attrs["calibration_period"] = f"{ref_start_year}-{ref_end_year}"
+            _release(reference, remove=True)
+            return answer
+    except Exception:
+        for buffer in buffers:
+            _release(buffer, remove=True)
+        if strict:
+            raise
+        logger.exception("Bias correction failed")
+        return None
+
+
+def _release(values, *, remove=False):
+    while (
+        not isinstance(values, np.memmap) and getattr(values, "base", None) is not None
+    ):
+        values = values.base
+    if isinstance(values, np.memmap):
+        path = values.filename
+        if not values._mmap.closed:
+            values._mmap.close()
+        if remove:
+            Path(path).unlink(missing_ok=True)
+
+
+def close_result(result, *, remove_work_files=False):
+    if result:
+        for key in ("corrected_hist_pr", "corrected_fut_pr"):
+            if key in result:
+                _release(result[key].values, remove=remove_work_files)

klimapy/BCSD/BiasCorrectionSpatialDownscaling.py

@@ -0,0 +1,89 @@
+"""BCSD 一体化入口及旧导入路径；计算实现位于同一子包下的 BiasCorrection 和 SpatialDownscaling。"""
+
+from .BiasCorrection import BC, create_cdf_spline, close_result
+import logging
+
+logger = logging.getLogger(__name__)
+
+from .SpatialDownscaling import SD, area_weighted_regrid, fft_smooth_three_harmonics
+from .._core.common import (
+    WINDOW_SIZE,
+    N_QUANTILES,
+    MIN_PRECIP,
+    MAX_RATIO,
+    MIN_RATIO,
+    DEBUG_CHECK,
+    DEBUG_PRINT_SAMPLES,
+    calculate_bounds,
+    get_precip_var_name,
+    convert_precip_to_mm_day,
+    cftime_to_datetime,
+    parse_cftime_date,
+    detect_calendar_type,
+    calculate_doy,
+    _safe_min,
+    _safe_max,
+    _safe_mean,
+    print_array_diagnostics,
+)
+
+
+def BCSD(
+    obs_path,
+    gcm_hist_path,
+    gcm_fut_path,
+    model_name,
+    output_dir,
+    ref_start_year,
+    ref_end_year,
+    *,
+    strict=False,
+    bc_options=None,
+    sd_options=None,
+):
+    """
+    BCSD完整流程：先BC后SD
+    """
+    logger.info(f"\n处理模式: {model_name}")
+    logger.info("=" * 60)
+
+    # 第一步：偏差校正 (BC)
+    bc_result = BC(
+        obs_path=obs_path,
+        gcm_hist_path=gcm_hist_path,
+        gcm_fut_path=gcm_fut_path,
+        model_name=model_name,
+        ref_start_year=ref_start_year,
+        ref_end_year=ref_end_year,
+        **dict(bc_options or {}, strict=strict),
+    )
+
+    if not bc_result:
+        logger.info(f"✗ {model_name} BC失败，停止处理")
+        return None
+
+    # 第二步：空间降尺度 (SD)
+    try:
+        sd_result = SD(
+            obs_path=obs_path,
+            bc_result=bc_result,
+            output_dir=output_dir,
+            ref_start_year=ref_start_year,
+            ref_end_year=ref_end_year,
+            **dict(sd_options or {}, strict=strict),
+        )
+    finally:
+        close_result(bc_result, remove_work_files=True)
+
+    if not sd_result:
+        logger.info(f"✗ {model_name} SD失败")
+        return None
+
+    logger.info(f"✓ {model_name} BCSD完整流程完成！")
+    logger.info("=" * 60)
+
+    return {
+        "model_name": model_name,
+        "hist_bcsd_path": sd_result["hist_bcsd_path"],
+        "fut_bcsd_path": sd_result["fut_bcsd_path"],
+    }

klimapy/BCSD/SpatialDownscaling.py

@@ -0,0 +1,198 @@
+"""Native streamed spatial downscaling; climatology and harmonics run in Rust."""
+
+import logging
+from contextlib import ExitStack
+from pathlib import Path
+import tempfile
+
+import numpy as np
+
+from .._backends import rust as _accelerate
+from ..native.correction import smooth
+from .._core.common import MIN_PRECIP, MIN_RATIO, MAX_RATIO
+from .._core.native_adapter import open_precip_dataset
+from .._core.engine import (
+    read_precip,
+    mapped_days,
+    calendar_info,
+    GridWeights,
+    NetCDFWriter,
+    allocate,
+    check_reference,
+)
+
+logger = logging.getLogger(__name__)
+DEBUG_CHECK = False
+__all__ = ["SD", "area_weighted_regrid", "fft_smooth_three_harmonics"]
+
+
+def fft_smooth_three_harmonics(daily_series, n_days, n_harmonics=3):
+    if len(daily_series) != n_days:
+        raise ValueError("Series length differs from n_days")
+    return smooth(daily_series, n_harmonics)
+
+
+def _smooth_block(block, n_harmonics=3):
+    return smooth(block, n_harmonics)
+
+
+def area_weighted_regrid(source_lat, source_lon, source_data, target_lat, target_lon):
+    with GridWeights(source_lat, source_lon, target_lat, target_lon) as plan:
+        return plan.apply(source_data).astype("float32")
+
+
+def SD(
+    obs_path,
+    bc_result,
+    output_dir,
+    ref_start_year,
+    ref_end_year,
+    *,
+    strict=False,
+    time_chunk_size=32,
+    spatial_chunk_size=16,
+):
+    try:
+        if any(
+            isinstance(n, bool) or not isinstance(n, int) or n < 1
+            for n in (time_chunk_size, spatial_chunk_size)
+        ):
+            raise ValueError("Chunk sizes must be positive integers")
+        with ExitStack() as stack:
+            obs = stack.enter_context(open_precip_dataset(obs_path))
+            check_reference(obs.time, ref_start_year, ref_end_year)
+            history, future = [
+                stack.enter_context(open_precip_dataset(bc_result[key]))
+                for key in ("corrected_hist_pr", "corrected_fut_pr")
+            ]
+            for d in ("lat", "lon"):
+                if not np.array_equal(history.coords[d], future.coords[d]):
+                    raise ValueError("Historical/future model grids differ")
+            cal, ndays = calendar_info(history.time)
+            if calendar_info(future.time)[0] != cal:
+                raise ValueError("Historical/future calendars differ")
+            model_name = str(bc_result.get("model_name", "unknown"))
+            outdir = Path(output_dir) / model_name
+            outdir.mkdir(parents=True, exist_ok=True)
+            work = stack.enter_context(
+                tempfile.TemporaryDirectory(prefix=".climate-", dir=outdir)
+            )
+
+            def mapped(shape, dtype, prefix):
+                a = allocate(shape, dtype, work, prefix)
+                stack.callback(a._mmap.close)
+                return a
+
+            shape = (ndays, *obs.shape[1:])
+            sums, counts, climate = (
+                mapped(shape, np.float64, "sum"),
+                mapped(shape, np.int32, "count"),
+                mapped(shape, np.float64, "climatology"),
+            )
+            sums[:] = 0.0
+            counts[:] = 0
+            years = obs.time.labels()[0]
+            days = mapped_days(obs.time, cal)
+            chunk = min(
+                time_chunk_size,
+                max(1, 64 * 1024**2 // (obs.shape[1] * obs.shape[2] * 8)),
+            )
+            for start in range(0, obs.shape[0], chunk):
+                stop = min(start + chunk, obs.shape[0])
+                labels = days[start:stop].copy() - 1
+                labels[
+                    (years[start:stop] < ref_start_year)
+                    | (years[start:stop] > ref_end_year)
+                ] = -1
+                _accelerate.climate_accumulate(
+                    obs.read(time=slice(start, stop)), labels, sums, counts
+                )
+            for row in range(0, obs.shape[1], spatial_chunk_size):
+                sl = (slice(None), slice(row, row + spatial_chunk_size), slice(None))
+                # Group means are part of the Rust climate pipeline, not a NumPy
+                # fallback. Use the native normalizer before harmonic smoothing.
+                climate[sl] = smooth(_accelerate.climate_mean(sums[sl], counts[sl]))
+            coarse_weights = stack.enter_context(
+                GridWeights(
+                    obs.coords["lat"],
+                    obs.coords["lon"],
+                    history.coords["lat"],
+                    history.coords["lon"],
+                )
+            )
+            coarse = np.empty((ndays, *history.shape[1:]))
+            for start in range(0, ndays, chunk):
+                coarse[start : start + chunk] = coarse_weights.apply(
+                    climate[start : start + chunk]
+                )
+            downscaler = stack.enter_context(
+                _accelerate.Downscaler(
+                    history.coords["lat"],
+                    history.coords["lon"],
+                    obs.coords["lat"],
+                    obs.coords["lon"],
+                    coarse,
+                    climate,
+                    minimum=MIN_PRECIP,
+                    low=MIN_RATIO,
+                    high=MAX_RATIO,
+                )
+            )
+            paths, failures = [], []
+            for label, data in [("historical", history), ("future", future)]:
+                labels = mapped_days(data.time, cal) - 1
+                if np.any(labels < 0):
+                    raise ValueError("Unmappable calendar days in model output")
+                path = outdir / f"{model_name}_{label}_BCSD.nc"
+                attrs = {
+                    "units": "mm/day",
+                    "method": "Rust FITPACK BC + low-harmonic SD",
+                    "calendar_type": cal,
+                    "reference_period": f"{ref_start_year}-{ref_end_year}",
+                    "model_name": model_name,
+                    "failed_time_steps": 0,
+                }
+                with NetCDFWriter(
+                    path, data.time, obs.coords["lat"], obs.coords["lon"], attrs
+                ) as writer:
+                    for start in range(0, len(labels), chunk):
+                        stop = min(start + chunk, len(labels))
+                        try:
+                            scaled = downscaler.apply(
+                                data.read(time=slice(start, stop)), labels[start:stop]
+                            )
+                        except Exception as exc:
+                            if strict:
+                                raise
+                            failures.append(
+                                {
+                                    "period": label,
+                                    "start": start,
+                                    "stop": stop,
+                                    "error": str(exc),
+                                }
+                            )
+                            scaled = np.full(
+                                (stop - start, *obs.shape[1:]), np.nan, dtype="float32"
+                            )
+                        writer.write(start, scaled)
+                    attrs["failed_time_steps"] = sum(
+                        f["stop"] - f["start"] for f in failures if f["period"] == label
+                    )
+                paths.append(str(path))
+            return {
+                "hist_bcsd_path": paths[0],
+                "fut_bcsd_path": paths[1],
+                "time_coordinate": "time",
+                "diagnostics": {
+                    "failed_chunks": failures,
+                    "climatology_days": ndays,
+                    "time_chunk_size": chunk,
+                    "spatial_chunk_size": spatial_chunk_size,
+                },
+            }
+    except Exception:
+        if strict:
+            raise
+        logger.exception("Spatial downscaling failed")
+        return None

klimapy/BCSD/__init__.py

@@ -0,0 +1,20 @@
+"""Bias correction and spatial downscaling stages and combined workflow."""
+
+import sys
+
+from . import BiasCorrection, SpatialDownscaling, BiasCorrectionSpatialDownscaling
+from .BiasCorrection import BC
+from .SpatialDownscaling import SD
+
+# Preserve the existing combined-workflow module import.
+BCSD = BiasCorrectionSpatialDownscaling
+sys.modules[__name__ + ".BCSD"] = BCSD
+
+__all__ = [
+    "BiasCorrection",
+    "SpatialDownscaling",
+    "BiasCorrectionSpatialDownscaling",
+    "BC",
+    "SD",
+    "BCSD",
+]