Skip to content

ayghri/ssh

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

7 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

E6 sea-level trend forecast — browser demo

A static site that loads the E6 graph-attention SSH-trend forecaster as ONNX and runs it in the user's browser via ONNX Runtime Web + WebGPU. The user picks an anchor year and a horizon, the model predicts the regional pattern of the SSH trend over [anchor, anchor + horizon] from the last 10 calendar years of observations, and the result is rendered to a 360×180 equirectangular canvas.

No build step. No npm install. ESM imports from a CDN; vanilla JS and one CSS file. Works from file:// if your browser allows ESM imports from file:// (Chrome blocks this; use a local HTTP server).

File layout

artifacts/website/
├── README.md            (this file)
├── index.html           static markup
├── style.css            visual theme
├── config.js            asset URLs the user edits
├── app.js               main logic
└── build/               <-- precompute outputs the user uploads to S3
    ├── e6.onnx
    ├── meta.json
    ├── obs_yearly_stats.bin
    ├── keep_mask.bin
    └── cos_lat_per_cell.bin

The build/ directory is generated by two scripts in /sync/repos/nemulate/scripts/ and is not checked in.

What the site does

  1. On load, fetches meta.json, keep_mask.bin, cos_lat_per_cell.bin, obs_yearly_stats.bin, and e6.onnx in parallel with progress updates. The model is ~12 MB; the obs stats are ~28 MB for 30 years of data.
  2. Loads ONNX Runtime Web from jsDelivr and creates an InferenceSession preferring the WebGPU execution provider with a WASM fallback.
  3. Populates the anchor-year dropdown from meta.available_anchors (years for which the previous 10 calendar years of observations are available).
  4. On Run, slices 10 years of yearly stats from the obs cache, builds a (1, N_keep, 201) Float32 tensor — exactly matching the layout the training-time YearlyStatsDataset produces — runs the ONNX session, scales the selected-horizon column by CESM2_SSH_GLOBALSTD × 10 = 37.993 to recover mm/yr, area-weighted demeans, computes vmax = max(1.5 × area-weighted std, 0.05), and paints to canvas with an RdBu_r colormap and a colorbar.

Model + asset facts

  • Model: TrendForecaster from nemulate.models.trend_head, the legacy h=192 inline-spine architecture. Hidden 192, 6 layers, 6 heads, K=8 neighbours, pos_dim=16, RMSNorm. ~3.09M params. Source ckpt at /sync/repos/nemulate/artifacts/trend_e6_legacy_phase2/trend.ckpt.
  • Input (ONNX): (1, 43335, 201) float32. 201 = 4 stats × 5 vars × 10 past years + 1 year-fraction scalar.
  • Output (ONNX): (1, 43335, 5) float32. Horizons (years): [5, 10, 15, 20, 30].
  • Post-scale: pred_mm_per_yr = onnx_output × 37.993. The browser matches scripts/56_plot_e6_focused.py (which directly plots the pred_forced_anchor_*.npy files in mm/yr).
  • Grid: 1° equirectangular, lat-major C-order on a 180×360 image. keep_mask[i] == 1 iff cell i is ocean. N_keep = 43335.
  • Variable order: [SHF, SSH, SST, TAUX, TAUY]. Channel order in the feature tensor: (stat, var, year) outer-to-inner, then the year-scalar in channel 200.
  • Obs preprocessing: AVISO + HadISST + ERA5 from /media/data/cache/obs_cache_xesmf.npz. SSH per-timestep area-weighted demean, then per-variable divide by CESM2_GLOBALSTD (matches scripts/38_forced_inference_obs.py).

Regenerating the assets

Both scripts live in /sync/repos/nemulate/scripts/ and write into artifacts/website/build/ (this directory).

# 1. Export the model -> build/e6.onnx (~12 MB)
PYTHONPATH=/sync/repos/nemulate \
  /media/misc/envs/nemul/bin/python /sync/repos/nemulate/scripts/65_export_e6_onnx.py \
    --ckpt /sync/repos/nemulate/artifacts/trend_e6_legacy_phase2/trend.ckpt \
    --out  /media/sync/syncthing/papers/2_wip/PhD-Thesis/artifacts/website/build/e6.onnx

# 2. Build the obs feature cache + grid metadata
PYTHONPATH=/sync/repos/nemulate \
  /media/misc/envs/nemul/bin/python /sync/repos/nemulate/scripts/66_export_web_assets.py \
    --out /media/sync/syncthing/papers/2_wip/PhD-Thesis/artifacts/website/build

Expected outputs and approximate sizes:

file shape dtype size
e6.onnx mixed ~12 MB
meta.json json <8 KB
obs_yearly_stats.bin (n_years, 4, 5, 43335) fp16 ~28 MB at 32 yr
keep_mask.bin (64800,) uint8 64 KB
cos_lat_per_cell.bin (43335,) fp32 170 KB

Things to double-check before re-running script 65

The exporter assumes the legacy training-time defaults from scripts/20_train_trend.py for the E6 run. These are:

  • n_years_input = 10confirmed from node_embed.0.weight shape (384, 217) → input dim 217 = 201 + pos_dim 16, and 201 - 1 = 200 = 4 × 5 × 10.
  • hidden_dim = 192, mlp_hidden = 768, num_layers = 6, heads = 6, pos_dim = 16, norm_type = "rms"all confirmed from the saved tensor shapes.
  • knn = 8, knn_hops = 1assumed, not confirmed. The trend_e6_legacy_phase2 train log does not echo these. The legacy runner template (run_e8e9_realhop_pipeline.sh) uses knn=8 and varies hops only in the "realhop" sweep, so I went with the click-default of 1. If E6 used knn_hops=2 or 3 the K-NN edge set is different and the model will produce noticeably different predictions. If you can find the original trend_e6 runner script (it appears to have been deleted along with the _BACKUP_ dir), cross-check and pass KNN_HOPS = … accordingly at the top of 65_export_e6_onnx.py.

Uploading to S3 (or any static CDN)

  1. Upload every file in artifacts/website/build/ to a bucket. Keep the file names. Enable CORS for the bucket:
    {
      "AllowedOrigins": ["*"],
      "AllowedMethods": ["GET"],
      "AllowedHeaders": ["*"],
      "MaxAgeSeconds": 3600
    }
    The .onnx and .bin files are content-addressable; long Cache-Control: public, max-age=31536000, immutable is fine.
  2. Open config.js and replace the five ./build/... paths with the resulting absolute URLs:
    export const CONFIG = {
      MODEL_URL:    "https://your-bucket.example/path/e6.onnx",
      META_URL:     "https://your-bucket.example/path/meta.json",
      STATS_URL:    "https://your-bucket.example/path/obs_yearly_stats.bin",
      KEEP_MASK_URL:"https://your-bucket.example/path/keep_mask.bin",
      COS_LAT_URL:  "https://your-bucket.example/path/cos_lat_per_cell.bin",
      ORT_URL:      "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.20.1/dist/ort.webgpu.min.mjs",
      ORT_WASM_BASE:"https://cdn.jsdelivr.net/npm/onnxruntime-web@1.20.1/dist/",
    };
  3. Upload index.html, style.css, config.js, app.js to any static host (S3 + CloudFront, Netlify, GitHub Pages, …). No CORS config needed on the static host beyond the usual same-origin rules for the HTML.

Local testing

python3 -m http.server 8000 --directory artifacts/website

Then open http://localhost:8000/ in Chrome 113+ / Edge 113+ / Safari 18+ / Firefox 121+. The status bar should walk through "loading metadata… loading grid masks… loading obs feature cache… loading ONNX runtime… loading model… ready. EP: …". Pick an anchor and a horizon and click Run — first inference takes ~1–3 s on a mid-range GPU (WebGPU kernel compilation + initial DMA); subsequent runs should be ~50–300 ms.

Manual verification checklist

  • Status bar reaches "ready" with no red error.
  • Anchor dropdown is populated with years (default range 2003..~2024 depending on the obs cache).
  • Horizon dropdown shows 5/10/15/20 yr options with the in-range ones labelled anchor–anchor+H and out-of-range ones labelled (ends YYYY, past obs).
  • Clicking Run produces a map within a few seconds.
  • The map shows a recognisable Pacific/Atlantic SSH-trend pattern: e.g. anchor 2015, horizon 10 yr (matches focused_2015-01_h10yr_e6_pred.png in artifacts/trend_e6_legacy_phase2/figures/) should show the tropical Pacific signature with std ~1.5 mm/yr → vmax ~2.3 mm/yr.
  • Numbers in the "vmax: …" line look comparable to script 56's log output (it prints the same per-anchor s_pred and v_pred).

Caveats

  • WebGPU support: Chrome 113+, Edge 113+, Safari 18+ (macOS 15+), Firefox 121+ behind a flag. On hardware/OS combos without WebGPU, the session falls back to WASM, which is ~10× slower for this model but still usable (~3–5 s per run on a modern laptop).
  • Memory: the obs stats array is held as fp16 in JS memory (~28 MB) and the per-anchor fp32 feature buffer (~33 MB) is allocated fresh each Run. Total JS heap during inference: ~80 MB. Should be fine on any non-mobile device.
  • Time-axis edges: at the most recent anchors the model is being asked to predict beyond the AVISO record. That is fine — the model takes only past data as input — but the predicted pattern cannot be validated against observations for those windows.
  • GMSL is removed: the model is trained to predict the forced-response SSH pattern under the CESM2 Boussinesq convention (globally zero per timestep). The browser subtracts the area-weighted mean from the output for the same reason. To compare against AVISO observed trends you must demean AVISO over the same window, exactly as script 56 does.
  • The exported model is fixed-shape: batch=1, N=43335, channels=201. Re-export with dynamic_axes if you want to batch.

About

No description, website, or topics provided.

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Contributors