Add trained checkpoint and inference artifacts

Browse files

Files changed (5) hide show

best.pt +3 -0
forecast.json +74 -0
forecast_world_map.html +61 -0
requirements.txt +5 -0
run_inference.py +94 -0

best.pt ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:ea6e045bea4e354c36b7fc0ccb72de99efa5634ad813de96e545aed4d7c2840e
+size 1436940

forecast.json ADDED Viewed

	@@ -0,0 +1,74 @@

+{
+  "storm": "Bavi",
+  "source": "JTWC/TCGP fixes through 2026-07-10 0600 UTC",
+  "initial_time": "2026-07-10T06:00:00Z",
+  "initial_lat": 21.9,
+  "initial_lon": 126.9,
+  "points": [
+    {
+      "lead_hours": 6,
+      "lat": 22.466064453125,
+      "lon": 126.55152130126953,
+      "p10_lat": 22.288002014160156,
+      "p90_lat": 22.62434959411621,
+      "p10_lon": 126.4314956665039,
+      "p90_lon": 126.70477294921875
+    },
+    {
+      "lead_hours": 12,
+      "lat": 22.821269989013672,
+      "lon": 126.27210998535156,
+      "p10_lat": 22.475460052490234,
+      "p90_lat": 23.098590850830078,
+      "p10_lon": 126.02592468261719,
+      "p90_lon": 126.5489273071289
+    },
+    {
+      "lead_hours": 24,
+      "lat": 23.898107528686523,
+      "lon": 125.54524230957031,
+      "p10_lat": 23.513370513916016,
+      "p90_lat": 24.373027801513672,
+      "p10_lon": 124.97864532470703,
+      "p90_lon": 126.13536834716797
+    },
+    {
+      "lead_hours": 48,
+      "lat": 25.304624557495117,
+      "lon": 124.66926574707031,
+      "p10_lat": 24.34754180908203,
+      "p90_lat": 26.398544311523438,
+      "p10_lon": 123.34161376953125,
+      "p90_lon": 125.87503814697266
+    },
+    {
+      "lead_hours": 72,
+      "lat": 26.52523422241211,
+      "lon": 123.29611206054688,
+      "p10_lat": 24.960338592529297,
+      "p90_lat": 28.163558959960938,
+      "p10_lon": 120.4675064086914,
+      "p90_lon": 125.69815826416016
+    },
+    {
+      "lead_hours": 96,
+      "lat": 26.998743057250977,
+      "lon": 122.69322204589844,
+      "p10_lat": 24.666839599609375,
+      "p90_lat": 29.167266845703125,
+      "p10_lon": 119.0267562866211,
+      "p90_lon": 125.30994415283203
+    },
+    {
+      "lead_hours": 120,
+      "lat": 28.687833786010742,
+      "lon": 121.19312286376953,
+      "p10_lat": 26.278013229370117,
+      "p90_lat": 31.288469314575195,
+      "p10_lon": 115.32806396484375,
+      "p90_lon": 125.7118911743164
+    }
+  ],
+  "device": "mps",
+  "note": "Mac checkpoint inference using current track fixes and mean-normalized atmospheric input because 2026 ERA5 fields are unavailable locally. Not an operational forecast."
+}

forecast_world_map.html ADDED Viewed

	@@ -0,0 +1,61 @@

+<!doctype html>
+<html lang="en">
+<head>
+  <meta charset="utf-8">
+  <meta name="viewport" content="width=device-width, initial-scale=1">
+  <title>Bavi Mac Model Forecast</title>
+  <link rel="stylesheet" href="https://unpkg.com/leaflet@1.9.4/dist/leaflet.css">
+  <style>
+    html, body, #map { height: 100%; margin: 0; }
+    body { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; }
+    .panel { position: absolute; z-index: 1000; top: 16px; left: 16px; width: min(360px, calc(100vw - 32px)); padding: 14px 16px; background: rgba(255,255,255,.96); border: 1px solid #d5d9df; box-shadow: 0 3px 14px rgba(0,0,0,.18); }
+    h1 { margin: 0 0 7px; font-size: 20px; }
+    p { margin: 5px 0; font-size: 13px; line-height: 1.4; color: #374151; }
+    .legend { display: flex; gap: 12px; flex-wrap: wrap; margin-top: 9px; font-size: 12px; color: #374151; }
+    .key { display: inline-flex; align-items: center; gap: 5px; }
+    .swatch { width: 20px; height: 3px; display: inline-block; }
+    .mean { background: #d7263d; }
+    .actual { background: #111827; }
+    .corridor { width: 20px; height: 10px; background: rgba(245,158,11,.32); border: 1px solid #d97706; }
+  </style>
+</head>
+<body>
+  <div id="map"></div>
+  <section class="panel">
+    <h1>Bavi: Mac checkpoint prediction</h1>
+    <p><b>Initial fix:</b> 10 Jul 2026 06:00 UTC, 21.9°N 126.9°E</p>
+    <p><b>Run:</b> local Mac Apple MPS, 50 ensemble members, 120-hour horizon</p>
+    <p><b>Important:</b> this is a research proxy using a mean-normalized atmospheric field because the matching 2026 ERA5 field is not available locally. It is not an operational warning forecast.</p>
+    <div class="legend">
+      <span class="key"><i class="swatch actual"></i>official initial fix</span>
+      <span class="key"><i class="swatch mean"></i>model mean</span>
+      <span class="key"><i class="swatch corridor"></i>10–90% corridor</span>
+    </div>
+  </section>
+  <script src="https://unpkg.com/leaflet@1.9.4/dist/leaflet.js"></script>
+  <script>
+    const initial = [21.9, 126.9];
+    const forecast = [
+      {h: 6, lat: 22.4336, lon: 126.5818, la: 22.3279, lb: 22.5795, oa: 126.4390, ob: 126.7272},
+      {h: 12, lat: 22.7873, lon: 126.3022, la: 22.5105, lb: 23.0514, oa: 126.0263, ob: 126.6464},
+      {h: 24, lat: 23.8919, lon: 125.4258, la: 23.5613, lb: 24.3624, oa: 124.9874, ob: 125.9527},
+      {h: 48, lat: 25.6544, lon: 124.6313, la: 24.6905, lb: 26.3908, oa: 123.1265, ob: 126.0750},
+      {h: 72, lat: 26.4361, lon: 123.5473, la: 24.9631, lb: 28.2892, oa: 120.7869, ob: 125.8778},
+      {h: 96, lat: 26.6257, lon: 121.9343, la: 24.2808, lb: 28.4809, oa: 118.0243, ob: 126.7538},
+      {h: 120, lat: 28.6126, lon: 121.8454, la: 26.2118, lb: 30.8639, oa: 117.2112, ob: 126.6462}
+    ];
+    const map = L.map('map', {worldCopyJump: true}).setView([25, 126], 4);
+    L.tileLayer('https://{s}.tile.openstreetmap.org/{z}/{x}/{y}.png', {maxZoom: 18, attribution: '&copy; OpenStreetMap contributors'}).addTo(map);
+    L.marker(initial).addTo(map).bindPopup('<b>Official initial fix</b><br>21.9°N, 126.9°E<br>10 Jul 2026 06:00 UTC').openPopup();
+    const meanLine = [initial];
+    forecast.forEach(p => {
+      meanLine.push([p.lat, p.lon]);
+      L.rectangle([[p.la, p.oa], [p.lb, p.ob]], {color:'#d97706', weight:1, fillColor:'#f59e0b', fillOpacity:.18}).addTo(map).bindTooltip('+' + p.h + 'h model 10–90% box');
+      L.circleMarker([p.lat, p.lon], {radius: 5, color:'#d7263d', fillColor:'#d7263d', fillOpacity:1}).addTo(map).bindTooltip('+' + p.h + 'h: ' + p.lat.toFixed(2) + '°N, ' + p.lon.toFixed(2) + '°E');
+    });
+    L.polyline(meanLine, {color:'#d7263d', weight:4, opacity:.9}).addTo(map);
+    L.polyline([initial, [21.9,126.9]], {color:'#111827', weight:5}).addTo(map);
+    map.fitBounds(L.latLngBounds(meanLine), {padding: [45, 45]});
+  </script>
+</body>
+</html>

requirements.txt ADDED Viewed

	@@ -0,0 +1,5 @@

+numpy==2.5.1
+torch==2.13.0
+scikit-learn==1.9.0
+matplotlib==3.11.0
+folium==0.20.0

run_inference.py ADDED Viewed

	@@ -0,0 +1,94 @@

+import json
+import argparse
+from pathlib import Path
+import numpy as np
+import torch
+from torch import nn
+ROOT = Path(__file__).resolve().parent
+parser = argparse.ArgumentParser(description="Run the Bavi typhoon checkpoint locally.")
+parser.add_argument("--checkpoint", type=Path, default=ROOT / "best.pt")
+parser.add_argument("--output", type=Path, default=ROOT / "forecast.json")
+parser.add_argument("--device", choices=["auto", "mps", "cpu"], default="auto")
+args = parser.parse_args()
+CKPT = args.checkpoint.expanduser().resolve()
+OUT = args.output.expanduser().resolve()
+class FieldEncoder(nn.Module):
+    def __init__(self, channels, hidden=128):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Conv2d(channels,64,3,padding=1), nn.GELU(), nn.BatchNorm2d(64),
+            nn.Conv2d(64,96,3,stride=2,padding=1), nn.GELU(), nn.BatchNorm2d(96),
+            nn.Conv2d(96,hidden,3,stride=2,padding=1), nn.GELU(),
+            nn.AdaptiveAvgPool2d(1))
+    def forward(self, x): return self.net(x).flatten(1)
+class ERA5CycloneEnsemble(nn.Module):
+    def __init__(self, track_dim, channels, steps, output_dim, hidden=192, latent=32):
+        super().__init__()
+        self.steps, self.output_dim, self.latent = steps, output_dim, latent
+        self.field = FieldEncoder(channels, 128)
+        self.track_gru = nn.GRU(track_dim,96,batch_first=True,bidirectional=True)
+        self.fuse = nn.Sequential(nn.Linear(128*steps+192,hidden),nn.GELU(),nn.Dropout(.15),nn.Linear(hidden,hidden),nn.GELU())
+        self.mean = nn.Linear(hidden,output_dim)
+        self.log_scale = nn.Linear(hidden,output_dim)
+        self.latent_proj = nn.Sequential(nn.Linear(latent,hidden),nn.GELU(),nn.Linear(hidden,output_dim))
+    def encode(self, track_x, field_x):
+        b,t,c,h,w=field_x.shape
+        field_z=self.field(field_x.reshape(b*t,c,h,w)).reshape(b,t,-1).flatten(1)
+        _,state=self.track_gru(track_x)
+        return self.fuse(torch.cat([field_z,state.transpose(0,1).flatten(1)],dim=1))
+    def sample(self, track_x, field_x, n, temperature=1.0):
+        z=self.encode(track_x,field_x)
+        mean=self.mean(z); log_scale=self.log_scale(z).clamp(-5,2)
+        eps=torch.randn(n,z.shape[0],self.latent,device=z.device)
+        low=self.latent_proj(eps.reshape(-1,self.latent)).reshape(n,z.shape[0],-1)
+        independent=torch.randn_like(low)*torch.exp(log_scale).unsqueeze(0)
+        return mean.unsqueeze(0)+temperature*(low+independent)
+raw=torch.load(CKPT,map_location="cpu",weights_only=False)
+if args.device == "mps":
+    device = torch.device("mps")
+elif args.device == "cpu":
+    device = torch.device("cpu")
+else:
+    device=torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+config=raw["config"]
+track_scaler=raw["track_scaler"]; y_scaler=raw["y_scaler"]
+field_mean=np.asarray(raw["field_mean"],dtype="float32")
+field_std=np.asarray(raw["field_std"],dtype="float32")
+model=ERA5CycloneEnsemble(9,10,1,28).to(device)
+model.load_state_dict(raw["model_state"]); model.eval()
+fixes=[
+    ("2026-07-09T12:00:00",19.2,128.8,100,952),
+    ("2026-07-09T18:00:00",20.1,128.2,90,953),
+    ("2026-07-10T00:00:00",20.8,127.3,75,964),
+    ("2026-07-10T06:00:00",21.9,126.9,75,962),
+]
+track=[]
+for i,(stamp,lat,lon,wind,pres) in enumerate(fixes):
+    dlat=0 if i==0 else lat-fixes[i-1][1]
+    dlon=0 if i==0 else ((lon-fixes[i-1][2]+180)%360)-180
+    track.append([lat,lon,wind,pres,dlat,dlon,float(np.hypot(dlat,dlon)),np.sin(2*np.pi*191/366),np.cos(2*np.pi*191/366)])
+track=np.asarray(track,dtype="float32")
+track[:,:1]-=track[-1:,0:1]
+track[:,1:2]=((track[:,1:2]-track[-1:,1:2]+180)%360)-180
+track[:,2]/=100.0
+track[:,3]=(track[:,3]-950.0)/50.0
+xtrack=track_scaler.transform(track).astype("float32")[None]
+field=np.zeros((1,1,10,33,33),dtype="float32")
+with torch.no_grad():
+    ens=model.sample(torch.from_numpy(xtrack).to(device),torch.from_numpy(field).to(device),50,float(config.get("sample_temperature",1.0))).cpu().numpy()
+ens=ens.mean(0,keepdims=True)+(ens-ens.mean(0,keepdims=True))*1.1297996044158936
+pred=y_scaler.inverse_transform(ens[:,0,:])
+base_lat,base_lon=fixes[-1][1],fixes[-1][2]
+points=[]
+for k,lead in enumerate(config["lead_hours"]):
+    j=4*k
+    lat=base_lat+pred[:,j]; lon=(base_lon+pred[:,j+1])%360
+    points.append({"lead_hours":int(lead),"lat":float(lat.mean()),"lon":float(lon.mean()),"p10_lat":float(np.quantile(lat,.1)),"p90_lat":float(np.quantile(lat,.9)),"p10_lon":float(np.quantile(lon,.1)),"p90_lon":float(np.quantile(lon,.9))})
+result={"storm":"Bavi","source":"JTWC/TCGP fixes through 2026-07-10 0600 UTC","initial_time":"2026-07-10T06:00:00Z","initial_lat":base_lat,"initial_lon":base_lon,"points":points,"device":str(device),"note":"Mac checkpoint inference using current track fixes and mean-normalized atmospheric input because 2026 ERA5 fields are unavailable locally. Not an operational forecast."}
+OUT.write_text(json.dumps(result,indent=2))
+print(json.dumps(result,indent=2))