Learning the atmosphere before predicting it.
Weather-JEPA is an open research project investigating whether reconstruction-free joint-embedding predictive learning can produce transferable, dynamically useful, uncertainty-aware representations of the atmosphere under accessible computational constraints.
The project’s long-term purpose is to contribute knowledge that may improve the anticipation and characterization of hazardous weather and ultimately support forecasting systems that protect human life.
Important
Weather-JEPA is a research platform. It is not an operational forecasting or warning system and must not be used for emergency decisions, evacuation guidance, navigation, aviation, marine operations, or other life-safety decisions.
| Item | Current state |
|---|---|
| Project phase | Research specification and repository foundation |
| Scientific status | Pre-implementation |
| Governing plan | docs/project-management/north-star-plan.md |
| Primary framework | PyTorch |
| Primary accelerator | Apple Metal Performance Shaders (MPS) |
| Reference hardware | Apple M2 Mac Studio, 32 GB unified memory |
| Primary data | ERA5 and WeatherBench 2-compatible datasets |
| Initial temporal interval | 6 hours |
| Initial model scale | Micro and Small tiers |
| License target | Apache License 2.0 |
| Operational status | Research only |
The project will not advance to a new phase simply because the next phase is more interesting. Each phase has an explicit exit gate.
- Mission
- Why Weather-JEPA
- Research Questions
- Scientific Comparison
- Mathematical Core
- Proposed Architecture
- Data Plan
- Development Sequence
- Evaluation Framework
- Engineering Standards
- Repository Structure
- Getting Started
- Experiment Governance
- Documentation and GitHub Pages
- Contributing
- Safety and Responsible Use
- License
- Citation
Weather-JEPA will investigate whether a compute-efficient, multiscale, JEPA-style atmospheric world model can learn representations that are:
- transferable across atmospheric tasks;
- dynamically stable over multiple forecast steps;
- robust to missing or corrupted inputs;
- useful for deterministic and probabilistic prediction;
- sensitive to hazardous atmospheric structure;
- reproducible on prosumer hardware.
The scientific contribution is upstream. Weather-JEPA does not replace meteorologists, numerical weather prediction, observing systems, emergency managers, or warning agencies.
Modern machine-learning weather systems demonstrate that neural models can learn powerful atmospheric transition operators. That success does not answer a different question:
Does a model learn a broadly useful representation of atmospheric dynamics, or does it become an excellent emulator of the particular fields and objective used during training?
Weather-JEPA studies that question directly.
A conventional supervised model predicts future physical fields. A masked autoencoder reconstructs missing physical values. A JEPA predicts the latent representation of hidden or future atmospheric states without requiring exact raw-value reconstruction.
This distinction matters because future atmospheric detail is not equally predictable at every scale. A representation objective may preserve synoptic organization, cross-variable structure, vertical coupling, and regime information without spending all model capacity reproducing uncertain local texture. It may also discard precisely the localized structure that matters most. The project will test both possibilities rather than selecting one by architectural faith.
Under matched architecture, data, parameter count, and optimization budget, does JEPA pretraining produce atmospheric representations that transfer more effectively than supervised forecasting or masked reconstruction?
Does the JEPA representation improve deterministic forecast skill across variables, levels, regions, spatial scales, and lead times?
Do JEPA-trained latent states support more stable multistep rollouts as measured by error growth, variance behavior, spectral drift, and decoded physical plausibility?
Are JEPA representations more robust to missing variables, masked regions, noisy inputs, temporal gaps, seasonal shifts, geographic transfer, and held-out extreme events?
Do JEPA representations improve the detection, classification, tracking, intensity estimation, or uncertainty characterization of selected hazardous phenomena?
What representation and forecast performance can be achieved per parameter, training hour, gigabyte of peak memory, and unit of processed data?
After deterministic stability is demonstrated, can a probabilistic latent predictor generate calibrated, spatially coherent alternative futures without averaging away high-impact possibilities?
The primary experiment compares three learning objectives using the same backbone wherever technically possible.
| System | Primary objective | Predicted target |
|---|---|---|
| Supervised forecast baseline | Direct future-state prediction | Physical fields |
| Masked autoencoder baseline | Reconstruction | Masked physical values |
| Weather-JEPA | Joint-embedding prediction | Target latent embeddings |
The comparison must control:
- architecture;
- input variables;
- pressure levels;
- temporal context;
- spatial resolution;
- parameter count within a declared tolerance;
- optimizer and learning-rate search budget;
- training examples;
- update count or compute budget;
- evaluation data;
- downstream probe protocol;
- random-seed policy.
Changing several of these simultaneously may produce a stronger model, but it will not isolate the effect of the learning objective.
- JEPA will improve linear-probe and few-shot transfer performance.
- JEPA may preserve large-scale dynamical structure better during rollout.
- Masked reconstruction may remain stronger for exact fine-scale recovery.
- JEPA gains, if present, may be largest for tasks driven by cross-variable and synoptic organization.
- A deterministic JEPA will not solve uncertainty by itself.
- Compute-constrained experiments can still identify objective-level effects.
Hypotheses are allowed to fail. Negative results remain project deliverables.
At time (t), the pressure-level atmospheric state is represented as
[ \mathbf{X}_t \in \mathbb{R}^{C \times L \times H \times W}, ]
where:
- (C): dynamic variables;
- (L): pressure levels;
- (H), (W): latitude and longitude dimensions.
Surface and static fields are represented separately:
[ \mathbf{S}_t \in \mathbb{R}^{C_s \times H \times W}, \qquad \mathbf{B} \in \mathbb{R}^{C_b \times H \times W}. ]
A sample contains a context window
[ \mathcal{C}t = {\mathbf{X}{t-(K-1)\Delta t},\ldots,\mathbf{X}_t} ]
and one or more forecast horizons
[ \mathcal{H} = {h_1,h_2,\ldots,h_m}. ]
The reference temporal interval is initially six hours.
Training-split statistics are applied by variable and pressure level:
\frac{X_{t,c,l,i,j}-\mu_{c,l,s}} {\sigma_{c,l,s}+\epsilon}. ]
Validation and test data must never influence normalization statistics.
Global field losses and metrics account for grid-cell area using normalized cosine-latitude weights:
[ w_i = \frac{\cos(\phi_i)} {\frac{1}{H}\sum_{k=1}^{H}\cos(\phi_k)}. ]
The trainable context encoder produces visible-context tokens:
[ \mathbf{Z}^{(c)}t = f{\theta} (M_c \odot \mathcal{C}_t,\mathbf{S}_t,\mathbf{B},\mathbf{e}_t). ]
The target encoder produces target representations:
f_{\xi} (\mathbf{X}{t+h},\mathbf{S}{t+h},\mathbf{B},\mathbf{e}_{t+h}). ]
Target parameters are updated by exponential moving average:
[ \xi \leftarrow \tau\xi+(1-\tau)\theta. ]
The target encoder does not receive predictor-loss gradients.
The predictor estimates hidden or future target embeddings from the context representation, target geometry, and forecast horizon:
g_{\psi}(\mathbf{Z}^{(c)}_t,M_y,\mathbf{q}_h). ]
Target families include:
- spatial blocks;
- vertical slices;
- cross-variable targets;
- future-time targets;
- multiscale targets;
- combined targets.
For target token set (\mathcal{T}),
\operatorname{sg}(\mathbf{z}^{(y)}_k) \right|_p. ]
Collapse is monitored through per-dimension variance, covariance, effective rank, singular values, pairwise similarity, and downstream probe performance.
If diagnostics justify it, the governed total objective may include variance, covariance, and physical-diagnostic terms:
\lambda_J\mathcal{L}{\mathrm{JEPA}} + \lambda_V\mathcal{L}{\mathrm{var}} + \lambda_C\mathcal{L}{\mathrm{cov}} + \lambda_P\mathcal{L}{\mathrm{physical}}. ]
Regularization is added because evidence requires it, not because the training curve has offended aesthetic expectations.
Decoded and latent autoregression are evaluated separately. A latent transition takes the form
F_{\psi} (\widehat{\mathbf{Z}}{t+n\Delta t},\mathbf{q}{\Delta t}). ]
Rollout diagnostics include drift, variance collapse or explosion, spectral distortion, front and gradient loss, unphysical extrema, phase error, and decoder accumulation error.
Probabilistic modeling begins only after deterministic stability:
[ p_{\psi} (\mathbf{Z}{t+h}\mid\mathbf{Z}{\le t},h). ]
Candidate methods include Gaussian latent distributions, mixture density prediction, conditional diffusion, flow matching, ensemble perturbation, and hazard-specific quantile or expectile heads.
Dropout at inference is not, by itself, a probabilistic research program.
ERA5 / WeatherBench-compatible state
|
v
validation -> normalization -> temporal windowing
|
v
variable + level + static + calendar + horizon embeddings
|
v
spatial patch embedding
|
v
factorized vertical-variable mixer
|
v
multiscale context encoder --------------------------.
| |
v v
JEPA predictor + target queries EMA target encoder
| |
'--------------- latent loss -----------------'
|
+--> lightweight decoder --> physical fields
|
+--> probes / classifiers / hazard heads
- dynamic pressure-level fields;
- dynamic surface fields;
- orography and land-sea mask;
- latitude and cyclic longitude;
- time-of-year and time-of-day features;
- variable identity;
- pressure-level identity;
- forecast-horizon conditioning.
The reference encoder is a hierarchical transformer or hybrid convolution-transformer with:
- spatial patch embedding;
- local or windowed attention at fine scales;
- downsampling for synoptic and planetary context;
- factorized mixing across variables and levels;
- residual connections;
- pre-normalization;
- optional stochastic depth only when justified by evidence.
A full global-attention transformer is not the default because its memory bill would arrive before its scientific contribution.
The predictor is intentionally narrower than the encoder and uses:
- projected context tokens;
- learnable target-mask tokens;
- spatial, vertical, variable, scale, and horizon embeddings;
- transformer or cross-attention blocks;
- target-latent projection.
The decoder supports physical evaluation but is not the center of the JEPA objective. Candidate designs include:
- patch-expansion transformer;
- U-Net-like convolutional decoder;
- multiscale feature-pyramid decoder.
- linear probe;
- shallow nonlinear probe;
- atmospheric-regime classifier;
- cyclone-event classifier;
- anomaly score;
- missing-variable reconstruction;
- later probabilistic hazard heads.
| Tier | Approximate parameters | Purpose |
|---|---|---|
| Micro | 1–5 million | Pipeline, tensor, loss, and collapse verification |
| Small | 10–30 million | Controlled objective comparison and ablation |
| Medium | 40–100 million | Final workstation-scale experiments |
| Large | More than 100 million | Requires profiling and explicit approval |
Scaling occurs only after the smaller tier passes its scientific gate.
- ERA5: primary long-duration global reanalysis source.
- WeatherBench 2: benchmark-compatible datasets, evaluation patterns, baselines, and metrics.
- Independent observations: later event validation using best-track, station, radiosonde, satellite, radar-derived, or official event records as licensing and alignment permit.
ERA5 is a model-assimilation product rather than a direct observation of truth. That limitation must remain visible in model claims.
- geopotential (
z); - temperature (
t); - zonal wind (
u); - meridional wind (
v); - specific humidity (
q).
- 1000 hPa;
- 925 hPa;
- 850 hPa;
- 700 hPa;
- 500 hPa;
- 300 hPa;
- 250 hPa;
- 200 hPa;
- 100 hPa.
A reduced seven-level configuration may be used during early pipeline experiments. Expansion requires an explicit decision record.
- mean sea-level pressure;
- 2 m temperature;
- 10 m zonal wind;
- 10 m meridional wind;
- total-column water vapor;
- accumulated precipitation with explicit accumulation semantics;
- surface pressure.
- orography;
- land-sea mask;
- latitude;
- cyclic longitude;
- optional land-state climatologies in later phases.
- synthetic tensors;
- tiny geographic subset;
- regional coarse grid;
- regional moderate grid;
- global coarse grid;
- higher-resolution expansion after profiling.
The initial reference grid is expected to be 1.5° or 1.0°. The development-region boundary must be selected through an Architecture Decision Record.
| Split | Years |
|---|---|
| Training | 1979–2017 |
| Validation | 2018–2019 |
| Test | 2020–2022 |
| Untouched challenge | 2023 onward, subject to data maturity |
Mandatory controls:
- no random frame-level split;
- no context-window overlap across split boundaries;
- no validation or test influence on normalization;
- complete event lifecycles remain in one split;
- final challenge data are evaluated only at declared gates.
CDS / WeatherBench source
-> immutable raw archive
-> checksum and metadata validation
-> coordinate normalization
-> variable and unit validation
-> regridding
-> temporal alignment
-> derived-variable generation
-> training-statistic computation
-> normalized Zarr stores
-> sample index manifests
-> DataLoader / DataPipe
Each variable must declare:
- canonical and source names;
- units;
- coordinates and vertical coordinate;
- valid range;
- missing-value meaning;
- instantaneous or accumulated semantics;
- transformations;
- normalization statistics;
- provenance;
- license;
- known issues.
| Stage | Purpose | Exit condition |
|---|---|---|
| 0 | Repository and environment foundation | Clean install, CI, tests, MPS smoke test, no secrets |
| 1 | Research specification | Every question maps to data, model, metric, and decision criterion |
| 2 | Data prototype | Traceable batch, invertible normalization, no leakage, data card |
| 3 | Synthetic verification | Losses, EMA, masks, horizons, gradients, and collapse checks verified |
| 4 | Persistence and climatology baselines | Non-neural baseline report complete |
| 5 | Shared backbone | Same encoder trains under all three primary objectives |
| 6 | Supervised forecasting | Stable direct-forecast recipe beats simple baselines |
| 7 | Masked autoencoder | Reconstruction and shared probe protocol validated |
| 8 | Weather-JEPA v1 | Stable, reproducible latent prediction without collapse |
| 9 | Controlled comparison | Objective comparison under matched controls |
| 10 | Latent dynamics | Decoded and latent rollouts characterized |
| 11 | Transfer and robustness | Few-shot, missing-input, shift, and probe results complete |
| 12 | Hazard evaluation | Event definitions, metrics, and retrospective analyses complete |
| 13 | Probabilistic Weather-JEPA | Calibration and coherent alternative futures demonstrated |
| 14 | Scaling and optimization | Each scale increase tied to a scientific question |
| 15 | External reproduction and release | Flagship result recreated from released artifacts |
See the North-Star Project Management Plan for complete tasks, gates, risks, and deliverables.
- latitude-weighted RMSE;
- mean absolute error;
- anomaly correlation coefficient;
- bias;
- variable-, level-, region-, season-, and lead-time stratification.
- spatial spectra;
- kinetic-energy spectra;
- variance by scale;
- gradient and front sharpness;
- temporal tendency distributions;
- rollout stability;
- structural similarity where scientifically appropriate.
- frozen linear probes;
- shallow nonlinear probes;
- few-shot adaptation;
- atmospheric-regime classification;
- anomaly retrieval;
- missing-variable recovery;
- cross-region and cross-season transfer;
- latent effective rank and covariance structure.
Initial candidates include:
- tropical cyclone genesis, track, and intensity;
- atmospheric rivers;
- blocking regimes;
- extreme precipitation;
- strong-wind or frontal events where the data support defensible labels.
Hazard evaluation must be event based, not inferred from one global mean score.
- continuous ranked probability score;
- Brier score;
- reliability diagrams;
- rank histograms;
- spread-skill relationships;
- multivariate energy score where appropriate;
- threshold-event calibration.
- parameters;
- training steps and wall-clock time;
- peak memory;
- processed samples;
- loader throughput;
- inference latency;
- energy estimate where feasible;
- performance per resource unit.
- multiple seeds for confirmatory experiments;
- confidence intervals or resampling where appropriate;
- predeclared primary metrics;
- effect sizes, not only rank order;
- all primary outcomes reported;
- exploratory and confirmatory results clearly separated.
Weather-JEPA adapts the NIST AI Risk Management Framework functions:
- Govern: roles, claims, risks, decisions, and release controls;
- Map: intended use, stakeholders, limitations, and misuse;
- Measure: validity, reliability, robustness, security, transparency, and uncertainty;
- Manage: prioritize risks, implement controls, monitor releases, and retire unsafe artifacts.
Every release requires a model card describing intended use, prohibited use, data, architecture, objective, evaluation, uncertainty, limitations, compute, license, and version.
- immutable raw layer;
- validated intermediate layer;
- curated training layer;
- explicit schemas and contracts;
- idempotent, resumable processing;
- source checksums;
- semantic dataset versions;
- machine-readable quality reports;
- training-only statistics;
- automatic split-leakage tests;
- no hidden manual transformations.
- configuration-driven runs;
- resolved configuration saved with every experiment;
- typed public APIs and documented tensor shapes;
- automated formatting, linting, typing, and testing;
- unit, integration, regression, and scientific tests;
- protected
mainbranch; - short-lived feature branches;
- semantic release tags;
- exact commit hashes in experiments;
- no untrusted checkpoint deserialization;
- dependency, secret, and vulnerability scanning;
- performance optimization only after measurement.
Every reported result must identify:
- experiment ID;
- hypothesis;
- code commit and dirty-tree state;
- resolved configuration;
- data version and manifest;
- random seeds;
- environment;
- hardware;
- runtime;
- checkpoints;
- evaluation procedure;
- figures and tables;
- conclusion and limitations.
weather-jepa/
├── README.md
├── LICENSE
├── CITATION.cff
├── SECURITY.md
├── CONTRIBUTING.md
├── CHANGELOG.md
├── pyproject.toml
├── configs/
│ ├── data/
│ ├── model/
│ ├── objective/
│ ├── train/
│ └── experiment/
├── docs/
│ ├── index.md
│ ├── _config.yml
│ ├── architecture/
│ ├── adr/
│ ├── rdr/
│ ├── standards/
│ ├── tutorials/
│ └── project-management/
├── src/weather_jepa/
│ ├── data/
│ ├── models/
│ ├── objectives/
│ ├── training/
│ ├── evaluation/
│ ├── hazards/
│ ├── visualization/
│ └── utils/
├── scripts/
├── tests/
│ ├── unit/
│ ├── integration/
│ ├── regression/
│ └── fixtures/
├── experiments/
├── reports/
├── figures/
├── notebooks/
└── papers/
Notebooks may support exploration and communication. Production training logic may not exist only in notebooks.
Note
Weather-JEPA is currently in the repository-foundation phase. The commands below define the intended interface and should be updated as Stage 0 is completed. Do not present planned commands as verified commands until CI confirms them.
- macOS or Linux;
- Python version declared in
pyproject.toml; - PyTorch with CPU or MPS support;
- sufficient storage for selected ERA5 subsets;
- credentials for approved data sources where required.
git clone https://github.com/<kdsmithtexnc>/weather-jepa.git
cd weather-jepa
python -m venv venv
source venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -e ".[dev]"python - <<'PY'
import torch
print("PyTorch:", torch.__version__)
print("MPS built:", torch.backends.mps.is_built())
print("MPS available:", torch.backends.mps.is_available())
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
x = torch.randn(8, 8, device=device)
print("Device:", device)
print("Result:", (x @ x.T).mean().item())
PY# Validate a configuration
weather-jepa config validate configs/experiment/micro_jepa.yaml
# Build a small data release
weather-jepa data build --config configs/data/era5_regional_v1.yaml
# Run data-quality checks
weather-jepa data validate --dataset data/processed/era5-regional-v1
# Train
weather-jepa train --config configs/experiment/micro_jepa.yaml
# Evaluate
weather-jepa evaluate --experiment EXP-YYYYMMDD-NNNN
# Generate an experiment card
weather-jepa report experiment EXP-YYYYMMDD-NNNNUntil implemented and tested, this command surface is a design contract rather than a claim of working software.
Every experiment receives an immutable identifier and a versioned experiment card.
Minimum fields:
experiment_id: EXP-YYYYMMDD-NNNN
title:
status: planned
research_question:
hypothesis:
exploratory_or_confirmatory:
code_commit:
dirty_tree:
data_version:
configuration_hash:
hardware:
seeds: []
primary_metrics: []
secondary_metrics: []
controls: []
artifacts: []
result:
limitations:
decision:
next_action:- Register the primary question before confirmatory training.
- Preserve failed runs and failure causes.
- Do not choose favorable seeds after the fact.
- Do not tune against the final challenge set.
- Keep objective comparisons within equivalent resource budgets.
- Report all primary outcomes.
- Record deviations from protocol.
Weather-JEPA uses two coordinated landing documents:
README.md— the repository landing page for developers, reviewers, and contributors.docs/index.md— the public GitHub Pages homepage for broader scientific and educational communication.
The governing project plan is:
The two landing pages intentionally overlap on mission and status, but they serve different audiences. Material scientific requirements belong in the north-star plan rather than being independently redefined in either landing page.
Recommended GitHub Pages configuration:
- source: Deploy from a branch;
- branch:
main; - folder:
/docs.
The Pages entry document is docs/index.md, not the repository README.
See docs/adr/0001-documentation-structure.md for the documentation decision.
Contributions will be accepted only when they preserve the project’s scientific controls and reproducibility requirements.
Before opening a pull request:
- read the north-star plan;
- open or reference an issue;
- create an ADR for architecture changes;
- create an RDR for research-protocol changes;
- add or update tests;
- update documentation and experiment schemas;
- disclose data, compute, or evaluation changes;
- run the full local quality suite.
A contribution that improves a metric by changing the data split without documenting it has not improved the project.
Detailed contribution instructions will live in CONTRIBUTING.md.
- atmospheric representation-learning research;
- controlled forecasting experiments;
- reproducible benchmark studies;
- educational visualization;
- robustness, uncertainty, and failure analysis.
- official weather warnings;
- emergency response decisions;
- aviation, marine, transportation, or navigation decisions;
- autonomous evacuation recommendations;
- commercial operational forecasting without independent validation;
- claims that the model “understands weather” or “saves lives” without evidence.
Public statements must distinguish:
- training results;
- validation results;
- held-out test results;
- retrospective event analysis;
- prospective evaluation;
- operational readiness.
The goal is to protect life. That goal increases the burden of evidence; it does not reduce it.
The project is intended for release under the Apache License 2.0.
Data remain governed by their original providers and licenses. The project license does not relicense ERA5, WeatherBench 2, observations, third-party checkpoints, or external artifacts.
A CITATION.cff file will be added before the first research release.
Until then, cite the project by repository release, version tag, and archived DOI when available. Do not cite an unversioned main branch as though it were a stable scientific artifact.
Weather-JEPA begins with a deliberately constrained computer, a deliberately narrow causal question, and a deliberately demanding evidence standard.
It will not win by being the largest model. It will be useful if it produces a clean answer to a difficult question, exposes where that answer fails, and leaves enough evidence for someone else to reproduce the result.