idx-flow 0.1.0

Index-based spherical convolutions for HEALPix grids in PyTorch

idx-flow

Index-based Spherical Convolutions for HEALPix Grids in PyTorch

A PyTorch library for efficient neural network operations on spherical data using HEALPix tessellation. This library implements index-based spherical convolutions that achieve O(N) computational complexity while preserving the equal-area properties essential for atmospheric and geophysical data analysis.

Citation

This architecture is based on the following paper:

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 DOI: 10.1109/IJCNN64981.2025.11227156

If you use this library in your research, please cite the paper above.

Features

• Efficient O(N) Complexity: Precomputed neighbor indices enable linear-time convolutions
• HEALPix Native: Built for the Hierarchical Equal Area isoLatitude Pixelization scheme
• PyTorch Integration: Seamless integration with PyTorch models and training pipelines
• Flexible Architecture: Support for encoder-decoder networks, GANs, and custom architectures
• Multiple Layer Types: Convolution, transpose convolution, upsampling, MLP, and pooling layers

Installation

From PyPI (when published)

pip install idx-flow

From Source

git clone https://github.com/otavio-feitosa/idx-flow.git
cd idx-flow
pip install -e .

Dependencies

• Python >= 3.8
• PyTorch >= 1.9.0
• NumPy >= 1.19.0
• healpy >= 1.15.0
• scikit-learn >= 0.24.0

Quick Start

Basic Usage

import torch
from idx_flow import SpatialConv, SpatialTransposeConv, compute_connection_indices

# Compute connection indices for downsampling (nside 64 -> 32)
indices_down, distances_down = compute_connection_indices(
    nside_in=64, nside_out=32, k=4
)

# Compute connection indices for upsampling (nside 32 -> 64)
indices_up, distances_up, weights_up = compute_connection_indices(
    nside_in=32, nside_out=64, k=4, return_weights=True
)

# Create layers
conv = SpatialConv(
    output_points=12 * 32**2,  # 12288 pixels
    connection_indices=indices_down,
    filters=64
)

transpose_conv = SpatialTransposeConv(
    output_points=12 * 64**2,  # 49152 pixels
    connection_indices=indices_up,
    kernel_weights=weights_up,
    filters=32
)

# Forward pass
x = torch.randn(8, 12 * 64**2, 32)  # [batch, points, channels]
encoded = conv(x)                    # [8, 12288, 64]
decoded = transpose_conv(encoded)    # [8, 49152, 32]

Building an Encoder-Decoder Network

import torch
import torch.nn as nn
from idx_flow import (
    SpatialConv,
    SpatialTransposeConv,
    SpatialBatchNorm,
    compute_connection_indices
)

class SphericalAutoencoder(nn.Module):
    """Autoencoder for spherical data on HEALPix grids."""

    def __init__(self, in_channels: int = 5, latent_dim: int = 64):
        super().__init__()

        # Precompute connection indices for each resolution level
        # Encoder: 256 -> 128 -> 64 -> 32
        idx_256_128, _ = compute_connection_indices(256, 128, k=4)
        idx_128_64, _ = compute_connection_indices(128, 64, k=4)
        idx_64_32, _ = compute_connection_indices(64, 32, k=4)

        # Decoder: 32 -> 64 -> 128 -> 256
        idx_32_64, _, w_32_64 = compute_connection_indices(32, 64, k=4, return_weights=True)
        idx_64_128, _, w_64_128 = compute_connection_indices(64, 128, k=4, return_weights=True)
        idx_128_256, _, w_128_256 = compute_connection_indices(128, 256, k=4, return_weights=True)

        # Encoder layers
        self.enc1 = SpatialConv(12*128**2, idx_256_128, filters=32)
        self.enc2 = SpatialConv(12*64**2, idx_128_64, filters=64)
        self.enc3 = SpatialConv(12*32**2, idx_64_32, filters=latent_dim)

        # Decoder layers
        self.dec1 = SpatialTransposeConv(12*64**2, idx_32_64, w_32_64, filters=64)
        self.dec2 = SpatialTransposeConv(12*128**2, idx_64_128, w_64_128, filters=32)
        self.dec3 = SpatialTransposeConv(12*256**2, idx_128_256, w_128_256, filters=in_channels)

        # Batch normalization
        self.bn1 = SpatialBatchNorm(32)
        self.bn2 = SpatialBatchNorm(64)
        self.bn3 = SpatialBatchNorm(latent_dim)

        self.activation = nn.SELU()

    def encode(self, x):
        x = self.activation(self.bn1(self.enc1(x)))
        x = self.activation(self.bn2(self.enc2(x)))
        x = self.activation(self.bn3(self.enc3(x)))
        return x

    def decode(self, z):
        x = self.activation(self.dec1(z))
        x = self.activation(self.dec2(x))
        x = self.dec3(x)
        return x

    def forward(self, x):
        z = self.encode(x)
        return self.decode(z)

# Example usage
model = SphericalAutoencoder(in_channels=5, latent_dim=64)
x = torch.randn(4, 12*256**2, 5)  # Batch of 4, nside=256, 5 channels
reconstruction = model(x)
print(f"Input shape: {x.shape}")
print(f"Output shape: {reconstruction.shape}")

API Reference

Layers

SpatialConv

Spatial convolution for downsampling on HEALPix grids.

SpatialConv(
    output_points: int,           # Number of output spatial points
    connection_indices: ndarray,  # [output_points, kernel_size] neighbor indices
    kernel_weights: ndarray = None,  # Optional distance-based weights
    filters: int = 32,            # Number of output channels
    bias: bool = True             # Include bias term
)

Input shape: [B, N_in, C_in] Output shape: [B, output_points, filters]

SpatialTransposeConv

Transpose convolution for upsampling on HEALPix grids.

SpatialTransposeConv(
    output_points: int,           # Number of output spatial points (higher res)
    connection_indices: ndarray,  # [output_points, kernel_size] neighbor indices
    kernel_weights: ndarray = None,  # Distance-based weights (recommended)
    filters: int = 32,            # Number of output channels
    bias: bool = True             # Include bias term
)

Input shape: [B, N_in, C_in] (lower resolution) Output shape: [B, output_points, filters] (higher resolution)

SpatialUpsampling

Non-learnable upsampling using distance-based interpolation.

SpatialUpsampling(
    output_points: int,           # Number of output spatial points
    connection_indices: ndarray,  # [output_points, kernel_size] neighbor indices
    distances: ndarray,           # [output_points, kernel_size] distances
    interpolation: str = "linear",  # "linear", "idw", or "gaussian"
    kernel_radius: float = None   # Radius for interpolation kernel
)

Input shape: [B, N_in, C_in] Output shape: [B, output_points, C_in] (channels preserved)

SpatialMLP

Multi-layer perceptron for complex spatial transformations.

SpatialMLP(
    output_points: int,                    # Number of output spatial points
    connection_indices: ndarray,           # [output_points, kernel_size]
    hidden_units: tuple = (32, 32, 32),   # Hidden layer dimensions
    activations: tuple = ("linear", "linear", "linear")  # Activation per layer
)

Input shape: [B, N_in, C_in] Output shape: [B, output_points, hidden_units[-1]]

SpatialPooling

Pooling operations over local neighborhoods.

SpatialPooling(
    output_points: int,           # Number of output spatial points
    connection_indices: ndarray,  # [output_points, kernel_size]
    pool_type: str = "mean"       # "mean", "max", or "sum"
)

Input shape: [B, N_in, C_in] Output shape: [B, output_points, C_in]

SpatialBatchNorm

Batch normalization for spatial data.

SpatialBatchNorm(
    num_features: int,     # Number of channels
    eps: float = 1e-5,     # Numerical stability constant
    momentum: float = 0.1, # Running statistics momentum
    affine: bool = True    # Learnable affine parameters
)

Input/Output shape: [B, N, C]

Utility Functions

compute_connection_indices

Convenience function for computing connection indices and weights.

indices, distances = compute_connection_indices(
    nside_in=64,          # Input resolution
    nside_out=32,         # Output resolution
    k=4                   # Number of neighbors
)

# With weights
indices, distances, weights = compute_connection_indices(
    nside_in=32, nside_out=64, k=4,
    return_weights=True,
    weight_method="inverse_square"
)

hp_distance

Compute geodesic distances and neighbor indices between HEALPix grids.

indices, distances_km = hp_distance(
    nside_in=64,   # Input resolution
    nside_out=32,  # Output resolution
    k=4            # Number of neighbors
)

get_weights

Calculate interpolation weights from distances.

weights = get_weights(
    distances,                    # [N, k] distance array
    method="inverse_square",      # Weighting method
    sigma_factor=0.5,            # For gaussian/exponential
    epsilon=1e-10                # Numerical stability
)

Methods: "inverse_square", "gaussian", "exponential", "tricube"

get_healpix_resolution_info

Get information about a HEALPix resolution.

info = get_healpix_resolution_info(nside=256)
# Returns: {
#     'nside': 256,
#     'npix': 786432,
#     'resolution_deg': 0.229,
#     'resolution_km': 25.65,
#     'area_sr': 1.598e-05,
#     'area_km2': 649.17
# }

Mathematical Background

Index-Based Spherical Convolution

The spatial convolution operation is defined as:

Y[b,i,f] = sum_k sum_c N[b,i,k,c] * W[k,c,f] + b[f]

Where:

• Y[b,i,f]: Output for batch b, spatial point i, feature channel f
• N[b,i,k,c]: Neighbor features gathered using precomputed indices
• W[k,c,f]: Learnable kernel weights
• b[f]: Bias term

Computational Complexity

Operation	Traditional	idx-flow
Grid construction	O(N^2)	O(N log N)
Neighbor lookup	O(N)	O(1)
Convolution	O(N^2)	O(N)

HEALPix Grid Properties

• Equal-area pixels: Area(pixel_i) = 4*pi / N_pix
• Number of pixels: N_pix = 12 * nside^2
• Resolution: Approximately sqrt(4*pi / N_pix) radians

Development

Setup Development Environment

git clone https://github.com/otavio-feitosa/idx-flow.git
cd idx-flow
pip install -e ".[dev]"

Running Tests

pytest tests/ -v

Code Formatting

black src/ tests/
isort src/ tests/

License

This project is licensed under the MIT License - see the LICENSE file for details.

Acknowledgments

• Monan Project, CEMPA Project, LAMCAD, and PGMet
• CNPq grants (processes 422614/2021-1, and 315349/2023-9)
• National Institute for Space Research (INPE)

Contributing

Contributions are welcome! Please feel free to submit a Pull Request. For major changes, please open an issue first to discuss what you would like to change.

1. Fork the repository
2. Create your feature branch (git checkout -b feature/AmazingFeature)
3. Commit your changes (git commit -m 'Add some AmazingFeature')
4. Push to the branch (git push origin feature/AmazingFeature)
5. Open a Pull Request