524 lines
19 KiB
Python
524 lines
19 KiB
Python
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
from datetime import datetime
|
|
|
|
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.probability_calibration import ( # noqa: E402
|
|
DEFAULT_CALIBRATION_FILE,
|
|
default_calibration_payload,
|
|
fit_calibration,
|
|
)
|
|
from src.analysis.deb_algorithm import load_history # noqa: E402
|
|
from src.database.runtime_state import ( # noqa: E402
|
|
DailyRecordRepository,
|
|
ProbabilitySnapshotRepository,
|
|
STATE_STORAGE_FILE,
|
|
STATE_STORAGE_SQLITE,
|
|
TrainingFeatureRecordRepository,
|
|
TruthRecordRepository,
|
|
get_state_storage_mode,
|
|
)
|
|
|
|
|
|
def _sf(value):
|
|
if value is None:
|
|
return None
|
|
try:
|
|
return float(value)
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _env_int(name, default=None):
|
|
try:
|
|
value = os.getenv(name)
|
|
if value is None or str(value).strip() == "":
|
|
return default
|
|
return int(value)
|
|
except Exception:
|
|
return default
|
|
|
|
|
|
def _log(enabled, message):
|
|
if enabled:
|
|
print(f"[fit_probability_calibration] {message}", flush=True)
|
|
|
|
|
|
def _load_json_if_exists(path):
|
|
if not path or not os.path.exists(path):
|
|
return {}
|
|
with open(path, "r", encoding="utf-8") as fh:
|
|
data = json.load(fh)
|
|
return data if isinstance(data, dict) else {}
|
|
|
|
|
|
def _legacy_training_samples_path():
|
|
return os.path.join(
|
|
PROJECT_ROOT,
|
|
"artifacts",
|
|
"probability_calibration",
|
|
"training_samples.json",
|
|
)
|
|
|
|
|
|
def _legacy_history_path():
|
|
return os.path.join(PROJECT_ROOT, "data", "daily_records.json")
|
|
|
|
|
|
def _legacy_snapshot_path():
|
|
return os.path.join(PROJECT_ROOT, "data", "probability_training_snapshots.jsonl")
|
|
|
|
|
|
def _default_history_arg():
|
|
return _legacy_history_path() if get_state_storage_mode() == STATE_STORAGE_FILE else None
|
|
|
|
|
|
def _default_snapshot_arg():
|
|
return _legacy_snapshot_path() if get_state_storage_mode() == STATE_STORAGE_FILE else None
|
|
|
|
|
|
def _load_history_with_fallback(path):
|
|
if not path:
|
|
if get_state_storage_mode() == STATE_STORAGE_SQLITE:
|
|
return DailyRecordRepository().load_all()
|
|
return {}
|
|
data = _load_json_if_exists(path)
|
|
if data:
|
|
return data
|
|
return load_history(path)
|
|
|
|
|
|
def _load_truth_history():
|
|
if get_state_storage_mode() != STATE_STORAGE_SQLITE:
|
|
return {}
|
|
try:
|
|
return TruthRecordRepository().load_all()
|
|
except Exception:
|
|
return {}
|
|
|
|
|
|
def _load_training_feature_history():
|
|
if get_state_storage_mode() != STATE_STORAGE_SQLITE:
|
|
return {}
|
|
try:
|
|
return TrainingFeatureRecordRepository().load_all()
|
|
except Exception:
|
|
return {}
|
|
|
|
|
|
def _load_snapshot_rows(path, limit=None):
|
|
if get_state_storage_mode() == STATE_STORAGE_SQLITE:
|
|
repo = ProbabilitySnapshotRepository()
|
|
if limit is not None and int(limit) > 0:
|
|
with repo.db.connect() as conn:
|
|
rows = conn.execute(
|
|
"""
|
|
SELECT payload_json
|
|
FROM probability_training_snapshots_store
|
|
ORDER BY id DESC
|
|
LIMIT ?
|
|
""",
|
|
(int(limit),),
|
|
).fetchall()
|
|
out = []
|
|
for row in reversed(rows):
|
|
try:
|
|
out.append(json.loads(row["payload_json"]))
|
|
except Exception:
|
|
continue
|
|
return out
|
|
return repo.load_all_rows()
|
|
rows = []
|
|
if not path or not os.path.exists(path):
|
|
return rows
|
|
with open(path, "r", encoding="utf-8") as fh:
|
|
for line in fh:
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
row = json.loads(line)
|
|
except Exception:
|
|
continue
|
|
if isinstance(row, dict):
|
|
rows.append(row)
|
|
return rows
|
|
|
|
|
|
def _load_legacy_training_samples(path=None):
|
|
payload = _load_json_if_exists(path or _legacy_training_samples_path())
|
|
rows = payload.get("samples") if isinstance(payload, dict) else None
|
|
if not isinstance(rows, list):
|
|
return []
|
|
return [row for row in rows if isinstance(row, dict)]
|
|
|
|
|
|
def _actual_high_for(history, truth_history, settlement_history, city, date_str):
|
|
city_rows = (history or {}).get(city) or {}
|
|
record = city_rows.get(date_str) or {}
|
|
actual_high = _sf(record.get("actual_high")) if isinstance(record, dict) else None
|
|
truth_record = ((truth_history.get(city) or {}).get(date_str) or {})
|
|
if actual_high is None and isinstance(truth_record, dict):
|
|
actual_high = _sf(truth_record.get("actual_high"))
|
|
filled = False
|
|
if actual_high is None:
|
|
actual_high = _sf(((settlement_history.get(city) or {}).get(date_str) or {}).get("max_temp"))
|
|
filled = actual_high is not None
|
|
metadata = {
|
|
"settlement_source": truth_record.get("settlement_source"),
|
|
"settlement_station_code": truth_record.get("settlement_station_code"),
|
|
"truth_version": truth_record.get("truth_version"),
|
|
"truth_updated_by": truth_record.get("updated_by"),
|
|
"truth_updated_at": truth_record.get("truth_updated_at"),
|
|
}
|
|
return actual_high, filled, metadata
|
|
|
|
|
|
def _extract_snapshot_samples(history, truth_history=None, snapshot_rows=None, settlement_history=None):
|
|
samples = []
|
|
filled_actual_from_history = 0
|
|
today = datetime.utcnow().strftime("%Y-%m-%d")
|
|
settlement_history = settlement_history or {}
|
|
|
|
for row in snapshot_rows or []:
|
|
city = str(row.get("city") or "").strip().lower()
|
|
date_str = str(row.get("date") or "").strip()
|
|
if not city or not date_str or date_str == today:
|
|
continue
|
|
|
|
actual_high, filled, truth_meta = _actual_high_for(
|
|
history,
|
|
truth_history or {},
|
|
settlement_history,
|
|
city,
|
|
date_str,
|
|
)
|
|
if actual_high is None:
|
|
continue
|
|
if filled:
|
|
filled_actual_from_history += 1
|
|
|
|
raw_mu = _sf(row.get("raw_mu"))
|
|
raw_sigma = _sf(row.get("raw_sigma"))
|
|
deb_prediction = _sf(row.get("deb_prediction"))
|
|
ensemble = row.get("ensemble") or {}
|
|
if not isinstance(ensemble, dict):
|
|
ensemble = {}
|
|
ens_median = _sf(ensemble.get("median"))
|
|
ensemble_spread = None
|
|
ens_p10 = _sf(ensemble.get("p10"))
|
|
ens_p90 = _sf(ensemble.get("p90"))
|
|
if ens_p10 is not None and ens_p90 is not None and ens_p90 >= ens_p10:
|
|
ensemble_spread = max(0.1, (ens_p90 - ens_p10) / 2.56)
|
|
multi_model = row.get("multi_model") or {}
|
|
if not isinstance(multi_model, dict):
|
|
multi_model = {}
|
|
forecast_values = [val for val in (_sf(v) for v in multi_model.values()) if val is not None]
|
|
forecast_values.sort()
|
|
if ensemble_spread is None:
|
|
if len(forecast_values) >= 2:
|
|
ensemble_spread = max(0.6, (forecast_values[-1] - forecast_values[0]) / 2.0)
|
|
elif raw_sigma is not None:
|
|
ensemble_spread = raw_sigma
|
|
else:
|
|
ensemble_spread = 1.0
|
|
if raw_sigma is None:
|
|
raw_sigma = ensemble_spread
|
|
|
|
peak_status = str(row.get("peak_status") or "before").strip().lower()
|
|
peak_flag = 0.0
|
|
if peak_status == "in_window":
|
|
peak_flag = 0.5
|
|
elif peak_status == "past":
|
|
peak_flag = 1.0
|
|
|
|
max_so_far = _sf(row.get("max_so_far"))
|
|
max_so_far_gap = None
|
|
if deb_prediction is not None and max_so_far is not None:
|
|
max_so_far_gap = deb_prediction - max_so_far
|
|
|
|
if raw_mu is None:
|
|
continue
|
|
|
|
samples.append(
|
|
{
|
|
"city": city,
|
|
"date": date_str,
|
|
"timestamp": row.get("timestamp"),
|
|
"actual_high": actual_high,
|
|
"raw_mu": raw_mu,
|
|
"raw_sigma": raw_sigma or 1.0,
|
|
"deb_prediction": deb_prediction,
|
|
"ens_median": ens_median if ens_median is not None else raw_mu,
|
|
"ensemble_spread": ensemble_spread,
|
|
"max_so_far_gap": max_so_far_gap,
|
|
"peak_flag": peak_flag,
|
|
"sample_source": "snapshot",
|
|
**truth_meta,
|
|
}
|
|
)
|
|
|
|
return samples, filled_actual_from_history
|
|
|
|
|
|
def _extract_daily_record_samples(
|
|
history,
|
|
training_feature_history=None,
|
|
truth_history=None,
|
|
settlement_history=None,
|
|
excluded_keys=None,
|
|
):
|
|
samples = []
|
|
filled_actual_from_history = 0
|
|
today = datetime.utcnow().strftime("%Y-%m-%d")
|
|
settlement_history = settlement_history or {}
|
|
excluded_keys = excluded_keys or set()
|
|
for city, city_rows in (history or {}).items():
|
|
if not isinstance(city_rows, dict):
|
|
continue
|
|
city_settlement = settlement_history.get(city) or {}
|
|
for date_str, record in city_rows.items():
|
|
if date_str == today or not isinstance(record, dict):
|
|
continue
|
|
if (city, date_str) in excluded_keys:
|
|
continue
|
|
actual_high = _sf(record.get("actual_high"))
|
|
truth_meta = ((truth_history or {}).get(city) or {}).get(date_str) or {}
|
|
if actual_high is None:
|
|
actual_high = _sf(truth_meta.get("actual_high"))
|
|
if actual_high is None:
|
|
actual_high = _sf((city_settlement.get(date_str) or {}).get("max_temp"))
|
|
if actual_high is not None:
|
|
filled_actual_from_history += 1
|
|
feature_record = ((training_feature_history or {}).get(city) or {}).get(date_str) or {}
|
|
source_record = feature_record if isinstance(feature_record, dict) and feature_record else record
|
|
deb_prediction = _sf(source_record.get("deb_prediction"))
|
|
raw_mu = _sf(source_record.get("mu")) or deb_prediction
|
|
forecasts = source_record.get("forecasts") or {}
|
|
if not isinstance(forecasts, dict):
|
|
forecasts = {}
|
|
forecast_values = [val for val in (_sf(v) for v in forecasts.values()) if val is not None]
|
|
forecast_values.sort()
|
|
forecast_median = (
|
|
forecast_values[len(forecast_values) // 2] if forecast_values else None
|
|
)
|
|
feature_snapshot = source_record.get("probability_features") or {}
|
|
if not isinstance(feature_snapshot, dict):
|
|
feature_snapshot = {}
|
|
|
|
ens_median = _sf(feature_snapshot.get("ens_median")) or forecast_median or raw_mu
|
|
ensemble_spread = _sf(feature_snapshot.get("ensemble_spread"))
|
|
if ensemble_spread is None:
|
|
if len(forecast_values) >= 2:
|
|
ensemble_spread = max(0.6, (forecast_values[-1] - forecast_values[0]) / 2.0)
|
|
else:
|
|
ensemble_spread = 1.0
|
|
raw_sigma = _sf(feature_snapshot.get("raw_sigma")) or ensemble_spread or 1.0
|
|
peak_status = str(feature_snapshot.get("peak_status") or "before").strip().lower()
|
|
peak_flag = 0.0
|
|
if peak_status == "in_window":
|
|
peak_flag = 0.5
|
|
elif peak_status == "past":
|
|
peak_flag = 1.0
|
|
|
|
if actual_high is None or raw_mu is None:
|
|
continue
|
|
|
|
max_so_far = _sf(feature_snapshot.get("max_so_far"))
|
|
max_so_far_gap = _sf(feature_snapshot.get("max_so_far_gap"))
|
|
if max_so_far_gap is None and max_so_far is not None and deb_prediction is not None:
|
|
max_so_far_gap = deb_prediction - max_so_far
|
|
|
|
samples.append(
|
|
{
|
|
"city": city,
|
|
"date": date_str,
|
|
"actual_high": actual_high,
|
|
"raw_mu": raw_mu,
|
|
"raw_sigma": raw_sigma,
|
|
"deb_prediction": deb_prediction,
|
|
"ens_median": ens_median,
|
|
"ensemble_spread": ensemble_spread,
|
|
"max_so_far_gap": max_so_far_gap,
|
|
"peak_flag": peak_flag,
|
|
"sample_source": "daily_record",
|
|
"settlement_source": truth_meta.get("settlement_source"),
|
|
"settlement_station_code": truth_meta.get("settlement_station_code"),
|
|
"truth_version": truth_meta.get("truth_version"),
|
|
"truth_updated_by": truth_meta.get("updated_by"),
|
|
"truth_updated_at": truth_meta.get("truth_updated_at"),
|
|
}
|
|
)
|
|
return samples, filled_actual_from_history
|
|
|
|
|
|
def _extract_samples(history, training_feature_history=None, truth_history=None, settlement_history=None, snapshot_rows=None):
|
|
snapshot_samples, snapshot_filled = _extract_snapshot_samples(
|
|
history,
|
|
truth_history=truth_history,
|
|
snapshot_rows=snapshot_rows or [],
|
|
settlement_history=settlement_history,
|
|
)
|
|
excluded_keys = {
|
|
(sample["city"], sample["date"])
|
|
for sample in snapshot_samples
|
|
}
|
|
daily_samples, daily_filled = _extract_daily_record_samples(
|
|
history,
|
|
training_feature_history=training_feature_history,
|
|
truth_history=truth_history,
|
|
settlement_history=settlement_history,
|
|
excluded_keys=excluded_keys,
|
|
)
|
|
return snapshot_samples + daily_samples, snapshot_filled + daily_filled
|
|
|
|
|
|
def merge_samples_with_legacy_archive(samples, legacy_samples=None):
|
|
merged = []
|
|
seen = set()
|
|
for sample in samples or []:
|
|
if not isinstance(sample, dict):
|
|
continue
|
|
key = (
|
|
str(sample.get("city") or "").strip().lower(),
|
|
str(sample.get("date") or "").strip(),
|
|
str(sample.get("sample_source") or "").strip().lower(),
|
|
)
|
|
if not key[0] or not key[1]:
|
|
continue
|
|
if key in seen:
|
|
continue
|
|
merged.append(sample)
|
|
seen.add(key)
|
|
for sample in legacy_samples or []:
|
|
if not isinstance(sample, dict):
|
|
continue
|
|
key = (
|
|
str(sample.get("city") or "").strip().lower(),
|
|
str(sample.get("date") or "").strip(),
|
|
str(sample.get("sample_source") or "").strip().lower(),
|
|
)
|
|
if not key[0] or not key[1]:
|
|
continue
|
|
if key in seen:
|
|
continue
|
|
merged.append(sample)
|
|
seen.add(key)
|
|
return merged
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Fit PolyWeather probability calibration parameters.")
|
|
parser.add_argument(
|
|
"--history-file",
|
|
default=_default_history_arg(),
|
|
help="Optional legacy daily_records.json path. In sqlite mode this defaults to the runtime database.",
|
|
)
|
|
parser.add_argument(
|
|
"--output",
|
|
default=DEFAULT_CALIBRATION_FILE,
|
|
help="Output JSON file for fitted calibration parameters.",
|
|
)
|
|
parser.add_argument(
|
|
"--settlement-history",
|
|
default=os.path.join(
|
|
PROJECT_ROOT,
|
|
"artifacts",
|
|
"probability_calibration",
|
|
"settlement_history.json",
|
|
),
|
|
help="Optional daily settlement history JSON built from historical CSV files.",
|
|
)
|
|
parser.add_argument(
|
|
"--snapshot-file",
|
|
default=_default_snapshot_arg(),
|
|
help="Optional legacy JSONL snapshot archive path. In sqlite mode this defaults to the runtime database.",
|
|
)
|
|
parser.add_argument(
|
|
"--version",
|
|
default=None,
|
|
help="Optional explicit calibration version.",
|
|
)
|
|
parser.add_argument(
|
|
"--snapshot-limit",
|
|
type=int,
|
|
default=_env_int("POLYWEATHER_EMOS_TRAINING_SNAPSHOT_LIMIT"),
|
|
help="Optional max number of recent probability snapshots to load from SQLite.",
|
|
)
|
|
parser.add_argument(
|
|
"--verbose",
|
|
action="store_true",
|
|
help="Print data loading and fitting progress.",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
_log(args.verbose, "loading daily records")
|
|
history = _load_history_with_fallback(args.history_file)
|
|
_log(args.verbose, f"loaded daily record cities={len(history or {})}")
|
|
_log(args.verbose, "loading training feature history")
|
|
training_feature_history = _load_training_feature_history()
|
|
_log(args.verbose, f"loaded training feature cities={len(training_feature_history or {})}")
|
|
_log(args.verbose, "loading truth history")
|
|
truth_history = _load_truth_history()
|
|
_log(args.verbose, f"loaded truth cities={len(truth_history or {})}")
|
|
_log(args.verbose, "loading settlement history")
|
|
settlement_history = _load_json_if_exists(args.settlement_history)
|
|
_log(args.verbose, f"loaded settlement history cities={len(settlement_history or {})}")
|
|
_log(
|
|
args.verbose,
|
|
"loading probability snapshots"
|
|
+ (f" limit={args.snapshot_limit}" if args.snapshot_limit else ""),
|
|
)
|
|
snapshot_rows = _load_snapshot_rows(args.snapshot_file, limit=args.snapshot_limit)
|
|
_log(args.verbose, f"loaded probability snapshots={len(snapshot_rows or [])}")
|
|
_log(args.verbose, "loading legacy training archive")
|
|
legacy_training_samples = _load_legacy_training_samples()
|
|
_log(args.verbose, f"loaded legacy training samples={len(legacy_training_samples or [])}")
|
|
_log(args.verbose, "extracting EMOS samples")
|
|
samples, filled_actual_from_history = _extract_samples(
|
|
history,
|
|
training_feature_history=training_feature_history,
|
|
truth_history=truth_history,
|
|
settlement_history=settlement_history,
|
|
snapshot_rows=snapshot_rows,
|
|
)
|
|
samples = merge_samples_with_legacy_archive(samples, legacy_training_samples)
|
|
_log(args.verbose, f"fitting calibration samples={len(samples or [])}")
|
|
calibration = fit_calibration(samples, version=args.version)
|
|
if not samples:
|
|
calibration = default_calibration_payload(
|
|
version=args.version,
|
|
reason="no_samples",
|
|
)
|
|
calibration.setdefault("metrics", {})
|
|
calibration["metrics"]["filled_actual_from_history"] = filled_actual_from_history
|
|
calibration["metrics"]["settlement_history_city_count"] = len(settlement_history)
|
|
calibration["metrics"]["legacy_archive_samples"] = len(legacy_training_samples)
|
|
try:
|
|
calibration["source"] = os.path.relpath(args.output, PROJECT_ROOT)
|
|
except ValueError:
|
|
calibration["source"] = os.path.abspath(args.output)
|
|
|
|
output_dir = os.path.dirname(os.path.abspath(args.output))
|
|
if output_dir:
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
with open(args.output, "w", encoding="utf-8") as fh:
|
|
json.dump(calibration, fh, ensure_ascii=False, indent=2)
|
|
|
|
_log(args.verbose, "done")
|
|
print(
|
|
"saved calibration to {path} with {count} samples".format(
|
|
path=args.output,
|
|
count=calibration.get("metrics", {}).get("sample_count", 0),
|
|
)
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|