访问官网
Github
文档
论文
可微分和加速的球面变换
S2FFT是一个Python包，用于使用JAX或PyTorch计算球面和旋转群上的傅里叶变换(Price & McEwen 2023)。它利用自动微分提供可微分变换，这些变换也可以部署在硬件加速器（如GPU和TPU）上。
更具体地说，S2FFT支持自旋球谐和Wigner变换（适用于实信号和复信号），在需要的地方支持伴随变换，并提供不同的优化选项（预计算或不预计算），用户可以根据可用资源和所需的角分辨率$L$进行选择。
[!重要] HEALPix在CPU上的长时间JIT编译问题已修复！GPU版本的修复即将推出。
[!提示] 从1.0.2版本开始，
S2FFT还提供了底层预计算变换的PyTorch实现。在未来的版本中，这种支持将扩展到我们的即时算法。
[!提示] 从1.1.0版本开始，
S2FFT还为现有的C/C++包提供JAX支持，特别是HEALPix和SSHT。这是通过用自定义JAX前端包装Python绑定来实现的。请注意，目前这种C/C++到JAX的互操作性仅限于CPU。
算法 :zap:
S2FFT利用了可以高度并行化和分布式的新算法结构，因此非常适合硬件加速器（即GPU和TPU）的架构。特别是，这些算法基于新的Wigner-d递归，在高角分辨率$L$下稳定。下图说明了递归过程（更多详情请参见Price & McEwen，准备中）。
有了这个递归，可以通过两步过程计算等纬度采样地图的球谐系数。首先，对每个纬度环进行经度上的1D傅里叶变换。其次，投影到实极坐标d函数上。可以预计算并存储所有实极坐标d函数以实现极端加速，但这会带来同样极端的内存开销，在L ~ 1024时是不可行的。另一种选择是递归计算实极坐标d函数，每次只计算投影的一部分，从而产生可忽略的内存开销，但执行速度略慢。下图说明了可分离的球谐变换（更多详情请参见Price & McEwen，准备中）。
采样 :earth_africa:
S2FFT中实现的算法结构可以支持任何等纬度采样方案。目前支持多种采样方案。
支持McEwen & Wiaux (2012)、Driscoll & Healy (1995)和Gauss-Legendre (1986)的等角度采样方案，它们具有相关的采样定理，因此可以以机器精度计算谐波变换。请注意，McEwen & Wiaux采样定理将球面上的奈奎斯特率降低了一半，与Driscoll & Healy方法相比，将所需的球面样本数量减半。
还支持流行的HEALPix采样方案(Gorski et al. 2005)。HEALPix采样不具有采样定理，因此相应的谐波变换无法达到机器精度，但会有一些误差。然而，HEALPix采样提供了面积相等的像素，这有许多实际优势。
[!注意]
出于算法原因，在高带限时HEALPix变换的JIT编译可能会变慢，这是由于XLA展开循环，目前无法避免。编译完成后，HEALPix变换应该以相关论文中概述的效率执行，因此这个额外的时间开销只需要承担一次。我们知道这个问题，正在努力修复。CPU执行的修复已经实现（请参见示例notebook）。GPU执行的修复即将推出。
安装 :computer:
S2FFT包的Python依赖项列在requirements/requirements-core.txt文件中，当运行以下命令时，pip将自动将它们安装到活动的Python环境中：
pip install s2fft
这将安装所有核心功能，包括JAX支持（也包括PyTorch支持）。
或者，可以直接从GitHub安装S2FFT包，方法是克隆此仓库，然后在仓库的根目录下运行：
pip install .
安装完成后，可以执行单元测试以确保安装成功。首先安装测试依赖，然后运行pytest：
pip install -r requirements/requirements-tests.txt
pytest tests/
发布版本的文档可在这里查看。要在本地构建文档，请运行：
pip install -r requirements/requirements-docs.txt
cd docs
make html
open _build/html/index.html
[!注意]
对于在各种notebook中可以找到的绘图功能，必须安装requirements/requirements-plotting.txt中的依赖项。
使用 :rocket:
导入和使用S2FFT非常简单，如下所示：
对于球面上的信号：
# 计算谐波系数
flm = s2fft.forward_jax(f, L)
# 映射回像素空间信号
f = s2fft.inverse_jax(flm, L)
对于旋转群上的信号：
# 计算Wigner系数
flmn = s2fft.wigner.forward_jax(f, L, N)
# 映射回像素空间信号
f = fft.wigner.inverse_jax(flmn, L, N)
有关使用的更多详细信息，请参阅文档和相关的notebooks。
[!注意]
我们还为预计算版本的变换提供PyTorch支持。这些可以通过forward/inverse_torch()调用。完整的PyTorch支持将在未来版本中提供。
SSHT/HEALPix的C/C++ JAX前端 :bulb:
S2FFT还为现有的C/C++包提供JAX支持，特别是HEALPix和SSHT。这是通过将Python绑定与自定义JAX前端包装来实现的。请注意，这种C/C++到JAX的互操作性目前仅限于CPU。
例如，可以通过以下方式调用这些球谐变换的替代后端：
# 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为现有包增加了梯度功能，适用于现代科学计算或机器学习应用！
有关使用的更多详细信息，请参阅相关的notebooks。
贡献者 ✨
感谢以下这些wonderful贡献者（表情符号对照表）:
 Matt Price 💻 👀 🤔 |  Jason McEwen 💻 👀 🤔 |  Matt Graham 💻 👀 |  sfmig 💻 👀 |  Devaraj Gopinathan 💻 |  Francois Lanusse 💻 🐛 |  Ikko Eltociear Ashimine 📖 |
 Kevin Mulder 🐛 |  Philipp Misof 🐛 |
我们鼓励任何感兴趣的开发者做出贡献。一个简单的首次贡献可以是添加对更多球面采样模式的支持！
引用 :books:
如果以任何方式使用此代码，我们恳请引用以下文章。该参考文献的BibTeX条目可能如下所示：
@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"
}
你可能还想考虑引用我们的相关论文，这些论文是此代码的基础：
@article{mcewen:fssht,
author = "Jason D. McEwen and Yves Wiaux",
title = "A novel sampling theorem on the sphere",
journal = "IEEE Trans. Sig. Proc.",
year = "2011",
volume = "59",
number = "12",
pages = "5876--5887",
eprint = "arXiv:1110.6298",
doi = "10.1109/TSP.2011.2166394"
}
@article{mcewen:so3,
作者 = "Jason D. McEwen 和 Martin Büttner 和 Boris Leistedt 和 Hiranya V. Peiris 和 Yves Wiaux",
标题 = "旋转群上的一种新颖采样定理",
期刊 = "IEEE 信号处理快报",
年份 = "2015",
卷 = "22",
期 = "12",
页码 = "2425--2429",
预印本 = "arXiv:1508.03101",
DOI = "10.1109/LSP.2015.2490676"
}
许可证 :memo:
我们希望这段代码能够为更广泛的社区所用，因此我们以MIT开源许可证提供此代码。
版权所有 2023 Matthew Price、Jason McEwen 及贡献者。
S2FFT是根据MIT许可证提供的免费软件。详情请参阅LICENSE文件。
