科学和工程的许多领域都会遇到在球体上定义的数据。对此类数据进行建模和分析通常需要傅里叶变换的球面对应物，即球面谐波变换。我们简要概述了球谐变换，并提出了一种新的可微分算法，该算法专为GPU上的加速而定制[1]。该算法在最近发布的 S2FFT python 包中实现，该包同时支持 JAX 和 PyTorch。

        我我们经常对分析存在于球体上的数据感兴趣。从量子化学、生物医学成像、气候物理学和地球物理学到更广阔的宇宙，应用的多样性是显着的。

        人们最熟悉的领域是在物理科学中，特别是在大气科学、地球物理建模和天体物理学中。

        最广为人知的球形数据案例的例子，如地球（左）和艺术家对天文观测的印象（右）。[地球图像来源 维基百科;天体物理学图片来自维基百科。

        这些问题自然是球形的，因为观测是在球体表面的每个点进行的：地球物理学的地球表面和天体物理学的天空。其他例子来自计算机图形学和视觉等应用，其中 360° 全景相机从各个方向捕捉您周围的世界。

        在许多情况下，手头问题的球形性质很容易看出;然而，情况并非总是如此。也许令人惊讶的是，球形数据在生物学学科中经常遇到，尽管球形方面通常不那么明显！由于我们经常关注生物学研究中的局部方向，例如水在大脑内扩散的方向，因此我们会遇到球形数据。

人脑中神经元连接的扩散张量成像。在每个体素中，神经元可以自由地向任何方向移动，因此问题自然是球形的。[阿尔弗雷德·安万德（Alfred Anwander）的动画，CC-BY许可。

鉴于此类数据的普遍性，开发出许多球形分析技术也就不足为奇了。对数据的频率分析可能很有见地，通常可以为进一步的分析或建模提供统计摘要或有效的表示。最近，几何深度学习技术已被证明对复杂域的数据分析非常有效[2–6]，特别是对于高度复杂的问题，如分子建模和蛋白质相互作用（参见我们之前的文章《几何深度学习简介》）。

二、傅立叶遇见勒让德

        因此，我们有关于球体的数据和各种可以分析球面数据的技术，但我们需要数学工具来做到这一点。具体来说，我们需要知道如何有效地将球形数据分解为频率。

        傅里叶变换提供了频率分解，通常用于计算数据中的统计相关性。许多物理系统也可以在频率空间中更直接地描述，因为每个频率可以独立演化。

        为了将标准的傅里叶变换扩展到球体，我们需要两位 17 世纪法国数学家的共识：约瑟夫·傅立叶和阿德里安-玛丽·勒让德。

        约瑟夫·傅立叶（左）和阿德里安-玛丽·勒让德（右）。可悲的是，勒让德的漫画是他唯一已知的形象。[傅里叶图片来源 维基百科。勒让德图片来源于维基百科。

        首先，让我们考虑如何将欧几里得数据分解为不同的频率。这种数据转换首先由约瑟夫·傅立叶（Joseph Fourier）推导，由下式给出

        它几乎随处可见，并且是本科物理的主要内容，这是有原因的！这是通过将我们的数据 f（x） 投影到一组三角函数（称为基）上来工作的。人们可以在球体上有效地做同样的事情，但基函数现在由球谐 Ylm 给出：

        其中 （θ， φ） 是通常的球面极坐标。

球谐基函数（实分量）。[来源于：Source： 维基百科。

        球谐波（如上图所示）可以进一步分解为指数和勒让德多项式的乘积——à la Adrien-Marie Legendre——为

        因此，球谐变换可以写为傅里叶变换，然后是相关的勒让德变换。真正的困难在于评估变换的勒让德部分：它要么计算成本高昂，要么内存耗大，这取决于人们选择的方法。

三、可区分性的重要性

可微分编程的发展开辟了许多新的分析类型。特别是，许多应用需要可微分的球面变换。

        可微分编程是一种编程范式，其中数字计算机程序可以通过自动微分进行整体微分。[1][2][3][4][5]这允许对程序中的参数进行基于梯度的优化，通常通过梯度下降，以及基于高阶导数信息的其他学习方法。可微分编程已广泛应用于各个领域，尤其是科学计算和机器学习。[5] 2016年初，欧洲航天局的高级概念团队提出了以系统的方式采用这种框架来改进学习算法的早期建议之一。

        球体上的机器学习模型需要可微变换，以便模型可以通过基于梯度的优化算法进行训练，即通过反向传播。

        用于混合数据驱动和基于模型的方法[8]的新兴物理增强机器学习方法[7]也需要可微分物理模型，在许多情况下，这些模型本身需要可微球变换。

        考虑到这一点，很明显，对于现代应用来说，一个有效的球谐变换算法是必要的，但还不够。可区分性是关键。

四、一个名为 S2FFT 的软件包

        这一切都很好，但如何有效地评估球谐变换呢？已经开发了各种算法，以及一些很棒的软件包。然而，对于现代应用程序，我们需要一个可区分的，可以在GPU等硬件加速器上运行，并且计算可扩展的应用程序。

        通过从头开始重新设计核心算法（如我们相应的论文 [1] 中深入描述的那样），我们最近开发了一个名为 S2FFT 的 python 包，它应该符合要求。

        S2FFT 是在 JAX 中实现的，JAX 是 Google 开发的一种可微分编程语言，还包括一个 PyTorch 前端。

S2FFT 是一个 Python 包，用于实现可微分和加速的球谐变换，具有 JAX 和 PyTorch 中的接口。[图片由作者创建。

        S2FFT提供两种操作模式：预先计算关联的勒让德函数，然后在运行时访问这些函数;或在转换过程中即时计算它们。预计算方法的速度几乎与你能得到的一样快，但存储所有勒让德函数值所需的内存会随着分辨率的立方变化而缩放，这可能是一个问题！相反，我们提供的第二种方法是递归地即时计算勒让德项，因此可以扩展到非常高的分辨率。

        此外，S2FFT还支持混合自动和手动微分方法，以便可以有效地计算梯度。

        该软件包旨在支持球体上的多种不同采样方案。在发布时，支持等角（McEwen & Wiaux [9]、Driscoll & Healy [10]）、Gauss-Legendre 和 HEALPix [11] 采样方案，尽管将来可以很容易地添加其他方案。

S2FFT支持的球体上的不同采样方案。[原创图由作者创作。

S2FFT 软件包在 PyPi 上可用，因此任何人都可以通过运行以下命令直接安装它：

pip install s2fft

        或者通过运行以下命令来获取 PyTorch 支持：

pip install "s2fft[torch]"

        从这里可以简单地调用顶级转换

import s2fft

# Compute forward spherical harmonic transform
flm = s2fft.forward_jax(f, L) 

# Compute inverse spherical harmonic transform
f = s2fft.inverse_jax(flm, L)  

        这些功能可以开箱即用，并作为层集成到现有模型中，无论是在 JAX 还是 PyTorch 中，都完全支持正向和反向模式区分。

五、未来展望

        随着研究人员对科学应用的可微分编程越来越感兴趣，迫切需要现代软件包来实现科学通常基于的基础数学方法，例如球谐变换。

        我们希望 S2FFT 在未来几年会有很大的用处，并很高兴看到人们会使用它！