idx-flow Documentation

Index-based Spherical Convolutions for HEALPix Grids in PyTorch

PyPI version Python 3.8+ License: MIT

PyTorch layers for O(N) spherical convolutions on HEALPix grids. Topology (connection indices) is precomputed once and stored as buffers; learnable weights are applied at runtime.

Citation

Important

If you use this library in your research, please cite:

Atmospheric Data Compression and Reconstruction Using Spherical GANs

Otavio Medeiros Feitosa, Haroldo F. de Campos Velho, Saulo R. Freitas, Juliana Aparecida Anochi, Angel Dominguez Chovert, Cesar M. L. de Oliveira Junior

International Joint Conference on Neural Networks (IJCNN), 2025

DOI: 10.1109/IJCNN64981.2025.11227156

@inproceedings{feitosa2025atmospheric,
   title={Atmospheric Data Compression and Reconstruction Using Spherical GANs},
   author={Feitosa, Otavio Medeiros and de Campos Velho, Haroldo F. and
           Freitas, Saulo R. and Anochi, Juliana Aparecida and
           Chovert, Angel Dominguez and de Oliveira Junior, Cesar M. L.},
   booktitle={International Joint Conference on Neural Networks (IJCNN)},
   year={2025},
   organization={IEEE},
   doi={10.1109/IJCNN64981.2025.11227156}
}

Quick Example

import torch
from idx_flow import SpatialConv, compute_connection_indices

indices, distances = compute_connection_indices(
    nside_in=64, nside_out=32, k=4
)

conv = SpatialConv(
    output_points=12 * 32**2,
    connection_indices=indices,
    filters=64,
    weight_init="kaiming_normal"
)

x = torch.randn(8, 12 * 64**2, 32)  # [batch, points, channels]
y = conv(x)
print(y.shape)  # torch.Size([8, 12288, 64])

Package Layout

All public names are re-exported from idx_flow directly:

from idx_flow import SpatialConv, SpatialViT, SpatialMLP

Internally the code is organized as:

  • conv – SpatialConv, SpatialTransposeConv, SpatialUpsampling

  • mlp – SpatialMLP, GlobalMLP

  • norm – SpatialBatchNorm, SpatialLayerNorm, SpatialInstanceNorm, SpatialGroupNorm

  • regularization – SpatialDropout, ChannelDropout

  • attention – SpatialSelfAttention

  • vit – SpatialPatchEmbedding, SpatialTransformerBlock, SpatialViT

  • pooling – SpatialPooling, Squeeze, Unsqueeze

  • functional – get_initializer, get_activation, type aliases

  • utils – hp_distance, get_weights, compute_connection_indices

Contents

User Guide

API Reference

Development

Indices and Tables

Acknowledgments

  • Monan Project, CEMPA Project, LAMCAD, PGMet

  • CNPq (processes 422614/2021-1 and 315349/2023-9)

  • National Institute for Space Research (INPE)