可微分和加速的球面变换

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贡献者（表情符号对照表）:

<table> <tbody> <tr> <td align="center" valign="top" width="14.28%"><a href="https://cosmomatt.github.io"><img src="https://avatars.githubusercontent.com/u/32554533?v=4?s=100" width="100px;" alt="Matt Price"/> Matt Price</a> <a href="https://github.com/astro-informatics/s2fft/commits?author=CosmoMatt" title="代码">💻</a> <a href="https://github.com/astro-informatics/s2fft/pulls?q=is%3Apr+reviewed-by%3ACosmoMatt" title="审核拉取请求">👀</a> <a href="#ideas-CosmoMatt" title="想法、规划与反馈">🤔</a></td> <td align="center" valign="top" width="14.28%"><a href="http://www.jasonmcewen.org"><img src="https://avatars.githubusercontent.com/u/3181701?v=4?s=100" width="100px;" alt="Jason McEwen "/> Jason McEwen </a> <a href="https://github.com/astro-informatics/s2fft/commits?author=jasonmcewen" title="代码">💻</a> <a href="https://github.com/astro-informatics/s2fft/pulls?q=is%3Apr+reviewed-by%3Ajasonmcewen" title="审核拉取请求">👀</a> <a href="#ideas-jasonmcewen" title="想法、规划与反馈">🤔</a></td> <td align="center" valign="top" width="14.28%"><a href="http://matt-graham.github.io"><img src="https://avatars.githubusercontent.com/u/6746980?v=4?s=100" width="100px;" alt="Matt Graham"/> Matt Graham</a> <a href="https://github.com/astro-informatics/s2fft/commits?author=matt-graham" title="代码">💻</a> <a href="https://github.com/astro-informatics/s2fft/pulls?q=is%3Apr+reviewed-by%3Amatt-graham" title="审核拉取请求">👀</a></td> <td align="center" valign="top" width="14.28%"><a href="https://sfmig.github.io/"><img src="https://avatars.githubusercontent.com/u/33267254?v=4?s=100" width="100px;" alt="sfmig"/> sfmig</a> <a href="https://github.com/astro-informatics/s2fft/commits?author=sfmig" title="代码">💻</a> <a href="https://github.com/astro-informatics/s2fft/pulls?q=is%3Apr+reviewed-by%3Asfmig" title="审核拉取请求">👀</a></td> <td align="center" valign="top" width="14.28%"><a href="https://github.com/Devaraj-G"><img src="https://avatars.githubusercontent.com/u/36169767?v=4?s=100" width="100px;" alt="Devaraj Gopinathan"/> Devaraj Gopinathan</a> <a href="https://github.com/astro-informatics/s2fft/commits?author=Devaraj-G" title="代码">💻</a></td> <td align="center" valign="top" width="14.28%"><a href="http://flanusse.net"><img src="https://avatars.githubusercontent.com/u/861591?v=4?s=100" width="100px;" alt="Francois Lanusse"/> Francois Lanusse</a> <a href="https://github.com/astro-informatics/s2fft/commits?author=EiffL" title="代码">💻</a> <a href="https://github.com/astro-informatics/s2fft/issues?q=author%3AEiffL" title="报告Bug">🐛</a></td> <td align="center" valign="top" width="14.28%"><a href="https://github.com/eltociear"><img src="https://avatars.githubusercontent.com/u/22633385?v=4?s=100" width="100px;" alt="Ikko Eltociear Ashimine"/> Ikko Eltociear Ashimine</a> <a href="https://github.com/astro-informatics/s2fft/commits?author=eltociear" title="文档">📖</a></td> </tr> <tr> <td align="center" valign="top" width="14.28%"><a href="https://github.com/kmulderdas"><img src="https://avatars.githubusercontent.com/u/33317219?v=4?s=100" width="100px;" alt="Kevin Mulder"/> Kevin Mulder</a> <a href="https://github.com/astro-informatics/s2fft/issues?q=author%3Akmulderdas" title="报告Bug">🐛</a></td> <td align="center" valign="top" width="14.28%"><a href="https://github.com/PhilippMisofCH"><img src="https://avatars.githubusercontent.com/u/142883157?v=4?s=100" width="100px;" alt="Philipp Misof"/> Philipp Misof</a> <a href="https://github.com/astro-informatics/s2fft/issues?q=author%3APhilippMisofCH" title="报告Bug">🐛</a></td> </tr> </tbody> </table>

我们鼓励任何感兴趣的开发者做出贡献。一个简单的首次贡献可以是添加对更多球面采样模式的支持！

引用 :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"    
}