S2FFT: 加速和可微分的球谐变换库

S2FFT简介

S2FFT是一个用于计算球面和旋转群上傅里叶变换的Python库,由Matthew A. Price和Jason D. McEwen于2023年开发。该库利用JAX和PyTorch实现,支持自动微分,可以进行可微分变换,并且可以在GPU和TPU等硬件加速器上部署。

S2FFT的主要特点包括:

• 支持自旋球谐变换和Wigner变换(适用于实信号和复信号)
• 提供所需的伴随变换
• 具有不同的优化选项(预计算或实时计算),可根据可用资源和所需的角分辨率进行选择
• 支持多种球面采样方案
• 提供JAX和PyTorch实现
• 为现有的C/C++库(如HEALPix和SSHT)提供JAX支持

算法设计

S2FFT采用了新的算法结构,可以高度并行化和分布式计算,非常适合在GPU和TPU等硬件加速器上运行。这些算法基于新的Wigner-d递归,可以稳定地处理高角分辨率L的数据。

球谐系数的计算主要分为两个步骤:

1. 对每个纬度环进行一维傅里叶变换
2. 投影到实极坐标d函数上

S2FFT提供了两种计算方式:

• 预计算并存储所有实极坐标d函数,可以极大加速计算,但会占用大量内存
• 递归计算实极坐标d函数,每次只计算投影的一部分,内存占用可以忽略不计,但执行速度稍慢

采样方案

S2FFT支持任何等纬度采样方案。目前支持的采样方案包括:

1. 等角度采样:

   • McEwen & Wiaux (2012)
   • Driscoll & Healy (1995)
   • Gauss-Legendre (1986)

   这些采样方案具有相关的采样定理,因此可以以机器精度计算谐波变换。其中McEwen & Wiaux采样定理将球面上的奈奎斯特采样率降低了一半。

2. HEALPix采样: 由Gorski等人于2005年提出,虽然不具有采样定理,但提供了等面积像素,在实际应用中有很多优势。

安装和使用

S2FFT可以通过pip安装:

pip install s2fft

这将安装包括JAX和PyTorch支持在内的所有核心功能。

基本使用示例:

import s2fft

# 球面上的信号
# 计算谐波系数
flm = s2fft.forward_jax(f, L)  
# 映射回像素空间信号
f = s2fft.inverse_jax(flm, L)

# 旋转群上的信号
# 计算Wigner系数
flmn = s2fft.wigner.forward_jax(f, L, N)
# 映射回像素空间信号
f = s2fft.wigner.inverse_jax(flmn, L, N)

S2FFT还为HEALPix和SSHT等现有C/C++包提供JAX支持:

# SSHT球谐变换
flm = s2fft.forward(f, L, sampling=["mw"], method="jax_ssht")  

# HEALPix球谐变换
flm = s2fft.forward(f, L, nside=nside, sampling="healpix", method="jax_healpy")  

这些JAX前端支持开箱即用的反向模式自动微分,底层链接到熟悉的C/C++包。这样,S2FFT为现有包增加了梯度功能,适用于现代科学计算或机器学习应用。

贡献者

S2FFT的开发得到了许多贡献者的支持,包括Matt Price、Jason McEwen、Matt Graham等。项目欢迎更多开发者参与贡献,例如添加对更多球面采样模式的支持。

引用和许可

如果使用S2FFT,请引用以下文章:

@article{price:s2fft, 
   author      = "Matthew A. Price and Jason D. McEwen",
   title       = "Differentiable and accelerated spherical harmonic and Wigner transforms",
   journal     = "Journal of Computational Physics, submitted",
   year        = "2023",
   eprint      = "arXiv:2311.14670"        
}

S2FFT采用MIT开源许可证发布,希望能够为更广泛的社区提供帮助。

总结

S2FFT为球面和旋转群上的傅里叶变换提供了高效、可微分和硬件加速的实现。它支持多种采样方案,并与现有的C/C++库集成,为现代科学计算和机器学习应用中的球面数据分析提供了强大而灵活的工具。无论是在天文学、地球科学还是计算机视觉等领域,S2FFT都能为研究人员和开发者提供valuable支持。