首页 > 其他分享 >s2fft库介绍:可微分和加速球谐变换

s2fft库介绍:可微分和加速球谐变换

时间:2024-03-18 14:58:08浏览次数:19  
标签:s2fft 微分 变换 S2FFT 球体 球谐 数据

一、说明

        科学和工程的许多领域都会遇到在球体上定义的数据。对此类数据进行建模和分析通常需要傅里叶变换的球面对应物,即球面谐波变换。我们简要概述了球谐变换,并提出了一种新的可微分算法,该算法专为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 在未来几年会有很大的用处,并很高兴看到人们会使用它!

标签:s2fft,微分,变换,S2FFT,球体,球谐,数据
From: https://blog.csdn.net/gongdiwudu/article/details/136772819

相关文章

  • 双曲型偏微分方程的几个概念
    本章研究了一类双曲型偏微分方程的一些基本性质。本书中研究的离散化技术主要基于偏微分方程的基本物理和数学特性。因此,有理由在偏微分方程的一些基础上投入一些精力。这里我们几乎只讨论双曲偏微分方程,特别是双曲守恒律。这主要有三个原因:(i)当忽略粘性和热传导的影响时,可压缩流......
  • 一阶微分方程
    常数变易法考虑以下一阶线性微分方程\[\dfrac{\text{d}y}{\text{d}x}=P(x)y+Q(x)\]先解齐次方程\[\begin{aligned}\dfrac{\text{d}y}{\text{d}x}&=P(x)y\\\dfrac{\text{d}y}{y}&=P(x)\text{d}x\\\lny&=P_1(x)+C\\y&=C\exp(P_1(x))\end{aligned}\]其......
  • 【习题】5.2 一阶常系数线性微分方程组
    [T050201]求方程组\(\begin{cases}x'+y'=y+z\\y'+z'=z+x\\z'+x'=x+y\end{cases}\)的通解.    解将方程组写成矩阵形式:\[\begin{aligned}&\begin{pmatrix}1&1&0\\0&1&1\\1&0&1\end{pmatrix}\begin{pmat......
  • 【习题】5.1 一阶线性微分方程的基本概念
    [T050101]设\(A\)为\(n\timesn\)常数矩阵,\(\Phi(t)\)是方程组\(X'=AX\)的标准基解矩阵\((\Phi(0)=E)\),证明\(\Phi(t)\Phi^{-1}(t_0)=\Phi(t-t_0)\),其中\(t_0\)是常数.    证由题设可知\(\Phi'(t)=A\Phi(t)\),将\(t\)换为\(t-t_0\),则\(\Phi......
  • 数值上如何对布洛赫函数进行微分?
    这个问题在计算陈数、贝里曲率等量时都会遇到。总结而言是两个问题:1.数值上产生的不同k点的布洛赫波是有不同的相位,不连续,所以此时怎么计算对布洛赫波的微分以及贝里曲率等物理量?2.当能带上有简并交叉点时,如何计算对布洛赫波的微分?对这两个问题,关老师有过很多博文:https://www.......
  • 分数阶微分 FOD(GL)
    分数阶微分方程讲义-华师大数院分数阶微分-清华大学出版社参考以上文件公式与代码,使用R代码编写,针对光谱数据,经对比alpha=1时与应用prospectr包函数计算结果一致fod<-function(spectrum,alpha){#initd<-dim(spectrum)FOD<-matrix(0,nrow=0,ncol=d[2])......
  • 初等函数微分和积分公式
         ......
  • 微分积分及其运算法则成对列举
         ......
  • 神经网络的反向传导的自动微分中前向模式和后向模式的适用情况
    神经网络的自动微分求导这里不做解释和介绍,自动微分求导中的前向模式和后向模式这里也不做解释和介绍。根据资料显示,如果一个神经网络的输入层维度为M,输出层维度为N,当M>N时,反向自动微分求导的后向模式计算效率高于前向模式,反之,如果M<N时,前向模式计算效率高于后向模式,下面给出自己......
  • 【常微分方程】2.1.6
    [T020101]设\(f(x)\)满足\(f(x+y)=\frac{f(x)+f(y)}{1-f(x)f(y)}\),且\(f'(0)\)存在,求\(f(x)\)的表达式.解令\(x=y=0\),则\(f(0)=\frac{2f(0)}{1-f(0)^2}\),得\(f(0)=0\).注意到\[f(x+\Deltax)=\frac{f(x)+f(\Deltax)}{1-f(x)f(\Deltax)}\Lo......