nifty-ls

一个快速的Lomb-Scargle周期图。它很棒，并使用NUFFT！

概述

用于识别不规则间距观测中周期性的Lomb-Scargle周期图，虽然有用但计算成本高昂。然而，它可以用数学方式表述为两个非均匀FFT（NUFFT）。这使我们能够利用Flatiron研究所的 finufft 软件包，它真的很快！它还支持GPU（CUDA）并比默认设置的Astropy的Lomb Scargle准确几个数量级。

背景

Press & Rybicki (1989) 的Lomb-Scargle方法将计算表述为四个加权三角函数和，并通过“抽提”到等间距网格来解决这些和。具体来说，和的形式为

\begin{align}
S_k &= \sum_{j=1}^M h_j \sin(2 \pi f_k t_j), \\
C_k &= \sum_{j=1}^M h_j \cos(2 \pi f_k t_j),
\end{align}

其中，$k$ 下标从0到 $N$，即频率分箱数，$f_k$ 是分箱 $k$ 的循环频率，$t_j$ 是观测时间（共有 $M$ 个），$h_j$ 是权重。

我们的关键观察是这恰好是一个非均匀FFT计算的内容！具体来说，在 finufft约定 中，“类型-1”（非均匀到均匀）复数NUFFT计算

g_k = \sum_{j=1}^M h_j e^{i k t_j}.

这个变换的复部和实部是Press & Rybicki的 $S_k$ 和 $C_k$，并对循环/角频率、$k$ 的域、实部与复数变换等进行了调整。finufft有一个特别快速且准确的扩散核（“半圆的指数”），它使用这个核而不是Press & Rybicki的抽提。

对于 $S_k$ 和 $C_k$ 的一些预处理和后处理来计算周期图，这可能会成为瓶颈，因为finufft非常快。这个软件包还优化并并行化了这些计算。

安装

来自PyPI

针对CPU支持

$ pip install nifty-ls

针对GPU（CUDA）支持

$ pip install nifty-ls[cuda]

默认情况下，安装支持CUDA 12；可以使用nifty-ls[cuda11]代替以支持CUDA 11（安装cupy-cuda11x）。

从源代码安装

首先，克隆仓库并将cd到仓库根目录

$ git clone https://www.github.com/flatironinstitute/nifty-ls
$ cd nifty-ls

然后，安装支持CPU

$ pip install .

安装支持GPU（CUDA）

$ pip install .[cuda]

或.[cuda11]以支持CUDA 11。

用于开发（默认在pyproject.toml中启用自动重建）

$ pip install nanobind scikit-build-core
$ pip install -e .[test] --no-build-isolation

开发人员还可能对在pyproject.toml中设置这些键感兴趣

[tool.scikit-build]
cmake.build-type = "Debug"
cmake.verbose = true
install.strip = false

为了最佳性能

您可能希望自行编译和安装finufft和cufinufft，以便它们可以针对您的硬件进行优化。为此，首先安装nifty-ls，然后按照finufft和cufinufft的Python安装说明进行操作，并按需配置库。

nifty-ls也可以根据上述说明从源代码构建以获得最佳性能，但由于大部分繁重计算都由(cu)finufft处理，因此性能提升很小。

用法

来自Astropy

导入nifty_ls将nifty-ls通过Astropy的LombScargle模块中的method="fastnifty"提供。名称前缀为"fast"，因为它属于假设具有均匀频率网格的快速方法系列。

import nifty_ls
from astropy.timeseries import LombScargle
frequency, power = LombScargle(t, y).autopower(method="fastnifty")

import matplotlib.pyplot as plt
import nifty_ls
import numpy as np
from astropy.timeseries import LombScargle

rng = np.random.default_rng(seed=123)
N = 1000
t = rng.uniform(0, 100, size=N)
y = np.sin(50 * t) + 1 + rng.poisson(size=N)

frequency, power = LombScargle(t, y).autopower(method='fastnifty')
plt.plot(frequency, power)
plt.xlabel('Frequency (cycles per unit time)')
plt.ylabel('Power')

要使用CUDA（cufinufft）后端，通过method_kws传递适当的参数

frequency, power = LombScargle(t, y).autopower(method="fastnifty", method_kws=dict(backend="cufinufft"))

在许多情况下，加速您的周期图只需在您的Astropy Lomb Scargle代码中设置method即可！更高级的使用，例如并行计算多个周期图，应直接通过nifty-ls接口进行。

来自nifty-ls（本地接口）

nifty-ls具有自己的接口，该接口提供的批量周期图的灵活性比Astropy接口更高。

单个周期图

单个周期图可以通过nifty-ls计算，如下所示

import nifty_ls
# with automatic frequency grid:
nifty_res = nifty_ls.lombscargle(t, y, dy)

# with user-specified frequency grid:
nifty_res = nifty_ls.lombscargle(t, y, dy, fmin=0.1, fmax=10, Nf=10**6)

import nifty_ls
import numpy as np

rng = np.random.default_rng(seed=123)
N = 1000
t = np.sort(rng.uniform(0, 100, size=N))
y = np.sin(50 * t) + 1 + rng.poisson(size=N)

# with automatic frequency grid:
nifty_res = nifty_ls.lombscargle(t, y)

# with user-specified frequency grid:
nifty_res = nifty_ls.lombscargle(t, y, fmin=0.1, fmax=10, Nf=10**6)

plt.plot(nifty_res.freq(), nifty_res.power)
plt.xlabel('Frequency (cycles per unit time)')
plt.ylabel('Power')

批量周期图

批量周期图（具有相同观测时间的多个对象）可以按如下方式计算

import nifty_ls
import numpy as np

N_t = 100
N_obj = 10
Nf = 200

rng = np.random.default_rng()
t = np.sort(rng.random(N_t))
obj_freqs = rng.random(N_obj).reshape(-1,1)
y_batch = np.sin(obj_freqs * t)
dy_batch = rng.random(y_batch.shape)

batched = nifty_ls.lombscargle(t, y_batch, dy_batch, Nf=Nf)
print(batched.power.shape)  # (10, 200)

请注意，此方法在具有相同观测时间的时间序列集合上同时计算多个周期图。此方法对于短时间序列特别有效，并且/或者当使用GPU时。

目前尚未实现支持具有不同观测时间的时间序列的批量，但计划实现。

限制

该代码仅支持固定间隔的频率网格；然而，finufft支持类型3 NUFFTs（非均匀到非均匀），这将启用任意频率网格。尚不清楚这有多有用，因此尚未实现，但如果您对此感兴趣，请打开GitHub问题。

性能

使用英特尔Icelake CPU的16个核心和NVIDIA A100 GPU，我们获得了以下性能。首先，我们将查看单个周期图的结果（即非批量）

在这种情况下，finufft比Astropy大5倍（带线程为11倍）的速度快，对于（非常）小的转换速度为2倍。相对于Astropy，小转换在更多频率分箱中进一步改进。（动态多线程调度转换是计划中的未来功能，特别有利于小的$N$。）

cufinufft对于大的$N$比Astropy快200倍！性能在小的$N$时趋于平稳，这主要是由于将数据发送到GPU和获取结果的开销。（在GPU上并发作业执行是计划中的另一个功能，将特别帮助小的$N$。）

以下演示了"批量模式"，其中从10个具有相同观测时间的不同时间序列中计算了10个周期图

在这里，finufft单线程的优势在问题规模上始终为6倍，而多线程优势在大型变换中可达30倍。

在这个案例中，GPU的200倍优势扩展到了更小的$N$，因为我们一次发送和接收更多数据。

我们注意到，多线程的finufft和cufinufft特别受益于批量变换，因为这揭示了更多的并行性，并分摊了固定的延迟。

在这些基准测试中，我们对finufft使用了FFTW_MEASURE，这提高了性能几个百分点。

多线程对小问题规模的影响很大；nifty-ls在这种情况下默认使用更少的线程。 “多线程”线使用1到16个线程。

在CPU上，nifty-ls不仅通过使用finufft获得性能，而且通过将预处理和后处理步骤卸载到编译扩展来提高性能。这些扩展使我们能够对元素进行更多处理，而不是对数组进行处理。换句话说，它们使“内核融合”（借用GPU计算中的一个术语）成为可能，从而提高了计算密度。

精度

虽然我们与Astropy的fast方法进行了性能比较，但这并不完全公平。nifty-ls比Astropy fast更准确！Astropy fast使用Press & Rybicki的外推近似，以速度换取精度，但多亏了finufft，nifty-ls可以两者兼得。

在下图中，我们用圆圈绘制了不同$N$（和默认$N_F \approx 12N$）的astropy、finufft和cufinufft的中值频谱误差和99百分位误差。

astropy的结果提供了两种情况：一种是标准情况，另一种是“最坏情况”。内部，astropy使用一个FFT网格，其大小是目标过采样率以上的下一个2的幂。通常，每次跳到新的2的幂都会导致精度提高。因此，“最坏情况”是那些不会导致这种跳跃的最高频率。

最坏情况评估中常见的误差为$\mathcal{O}(10\%)$或更大。典型评估中常见的误差为$\mathcal{O}(1\%)$或更大。nifty-ls的精度保守估计比这高6个数量级。

上图中的参考结果来自“相绕行”方法，该方法使用三角恒等式来避免昂贵的正弦和余弦评估。也可以使用astropy的fast方法作为参考，通过use_fft=False启用精确评估。发现相同的结果，但相绕行要快几个数量级（但仍然与finufft不相竞争）。

总之，nifty-ls在保持高性能的同时，也具有高度的准确性。

float32 vs float64

虽然32位浮点数可以为finufft和cufinufft提供实质性的速度提升，但我们通常不建议在Lomb-Scargle中使用它们。原因是问题的挑战性条件数。条件数是输出对输入微小扰动的响应——换句话说，是导数。 可以很容易地证明 NUFFT相对于非均匀点的导数与$N$成正比，即变换长度（即模式数）。换句话说，观测时间误差被放大$\mathcal{O}(N)$。由于float32的相对误差为$\mathcal{O}(10^{-7})$，长度为$10^5$的变换已经受到$\mathcal{O}(1\%)$误差的影响。因此，我们在nifty-ls中关注float64，但float32也由所有后端原生支持，供有冒险心的用户使用。

条件数也是上述图中误差随$N$增加的温和上升趋势的可能贡献者，至少对于finufft/cufinufft。对于float64的相对误差为$\mathcal{O}(10^{-16})$和变换长度为$\mathcal{O}(10^6)$，最小误差为$\mathcal{O}(10^{-10})$。

测试

首先，从源代码安装（pip install .[test]）。然后，从仓库根目录运行

$ pytest

测试定义在 tests/ 目录中，包括对 nifty-ls 和 Astropy 的迷你基准测试，如下所示

$ pytest
======================================================== test session starts =========================================================
platform linux -- Python 3.10.13, pytest-8.1.1, pluggy-1.4.0
benchmark: 4.0.0 (defaults: timer=time.perf_counter disable_gc=True min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000)
rootdir: /mnt/home/lgarrison/nifty-ls
configfile: pyproject.toml
plugins: benchmark-4.0.0, asdf-2.15.0, anyio-3.6.2, hypothesis-6.23.1
collected 36 items                                                                                                                   

tests/test_ls.py ......................                                                                                        [ 61%]
tests/test_perf.py ..............                                                                                              [100%]


----------------------------------------- benchmark 'Nf=1000': 5 tests ----------------------------------------
Name (time in ms)                       Min                Mean            StdDev            Rounds  Iterations
---------------------------------------------------------------------------------------------------------------
test_batched[finufft-1000]           6.8418 (1.0)        7.1821 (1.0)      0.1831 (1.32)         43           1
test_batched[cufinufft-1000]         7.7027 (1.13)       8.6634 (1.21)     0.9555 (6.89)         74           1
test_unbatched[finufft-1000]       110.7541 (16.19)    111.0603 (15.46)    0.1387 (1.0)          10           1
test_unbatched[astropy-1000]       441.2313 (64.49)    441.9655 (61.54)    1.0732 (7.74)          5           1
test_unbatched[cufinufft-1000]     488.2630 (71.36)    496.0788 (69.07)    6.1908 (44.63)         5           1
---------------------------------------------------------------------------------------------------------------

--------------------------------- benchmark 'Nf=10000': 3 tests ----------------------------------
Name (time in ms)            Min              Mean            StdDev            Rounds  Iterations
--------------------------------------------------------------------------------------------------
test[finufft-10000]       1.8481 (1.0)      1.8709 (1.0)      0.0347 (1.75)        507           1
test[cufinufft-10000]     5.1269 (2.77)     5.2052 (2.78)     0.3313 (16.72)       117           1
test[astropy-10000]       8.1725 (4.42)     8.2176 (4.39)     0.0198 (1.0)         113           1
--------------------------------------------------------------------------------------------------

----------------------------------- benchmark 'Nf=100000': 3 tests ----------------------------------
Name (time in ms)              Min               Mean            StdDev            Rounds  Iterations
-----------------------------------------------------------------------------------------------------
test[cufinufft-100000]      5.8566 (1.0)       6.0411 (1.0)      0.7407 (10.61)       159           1
test[finufft-100000]        6.9766 (1.19)      7.1816 (1.19)     0.0748 (1.07)        132           1
test[astropy-100000]       47.9246 (8.18)     48.0828 (7.96)     0.0698 (1.0)          19           1
-----------------------------------------------------------------------------------------------------

------------------------------------- benchmark 'Nf=1000000': 3 tests --------------------------------------
Name (time in ms)                  Min                  Mean            StdDev            Rounds  Iterations
------------------------------------------------------------------------------------------------------------
test[cufinufft-1000000]         8.0038 (1.0)          8.5193 (1.0)      1.3245 (1.62)         84           1
test[finufft-1000000]          74.9239 (9.36)        76.5690 (8.99)     0.8196 (1.0)          10           1
test[astropy-1000000]       1,430.4282 (178.72)   1,434.7986 (168.42)   5.5234 (6.74)          5           1
------------------------------------------------------------------------------------------------------------

Legend:
  Outliers: 1 Standard Deviation from Mean; 1.5 IQR (InterQuartile Range) from 1st Quartile and 3rd Quartile.
  OPS: Operations Per Second, computed as 1 / Mean
======================================================== 36 passed in 30.81s =========================================================

结果是在 16 核的英特尔 Icelake CPU 和 1 个 NVIDIA A100 GPU 上获得的。与最快的运行时间相比的比率显示在括号中。您可能在您的平台上获得非常不同的性能！特别是最慢的 Astropy 结果可能取决于您安装的 Numpy 发行版及其三角函数性能。

作者

nifty-ls 最初由 Lehman Garrison 实现，基于 Dan Foreman-Mackey 在 dfm/nufft-ls 仓库中的工作，并在 Alex Barnett 的咨询下完成。

致谢

nifty-ls 直接建立在 Alex Barnett 和其他人出色的 finufft 包之上（参见 finufft 致谢）。

本包的许多部分是对 Astropy LombScargle 的改编，特别是 Press & Rybicki (1989) 方法。