From 459909539ced012d1a255976978472d47dfde320 Mon Sep 17 00:00:00 2001 From: "2569718930@qq.com" <2569718930@qq.com> Date: Thu, 2 Jul 2026 20:16:24 +0800 Subject: [PATCH] Add WeatherNext2 worker and prevent empty scan cache --- .gitignore | 3 + deploy.sh | 3 + docker-compose.yml | 53 ++ docs/weathernext2-integration.md | 97 +++ .../__tests__/modelSummaryDashboard.test.ts | 7 +- .../__tests__/scanTerminalCache.test.ts | 58 ++ .../scan-terminal/use-scan-terminal-query.ts | 20 +- frontend/lib/model-summary.ts | 1 + requirements.lock | 153 ++++- requirements.txt | 9 + src/analysis/deb_algorithm.py | 1 + src/analysis/trend_engine.py | 78 ++- src/analysis/weathernext2_calibration.py | 325 ++++++++++ src/data_collection/weather_sources.py | 41 +- src/data_collection/weathernext2_fetcher.py | 243 ++++++++ src/data_collection/weathernext2_sources.py | 581 ++++++++++++++++++ tests/test_deployment_runtime_config.py | 30 + tests/test_trend_engine.py | 58 ++ tests/test_weathernext2_calibration.py | 92 +++ tests/test_weathernext2_sources.py | 259 ++++++++ tests/test_weathernext2_worker_service.py | 108 ++++ web/scan_terminal_city_row.py | 43 +- web/weathernext2_worker.py | 91 +++ web/weathernext2_worker_service.py | 288 +++++++++ 24 files changed, 2620 insertions(+), 22 deletions(-) create mode 100644 docs/weathernext2-integration.md create mode 100644 frontend/components/dashboard/scan-terminal/__tests__/scanTerminalCache.test.ts create mode 100644 src/analysis/weathernext2_calibration.py create mode 100644 src/data_collection/weathernext2_fetcher.py create mode 100644 src/data_collection/weathernext2_sources.py create mode 100644 tests/test_weathernext2_calibration.py create mode 100644 tests/test_weathernext2_sources.py create mode 100644 tests/test_weathernext2_worker_service.py create mode 100644 web/weathernext2_worker.py create mode 100644 web/weathernext2_worker_service.py diff --git a/.gitignore b/.gitignore index 120355c1..2352188f 100644 --- a/.gitignore +++ b/.gitignore @@ -1,5 +1,8 @@ # Secrets .env +secrets/ +*.service-account.json +gcp-sa.json # Scratch / temp scripts scratch/ diff --git a/deploy.sh b/deploy.sh index 12f67ab7..9213a03f 100644 --- a/deploy.sh +++ b/deploy.sh @@ -380,6 +380,9 @@ compose_up_retry "cache warmer" -d --no-deps polyweather_warmer echo "Updating training settlement worker..." compose_up_retry "training settlement" -d --no-deps polyweather_training_settlement +echo "Updating WeatherNext2 worker..." +compose_up_retry "WeatherNext2 worker" -d --no-deps polyweather_weathernext2_worker + echo "Updating frontend..." compose_up_retry "frontend" -d --no-deps polyweather_frontend diff --git a/docker-compose.yml b/docker-compose.yml index 5fef525a..0c779826 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -139,6 +139,9 @@ services: SUPABASE_SERVICE_ROLE_KEY: ${SUPABASE_SERVICE_ROLE_KEY} SUPABASE_URL: ${SUPABASE_URL} UVICORN_WORKERS: ${UVICORN_WORKERS:-2} + WEATHERNEXT2_CITY_HIGHS_PATH: ${WEATHERNEXT2_CITY_HIGHS_PATH:-/app/data/weathernext2_city_highs.json} + WEATHERNEXT2_ENABLED: ${WEATHERNEXT2_ENABLED:-1} + WEATHERNEXT2_MODEL_DIR: ${WEATHERNEXT2_MODEL_DIR:-/app/data/models/weathernext2_calibrator} healthcheck: interval: 30s retries: 3 @@ -296,6 +299,56 @@ services: volumes: - ${POLYWEATHER_RUNTIME_DATA_DIR:-/var/lib/polyweather}:/var/lib/polyweather - ${POLYWEATHER_RUNTIME_DATA_DIR:-/var/lib/polyweather}:/app/data + polyweather_weathernext2_worker: + command: python -m web.weathernext2_worker + container_name: polyweather_weathernext2_worker + depends_on: + polyweather_redis: + condition: service_healthy + polyweather_web: + condition: service_started + logging: + driver: "json-file" + options: + max-size: "50m" + max-file: "3" + cpus: ${POLYWEATHER_WEATHERNEXT2_CPUS:-1.50} + env_file: *id001 + environment: + GOOGLE_APPLICATION_CREDENTIALS: ${GOOGLE_APPLICATION_CREDENTIALS:-/app/secrets/gcp-sa.json} + POLYWEATHER_EVENT_STORE: ${POLYWEATHER_EVENT_STORE:-redis} + POLYWEATHER_OBSERVATION_COLLECTOR_ENABLED: 'false' + POLYWEATHER_REDIS_URL: ${POLYWEATHER_REDIS_URL:-redis://polyweather_redis:6379/0} + POLYWEATHER_SCAN_TERMINAL_PREWARM_ENABLED: 'false' + POLYWEATHER_SERVICE_ROLE: weathernext2_worker + WEATHERNEXT2_BACKEND: ${WEATHERNEXT2_BACKEND:-gcs_zarr} + WEATHERNEXT2_CITY_HIGHS_PATH: ${WEATHERNEXT2_CITY_HIGHS_PATH:-/app/data/weathernext2_city_highs.json} + WEATHERNEXT2_DATA_ROOT: ${WEATHERNEXT2_DATA_ROOT:-/app/data/weathernext2} + WEATHERNEXT2_ENABLED: ${WEATHERNEXT2_ENABLED:-1} + WEATHERNEXT2_GCS_ZARR_URI: ${WEATHERNEXT2_GCS_ZARR_URI:-gs://weathernext/weathernext_2_0_0/zarr} + WEATHERNEXT2_MODEL_DIR: ${WEATHERNEXT2_MODEL_DIR:-/app/data/models/weathernext2_calibrator} + WEATHERNEXT2_WORKER_INITIAL_DELAY_SEC: ${WEATHERNEXT2_WORKER_INITIAL_DELAY_SEC:-90} + WEATHERNEXT2_WORKER_INTERVAL_SEC: ${WEATHERNEXT2_WORKER_INTERVAL_SEC:-21600} + healthcheck: + interval: 60s + retries: 3 + test: + - CMD + - python + - -c + - import sqlite3; c=sqlite3.connect('/var/lib/polyweather/polyweather.db'); + c.execute('SELECT 1'); c.close() + timeout: 10s + image: ghcr.io/yangyuan-zhen/polyweather-backend:${IMAGE_TAG:-latest} + mem_limit: ${POLYWEATHER_WEATHERNEXT2_MEM_LIMIT:-2g} + memswap_limit: ${POLYWEATHER_WEATHERNEXT2_MEMSWAP_LIMIT:-3g} + pids_limit: 512 + restart: unless-stopped + user: ${UID:-1000}:${GID:-1000} + volumes: + - ${POLYWEATHER_RUNTIME_DATA_DIR:-/var/lib/polyweather}:/var/lib/polyweather + - ${POLYWEATHER_RUNTIME_DATA_DIR:-/var/lib/polyweather}:/app/data + - ./secrets:/app/secrets:ro x-polyweather-base: env_file: *id001 image: ghcr.io/yangyuan-zhen/polyweather-backend:${IMAGE_TAG:-latest} diff --git a/docs/weathernext2-integration.md b/docs/weathernext2-integration.md new file mode 100644 index 00000000..834243a5 --- /dev/null +++ b/docs/weathernext2-integration.md @@ -0,0 +1,97 @@ +# WeatherNext 2 接入说明 + +PolyWeather 已新增 WeatherNext 2 的内部接入层，当前定位是“概率底座”和“DEB 参考模型”，不是公开 API。 + +## 已落地 + +- `src.data_collection.weathernext2_sources` + - 将 64 成员最高温预报聚合成 Polymarket 口径概率桶。 + - 摄氏城市按单温度选项，例如 `33°C`。 + - 华氏城市按两度市场选项，例如 `94-95°F`。 + - 支持从 hourly 成员序列按城市当地日期切出今日最高温。 +- `WeatherDataCollector.fetch_weathernext2_probability` + - 默认关闭，不影响现有生产采集。 + - `WEATHERNEXT2_ENABLED=1` 后启用。 + - 优先读取 worker 生成的 `/app/data/weathernext2_city_highs.json`，再降级读取 fixture。 + - 若存在 LightGBM 校准模型，会对 raw 成员分布做 quantile residual 校准后再生成市场桶概率。 +- `web.weathernext2_worker` + - 独立低频 worker，不在 API/Telegram 请求中训练或访问 GCS。 + - 从 GCS Zarr 对城市经纬度最近格点取 2m temperature hourly ensemble。 + - 按城市当地日期切今日最高温，写入离线 JSON 产物。 + - 使用 `training_feature_records_store` + `truth_records_store` 训练 LightGBM q10/q50/q90 residual calibrator。 +- `trend_engine` + - 若 `weather_data["weathernext2"]` 存在，`distribution_full` 使用 WeatherNext 2 概率桶，`probability_engine=weathernext2`。 + - WeatherNext 2 ensemble median/mean 会作为 `WeatherNext 2` 进入 `current_forecasts`，供 DEB 融合参考。 + +## 环境变量 + +```bash +WEATHERNEXT2_ENABLED=1 +WEATHERNEXT2_BACKEND=gcs_zarr +WEATHERNEXT2_DATA_ROOT=/app/data/weathernext2 +WEATHERNEXT2_CITY_HIGHS_PATH=/app/data/weathernext2_city_highs.json +WEATHERNEXT2_MODEL_DIR=/app/data/models/weathernext2_calibrator +WEATHERNEXT2_WORKER_INTERVAL_SEC=21600 +WEATHERNEXT2_CACHE_TTL_SEC=21600 +WEATHERNEXT2_GCS_ZARR_URI=gs://weathernext/weathernext_2_0_0/zarr +WEATHERNEXT2_MEAN_GCS_ZARR_URI=gs://weathernext/weathernext_2_0_0_mean/zarr +GOOGLE_APPLICATION_CREDENTIALS=/app/secrets/gcp-sa.json +``` + +Google 官方说明中 WeatherNext 2 GCS Zarr 路径为 `gs://weathernext/weathernext_2_0_0/zarr`，访问前需要 WeatherNext Data Request 权限。缺 GCP 凭据或无访问权限时，worker 会失败退出本轮，不覆盖旧产物；Web/API 仍按现有数据源运行。 + +## Artifact / Fixture 格式 + +```json +{ + "schema_version": 1, + "source": "weathernext2", + "backend": "gcs_zarr", + "generated_at": "2026-07-01T10:00:00+00:00", + "cities": { + "houston": { + "target_date": "2026-06-29", + "source_run": "2026-06-29T00:00:00Z", + "member_highs": { + "member_00": 94.1, + "member_01": 94.8, + "member_02": 95.2, + "member_03": 96.7 + }, + "summary": { + "members": 4, + "median": 95.0 + }, + "buckets": [] + } + } +} +``` + +也可以提供 hourly 成员序列： + +```json +{ + "shanghai": { + "target_date": "2026-06-29", + "timezone_offset_seconds": 28800, + "utc_times": ["2026-06-28T16:00:00Z", "2026-06-29T06:00:00Z"], + "member_hourly": { + "member_00": [31.2, 33.0], + "member_01": [30.8, 32.6] + } + } +} +``` + +## 后续真实拉取 + +Google 官方 WeatherNext 2 数据源包括 GCS Zarr、BigQuery、Earth Engine 和 Vertex AI。当前第一版已经按独立离线 job 落地： + +- 每 6 小时拉取最近一次 WeatherNext 2 run。 +- 对 50 个城市插值到机场/结算点。 +- 切城市当地日期今日最高温。 +- 写入 `/app/data/weathernext2_city_highs.json` 或数据库表。 +- Web/API 只读本地结果，避免实时请求 Google 数据集影响响应时间。 + +后续可以评估 BigQuery 后端作为替代路径，但不应让 API 请求直接依赖外部 WeatherNext 数据集。 diff --git a/frontend/components/dashboard/scan-terminal/__tests__/modelSummaryDashboard.test.ts b/frontend/components/dashboard/scan-terminal/__tests__/modelSummaryDashboard.test.ts index 0da66584..2a88fdf8 100644 --- a/frontend/components/dashboard/scan-terminal/__tests__/modelSummaryDashboard.test.ts +++ b/frontend/components/dashboard/scan-terminal/__tests__/modelSummaryDashboard.test.ts @@ -30,6 +30,7 @@ export function runTests() { model_cluster_sources: { ECMWF: 28.8, GFS: 29.4, + "WeatherNext 2": 29.8, }, }, { @@ -144,8 +145,9 @@ export function runTests() { assert( MODEL_SUMMARY_MODEL_COLUMNS.map((column) => column.key).includes("AROME HD") && MODEL_SUMMARY_MODEL_COLUMNS.map((column) => column.key).includes("HRRR") && - MODEL_SUMMARY_MODEL_COLUMNS.map((column) => column.key).includes("NAM"), - "model summary must expose the fixed model columns including optional short-range models", + MODEL_SUMMARY_MODEL_COLUMNS.map((column) => column.key).includes("NAM") && + MODEL_SUMMARY_MODEL_COLUMNS.map((column) => column.key).includes("WeatherNext 2"), + "model summary must expose the fixed model columns including optional short-range and WeatherNext 2 models", ); assert(summaryRows.length === 5, "model summary should keep one row per city"); assert(summaryRows[0].cityName === "Beijing", "model summary should sort by resolved region then city name"); @@ -179,6 +181,7 @@ export function runTests() { "model summary should aggregate Fahrenheit probabilities into two-degree market option labels", ); assert(parisRow.marketMatches.length === 3, "model summary should keep every Polymarket tradable bucket"); + assert(beijingRow?.models["WeatherNext 2"] === 29.8, "model summary should preserve WeatherNext 2 model representative"); assert( parisRow.marketMatches[0].label === "32°C" && parisRow.marketMatches[0].modelProbability === null && diff --git a/frontend/components/dashboard/scan-terminal/__tests__/scanTerminalCache.test.ts b/frontend/components/dashboard/scan-terminal/__tests__/scanTerminalCache.test.ts new file mode 100644 index 00000000..6b099231 --- /dev/null +++ b/frontend/components/dashboard/scan-terminal/__tests__/scanTerminalCache.test.ts @@ -0,0 +1,58 @@ +import { + readScanCache, + writeScanCache, +} from "@/components/dashboard/scan-terminal/use-scan-terminal-query"; + +function assert(condition: unknown, message: string) { + if (!condition) throw new Error(message); +} + +export function runTests() { + const originalLocalStorage = globalThis.localStorage; + const storage = new Map(); + const memoryStorage = { + getItem: (key: string) => storage.get(key) ?? null, + removeItem: (key: string) => { storage.delete(key); }, + setItem: (key: string, value: string) => { storage.set(key, value); }, + } as Storage; + + Object.defineProperty(globalThis, "localStorage", { + configurable: true, + value: memoryStorage, + }); + + try { + writeScanCache( + { + generated_at: "2026-07-02T12:00:00Z", + rows: [{ city: "ankara", deb_prediction: 20.3 } as any], + status: "ready", + } as any, + "all", + "model-summary", + ); + assert( + readScanCache("all", "model-summary", { allowStale: true })?.rows?.length === 1, + "non-empty scan cache should remain readable", + ); + + writeScanCache( + { + generated_at: "2026-07-02T12:05:00Z", + rows: [], + status: "ready", + } as any, + "all", + "model-summary", + ); + assert( + readScanCache("all", "model-summary", { allowStale: true }) === null, + "empty scan cache responses must be ignored so model summary can revalidate rows", + ); + } finally { + Object.defineProperty(globalThis, "localStorage", { + configurable: true, + value: originalLocalStorage, + }); + } +} diff --git a/frontend/components/dashboard/scan-terminal/use-scan-terminal-query.ts b/frontend/components/dashboard/scan-terminal/use-scan-terminal-query.ts index 4ab1f6aa..a8ee7664 100644 --- a/frontend/components/dashboard/scan-terminal/use-scan-terminal-query.ts +++ b/frontend/components/dashboard/scan-terminal/use-scan-terminal-query.ts @@ -36,7 +36,7 @@ function scanCacheKey(tradingRegion: string, cacheScope: string): string { return `${SCAN_CACHE_PREFIX}:${cacheScope || "terminal"}:${tradingRegion || "all"}`; } -function readScanCache( +export function readScanCache( tradingRegion: string, cacheScope: string, options?: { allowStale?: boolean }, @@ -47,15 +47,27 @@ function readScanCache( const cached = JSON.parse(raw); const age = Date.now() - Number(cached.ts || 0); const maxAge = options?.allowStale ? MAX_STALE_SCAN_CACHE_MS : SCAN_CACHE_TTL_MS; - if (cached.ts && age >= 0 && age < maxAge && cached.data?.rows) { + if ( + cached.ts && + age >= 0 && + age < maxAge && + Array.isArray(cached.data?.rows) && + cached.data.rows.length > 0 + ) { return cached.data; } } catch { /* ignore */ } return null; } -function writeScanCache(data: ScanTerminalResponse, tradingRegion: string, cacheScope: string) { - try { localStorage.setItem(scanCacheKey(tradingRegion, cacheScope), JSON.stringify({ ts: Date.now(), data })); } catch { /* ignore */ } +export function writeScanCache(data: ScanTerminalResponse, tradingRegion: string, cacheScope: string) { + try { + if (!Array.isArray(data.rows) || data.rows.length <= 0) { + localStorage.removeItem(scanCacheKey(tradingRegion, cacheScope)); + return; + } + localStorage.setItem(scanCacheKey(tradingRegion, cacheScope), JSON.stringify({ ts: Date.now(), data })); + } catch { /* ignore */ } } function normalizeCityKey(city: string | null | undefined) { diff --git a/frontend/lib/model-summary.ts b/frontend/lib/model-summary.ts index 61192780..aa1eaf8d 100644 --- a/frontend/lib/model-summary.ts +++ b/frontend/lib/model-summary.ts @@ -16,6 +16,7 @@ export const MODEL_SUMMARY_MODEL_COLUMNS = [ { key: "AROME HD", label: "AROME HD" }, { key: "HRRR", label: "HRRR" }, { key: "NAM", label: "NAM" }, + { key: "WeatherNext 2", label: "WeatherNext 2" }, ] as const; export type ModelSummaryColumnKey = (typeof MODEL_SUMMARY_MODEL_COLUMNS)[number]["key"]; diff --git a/requirements.lock b/requirements.lock index a6f3cd1d..5a5490d0 100644 --- a/requirements.lock +++ b/requirements.lock @@ -4,6 +4,7 @@ aiohappyeyeballs==2.6.2 # via aiohttp aiohttp==3.14.1 # via + # gcsfs # pytelegrambotapi # web3 aiosignal==1.4.0 @@ -28,6 +29,8 @@ certifi==2026.5.20 # httpx # netcdf4 # requests +cffi==2.0.0 ; platform_python_implementation != 'PyPy' + # via cryptography cftime==1.6.5 # via netcdf4 charset-normalizer==3.4.7 @@ -40,8 +43,14 @@ colorama==0.4.6 ; sys_platform == 'win32' # via # click # loguru +cryptography==49.0.0 + # via google-auth cytoolz==1.1.0 ; implementation_name == 'cpython' # via eth-utils +decorator==5.3.1 + # via gcsfs +donfig==0.8.1.post1 + # via zarr eth-abi==5.2.0 # via # eth-account @@ -81,6 +90,58 @@ frozenlist==1.8.0 # via # aiohttp # aiosignal +fsspec==2026.6.0 + # via + # -r requirements.txt + # gcsfs +gcsfs==2026.6.0 + # via -r requirements.txt +google-api-core==2.31.0 + # via + # google-cloud-core + # google-cloud-storage + # google-cloud-storage-control +google-auth==2.55.1 + # via + # gcsfs + # google-api-core + # google-auth-oauthlib + # google-cloud-core + # google-cloud-storage + # google-cloud-storage-control +google-auth-oauthlib==1.4.0 + # via gcsfs +google-cloud-core==2.6.0 + # via google-cloud-storage +google-cloud-storage==3.12.0 + # via + # -r requirements.txt + # gcsfs +google-cloud-storage-control==1.12.0 + # via gcsfs +google-crc32c==1.8.0 + # via + # google-cloud-storage + # google-resumable-media + # zarr +google-resumable-media==2.10.0 + # via google-cloud-storage +googleapis-common-protos==1.75.0 + # via + # google-api-core + # grpc-google-iam-v1 + # grpcio-status +grpc-google-iam-v1==0.14.4 + # via google-cloud-storage-control +grpcio==1.81.1 + # via + # google-api-core + # google-cloud-storage-control + # googleapis-common-protos + # grpc-google-iam-v1 + # grpcio-status +grpcio-status==1.81.1 + # via google-api-core h11==0.16.0 # via # httpcore @@ -100,33 +161,72 @@ idna==3.18 # httpx # requests # yarl +joblib==1.5.3 + # via + # -r requirements.txt + # scikit-learn +lightgbm==4.6.0 + # via -r requirements.txt loguru==0.7.3 # via -r requirements.txt multidict==6.7.1 # via # aiohttp # yarl +narwhals==2.23.0 + # via scikit-learn netcdf4==1.7.4 # via -r requirements.txt +numcodecs==0.16.5 + # via zarr numpy==2.4.6 - -torch -typing_extensions==4.15.0 - # via torch # via # -r requirements.txt # cftime + # lightgbm # netcdf4 + # numcodecs + # pandas + # scikit-learn + # scipy + # xarray + # zarr +oauthlib==3.3.1 + # via requests-oauthlib +packaging==26.2 + # via + # xarray + # zarr +pandas==3.0.3 + # via xarray parsimonious==0.10.0 # via eth-abi propcache==0.5.2 # via # aiohttp # yarl +proto-plus==1.28.0 + # via + # google-api-core + # google-cloud-storage-control +protobuf==7.35.1 + # via + # google-api-core + # google-cloud-storage-control + # googleapis-common-protos + # grpc-google-iam-v1 + # grpcio-status + # proto-plus pyaes==1.6.1 # via telethon pyasn1==0.6.3 - # via rsa + # via + # pyasn1-modules + # rsa +pyasn1-modules==0.4.2 + # via google-auth +pycparser==3.0 ; implementation_name != 'PyPy' and platform_python_implementation != 'PyPy' + # via cffi pycryptodome==3.23.0 # via # eth-hash @@ -141,6 +241,8 @@ pydantic-core==2.46.4 # via pydantic pytelegrambotapi==4.34.0 # via -r requirements.txt +python-dateutil==2.9.0.post0 + # via pandas python-dotenv==1.2.2 # via -r requirements.txt pyunormalize==17.0.0 @@ -148,7 +250,9 @@ pyunormalize==17.0.0 pywin32==312 ; sys_platform == 'win32' # via web3 pyyaml==6.0.3 - # via -r requirements.txt + # via + # -r requirements.txt + # donfig redis==6.4.0 # via -r requirements.txt regex==2026.5.9 @@ -156,18 +260,40 @@ regex==2026.5.9 requests==2.34.2 # via # -r requirements.txt + # gcsfs + # google-api-core + # google-cloud-storage # pytelegrambotapi + # requests-oauthlib # web3 +requests-oauthlib==2.0.0 + # via google-auth-oauthlib rlp==4.1.0 # via # eth-account # eth-rlp rsa==4.9.1 # via telethon +scikit-learn==1.9.0 + # via -r requirements.txt +scipy==1.17.1 ; python_full_version < '3.12' + # via + # -r requirements.txt + # lightgbm + # scikit-learn +scipy==1.18.0 ; python_full_version >= '3.12' + # via + # -r requirements.txt + # lightgbm + # scikit-learn +six==1.17.0 + # via python-dateutil starlette==1.3.1 # via fastapi telethon==1.43.2 # via -r requirements.txt +threadpoolctl==3.6.0 + # via scikit-learn toolz==1.1.0 ; implementation_name == 'cpython' or implementation_name == 'pypy' # via # cytoolz @@ -181,17 +307,22 @@ typing-extensions==4.15.0 # anyio # eth-typing # fastapi + # grpcio + # numcodecs # pydantic # pydantic-core # starlette # typing-inspection # web3 + # zarr typing-inspection==0.4.2 # via # fastapi # pydantic -tzdata==2026.2 ; sys_platform == 'win32' - # via -r requirements.txt +tzdata==2026.2 ; sys_platform == 'emscripten' or sys_platform == 'win32' + # via + # -r requirements.txt + # pandas urllib3==2.7.0 # via # requests @@ -206,5 +337,11 @@ websockets==15.0.1 # web3 win32-setctime==1.2.0 ; sys_platform == 'win32' # via loguru +xarray==2026.4.0 + # via -r requirements.txt yarl==1.24.2 # via aiohttp +zarr==3.1.6 ; python_full_version < '3.12' + # via -r requirements.txt +zarr==3.2.1 ; python_full_version >= '3.12' + # via -r requirements.txt diff --git a/requirements.txt b/requirements.txt index fe0e1854..2c51f888 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,6 +7,15 @@ python-dotenv>=1.1,<2 PyYAML>=6.0,<7 netCDF4>=1.7,<2 numpy>=2.2,<3 +scipy>=1.10,<2 +scikit-learn>=1.4,<2 +lightgbm>=4.6,<5 +joblib>=1.3,<2 +xarray>=2024.7,<2027 +zarr>=2.18,<4 +fsspec>=2024.10,<2027 +gcsfs>=2024.10,<2027 +google-cloud-storage>=2.18,<4 web3>=7.13,<8 fastapi>=0.115,<1 uvicorn>=0.34,<1 diff --git a/src/analysis/deb_algorithm.py b/src/analysis/deb_algorithm.py index 1e62c7b4..7a59a1c8 100644 --- a/src/analysis/deb_algorithm.py +++ b/src/analysis/deb_algorithm.py @@ -110,6 +110,7 @@ def _deb_model_priority(model_name: str) -> int: "hko": 45, "lgbm": 50, "openmeteo": 15, + "weathernext2": 30, }.get(normalized, 10) diff --git a/src/analysis/trend_engine.py b/src/analysis/trend_engine.py index a47bc1df..3f9a0f5f 100644 --- a/src/analysis/trend_engine.py +++ b/src/analysis/trend_engine.py @@ -50,6 +50,49 @@ def _sf(v): return None +def _weathernext2_probability_payload( + weather_data: Dict[str, Any], +) -> Optional[Dict[str, Any]]: + source = weather_data.get("weathernext2") + if not isinstance(source, dict): + return None + raw_buckets = source.get("buckets") + if not isinstance(raw_buckets, list) or not raw_buckets: + return None + + buckets = [] + for bucket in raw_buckets: + if not isinstance(bucket, dict): + continue + probability = _sf(bucket.get("probability")) + if probability is None or probability <= 0: + continue + copied = dict(bucket) + copied["probability"] = round(probability, 3) + buckets.append(copied) + if not buckets: + return None + + summary = source.get("summary") if isinstance(source.get("summary"), dict) else {} + mu = _sf(summary.get("median")) + if mu is None: + mu = _sf(summary.get("mean")) + if mu is None: + top_bucket = max(buckets, key=lambda item: _sf(item.get("probability")) or 0) + mu = _sf(top_bucket.get("value")) + + return { + "engine": "weathernext2", + "mu": mu, + "probabilities": sorted( + buckets, + key=lambda item: _sf(item.get("probability")) or 0, + reverse=True, + )[:4], + "probabilities_all": buckets, + } + + def _median(values: List[float]) -> Optional[float]: if not values: return None @@ -548,6 +591,18 @@ def analyze_weather_trend( for m_name, m_val in mm_forecasts.items(): if m_val is not None and not _is_excluded_model_name(m_name): current_forecasts[m_name] = _sf(m_val) + weathernext2 = weather_data.get("weathernext2") + if isinstance(weathernext2, dict): + weathernext2_summary = ( + weathernext2.get("summary") + if isinstance(weathernext2.get("summary"), dict) + else {} + ) + weathernext2_median = _sf(weathernext2_summary.get("median")) + if weathernext2_median is None: + weathernext2_median = _sf(weathernext2_summary.get("mean")) + if weathernext2_median is not None: + current_forecasts["WeatherNext 2"] = weathernext2_median forecast_highs = [h for h in current_forecasts.values() if h is not None] forecast_high = max(forecast_highs) if forecast_highs else None forecast_median = ( @@ -865,9 +920,12 @@ def analyze_weather_trend( # === Probability Engine === probabilities: List[Dict[str, Any]] = [] probabilities_all: List[Dict[str, Any]] = [] + probability_engine = "legacy" forecast_miss_deg = 0.0 + weathernext2_probs = _weathernext2_probability_payload(weather_data) if is_dead_market: + probability_engine = "dead_market" settled_wu = apply_city_settlement(city_name, max_so_far) if max_so_far is not None else 0 dead_msg = ( f"🎲 结算预测：已锁定 {settled_wu}{temp_symbol} " @@ -881,6 +939,24 @@ def analyze_weather_trend( {"value": settled_wu, "range": f"[{settled_wu-0.5}~{settled_wu+0.5})", "probability": 1.0} ] probabilities_all = probabilities + elif weathernext2_probs: + if max_so_far is not None and forecast_median is not None: + forecast_miss_deg = round(forecast_median - max_so_far, 1) + probability_engine = "weathernext2" + mu = weathernext2_probs.get("mu") or mu + probabilities = weathernext2_probs.get("probabilities", []) + probabilities_all = weathernext2_probs.get("probabilities_all", probabilities) + prob_parts = [] + for bucket in probabilities[:4]: + label = str(bucket.get("label") or bucket.get("range") or bucket.get("value") or "").strip() + probability = _sf(bucket.get("probability")) + if label and probability is not None: + prob_parts.append(f"{label} {probability * 100:.0f}%") + if prob_parts: + mu_label = f"μ={mu:.1f}" if mu is not None else "μ=--" + prob_str = " | ".join(prob_parts) + insights.append(f"🎲 WeatherNext 2 概率 ({mu_label})：{prob_str}") + ai_features.append(f"🎲 WeatherNext 2 概率分布：{prob_str}") elif (ens_p10 is not None and ens_p90 is not None) or fallback_sigma: # Forecast miss magnitude if max_so_far is not None and forecast_median is not None: @@ -1117,7 +1193,7 @@ def analyze_weather_trend( "mu": mu, "probabilities": probabilities, "probabilities_all": probabilities_all or probabilities, - "probability_engine": "legacy", + "probability_engine": probability_engine, "trend_info": { "direction": trend_direction if 'trend_direction' in dir() else "unknown", "recent": recent_list, diff --git a/src/analysis/weathernext2_calibration.py b/src/analysis/weathernext2_calibration.py new file mode 100644 index 00000000..d92938d0 --- /dev/null +++ b/src/analysis/weathernext2_calibration.py @@ -0,0 +1,325 @@ +from __future__ import annotations + +import json +import math +import os +import time +from pathlib import Path +from typing import Any, Dict, Iterable, Optional + +from src.data_collection.weathernext2_sources import build_weathernext2_city_probability + + +QUANTILES = {"q10": 0.10, "q50": 0.50, "q90": 0.90} +FEATURE_NAMES = [ + "city_code", + "wn2_mean", + "wn2_median", + "wn2_p10", + "wn2_p25", + "wn2_p75", + "wn2_p90", + "wn2_spread", + "deb_prediction_c", + "model_median_c", + "model_spread", + "current_max_so_far_c", + "local_hour", + "month", + "day_of_year", + "observation_progress", +] + + +def _sf(value: Any) -> Optional[float]: + try: + if value is None or value == "": + return None + parsed = float(value) + except (TypeError, ValueError): + return None + return parsed if math.isfinite(parsed) else None + + +def _date_parts(value: Any) -> tuple[float, float]: + text = str(value or "").strip() + try: + parsed = time.strptime(text[:10], "%Y-%m-%d") + return float(parsed.tm_mon), float(parsed.tm_yday) + except Exception: + return 0.0, 0.0 + + +def _summary_from_record(record: Dict[str, Any]) -> Dict[str, Any]: + wn2 = record.get("weathernext2") if isinstance(record.get("weathernext2"), dict) else {} + summary = wn2.get("summary") if isinstance(wn2.get("summary"), dict) else {} + return summary if isinstance(summary, dict) else {} + + +def _city_key(value: Any) -> str: + return str(value or "").strip().lower() + + +def _build_city_index(records: Iterable[Dict[str, Any]]) -> Dict[str, int]: + cities = sorted({_city_key(record.get("city")) for record in records if _city_key(record.get("city"))}) + return {city: idx for idx, city in enumerate(cities)} + + +def _feature_row(record: Dict[str, Any], city_index: Dict[str, int]) -> Optional[list[float]]: + summary = _summary_from_record(record) + median = _sf(summary.get("median")) + if median is None: + return None + target_date = record.get("target_date") or record.get("date") + month, day_of_year = _date_parts(target_date) + city = _city_key(record.get("city")) + current = _sf(record.get("current_max_so_far_c")) + progress = _sf(record.get("observation_progress")) + if progress is None: + local_hour = _sf(record.get("local_hour")) + progress = min(max((local_hour or 0.0) / 24.0, 0.0), 1.0) + + def fallback(name: str, default: float) -> float: + parsed = _sf(summary.get(name)) + return default if parsed is None else parsed + + return [ + float(city_index.get(city, -1)), + fallback("mean", median), + median, + fallback("p10", median), + fallback("p25", median), + fallback("p75", median), + fallback("p90", median), + fallback("spread", 0.0), + _sf(record.get("deb_prediction_c")) or median, + _sf(record.get("model_median_c")) or median, + _sf(record.get("model_spread")) or 0.0, + current if current is not None else median, + _sf(record.get("local_hour")) or 0.0, + month, + day_of_year, + progress, + ] + + +def _training_xy( + records: Iterable[Dict[str, Any]], + city_index: Dict[str, int], +) -> tuple[list[list[float]], list[float], list[Dict[str, Any]]]: + features: list[list[float]] = [] + residuals: list[float] = [] + kept: list[Dict[str, Any]] = [] + for record in records: + summary = _summary_from_record(record) + median = _sf(summary.get("median")) + actual = _sf(record.get("actual_high_c", record.get("actual_high"))) + row = _feature_row(record, city_index) + if median is None or actual is None or row is None: + continue + features.append(row) + residuals.append(actual - median) + kept.append(record) + return features, residuals, kept + + +def train_lightgbm_quantile_calibrator( + records: Iterable[Dict[str, Any]], + *, + model_dir: os.PathLike[str] | str, + min_global_samples: int = 150, + min_city_samples: int = 5, +) -> Dict[str, Any]: + rows = [record for record in records if isinstance(record, dict)] + city_index = _build_city_index(rows) + features, residuals, kept = _training_xy(rows, city_index) + if len(features) < int(min_global_samples): + return { + "trained": False, + "reason": "insufficient_global_samples", + "samples": len(features), + } + + city_counts: Dict[str, int] = {} + for record in kept: + city_counts[_city_key(record.get("city"))] = city_counts.get(_city_key(record.get("city")), 0) + 1 + if not any(count >= int(min_city_samples) for count in city_counts.values()): + return { + "trained": False, + "reason": "insufficient_city_samples", + "samples": len(features), + } + + try: + import joblib # type: ignore + from lightgbm import LGBMRegressor # type: ignore + except Exception as exc: + return { + "trained": False, + "reason": "missing_lightgbm", + "samples": len(features), + "error": str(exc), + } + + target_dir = Path(model_dir) + target_dir.mkdir(parents=True, exist_ok=True) + models = {} + for key, alpha in QUANTILES.items(): + model = LGBMRegressor( + objective="quantile", + alpha=alpha, + n_estimators=45, + learning_rate=0.08, + num_leaves=15, + min_child_samples=5, + random_state=42, + n_jobs=2, + verbosity=-1, + ) + model.fit(features, residuals) + models[key] = model + joblib.dump(model, target_dir / f"{key}.pkl") + + metadata = { + "model_version": f"weathernext2_lightgbm_quantile_{int(time.time())}", + "engine": "lightgbm_quantile", + "samples": len(features), + "feature_names": FEATURE_NAMES, + "city_index": city_index, + "city_counts": city_counts, + "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), + } + (target_dir / "metadata.json").write_text( + json.dumps(metadata, ensure_ascii=False, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + ordered = True + for feature in features[: min(len(features), 50)]: + preds = sorted(float(models[key].predict([feature])[0]) for key in ("q10", "q50", "q90")) + if preds[0] > preds[1] or preds[1] > preds[2]: + ordered = False + break + return { + "trained": True, + "samples": len(features), + "model_dir": str(target_dir), + "model_version": metadata["model_version"], + "validation": {"ordered_quantiles": ordered}, + } + + +def _load_model_bundle(model_dir: os.PathLike[str] | str) -> Optional[Dict[str, Any]]: + target_dir = Path(model_dir) + metadata_path = target_dir / "metadata.json" + if not metadata_path.is_file(): + return None + try: + import joblib # type: ignore + + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + return { + "metadata": metadata, + "models": { + key: joblib.load(target_dir / f"{key}.pkl") + for key in QUANTILES + }, + } + except Exception: + return None + + +def _predict_residual_quantiles(record: Dict[str, Any], bundle: Dict[str, Any]) -> Optional[Dict[str, float]]: + metadata = bundle.get("metadata") or {} + city_index = metadata.get("city_index") if isinstance(metadata.get("city_index"), dict) else {} + feature = _feature_row(record, city_index) + if feature is None: + return None + raw = { + key: float(model.predict([feature])[0]) + for key, model in (bundle.get("models") or {}).items() + } + values = sorted([raw.get("q10", 0.0), raw.get("q50", 0.0), raw.get("q90", 0.0)]) + return {"q10": values[0], "q50": values[1], "q90": values[2]} + + +def _calibrate_members( + member_highs: Dict[str, Any], + raw_summary: Dict[str, Any], + residuals: Dict[str, float], +) -> Dict[str, float]: + raw_median = _sf(raw_summary.get("median")) + raw_p10 = _sf(raw_summary.get("p10")) + raw_p90 = _sf(raw_summary.get("p90")) + if raw_median is None: + return {} + target_p10 = (raw_p10 if raw_p10 is not None else raw_median) + residuals["q10"] + target_median = raw_median + residuals["q50"] + target_p90 = (raw_p90 if raw_p90 is not None else raw_median) + residuals["q90"] + + calibrated = {} + for member_id, value in (member_highs or {}).items(): + parsed = _sf(value) + if parsed is None: + continue + if parsed <= raw_median: + raw_span = max(raw_median - (raw_p10 if raw_p10 is not None else raw_median), 0.1) + target_span = max(target_median - target_p10, 0.1) + adjusted = target_median - (raw_median - parsed) / raw_span * target_span + else: + raw_span = max((raw_p90 if raw_p90 is not None else raw_median) - raw_median, 0.1) + target_span = max(target_p90 - target_median, 0.1) + adjusted = target_median + (parsed - raw_median) / raw_span * target_span + calibrated[str(member_id)] = round(adjusted, 1) + return calibrated + + +def apply_quantile_calibration_to_payload( + payload: Dict[str, Any], + *, + model_dir: os.PathLike[str] | str, +) -> Dict[str, Any]: + bundle = _load_model_bundle(model_dir) + if not bundle: + return dict(payload) + raw_summary = payload.get("summary") if isinstance(payload.get("summary"), dict) else {} + member_highs = payload.get("member_highs") if isinstance(payload.get("member_highs"), dict) else {} + if not raw_summary or not member_highs: + return dict(payload) + + record = { + "city": payload.get("city"), + "target_date": payload.get("target_date"), + "weathernext2": {"summary": raw_summary}, + } + residuals = _predict_residual_quantiles(record, bundle) + if residuals is None: + return dict(payload) + + calibrated_members = _calibrate_members(member_highs, raw_summary, residuals) + if not calibrated_members: + return dict(payload) + + calibrated = build_weathernext2_city_probability( + city=str(payload.get("city") or ""), + member_highs=calibrated_members, + temp_symbol=str(payload.get("temp_symbol") or "°C"), + target_date=payload.get("target_date"), + source_run=payload.get("source_run"), + generated_at=payload.get("generated_at"), + ) + metadata = bundle.get("metadata") or {} + calibrated["calibration"] = { + "engine": "lightgbm_quantile", + "model_version": metadata.get("model_version"), + "samples": metadata.get("samples"), + "residual_quantiles": {key: round(value, 3) for key, value in residuals.items()}, + "raw_summary": raw_summary, + "calibrated_summary": calibrated.get("summary"), + } + calibrated["raw_weathernext2"] = { + "summary": raw_summary, + "buckets": payload.get("buckets") or [], + "top_bucket": payload.get("top_bucket"), + } + return calibrated diff --git a/src/data_collection/weather_sources.py b/src/data_collection/weather_sources.py index 3541a75c..474d744b 100644 --- a/src/data_collection/weather_sources.py +++ b/src/data_collection/weather_sources.py @@ -18,6 +18,7 @@ from src.data_collection.metar_sources import MetarSourceMixin from src.data_collection.mgm_sources import MgmSourceMixin from src.data_collection.jma_amedas_sources import JmaAmedasSourceMixin from src.data_collection.nws_open_meteo_sources import NwsOpenMeteoSourceMixin +from src.data_collection.weathernext2_sources import WeatherNext2SourceMixin from src.data_collection.amos_station_sources import AmosStationSourceMixin from src.data_collection.amsc_awos_sources import AmscAwosSourceMixin from src.data_collection.fmi_sources import FmiSourceMixin @@ -35,7 +36,7 @@ from src.data_collection.forecast_source_bundle import fetch_open_meteo_forecast from src.database.db_manager import DBManager -class WeatherDataCollector(OpenMeteoCacheMixin, SettlementSourceMixin, MetarSourceMixin, MgmSourceMixin, JmaAmedasSourceMixin, NwsOpenMeteoSourceMixin, AmosStationSourceMixin, AmscAwosSourceMixin, FmiSourceMixin, KnmiSourceMixin, HkoObsSourceMixin, CowinSourceMixin, MadisSourceMixin, SingaporeMssSourceMixin, ImsSourceMixin, NcmSourceMixin, AerowebSourceMixin, WundergroundHistoricalMixin): +class WeatherDataCollector(OpenMeteoCacheMixin, SettlementSourceMixin, MetarSourceMixin, MgmSourceMixin, JmaAmedasSourceMixin, NwsOpenMeteoSourceMixin, WeatherNext2SourceMixin, AmosStationSourceMixin, AmscAwosSourceMixin, FmiSourceMixin, KnmiSourceMixin, HkoObsSourceMixin, CowinSourceMixin, MadisSourceMixin, SingaporeMssSourceMixin, ImsSourceMixin, NcmSourceMixin, AerowebSourceMixin, WundergroundHistoricalMixin): """ Multi-source weather data collector @@ -189,9 +190,11 @@ class WeatherDataCollector(OpenMeteoCacheMixin, SettlementSourceMixin, MetarSour self._open_meteo_cache: Dict[str, Dict] = {} self._ensemble_cache: Dict[str, Dict] = {} self._multi_model_cache: Dict[str, Dict] = {} + self._weathernext2_cache: Dict[str, Dict] = {} self._open_meteo_cache_lock = threading.Lock() self._ensemble_cache_lock = threading.Lock() self._multi_model_cache_lock = threading.Lock() + self._weathernext2_cache_lock = threading.Lock() # Open-Meteo 共享 429 冷却计时器：触发限流后所有 OM 端点暂停请求 self._open_meteo_rate_limit_until: float = 0.0 self._open_meteo_rl_cooldown: int = int( @@ -1740,6 +1743,26 @@ class WeatherDataCollector(OpenMeteoCacheMixin, SettlementSourceMixin, MetarSour if multi_model_data: results["multi_model"] = multi_model_data + def _attach_weathernext2_model( + self, + results: Dict, + city: str, + lat: float, + lon: float, + use_fahrenheit: bool, + *, + timezone_offset_seconds: Optional[int] = None, + ) -> None: + payload = self.fetch_weathernext2_probability( + city, + lat, + lon, + use_fahrenheit=use_fahrenheit, + timezone_offset_seconds=timezone_offset_seconds, + ) + if payload: + results["weathernext2"] = payload + def fetch_all_sources( self, city: str, @@ -1794,6 +1817,14 @@ class WeatherDataCollector(OpenMeteoCacheMixin, SettlementSourceMixin, MetarSour results["open-meteo"] = open_meteo # 获取时区偏移以过滤 METAR utc_offset = open_meteo.get("utc_offset", 0) + self._attach_weathernext2_model( + results, + city_lower, + lat, + lon, + use_fahrenheit, + timezone_offset_seconds=utc_offset, + ) if supports_aviationweather: metar_data = self.fetch_metar( city, use_fahrenheit=use_fahrenheit, utc_offset=utc_offset @@ -1844,6 +1875,14 @@ class WeatherDataCollector(OpenMeteoCacheMixin, SettlementSourceMixin, MetarSour fallback_utc_offset = int( self.CITY_REGISTRY.get(city_lower, {}).get("tz_offset", 0) ) + self._attach_weathernext2_model( + results, + city_lower, + lat, + lon, + use_fahrenheit, + timezone_offset_seconds=fallback_utc_offset, + ) if supports_aviationweather: metar_data = self.fetch_metar( city, diff --git a/src/data_collection/weathernext2_fetcher.py b/src/data_collection/weathernext2_fetcher.py new file mode 100644 index 00000000..47230cbc --- /dev/null +++ b/src/data_collection/weathernext2_fetcher.py @@ -0,0 +1,243 @@ +from __future__ import annotations + +import math +import os +from datetime import datetime, timezone +from typing import Any, Dict, Mapping, Optional, Sequence + + +TEMPERATURE_VARIABLE_CANDIDATES = ( + "temperature_2m", + "2m_temperature", + "t2m", + "air_temperature_2m", + "temperature", +) + + +def _mapping_keys(dataset: Any) -> list[str]: + if isinstance(dataset, Mapping): + return [str(key) for key in dataset.keys()] + data_vars = getattr(dataset, "data_vars", None) + if data_vars is not None: + try: + return [str(key) for key in data_vars.keys()] + except Exception: + return [] + return [] + + +def select_temperature_variable(dataset: Any, preferred: Optional[str] = None) -> str: + if preferred: + keys = set(_mapping_keys(dataset)) + if preferred in keys: + return preferred + raise KeyError(f"WeatherNext2 temperature variable not found: {preferred}") + + keys = _mapping_keys(dataset) + lowered = {key.lower(): key for key in keys} + for candidate in TEMPERATURE_VARIABLE_CANDIDATES: + if candidate.lower() in lowered: + return lowered[candidate.lower()] + + for key in keys: + normalized = key.lower().replace("-", "_") + if "temperature" in normalized and ("2m" in normalized or "two_meter" in normalized): + return key + raise KeyError("WeatherNext2 2m temperature variable was not found") + + +def normalize_temperature_value( + value: Any, + units: str = "", + *, + use_fahrenheit: bool = False, +) -> Optional[float]: + try: + parsed = float(value) + except (TypeError, ValueError): + return None + if not math.isfinite(parsed): + return None + + unit = str(units or "").strip().lower() + if unit in {"k", "kelvin"} or parsed > 150: + celsius = parsed - 273.15 + elif unit in {"f", "fahrenheit", "degf", "degree_fahrenheit"}: + celsius = (parsed - 32.0) * 5.0 / 9.0 + else: + celsius = parsed + + if use_fahrenheit: + return round(celsius * 9.0 / 5.0 + 32.0, 1) + return round(celsius, 1) + + +def _get_dataset_value(dataset: Any, key: str) -> Any: + if isinstance(dataset, Mapping): + return dataset[key] + return dataset[key] + + +def _as_list(values: Any) -> list[Any]: + if hasattr(values, "values"): + try: + values = values.values + except Exception: + pass + if hasattr(values, "tolist"): + try: + return values.tolist() + except Exception: + pass + return list(values or []) + + +def _nearest_index(values: Sequence[Any], target: float) -> int: + numeric_values = [float(value) for value in values] + return min(range(len(numeric_values)), key=lambda idx: abs(numeric_values[idx] - target)) + + +def _variable_units(dataset: Any, temp_var: str) -> str: + if isinstance(dataset, Mapping): + units = dataset.get("units") + if isinstance(units, Mapping): + return str(units.get(temp_var) or "") + return "" + variable = dataset[temp_var] + attrs = getattr(variable, "attrs", {}) or {} + return str(attrs.get("units") or attrs.get("unit") or "") + + +def _variable_values(dataset: Any, temp_var: str) -> Any: + variable = _get_dataset_value(dataset, temp_var) + if hasattr(variable, "values"): + return variable.values + return variable + + +def _index_nested(values: Any, indexes: Sequence[int]) -> Any: + current = values + for idx in indexes: + current = current[idx] + return current + + +def _dataset_dimension_values(dataset: Any, candidates: Sequence[str]) -> tuple[str, list[Any]]: + keys = _mapping_keys(dataset) + lowered = {key.lower(): key for key in keys} + for candidate in candidates: + key = lowered.get(candidate.lower()) + if key is not None: + return key, _as_list(_get_dataset_value(dataset, key)) + + coords = getattr(dataset, "coords", None) + if coords is not None: + for candidate in candidates: + if candidate in coords: + return candidate, _as_list(coords[candidate]) + raise KeyError(f"WeatherNext2 dimension not found: {candidates}") + + +def _parse_time(value: Any) -> str: + if isinstance(value, datetime): + dt = value + else: + text = str(value or "").strip() + if text.endswith("Z"): + text = f"{text[:-1]}+00:00" + try: + dt = datetime.fromisoformat(text) + except ValueError: + return str(value) + if dt.tzinfo is None: + dt = dt.replace(tzinfo=timezone.utc) + return dt.astimezone(timezone.utc).isoformat().replace("+00:00", "Z") + + +def extract_member_hourly_from_grid_dataset( + dataset: Any, + *, + lat: float, + lon: float, + temp_var: Optional[str] = None, + use_fahrenheit: bool = False, +) -> Dict[str, Any]: + """Extract all member hourly temperatures for the nearest grid point.""" + selected_var = select_temperature_variable(dataset, temp_var or os.getenv("WEATHERNEXT2_TEMP_VAR") or None) + lat_name, lat_values = _dataset_dimension_values(dataset, ("lat", "latitude")) + lon_name, lon_values = _dataset_dimension_values(dataset, ("lon", "longitude")) + member_name, member_values = _dataset_dimension_values(dataset, ("member", "realization", "ensemble_member")) + time_name, time_values = _dataset_dimension_values(dataset, ("time", "valid_time")) + + lat_idx = _nearest_index(lat_values, lat) + lon_idx = _nearest_index(lon_values, lon) + units = _variable_units(dataset, selected_var) + + variable = _get_dataset_value(dataset, selected_var) + dims = list(getattr(variable, "dims", []) or [member_name, time_name, lat_name, lon_name]) + values = _variable_values(dataset, selected_var) + dim_indexes = { + member_name: None, + time_name: None, + lat_name: lat_idx, + lon_name: lon_idx, + } + + member_hourly: Dict[str, list[Optional[float]]] = {} + for member_idx, member in enumerate(member_values): + hourly = [] + for time_idx, _time_value in enumerate(time_values): + dim_indexes[member_name] = member_idx + dim_indexes[time_name] = time_idx + indexes = [int(dim_indexes[dim]) for dim in dims] + hourly.append( + normalize_temperature_value( + _index_nested(values, indexes), + units, + use_fahrenheit=use_fahrenheit, + ) + ) + try: + member_label = f"member_{int(member):02d}" + except Exception: + member_label = f"member_{member_idx:02d}" + member_hourly[member_label] = hourly + + return { + "temp_var": selected_var, + "units": units, + "lat_dim": lat_name, + "lon_dim": lon_name, + "member_dim": member_name, + "time_dim": time_name, + "nearest_lat": float(lat_values[lat_idx]), + "nearest_lon": float(lon_values[lon_idx]), + "utc_times": [_parse_time(value) for value in time_values], + "member_hourly": member_hourly, + } + + +def open_weathernext2_zarr_dataset(uri: str): + """Open a WeatherNext 2 Zarr dataset lazily with optional runtime dependencies.""" + try: + import xarray as xr # type: ignore + except Exception as exc: # pragma: no cover - depends on deployment extras + raise RuntimeError("xarray is required for WeatherNext2 GCS/Zarr access") from exc + + storage_options = {} + credentials = str(os.getenv("GOOGLE_APPLICATION_CREDENTIALS", "") or "").strip() + if credentials: + storage_options["token"] = "google_default" + try: + return xr.open_zarr( + uri, + storage_options=storage_options or None, + consolidated=True, + ) + except Exception: + return xr.open_zarr( + uri, + storage_options=storage_options or None, + consolidated=False, + ) diff --git a/src/data_collection/weathernext2_sources.py b/src/data_collection/weathernext2_sources.py new file mode 100644 index 00000000..a097fd12 --- /dev/null +++ b/src/data_collection/weathernext2_sources.py @@ -0,0 +1,581 @@ +from __future__ import annotations + +import math +import os +import json +import threading +import time +from datetime import date, datetime, timedelta, timezone +from typing import Any, Dict, Iterable, Mapping, Optional, Sequence, Union + +from loguru import logger + + +WEATHERNEXT2_SOURCE = "weathernext2" +WEATHERNEXT2_PROVIDER = "google_deepmind" +WEATHERNEXT2_GCS_ZARR_URI = "gs://weathernext/weathernext_2_0_0/zarr" +WEATHERNEXT2_MEAN_GCS_ZARR_URI = "gs://weathernext/weathernext_2_0_0_mean/zarr" + +Number = Union[int, float] + + +def _numeric(value: Any) -> Optional[float]: + if value is None: + return None + try: + parsed = float(value) + except (TypeError, ValueError): + return None + return parsed if math.isfinite(parsed) else None + + +def _round1(value: Number) -> float: + return round(float(value), 1) + + +def _round3(value: Number) -> float: + return round(float(value), 3) + + +def _settle_integer(value: Number) -> int: + parsed = float(value) + if parsed >= 0: + return int(math.floor(parsed + 0.5)) + return int(math.ceil(parsed - 0.5)) + + +def _is_fahrenheit(temp_symbol: str) -> bool: + return "F" in str(temp_symbol or "").upper() + + +def _unit(temp_symbol: str) -> str: + return "°F" if _is_fahrenheit(temp_symbol) else "°C" + + +def market_bucket_for_temperature(value: Number, temp_symbol: str = "°C") -> Dict[str, Any]: + """Return the tradable market option bucket for a single member temperature.""" + unit = _unit(temp_symbol) + settled = _settle_integer(value) + if _is_fahrenheit(unit): + lower_value = settled if settled % 2 == 0 else settled - 1 + upper_value = lower_value + 1 + return { + "key": f"{lower_value}-{upper_value}{unit}", + "label": f"{lower_value}-{upper_value}{unit}", + "lower": float(lower_value) - 0.5, + "upper": float(upper_value) + 0.5, + "sort_value": float(lower_value), + } + return { + "key": f"{settled}{unit}", + "label": f"{settled}{unit}", + "lower": float(settled) - 0.5, + "upper": float(settled) + 0.5, + "sort_value": float(settled), + } + + +def _member_items(member_highs: Union[Mapping[str, Any], Sequence[Any]]) -> Iterable[tuple[str, float]]: + if isinstance(member_highs, Mapping): + iterable = member_highs.items() + else: + iterable = ((f"member_{idx:02d}", value) for idx, value in enumerate(member_highs)) + for member_id, value in iterable: + parsed = _numeric(value) + if parsed is not None: + yield str(member_id), parsed + + +def summarize_member_highs(member_highs: Union[Mapping[str, Any], Sequence[Any]]) -> Dict[str, Any]: + values = sorted(value for _member_id, value in _member_items(member_highs)) + if not values: + return { + "members": 0, + "mean": None, + "median": None, + "p10": None, + "p25": None, + "p75": None, + "p90": None, + "min": None, + "max": None, + "spread": None, + } + + def percentile(q: float) -> float: + if len(values) == 1: + return values[0] + position = (len(values) - 1) * q + lower_idx = int(math.floor(position)) + upper_idx = int(math.ceil(position)) + if lower_idx == upper_idx: + return values[lower_idx] + weight = position - lower_idx + return values[lower_idx] * (1 - weight) + values[upper_idx] * weight + + return { + "members": len(values), + "mean": _round1(sum(values) / len(values)), + "median": _round1(percentile(0.5)), + "p10": _round1(percentile(0.1)), + "p25": _round1(percentile(0.25)), + "p75": _round1(percentile(0.75)), + "p90": _round1(percentile(0.9)), + "min": _round1(values[0]), + "max": _round1(values[-1]), + "spread": _round1(values[-1] - values[0]), + } + + +def build_market_bucket_probabilities( + member_highs: Union[Mapping[str, Any], Sequence[Any]], + temp_symbol: str = "°C", +) -> list[Dict[str, Any]]: + grouped: Dict[str, Dict[str, Any]] = {} + total_members = 0 + + for _member_id, value in _member_items(member_highs): + total_members += 1 + option = market_bucket_for_temperature(value, temp_symbol) + bucket = grouped.setdefault( + option["key"], + { + "key": option["key"], + "label": option["label"], + "lower": option["lower"], + "upper": option["upper"], + "sort_value": option["sort_value"], + "weighted_sum": 0.0, + "member_count": 0, + }, + ) + bucket["weighted_sum"] += value + bucket["member_count"] += 1 + + if total_members <= 0: + return [] + + buckets = [] + for bucket in grouped.values(): + member_count = int(bucket["member_count"]) + buckets.append( + { + "key": bucket["key"], + "label": bucket["label"], + "lower": bucket["lower"], + "upper": bucket["upper"], + "value": _round1(bucket["weighted_sum"] / member_count), + "probability": _round3(member_count / total_members), + "member_count": member_count, + "total_members": total_members, + } + ) + + return sorted(buckets, key=lambda item: (float(item["lower"]), float(item["upper"]))) + + +def _parse_utc_time(value: Any) -> Optional[datetime]: + if isinstance(value, datetime): + parsed = value + else: + text = str(value or "").strip() + if not text: + return None + if text.endswith("Z"): + text = f"{text[:-1]}+00:00" + try: + parsed = datetime.fromisoformat(text) + except ValueError: + return None + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + return parsed.astimezone(timezone.utc) + + +def _parse_date(value: Any) -> Optional[date]: + if isinstance(value, date) and not isinstance(value, datetime): + return value + text = str(value or "").strip() + if not text: + return None + try: + return datetime.fromisoformat(text[:10]).date() + except ValueError: + return None + + +def build_city_local_daily_highs_from_hourly( + member_hourly: Mapping[str, Sequence[Any]], + utc_times: Sequence[Any], + timezone_offset_seconds: int, + target_local_date: Any, +) -> Dict[str, float]: + """Reduce hourly ensemble member temperatures to local-date daily highs.""" + target_date = _parse_date(target_local_date) + if target_date is None: + return {} + + parsed_times = [_parse_utc_time(value) for value in utc_times] + highs: Dict[str, float] = {} + offset = timedelta(seconds=int(timezone_offset_seconds or 0)) + + for member_id, values in member_hourly.items(): + member_values = [] + for idx, utc_time in enumerate(parsed_times): + if utc_time is None or idx >= len(values): + continue + local_date = (utc_time + offset).date() + if local_date != target_date: + continue + parsed = _numeric(values[idx]) + if parsed is not None: + member_values.append(parsed) + if member_values: + highs[str(member_id)] = _round1(max(member_values)) + + return highs + + +def build_weathernext2_city_probability( + *, + city: str, + member_highs: Union[Mapping[str, Any], Sequence[Any]], + temp_symbol: str = "°C", + target_date: Optional[str] = None, + source_run: Optional[str] = None, + generated_at: Optional[str] = None, +) -> Dict[str, Any]: + """Build a WeatherNext 2 probability payload aligned with tradable market options.""" + buckets = build_market_bucket_probabilities(member_highs, temp_symbol=temp_symbol) + summary = summarize_member_highs(member_highs) + normalized_member_highs = { + member_id: _round1(value) + for member_id, value in _member_items(member_highs) + } + top_bucket = ( + max(buckets, key=lambda item: (float(item["probability"]), float(item["lower"]))) + if buckets + else None + ) + + return { + "source": WEATHERNEXT2_SOURCE, + "provider": WEATHERNEXT2_PROVIDER, + "city": str(city or "").strip(), + "target_date": target_date, + "source_run": source_run, + "generated_at": generated_at or datetime.now(timezone.utc).isoformat(), + "temp_symbol": _unit(temp_symbol), + "members": int(summary["members"] or 0), + "member_highs": normalized_member_highs, + "summary": summary, + "buckets": buckets, + "top_bucket": top_bucket, + "bucket_policy": "celsius_single_fahrenheit_two_degree_market_options", + "gcs_zarr_uri": os.getenv("WEATHERNEXT2_GCS_ZARR_URI", WEATHERNEXT2_GCS_ZARR_URI), + "mean_gcs_zarr_uri": os.getenv("WEATHERNEXT2_MEAN_GCS_ZARR_URI", WEATHERNEXT2_MEAN_GCS_ZARR_URI), + } + + +def _env_enabled(name: str, default: str = "0") -> bool: + return str(os.getenv(name, default) or "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + + +def _city_key(city: str) -> str: + return str(city or "").strip().lower() + + +def _load_fixture_payload(path: str, city: str) -> Optional[Dict[str, Any]]: + if not path: + return None + try: + with open(path, "r", encoding="utf-8") as fh: + payload = json.load(fh) + except Exception as exc: + logger.warning("WeatherNext2 fixture load failed path={}: {}", path, exc) + return None + if not isinstance(payload, dict): + return None + + key = _city_key(city) + if isinstance(payload.get(key), dict): + return dict(payload[key]) + for candidate, value in payload.items(): + if _city_key(candidate) == key and isinstance(value, dict): + return dict(value) + if "member_highs" in payload or "member_hourly" in payload: + return dict(payload) + return None + + +def _weathernext2_data_root() -> str: + return str(os.getenv("WEATHERNEXT2_DATA_ROOT", "/app/data/weathernext2") or "").strip() + + +def _weathernext2_artifact_path() -> str: + configured = str(os.getenv("WEATHERNEXT2_CITY_HIGHS_PATH", "") or "").strip() + if configured: + return configured + return os.path.join(_weathernext2_data_root(), "weathernext2_city_highs.json") + + +def _weathernext2_model_dir() -> str: + return str( + os.getenv( + "WEATHERNEXT2_MODEL_DIR", + "/app/data/models/weathernext2_calibrator", + ) + or "" + ).strip() + + +def _load_artifact_payload(path: str, city: str) -> Optional[Dict[str, Any]]: + if not path or not os.path.isfile(path): + return None + try: + with open(path, "r", encoding="utf-8") as fh: + payload = json.load(fh) + except Exception as exc: + logger.warning("WeatherNext2 artifact load failed path={}: {}", path, exc) + return None + if not isinstance(payload, dict): + return None + + key = _city_key(city) + city_payloads: Any = payload.get("cities") + if isinstance(city_payloads, dict): + if isinstance(city_payloads.get(key), dict): + return dict(city_payloads[key]) + for candidate, value in city_payloads.items(): + if _city_key(candidate) == key and isinstance(value, dict): + return dict(value) + if isinstance(city_payloads, list): + for value in city_payloads: + if not isinstance(value, dict): + continue + if _city_key(value.get("city")) == key: + return dict(value) + + if isinstance(payload.get(key), dict): + return dict(payload[key]) + for candidate, value in payload.items(): + if _city_key(candidate) == key and isinstance(value, dict): + return dict(value) + return None + + +class WeatherNext2SourceMixin: + """Optional WeatherNext 2 probability source. + + The live Google datasets require project access and optional heavy client + dependencies. This mixin keeps production safe by supporting a fixture or + prepared payload first, while reserving backend configuration for the real + GCS/BigQuery fetcher. + """ + + def _weathernext2_enabled(self) -> bool: + return _env_enabled("WEATHERNEXT2_ENABLED") + + def _weathernext2_cache_key( + self, + city: str, + lat: float, + lon: float, + use_fahrenheit: bool, + target_date: Optional[str], + ) -> str: + return ( + f"{_city_key(city)}:{round(float(lat), 4)}:{round(float(lon), 4)}:" + f"{'f' if use_fahrenheit else 'c'}:{target_date or ''}" + ) + + def _weathernext2_cache_state(self) -> tuple[Dict[str, Dict[str, Any]], threading.Lock, int]: + if not hasattr(self, "_weathernext2_cache"): + self._weathernext2_cache = {} + if not hasattr(self, "_weathernext2_cache_lock"): + self._weathernext2_cache_lock = threading.Lock() + ttl_sec = int(os.getenv("WEATHERNEXT2_CACHE_TTL_SEC", "21600")) + return self._weathernext2_cache, self._weathernext2_cache_lock, ttl_sec + + def _weathernext2_target_date( + self, + timezone_offset_seconds: Optional[int], + target_date: Optional[str], + ) -> str: + if target_date: + return str(target_date) + offset = timedelta(seconds=int(timezone_offset_seconds or 0)) + return (datetime.now(timezone.utc) + offset).date().isoformat() + + def _weathernext2_from_fixture( + self, + *, + city: str, + use_fahrenheit: bool, + target_date: str, + timezone_offset_seconds: Optional[int], + ) -> Optional[Dict[str, Any]]: + fixture_path = str(os.getenv("WEATHERNEXT2_FIXTURE_PATH", "") or "").strip() + fixture = _load_fixture_payload(fixture_path, city) + if not fixture: + return None + + temp_symbol = "°F" if use_fahrenheit else "°C" + fixture_target_date = str(fixture.get("target_date") or target_date) + member_highs = fixture.get("member_highs") + if member_highs is None and isinstance(fixture.get("member_hourly"), dict): + member_highs = build_city_local_daily_highs_from_hourly( + fixture["member_hourly"], + fixture.get("utc_times") or fixture.get("times") or [], + int(timezone_offset_seconds or fixture.get("timezone_offset_seconds") or 0), + fixture_target_date, + ) + if member_highs is None: + return None + + return build_weathernext2_city_probability( + city=city, + member_highs=member_highs, + temp_symbol=temp_symbol, + target_date=fixture_target_date, + source_run=fixture.get("source_run"), + generated_at=fixture.get("generated_at"), + ) + + def _weathernext2_normalize_prepared_payload( + self, + *, + payload: Mapping[str, Any], + city: str, + use_fahrenheit: bool, + target_date: str, + timezone_offset_seconds: Optional[int], + ) -> Optional[Dict[str, Any]]: + temp_symbol = "°F" if use_fahrenheit else "°C" + payload_target_date = str(payload.get("target_date") or target_date) + member_highs = payload.get("member_highs") + if member_highs is None and isinstance(payload.get("member_hourly"), dict): + member_highs = build_city_local_daily_highs_from_hourly( + payload["member_hourly"], + payload.get("utc_times") or payload.get("times") or [], + int(timezone_offset_seconds or payload.get("timezone_offset_seconds") or 0), + payload_target_date, + ) + if member_highs is None and isinstance(payload.get("buckets"), list): + normalized = dict(payload) + normalized.setdefault("source", WEATHERNEXT2_SOURCE) + normalized.setdefault("provider", WEATHERNEXT2_PROVIDER) + normalized.setdefault("city", city) + normalized.setdefault("target_date", payload_target_date) + normalized.setdefault("temp_symbol", temp_symbol) + normalized.setdefault( + "gcs_zarr_uri", + os.getenv("WEATHERNEXT2_GCS_ZARR_URI", WEATHERNEXT2_GCS_ZARR_URI), + ) + return normalized + if member_highs is None: + return None + return build_weathernext2_city_probability( + city=city, + member_highs=member_highs, + temp_symbol=temp_symbol, + target_date=payload_target_date, + source_run=payload.get("source_run"), + generated_at=payload.get("generated_at"), + ) + + def _weathernext2_from_artifact( + self, + *, + city: str, + use_fahrenheit: bool, + target_date: str, + timezone_offset_seconds: Optional[int], + ) -> Optional[Dict[str, Any]]: + artifact = _load_artifact_payload(_weathernext2_artifact_path(), city) + if not artifact: + return None + payload = self._weathernext2_normalize_prepared_payload( + payload=artifact, + city=city, + use_fahrenheit=use_fahrenheit, + target_date=target_date, + timezone_offset_seconds=timezone_offset_seconds, + ) + if payload is None: + return None + try: + from src.analysis.weathernext2_calibration import apply_quantile_calibration_to_payload + + return apply_quantile_calibration_to_payload( + payload, + model_dir=_weathernext2_model_dir(), + ) + except Exception as exc: + logger.warning("WeatherNext2 calibration apply failed city={}: {}", city, exc) + return payload + + def fetch_weathernext2_probability( + self, + city: str, + lat: float, + lon: float, + *, + use_fahrenheit: bool = False, + target_date: Optional[str] = None, + timezone_offset_seconds: Optional[int] = None, + ) -> Optional[Dict[str, Any]]: + if not self._weathernext2_enabled(): + return None + if lat is None or lon is None: + return None + + resolved_target_date = self._weathernext2_target_date( + timezone_offset_seconds, + target_date, + ) + cache_key = self._weathernext2_cache_key( + city, + lat, + lon, + use_fahrenheit, + resolved_target_date, + ) + cache, lock, ttl_sec = self._weathernext2_cache_state() + now_ts = time.time() + with lock: + cached = cache.get(cache_key) + cached_data = cached.get("data") if isinstance(cached, dict) else None + if isinstance(cached_data, dict) and now_ts - float(cached.get("t", 0)) < ttl_sec: + return dict(cached_data) + + payload = self._weathernext2_from_artifact( + city=city, + use_fahrenheit=use_fahrenheit, + target_date=resolved_target_date, + timezone_offset_seconds=timezone_offset_seconds, + ) + if payload is None: + payload = self._weathernext2_from_fixture( + city=city, + use_fahrenheit=use_fahrenheit, + target_date=resolved_target_date, + timezone_offset_seconds=timezone_offset_seconds, + ) + if payload is None: + backend = str(os.getenv("WEATHERNEXT2_BACKEND", "") or "").strip().lower() + if backend: + logger.warning( + "WeatherNext2 backend={} configured but live fetcher is not installed; use fixture/prepared payload first", + backend, + ) + return None + + with lock: + cache[cache_key] = {"t": now_ts, "data": dict(payload)} + return payload diff --git a/tests/test_deployment_runtime_config.py b/tests/test_deployment_runtime_config.py index f7b61a84..8be9e79e 100644 --- a/tests/test_deployment_runtime_config.py +++ b/tests/test_deployment_runtime_config.py @@ -100,6 +100,13 @@ def test_docker_compose_isolates_collector_from_web_and_bot_services(): training_settlement_block = compose.split( " polyweather_training_settlement:", 1, + )[1].split( + "\n polyweather_weathernext2_worker:", + 1, + )[0] + weathernext2_block = compose.split( + " polyweather_weathernext2_worker:", + 1, )[1].split( "\nx-polyweather-base:", 1, @@ -110,6 +117,7 @@ def test_docker_compose_isolates_collector_from_web_and_bot_services(): assert "POLYWEATHER_SERVICE_ROLE: collector" in collector_block assert "POLYWEATHER_SERVICE_ROLE: warmer" in warmer_block assert "POLYWEATHER_SERVICE_ROLE: training_settlement" in training_settlement_block + assert "POLYWEATHER_SERVICE_ROLE: weathernext2_worker" in weathernext2_block assert "redis-server --appendonly yes --maxmemory ${POLYWEATHER_REDIS_MAXMEMORY:-512mb} --maxmemory-policy noeviction" in compose assert "POLYWEATHER_SCAN_TERMINAL_PREWARM_ENABLED: 'false'" in bot_block assert "POLYWEATHER_EVENT_STORE: ${POLYWEATHER_EVENT_STORE:-redis}" in web_block @@ -126,7 +134,9 @@ def test_docker_compose_isolates_collector_from_web_and_bot_services(): assert "POLYWEATHER_OBSERVATION_COLLECTOR_ENABLED: 'true'" in collector_block assert "POLYWEATHER_OBSERVATION_COLLECTOR_ENABLED: 'false'" in warmer_block assert "POLYWEATHER_OBSERVATION_COLLECTOR_ENABLED: 'false'" in training_settlement_block + assert "POLYWEATHER_OBSERVATION_COLLECTOR_ENABLED: 'false'" in weathernext2_block assert "command: python -m web.training_settlement_worker" in training_settlement_block + assert "command: python -m web.weathernext2_worker" in weathernext2_block assert ( "POLYWEATHER_TRAINING_SETTLEMENT_INTERVAL_SEC: " "${POLYWEATHER_TRAINING_SETTLEMENT_INTERVAL_SEC:-21600}" @@ -142,6 +152,25 @@ def test_docker_compose_isolates_collector_from_web_and_bot_services(): assert "POLYWEATHER_CITY_DETAIL_BATCH_QUEUE_WAIT_MS: ${POLYWEATHER_CITY_DETAIL_BATCH_QUEUE_WAIT_MS:-3000}" in web_block assert "POLYWEATHER_CITY_DETAIL_BATCH_PARTIAL_TIMEOUT_MS: ${POLYWEATHER_CITY_DETAIL_BATCH_PARTIAL_TIMEOUT_MS:-8000}" in web_block assert "UVICORN_WORKERS: ${UVICORN_WORKERS:-2}" in web_block + assert "WEATHERNEXT2_ENABLED: ${WEATHERNEXT2_ENABLED:-1}" in web_block + assert "WEATHERNEXT2_ENABLED: ${WEATHERNEXT2_ENABLED:-1}" in weathernext2_block + assert "WEATHERNEXT2_BACKEND: ${WEATHERNEXT2_BACKEND:-gcs_zarr}" in weathernext2_block + assert ( + "WEATHERNEXT2_GCS_ZARR_URI: " + "${WEATHERNEXT2_GCS_ZARR_URI:-gs://weathernext/weathernext_2_0_0/zarr}" + in weathernext2_block + ) + assert ( + "WEATHERNEXT2_MODEL_DIR: " + "${WEATHERNEXT2_MODEL_DIR:-/app/data/models/weathernext2_calibrator}" + in weathernext2_block + ) + assert ( + "GOOGLE_APPLICATION_CREDENTIALS: " + "${GOOGLE_APPLICATION_CREDENTIALS:-/app/secrets/gcp-sa.json}" + in weathernext2_block + ) + assert "./secrets:/app/secrets:ro" in weathernext2_block assert "POLYWEATHER_COLLECTOR_PATCH_ENDPOINT: ''" in bot_block assert "POLYWEATHER_COLLECTOR_PATCH_ENDPOINT: ''" in web_block assert ( @@ -312,6 +341,7 @@ def test_deploy_script_retries_compose_recreate_races(): assert 'compose_up_retry "observation collector" -d --no-deps polyweather_collector' in script assert 'compose_up_retry "cache warmer" -d --no-deps polyweather_warmer' in script assert 'compose_up_retry "training settlement" -d --no-deps polyweather_training_settlement' in script + assert 'compose_up_retry "WeatherNext2 worker" -d --no-deps polyweather_weathernext2_worker' in script assert 'compose_up_retry "frontend" -d --no-deps polyweather_frontend' in script diff --git a/tests/test_trend_engine.py b/tests/test_trend_engine.py index 592aec99..78db712e 100644 --- a/tests/test_trend_engine.py +++ b/tests/test_trend_engine.py @@ -241,6 +241,64 @@ class TestMuCalculation: assert sd["peak_hours"] == ["15:00", "16:00"] assert sd["peak_status"] == "before" + @patch("src.analysis.trend_engine.calculate_dynamic_weights", return_value=(None, "")) + @patch("src.analysis.trend_engine.get_deb_accuracy", return_value=None) + @patch("src.analysis.trend_engine.update_daily_record") + def test_weathernext2_probability_replaces_legacy_distribution( + self, _udr, _deb_acc, _dw + ): + data = _make_weather_data( + cur_temp=30.0, + max_so_far=31.0, + om_today_high=33.0, + ens_median=32.0, + local_time="2026-03-04 11:00", + ) + data["weathernext2"] = { + "summary": {"median": 33.1}, + "buckets": [ + { + "label": "32°C", + "value": 32.1, + "lower": 31.5, + "upper": 32.5, + "probability": 0.41, + }, + { + "label": "33°C", + "value": 33.0, + "lower": 32.5, + "upper": 33.5, + "probability": 0.56, + }, + ], + } + + _, ai_context, sd = analyze_weather_trend(data, "°C", "test_city") + + assert sd["probability_engine"] == "weathernext2" + assert sd["probabilities_all"] == data["weathernext2"]["buckets"] + assert sd["mu"] == 33.1 + assert "WeatherNext 2 概率分布" in ai_context + + @patch("src.analysis.trend_engine.calculate_dynamic_weights", return_value=(None, "")) + @patch("src.analysis.trend_engine.get_deb_accuracy", return_value=None) + @patch("src.analysis.trend_engine.update_daily_record") + def test_weathernext2_median_enters_current_forecasts( + self, _udr, _deb_acc, _dw + ): + data = _make_weather_data(local_time="2026-03-04 11:00") + data["weathernext2"] = { + "summary": {"median": 33.1}, + "buckets": [ + {"label": "33°C", "value": 33.0, "probability": 0.56}, + ], + } + + _, _, sd = analyze_weather_trend(data, "°C", "test_city") + + assert sd["current_forecasts"]["WeatherNext 2"] == 33.1 + class TestDebEnsembleSignal: @patch( diff --git a/tests/test_weathernext2_calibration.py b/tests/test_weathernext2_calibration.py new file mode 100644 index 00000000..4bae6b8d --- /dev/null +++ b/tests/test_weathernext2_calibration.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from src.analysis.weathernext2_calibration import ( + apply_quantile_calibration_to_payload, + train_lightgbm_quantile_calibrator, +) +from src.data_collection.weathernext2_sources import build_weathernext2_city_probability + + +def _synthetic_training_rows(count: int = 170): + rows = [] + for idx in range(count): + city = "houston" if idx % 2 == 0 else "shanghai" + median = 30.0 + (idx % 7) * 0.2 + residual = 1.0 if city == "houston" else -0.5 + rows.append( + { + "city": city, + "target_date": f"2026-05-{idx % 28 + 1:02d}", + "actual_high_c": median + residual, + "weathernext2": { + "summary": { + "mean": median, + "median": median, + "p10": median - 1.0, + "p25": median - 0.5, + "p75": median + 0.5, + "p90": median + 1.0, + "spread": 2.0, + } + }, + "deb_prediction_c": median + 0.2, + "model_median_c": median + 0.1, + "model_spread": 1.4, + "current_max_so_far_c": median - 2.0, + "local_hour": 12, + } + ) + return rows + + +def test_lightgbm_quantile_calibrator_trains_and_saves_ordered_quantiles(tmp_path): + result = train_lightgbm_quantile_calibrator( + _synthetic_training_rows(), + model_dir=tmp_path, + min_global_samples=150, + min_city_samples=5, + ) + + assert result["trained"] is True + assert result["samples"] == 170 + assert (tmp_path / "metadata.json").is_file() + assert (tmp_path / "q10.pkl").is_file() + assert (tmp_path / "q50.pkl").is_file() + assert (tmp_path / "q90.pkl").is_file() + assert result["validation"]["ordered_quantiles"] is True + + +def test_lightgbm_quantile_calibrator_skips_when_samples_are_insufficient(tmp_path): + result = train_lightgbm_quantile_calibrator( + _synthetic_training_rows(20), + model_dir=tmp_path, + min_global_samples=150, + min_city_samples=5, + ) + + assert result["trained"] is False + assert result["reason"] == "insufficient_global_samples" + + +def test_calibrated_distribution_rebuilds_market_buckets_from_shifted_members(tmp_path): + train_lightgbm_quantile_calibrator( + _synthetic_training_rows(), + model_dir=tmp_path, + min_global_samples=150, + min_city_samples=5, + ) + raw = build_weathernext2_city_probability( + city="houston", + member_highs=[30.0, 30.2, 30.4, 30.6], + temp_symbol="°C", + target_date="2026-06-29", + ) + + calibrated = apply_quantile_calibration_to_payload(raw, model_dir=tmp_path) + + assert calibrated["calibration"]["engine"] == "lightgbm_quantile" + assert calibrated["calibration"]["samples"] == 170 + assert calibrated["calibration"]["raw_summary"]["median"] == raw["summary"]["median"] + assert calibrated["calibration"]["calibrated_summary"]["median"] > raw["summary"]["median"] + assert calibrated["buckets"] + assert calibrated["top_bucket"]["label"].endswith("°C") diff --git a/tests/test_weathernext2_sources.py b/tests/test_weathernext2_sources.py new file mode 100644 index 00000000..d60203aa --- /dev/null +++ b/tests/test_weathernext2_sources.py @@ -0,0 +1,259 @@ +from datetime import datetime, timezone +import json + +from src.data_collection.weather_sources import WeatherDataCollector +from src.data_collection.weathernext2_fetcher import ( + extract_member_hourly_from_grid_dataset, + normalize_temperature_value, + open_weathernext2_zarr_dataset, + select_temperature_variable, +) +from src.data_collection.weathernext2_sources import ( + build_city_local_daily_highs_from_hourly, + build_weathernext2_city_probability, + market_bucket_for_temperature, + summarize_member_highs, +) + + +def test_market_bucket_for_temperature_uses_single_celsius_options(): + bucket = market_bucket_for_temperature(32.6, "°C") + + assert bucket["label"] == "33°C" + assert bucket["lower"] == 32.5 + assert bucket["upper"] == 33.5 + + +def test_market_bucket_for_temperature_groups_fahrenheit_by_two_degree_market_option(): + assert market_bucket_for_temperature(94.2, "°F")["label"] == "94-95°F" + assert market_bucket_for_temperature(95.4, "°F")["label"] == "94-95°F" + assert market_bucket_for_temperature(96.1, "°F")["label"] == "96-97°F" + + +def test_weathernext2_probability_aggregates_member_highs_to_market_buckets(): + probability = build_weathernext2_city_probability( + city="Houston", + member_highs={ + "member_00": 94.1, + "member_01": 94.8, + "member_02": 95.2, + "member_03": 96.7, + }, + temp_symbol="°F", + target_date="2026-06-29", + source_run="2026-06-29T00:00:00Z", + ) + + assert probability["source"] == "weathernext2" + assert probability["members"] == 4 + assert probability["top_bucket"]["label"] == "94-95°F" + assert probability["top_bucket"]["probability"] == 0.75 + assert probability["buckets"] == [ + { + "key": "94-95°F", + "label": "94-95°F", + "lower": 93.5, + "upper": 95.5, + "value": 94.7, + "probability": 0.75, + "member_count": 3, + "total_members": 4, + }, + { + "key": "96-97°F", + "label": "96-97°F", + "lower": 95.5, + "upper": 97.5, + "value": 96.7, + "probability": 0.25, + "member_count": 1, + "total_members": 4, + }, + ] + + +def test_weathernext2_probability_keeps_celsius_market_option_labels(): + probability = build_weathernext2_city_probability( + city="Shanghai", + member_highs=[31.6, 32.2, 32.8, 33.1], + temp_symbol="°C", + target_date="2026-06-29", + ) + + labels = [bucket["label"] for bucket in probability["buckets"]] + + assert labels == ["32°C", "33°C"] + assert probability["top_bucket"]["label"] == "33°C" + assert probability["top_bucket"]["probability"] == 0.5 + + +def test_city_local_daily_highs_from_hourly_uses_target_local_date(): + highs = build_city_local_daily_highs_from_hourly( + member_hourly={ + "member_00": [21.0, 26.0, 28.0, 24.0], + "member_01": [20.0, 25.0, None, 27.0], + }, + utc_times=[ + datetime(2026, 6, 28, 15, tzinfo=timezone.utc), # local 2026-06-28 23:00 + datetime(2026, 6, 28, 16, tzinfo=timezone.utc), # local 2026-06-29 00:00 + datetime(2026, 6, 29, 6, tzinfo=timezone.utc), # local 2026-06-29 14:00 + datetime(2026, 6, 29, 16, tzinfo=timezone.utc), # local 2026-06-30 00:00 + ], + timezone_offset_seconds=8 * 3600, + target_local_date="2026-06-29", + ) + + assert highs == {"member_00": 28.0, "member_01": 25.0} + + +def test_summarize_member_highs_reports_ensemble_shape(): + summary = summarize_member_highs([30, 32, 34, 36, 38]) + + assert summary == { + "members": 5, + "mean": 34.0, + "median": 34.0, + "p10": 30.8, + "p25": 32.0, + "p75": 36.0, + "p90": 37.2, + "min": 30.0, + "max": 38.0, + "spread": 8.0, + } + + +def test_collector_reads_weathernext2_fixture_when_enabled(monkeypatch, tmp_path): + fixture_path = tmp_path / "weathernext2.json" + fixture_path.write_text( + json.dumps( + { + "houston": { + "target_date": "2026-06-29", + "source_run": "2026-06-29T00:00:00Z", + "member_highs": [94.1, 94.8, 95.2, 96.7], + } + } + ), + encoding="utf-8", + ) + monkeypatch.setenv("WEATHERNEXT2_ENABLED", "1") + monkeypatch.setenv("WEATHERNEXT2_FIXTURE_PATH", str(fixture_path)) + + collector = WeatherDataCollector({}) + + payload = collector.fetch_weathernext2_probability( + "Houston", + lat=29.7604, + lon=-95.3698, + use_fahrenheit=True, + target_date="2026-06-29", + timezone_offset_seconds=-5 * 3600, + ) + + assert payload["source"] == "weathernext2" + assert payload["top_bucket"]["label"] == "94-95°F" + assert payload["buckets"][0]["probability"] == 0.75 + + +def test_collector_reads_weathernext2_worker_artifact_before_fixture(monkeypatch, tmp_path): + artifact_path = tmp_path / "weathernext2_city_highs.json" + fixture_path = tmp_path / "fixture.json" + artifact_path.write_text( + json.dumps( + { + "schema_version": 1, + "cities": { + "houston": { + "target_date": "2026-06-29", + "source_run": "2026-06-29T00:00:00Z", + "member_highs": [94.1, 94.8, 95.2, 96.7], + } + }, + } + ), + encoding="utf-8", + ) + fixture_path.write_text( + json.dumps({"houston": {"member_highs": [80.0, 80.2]}}), + encoding="utf-8", + ) + monkeypatch.setenv("WEATHERNEXT2_ENABLED", "1") + monkeypatch.setenv("WEATHERNEXT2_CITY_HIGHS_PATH", str(artifact_path)) + monkeypatch.setenv("WEATHERNEXT2_FIXTURE_PATH", str(fixture_path)) + monkeypatch.setenv("WEATHERNEXT2_MODEL_DIR", str(tmp_path / "missing-model")) + + collector = WeatherDataCollector({}) + + payload = collector.fetch_weathernext2_probability( + "Houston", + lat=29.7604, + lon=-95.3698, + use_fahrenheit=True, + target_date="2026-06-29", + timezone_offset_seconds=-5 * 3600, + ) + + assert payload["source"] == "weathernext2" + assert payload["top_bucket"]["label"] == "94-95°F" + assert payload["buckets"][0]["probability"] == 0.75 + + +def test_weathernext2_grid_extractor_selects_temp_var_nearest_point_and_converts_kelvin(): + dataset = { + "lat": [10.0, 20.0], + "lon": [30.0, 40.0], + "member": [0, 1], + "time": [ + "2026-06-29T00:00:00Z", + "2026-06-29T06:00:00Z", + "2026-06-29T12:00:00Z", + ], + "temperature_2m": [ + [ + [[290.0, 291.0], [292.0, 293.0]], + [[294.0, 295.0], [296.0, 297.0]], + [[298.0, 299.0], [300.0, 301.0]], + ], + [ + [[289.0, 290.0], [291.0, 292.0]], + [[293.0, 294.0], [295.0, 296.0]], + [[297.0, 298.0], [299.0, 300.0]], + ], + ], + "units": {"temperature_2m": "K"}, + } + + assert select_temperature_variable(dataset) == "temperature_2m" + assert normalize_temperature_value(300.15, "K") == 27.0 + + extracted = extract_member_hourly_from_grid_dataset( + dataset, + lat=18.0, + lon=38.0, + ) + + assert extracted["temp_var"] == "temperature_2m" + assert extracted["nearest_lat"] == 20.0 + assert extracted["nearest_lon"] == 40.0 + assert extracted["member_hourly"]["member_00"] == [19.9, 23.9, 27.9] + assert extracted["member_hourly"]["member_01"] == [18.9, 22.9, 26.9] + + +def test_weathernext2_zarr_uses_google_default_credentials(monkeypatch): + calls = [] + + class FakeXarray: + @staticmethod + def open_zarr(uri, **kwargs): + calls.append((uri, kwargs)) + return {"ok": True} + + import src.data_collection.weathernext2_fetcher as fetcher + + monkeypatch.setattr(fetcher, "xr", FakeXarray, raising=False) + monkeypatch.setitem(__import__("sys").modules, "xarray", FakeXarray) + monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", "/app/secrets/gcp-sa.json") + + assert open_weathernext2_zarr_dataset("gs://weathernext/example") == {"ok": True} + assert calls[0][1]["storage_options"]["token"] == "google_default" diff --git a/tests/test_weathernext2_worker_service.py b/tests/test_weathernext2_worker_service.py new file mode 100644 index 00000000..c162b28f --- /dev/null +++ b/tests/test_weathernext2_worker_service.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +import json +from datetime import date, timedelta + +from src.database.runtime_state import ( + RuntimeStateDB, + TrainingFeatureRecordRepository, + TruthRecordRepository, +) +from web.weathernext2_worker_service import run_weathernext2_cycle + + +def _seed_training_rows(feature_repo, truth_repo, count: int = 170): + for idx in range(count): + city = "houston" if idx % 2 == 0 else "shanghai" + target_date = (date(2026, 1, 1) + timedelta(days=idx)).isoformat() + median = 30.0 + (idx % 7) * 0.2 + residual = 1.0 if city == "houston" else -0.5 + feature_repo.upsert_record( + city, + target_date, + { + "weathernext2": { + "summary": { + "mean": median, + "median": median, + "p10": median - 1.0, + "p25": median - 0.5, + "p75": median + 0.5, + "p90": median + 1.0, + "spread": 2.0, + } + }, + "deb_prediction": median + 0.2, + "model_median_c": median + 0.1, + "model_spread": 1.4, + "current_max_so_far_c": median - 2.0, + "local_hour": 12, + }, + ) + truth_repo.upsert_truth( + city=city, + target_date=target_date, + actual_high=median + residual, + settlement_source="metar", + settlement_station_code="TEST", + settlement_station_label="TEST", + truth_version="test", + updated_by="test", + is_final=True, + ) + + +def test_weathernext2_worker_writes_city_artifact_and_trains_calibrator(tmp_path): + db = RuntimeStateDB(str(tmp_path / "polyweather.db")) + feature_repo = TrainingFeatureRecordRepository(db) + truth_repo = TruthRecordRepository(db) + _seed_training_rows(feature_repo, truth_repo) + + def fake_fetcher(city, meta, target_date): + return { + "city": city, + "target_date": target_date, + "source_run": "2026-06-29T00:00:00Z", + "member_highs": [30.0, 30.2, 30.4, 30.6], + } + + output_path = tmp_path / "weathernext2_city_highs.json" + model_dir = tmp_path / "model" + result = run_weathernext2_cycle( + city_registry={ + "houston": {"name": "Houston", "lat": 29.7, "lon": -95.3, "use_fahrenheit": False, "tz_offset": -5 * 3600}, + }, + output_path=str(output_path), + model_dir=str(model_dir), + fetcher=fake_fetcher, + feature_repo=feature_repo, + truth_repo=truth_repo, + min_global_samples=150, + min_city_samples=5, + ) + + assert result["status"] == "written" + assert result["city_count"] == 1 + assert result["calibration"]["trained"] is True + assert (model_dir / "metadata.json").is_file() + payload = json.loads(output_path.read_text(encoding="utf-8")) + assert payload["cities"]["houston"]["member_highs"] == [30.0, 30.2, 30.4, 30.6] + assert payload["calibration_training"]["samples"] == 170 + + +def test_weathernext2_worker_does_not_overwrite_artifact_when_no_city_payloads(tmp_path): + db = RuntimeStateDB(str(tmp_path / "polyweather.db")) + output_path = tmp_path / "weathernext2_city_highs.json" + output_path.write_text('{"cities":{"old":{"member_highs":[1]}}}', encoding="utf-8") + + result = run_weathernext2_cycle( + city_registry={"houston": {"name": "Houston", "lat": 29.7, "lon": -95.3}}, + output_path=str(output_path), + model_dir=str(tmp_path / "model"), + fetcher=lambda _city, _meta, _target_date: None, + feature_repo=TrainingFeatureRecordRepository(db), + truth_repo=TruthRecordRepository(db), + ) + + assert result["status"] == "no_city_payloads" + assert json.loads(output_path.read_text(encoding="utf-8"))["cities"]["old"]["member_highs"] == [1] diff --git a/web/scan_terminal_city_row.py b/web/scan_terminal_city_row.py index 6af494f2..77d80cd8 100644 --- a/web/scan_terminal_city_row.py +++ b/web/scan_terminal_city_row.py @@ -78,6 +78,34 @@ def _city_local_time(city: str, utc_offset_seconds: Optional[int] = None) -> str return _city_local_now(city, utc_offset_seconds).strftime("%H:%M") +def _model_sources_with_weathernext2( + base_models: Any, + data: Dict[str, Any], +) -> Dict[str, Any]: + models = ( + { + str(k): v + for k, v in (base_models or {}).items() + if v is not None + } + if isinstance(base_models, dict) + else {} + ) + weathernext2 = data.get("weathernext2") + if isinstance(weathernext2, dict): + summary = ( + weathernext2.get("summary") + if isinstance(weathernext2.get("summary"), dict) + else {} + ) + representative = _safe_float(summary.get("median")) + if representative is None: + representative = _safe_float(summary.get("mean")) + if representative is not None: + models["WeatherNext 2"] = round(float(representative), 1) + return models + + def _panel_cache_stale_reason(city: str, cached_entry: Dict[str, Any], payload: Dict[str, Any]) -> Optional[str]: updated_at_ts = _safe_float(cached_entry.get("updated_at_ts")) if updated_at_ts is None or time.time() - updated_at_ts > SCAN_PANEL_CACHE_MAX_AGE_SEC: @@ -611,7 +639,12 @@ def _build_terminal_row( "distribution_full": scan.get("distribution_full") or scan.get("distribution_preview") or row.get("distribution_preview") or [], "probability_engine": scan.get("probability_engine") or (data.get("probabilities") or {}).get("engine"), "probability_calibration_mode": scan.get("probability_calibration_mode") or (data.get("probabilities") or {}).get("calibration_mode"), - "model_cluster_sources": daily_entry.get("models") if isinstance(daily_entry.get("models"), dict) else data.get("multi_model", {}).get("forecasts"), + "model_cluster_sources": _model_sources_with_weathernext2( + daily_entry.get("models") + if isinstance(daily_entry.get("models"), dict) + else data.get("multi_model", {}).get("forecasts"), + data, + ), "window_phase": row.get("window_phase") or scan.get("window_phase"), "window_score": row.get("window_score") if row.get("window_score") is not None else scan.get("window_score"), "signal_status": scan.get("signal_status"), @@ -744,13 +777,11 @@ def _build_quick_row( "network_provider": data.get("official_network_source") or official_status.get("provider_code"), "network_provider_label": official_status.get("provider_label"), "deb_prediction": deb.get("prediction"), - "model_cluster_sources": ( + "model_cluster_sources": _model_sources_with_weathernext2( daily_entry.get("models") if isinstance(daily_entry.get("models"), dict) - else { - str(k): v for k, v in multi.get("forecasts", {}).items() - if v is not None - } + else multi.get("forecasts", {}), + data, ), "distribution_preview": distribution[:6] if distribution else [], "distribution_full": probs.get("distribution_all") or distribution, diff --git a/web/weathernext2_worker.py b/web/weathernext2_worker.py new file mode 100644 index 00000000..e1d6eb45 --- /dev/null +++ b/web/weathernext2_worker.py @@ -0,0 +1,91 @@ +"""Standalone WeatherNext 2 offline worker.""" + +from __future__ import annotations + +import argparse +import json +import os +import signal +import threading +import time +from types import FrameType +from typing import Optional + +from loguru import logger + +from web.weathernext2_worker_service import run_weathernext2_cycle + + +_STOP_EVENT = threading.Event() + + +def _env_int(name: str, default: int) -> int: + raw = os.getenv(name) + if raw is None: + return default + try: + return int(raw) + except Exception: + return default + + +def _handle_stop_signal(signum: int, _frame: Optional[FrameType]) -> None: + logger.info("WeatherNext2 worker stopping signal={}", signum) + _STOP_EVENT.set() + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Fetch WeatherNext 2 city highs and train LightGBM calibration." + ) + parser.add_argument("--once", action="store_true", help="Run one cycle and exit.") + parser.add_argument( + "--interval-sec", + type=int, + default=_env_int("WEATHERNEXT2_WORKER_INTERVAL_SEC", 21600), + ) + parser.add_argument( + "--initial-delay-sec", + type=int, + default=_env_int("WEATHERNEXT2_WORKER_INITIAL_DELAY_SEC", 90), + ) + parser.add_argument("--cities", nargs="*", default=None) + return parser.parse_args() + + +def _run_once(*, cities: Optional[list[str]]) -> dict: + result = run_weathernext2_cycle(cities=cities) + logger.info("WeatherNext2 worker result={}", json.dumps(result, ensure_ascii=False)) + return result + + +def main() -> None: + signal.signal(signal.SIGINT, _handle_stop_signal) + signal.signal(signal.SIGTERM, _handle_stop_signal) + args = _parse_args() + + if args.once: + result = _run_once(cities=args.cities) + print(json.dumps(result, ensure_ascii=False)) + return + + interval_sec = max(1800, int(args.interval_sec or 21600)) + initial_delay_sec = max(0, int(args.initial_delay_sec or 0)) + logger.info("WeatherNext2 worker started interval={}s", interval_sec) + if initial_delay_sec and _STOP_EVENT.wait(initial_delay_sec): + return + + while not _STOP_EVENT.is_set(): + started = time.time() + try: + _run_once(cities=args.cities) + except Exception as exc: + logger.exception("WeatherNext2 cycle failed: {}", exc) + elapsed = time.time() - started + wait_for = max(5.0, interval_sec - elapsed) + if _STOP_EVENT.wait(wait_for): + break + + +if __name__ == "__main__": + main() diff --git a/web/weathernext2_worker_service.py b/web/weathernext2_worker_service.py new file mode 100644 index 00000000..8da6d310 --- /dev/null +++ b/web/weathernext2_worker_service.py @@ -0,0 +1,288 @@ +"""Offline WeatherNext 2 fetch and calibration worker service.""" + +from __future__ import annotations + +import json +import os +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Callable, Dict, Iterable, Mapping, Optional, Sequence + +from loguru import logger + +from src.analysis.weathernext2_calibration import train_lightgbm_quantile_calibrator +from src.data_collection.city_registry import CITY_REGISTRY +from src.data_collection.weathernext2_fetcher import ( + extract_member_hourly_from_grid_dataset, + open_weathernext2_zarr_dataset, +) +from src.data_collection.weathernext2_sources import ( + WEATHERNEXT2_GCS_ZARR_URI, + build_city_local_daily_highs_from_hourly, + build_weathernext2_city_probability, +) +from src.database.runtime_state import ( + TrainingFeatureRecordRepository, + TruthRecordRepository, +) + + +CityFetcher = Callable[[str, Mapping[str, Any], str], Optional[Mapping[str, Any]]] + + +def _env_int(name: str, default: int) -> int: + raw = os.getenv(name) + if raw is None: + return default + try: + return int(raw) + except Exception: + return default + + +def _normalize_city(city: str) -> str: + return str(city or "").strip().lower().replace("-", " ") + + +def _selected_city_names( + city_registry: Mapping[str, Mapping[str, Any]], + cities: Optional[Iterable[str]], +) -> Sequence[str]: + selected = {_normalize_city(city) for city in (cities or []) if _normalize_city(city)} + names = [] + for city in sorted(city_registry.keys()): + normalized = _normalize_city(city) + if selected and normalized not in selected: + continue + names.append(normalized) + return tuple(names) + + +def _city_target_date(city_meta: Mapping[str, Any]) -> str: + offset_seconds = int(city_meta.get("tz_offset") or 0) + now_utc = datetime.now(timezone.utc).timestamp() + return datetime.fromtimestamp(now_utc + offset_seconds, tz=timezone.utc).date().isoformat() + + +def _safe_float(value: Any) -> Optional[float]: + try: + parsed = float(value) + except (TypeError, ValueError): + return None + return parsed if parsed == parsed else None + + +def _dataset_source_run(dataset: Any) -> Optional[str]: + attrs = getattr(dataset, "attrs", None) + if isinstance(attrs, Mapping): + for key in ("source_run", "forecast_reference_time", "init_time", "run_time"): + value = attrs.get(key) + if value: + return str(value) + for key in ("init_time", "forecast_reference_time", "time"): + try: + values = dataset[key].values + except Exception: + try: + values = dataset[key] + except Exception: + continue + try: + first = values[-1] + except Exception: + first = values + if first is not None: + return str(first) + return None + + +def _default_city_fetcher(dataset: Any, temp_var: Optional[str]) -> CityFetcher: + source_run = _dataset_source_run(dataset) + + def fetch(city: str, meta: Mapping[str, Any], target_date: str) -> Optional[Mapping[str, Any]]: + lat = _safe_float(meta.get("lat")) + lon = _safe_float(meta.get("lon")) + if lat is None or lon is None: + return None + use_fahrenheit = bool(meta.get("use_fahrenheit") or meta.get("f")) + extracted = extract_member_hourly_from_grid_dataset( + dataset, + lat=lat, + lon=lon, + temp_var=temp_var, + use_fahrenheit=use_fahrenheit, + ) + member_highs = build_city_local_daily_highs_from_hourly( + extracted["member_hourly"], + extracted["utc_times"], + int(meta.get("tz_offset") or 0), + target_date, + ) + if not member_highs: + return None + payload = build_weathernext2_city_probability( + city=city, + member_highs=member_highs, + temp_symbol="°F" if use_fahrenheit else "°C", + target_date=target_date, + source_run=source_run, + ) + payload["fetch_metadata"] = { + "backend": "gcs_zarr", + "temp_var": extracted.get("temp_var"), + "units": extracted.get("units"), + "nearest_lat": extracted.get("nearest_lat"), + "nearest_lon": extracted.get("nearest_lon"), + "requested_lat": lat, + "requested_lon": lon, + } + return payload + + return fetch + + +def _atomic_write_json(path: Path, payload: Mapping[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp_path = path.with_suffix(f"{path.suffix}.tmp") + tmp_path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8") + os.replace(str(tmp_path), str(path)) + + +def _flatten_training_records_for_weathernext2( + *, + feature_repo: Optional[TrainingFeatureRecordRepository] = None, + truth_repo: Optional[TruthRecordRepository] = None, +) -> list[Dict[str, Any]]: + features = (feature_repo or TrainingFeatureRecordRepository()).load_all() + truth = (truth_repo or TruthRecordRepository()).load_all() + rows: list[Dict[str, Any]] = [] + for city, by_date in features.items(): + city_truth = truth.get(city, {}) + for target_date, payload in by_date.items(): + if not isinstance(payload, dict): + continue + actual = payload.get("actual_high") + if actual is None and isinstance(city_truth.get(target_date), dict): + actual = city_truth[target_date].get("actual_high") + if actual is None: + continue + row = dict(payload) + row["city"] = city + row["target_date"] = target_date + row["actual_high_c"] = actual + if "deb_prediction_c" not in row and row.get("deb_prediction") is not None: + row["deb_prediction_c"] = row.get("deb_prediction") + if "weathernext2" not in row and isinstance(row.get("weather_data"), dict): + wn2 = row["weather_data"].get("weathernext2") + if isinstance(wn2, dict): + row["weathernext2"] = wn2 + rows.append(row) + return rows + + +def run_weathernext2_cycle( + *, + city_registry: Optional[Mapping[str, Mapping[str, Any]]] = None, + cities: Optional[Iterable[str]] = None, + output_path: Optional[str] = None, + model_dir: Optional[str] = None, + fetcher: Optional[CityFetcher] = None, + dataset: Any = None, + feature_repo: Optional[TrainingFeatureRecordRepository] = None, + truth_repo: Optional[TruthRecordRepository] = None, + min_global_samples: Optional[int] = None, + min_city_samples: Optional[int] = None, +) -> Dict[str, Any]: + registry = city_registry or CITY_REGISTRY + data_root = Path(os.getenv("WEATHERNEXT2_DATA_ROOT", "/app/data/weathernext2")) + out_path = Path(output_path or os.getenv("WEATHERNEXT2_CITY_HIGHS_PATH", "") or data_root / "weathernext2_city_highs.json") + calibrator_dir = Path(model_dir or os.getenv("WEATHERNEXT2_MODEL_DIR", "/app/data/models/weathernext2_calibrator")) + selected = _selected_city_names(registry, cities) + backend = str(os.getenv("WEATHERNEXT2_BACKEND", "gcs_zarr") or "gcs_zarr").strip().lower() + + generated_at = datetime.now(timezone.utc).isoformat() + opened_dataset = dataset + if fetcher is None: + if opened_dataset is None: + if backend != "gcs_zarr": + return {"ok": False, "status": "skipped", "reason": "unsupported_backend", "backend": backend} + uri = os.getenv("WEATHERNEXT2_GCS_ZARR_URI", WEATHERNEXT2_GCS_ZARR_URI) + opened_dataset = open_weathernext2_zarr_dataset(uri) + fetcher = _default_city_fetcher( + opened_dataset, + str(os.getenv("WEATHERNEXT2_TEMP_VAR", "") or "").strip() or None, + ) + + city_payloads: Dict[str, Dict[str, Any]] = {} + failed = 0 + items = [] + for city in selected: + meta = registry.get(city) or {} + target_date = _city_target_date(meta) + try: + payload = fetcher(city, meta, target_date) + if not isinstance(payload, Mapping): + items.append({"city": city, "ok": True, "status": "empty", "target_date": target_date}) + continue + city_payloads[city] = dict(payload) + items.append({"city": city, "ok": True, "status": "written", "target_date": target_date}) + except Exception as exc: + failed += 1 + logger.warning("WeatherNext2 city fetch failed city={}: {}", city, exc) + items.append({"city": city, "ok": False, "status": "failed", "target_date": target_date, "error": str(exc)}) + + training_rows = _flatten_training_records_for_weathernext2( + feature_repo=feature_repo, + truth_repo=truth_repo, + ) + calibration = train_lightgbm_quantile_calibrator( + training_rows, + model_dir=calibrator_dir, + min_global_samples=min_global_samples + if min_global_samples is not None + else _env_int("WEATHERNEXT2_MIN_GLOBAL_SAMPLES", 150), + min_city_samples=min_city_samples + if min_city_samples is not None + else _env_int("WEATHERNEXT2_MIN_CITY_SAMPLES", 5), + ) + + if not city_payloads: + return { + "ok": failed == 0, + "status": "no_city_payloads", + "backend": backend, + "generated_at": generated_at, + "city_count": 0, + "failed": failed, + "items": items, + "calibration": calibration, + "output_path": str(out_path), + } + + artifact = { + "schema_version": 1, + "source": "weathernext2", + "backend": backend, + "gcs_zarr_uri": os.getenv("WEATHERNEXT2_GCS_ZARR_URI", WEATHERNEXT2_GCS_ZARR_URI), + "generated_at": generated_at, + "city_count": len(city_payloads), + "cities": city_payloads, + "fetch_metadata": { + "selected_city_count": len(selected), + "failed": failed, + }, + "calibration_training": calibration, + } + _atomic_write_json(out_path, artifact) + + return { + "ok": failed == 0, + "status": "written", + "backend": backend, + "generated_at": generated_at, + "city_count": len(city_payloads), + "failed": failed, + "items": items, + "calibration": calibration, + "output_path": str(out_path), + }