Files
PolyWeather/scripts/backfill_probability_shadow_history.py
T
2026-04-02 23:24:38 +08:00

154 lines
4.8 KiB
Python

import argparse
import json
import os
import sys
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if PROJECT_ROOT not in sys.path:
sys.path.insert(0, PROJECT_ROOT)
from src.analysis.deb_algorithm import load_history, save_history # noqa: E402
from src.analysis.probability_calibration import ( # noqa: E402
ENGINE_MODE_EMOS_SHADOW,
apply_probability_calibration,
build_probability_features,
)
from scripts.fit_probability_calibration import _default_history_arg # noqa: E402
def _sample_to_features(sample):
peak_flag = sample.get("peak_flag")
if peak_flag == 1.0:
peak_status = "past"
elif peak_flag == 0.5:
peak_status = "in_window"
else:
peak_status = "before"
return build_probability_features(
city_name=sample.get("city") or "",
raw_mu=sample.get("raw_mu"),
raw_sigma=sample.get("raw_sigma"),
deb_prediction=sample.get("deb_prediction"),
ens_data={
"median": sample.get("ens_median"),
"p10": None,
"p90": None,
},
current_forecasts={},
max_so_far=None,
peak_status=peak_status,
local_hour_frac=None,
)
def main():
parser = argparse.ArgumentParser(description="Backfill shadow probability snapshots into daily records.")
parser.add_argument(
"--history-file",
default=_default_history_arg(),
)
parser.add_argument(
"--training-samples",
default=os.path.join(
PROJECT_ROOT,
"artifacts",
"probability_calibration",
"training_samples.json",
),
)
parser.add_argument(
"--calibration-file",
default=os.path.join(
PROJECT_ROOT,
"artifacts",
"probability_calibration",
"default.json",
),
)
args = parser.parse_args()
history = load_history(args.history_file)
with open(args.training_samples, "r", encoding="utf-8") as fh:
training_payload = json.load(fh)
updated = 0
touched = 0
for sample in training_payload.get("samples") or []:
city = str(sample.get("city") or "").strip().lower()
date_str = str(sample.get("date") or "").strip()
if not city or not date_str:
continue
record = ((history.get(city) or {}).get(date_str) or {})
if not isinstance(record, dict) or not record:
continue
legacy_distribution = [
{"value": row.get("v"), "probability": row.get("p")}
for row in (record.get("prob_snapshot") or [])
if isinstance(row, dict) and row.get("v") is not None
]
if not legacy_distribution:
continue
calibration = apply_probability_calibration(
city_name=city,
temp_symbol="°F" if city in {"atlanta", "chicago", "dallas", "miami", "new york", "seattle"} else "°C",
raw_mu=sample.get("raw_mu"),
raw_sigma=sample.get("raw_sigma"),
max_so_far=None,
legacy_distribution=legacy_distribution,
features=_sample_to_features(sample),
calibration_path=args.calibration_file,
mode=ENGINE_MODE_EMOS_SHADOW,
)
shadow_distribution = calibration.get("shadow_distribution") or []
compact_shadow = [
{
"v": int(row.get("value")),
"p": round(float(row.get("probability") or 0.0), 3),
}
for row in shadow_distribution[:4]
if row.get("value") is not None
]
compact_calibration = {
"mode": calibration.get("mode"),
"engine": calibration.get("engine"),
"version": calibration.get("calibration_version"),
"source": calibration.get("calibration_source"),
"raw_mu": calibration.get("raw_mu"),
"raw_sigma": calibration.get("raw_sigma"),
"calibrated_mu": calibration.get("calibrated_mu"),
"calibrated_sigma": calibration.get("calibrated_sigma"),
}
touched += 1
if (
record.get("shadow_prob_snapshot") == compact_shadow
and record.get("probability_calibration") == compact_calibration
):
continue
record["shadow_prob_snapshot"] = compact_shadow
record["probability_calibration"] = compact_calibration
history[city][date_str] = record
updated += 1
save_history(args.history_file, history)
print(
json.dumps(
{
"samples_seen": len(training_payload.get("samples") or []),
"records_considered": touched,
"records_updated": updated,
},
ensure_ascii=False,
indent=2,
)
)
if __name__ == "__main__":
main()