File size: 927 Bytes
2385fd7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | """Public Trackformer1.1 loading and inference API."""
from __future__ import annotations
from pathlib import Path
from trackformer_1_1_intensity import Trackformer11IntensityEnsemble
from trackformer_1_1_route import (
LEAD_HOURS,
build_pacific_route,
detect_pressure_systems,
forecast_pacific_state,
)
PACKAGE_ROOT = Path(__file__).resolve().parent
MODEL_ROOT = PACKAGE_ROOT / "models" / "trackformer_1_1"
CALIBRATION_PATH = MODEL_ROOT / "trackformer_1_1_calibration.json"
def load_intensity(device: str | None = None) -> Trackformer11IntensityEnsemble:
"""Load the frozen Trackformer1.1 intensity and structure experts."""
return Trackformer11IntensityEnsemble(MODEL_ROOT, CALIBRATION_PATH, device=device)
__all__ = [
"CALIBRATION_PATH",
"LEAD_HOURS",
"MODEL_ROOT",
"build_pacific_route",
"detect_pressure_systems",
"forecast_pacific_state",
"load_intensity",
]
|