Software
HEAL-SWIN is an extension of the popular vision-transformer SWIN, capable of operating on high-dimensional spherical data. HEAL-SWIN uses the HEALPix grid and has only minimal computational overhead compared to the baseline SWIN model.
We provide a reference implementation for HEAL-SWIN focused on applications for spherical images, in particular those taken with fisheye lenses. We include a containerized training framework based on PyTorch Lightning and parameter logging with MLFlow. SLURM-based jobscripts for published reference runs and baselines are also available.
This package implements the NTK and NNGP recursions of several group convolutional layers enabling the analysis of the infinite-width training dynamics of equivariant neural networks. The package is written as an extension to the neural-tangents library and is implemented in JAX. Supported symmetry groups are the planar roto-translation group \(C_4 \ltimes \mathbb{R}^2\) and the group of 3D rotations \(\mathrm{SO}(3)\).
We use an equal area gridding (HEALPix) of the sphere to perform global weather forecasting with a volumetric transformer architecture, outperforming baselines on an equiangular grid.
We provide an implementation in PyTorch including training and evaluation pipelines, ERA5 data ingestion via CDSAPI, and metrics consolidation with DuckDB. SLURM-based run scripts and persisted experiment configurations for reproducing published results are also included.
