这种修复插值可以在 Python 中变得更快吗？

Question

根据本文 Garcia 等人，Matlab 中有一个使用离散余弦变换编写的 inpaint 函数（Inpaintn）来填充多维数据集中的缺失值。等人。 （2012）。我尝试将此代码（inpaintn.m）移植到Python中，如下所示，

import numpy as np
from scipy.ndimage import distance_transform_edt
from scipy.fft import idctn, dctn
from tqdm import tqdm

def fill_nd(data, invalid=None):
    if invalid is None: invalid = np.isnan(data)

    ind = distance_transform_edt(invalid, return_distances=False, return_indices=True)
    return data[tuple(ind)]


def InitialGuess(y, I):
    z = fill_nd(y)
    s0 = 3
    return z, s0


def idctnn(y):
    return idctn(y, norm='ortho')


def dctnn(y):
    return dctn(y, norm='ortho')


def inpaint(xx, y0=[], n=100, m=2, verbose=False):
    x = xx.copy() #as it changes x itself, so copying it to another variable.

    sizx = np.shape(x)
    d = np.ndim(x)
    Lambda = np.zeros(sizx, dtype='float')

    for i in range(0, d):
        siz0 = np.ones(d, dtype='int')
        siz0[i] = sizx[i]
        Lambda = Lambda + np.cos(np.pi * np.reshape(np.arange(1, sizx[i] + 0.1) - 1, siz0) / sizx[i])

    Lambda = 2 * (d - Lambda)

    # Initial condition
    W = np.isfinite(x)
    if len(y0) == len(x):
        y = y0
        s0 = 3  # note: s = 10 ^ s0
    else:
        if np.any(~W):
            if verbose: print('Initial Guess as Nearest Neighbors')
            y, s0 = InitialGuess(x, np.isfinite(x).astype('bool'))
        else:
            y = x
            s0 = 3
            # return x
    x[~W] = 0.

    # Smoothness parameters: from high to negligible
    s = np.logspace(s0, -6, n)

    RF = 2.  # Relaxation Factor
    Lambda = Lambda ** m

    if verbose: print('Inpainting .......')

    for i in tqdm(range(n)):
        Gamma = 1. / (1 + s[i] * Lambda)
        y = RF * idctnn(Gamma * dctnn((W * (x - y)) + y)) + (1 - RF) * y
        
    y[W] = x[W]

    return y

代码运行良好，但我一直在努力寻找使代码运行得更快的方法，特别是因为我的数据集很大。使用这种类型插值的优点是，我可以提供整个 3D 数据集（包含时间和网格坐标）来填充缺失值，而不是为每个时间坐标执行此操作。

这是一个使用 python 的示例数据集

import numpy as np

#A 3D dataset with dimensions (time, latitude, longitude)
X = np.random.randn(1000,180,360)

# Randomly choosing indices to insert 64800 NaN values (say). 
#NaNs can also be present as blocks in the data, not randomly dispersed as below.
index_nan = np.random.choice(X.size, 64800, replace=False)

#Inserting NaNs. 
X.ravel()[index_nan] = np.nan

我尝试过一些方法，但没有成功

1. 使用 Numba

jit 装饰器让它变得更慢，即使有像

parallel/fastmath/vectorize,nopython=True 

这样的选项。

2. 使用 Cython

我尝试排版这些函数中使用的所有变量，但它仍然比原生 python 实现慢。而且，在我的机器上编译 Cython 代码很麻烦。

3. 使用 Numpy 向量化

我已经用

scipy

函数替换了离散余弦变换函数及其反函数，但我似乎无法想到如何对内部 for 循环进行矢量化以使其快速，以及是否可能。 我尝试过分析我的代码，瓶颈似乎在于使用

scipy

进行离散余弦变换。还有其他瓶颈，但对我来说没有意义。我还附上了一张用于分析的图像。

如果有可行的方法来加速这段代码，那确实会有很大帮助。我在Python方面并不是很先进，但是我可以从中学到很多东西，特别是我的问题的可行性。

Answer 1

该算法适用于相当大的数组（不适合 CPU 缓存），部分解释了为什么它有点慢。此外，众所周知，DCT/IDCT 是昂贵的操作。话虽这么说，您可以使用 Numba 的 JIT 和 scipy 函数的

workers=-1

选项来并行化算法。此外，您可以通过就地工作避免创建许多昂贵的临时数组。这是未经测试的结果代码：

# In-place computation
def idctnn(y):
    return idctn(y, norm='ortho', workers=-1, overwrite_x=True)


# In-place computation
def dctnn(y):
    return dctn(y, norm='ortho', workers=-1, overwrite_x=True)


# In-place computation (writes in `Transformed`)
@nb.njit('void(float64[:,:,::1], float64[:,:,::1], float64)', parallel=True)
def ComputeGammaTransform(Transformed, Lambda, sVal):
    for i in nb.prange(Transformed.shape[0]):
        for j in range(Transformed.shape[1]):
            for k in range(Transformed.shape[2]):
                Transformed[i, j, k] /= (1. + sVal * Lambda[i, j, k])


# Out-of-place computation (writes in `out`)
@nb.njit('void(float64[:,:,::1], float64[:,:,::1], float64[:,:,::1], boolean[:,:,::1])', parallel=True)
def ComputeDctInput(out, x, y, W):
    for i in nb.prange(out.shape[0]):
        for j in range(out.shape[1]):
            for k in range(out.shape[2]):
                out[i, j, k] = W[i, j, k] * (x[i, j, k] - y[i, j, k]) + y[i, j, k]


# In-place computation (writes in `y`)
@nb.njit('void(float64[:,:,::1], float64[:,:,::1], float64)', parallel=True)
def ComputeDctOutput(dctResult, y, RF):
    for i in nb.prange(y.shape[0]):
        for j in range(y.shape[1]):
            for k in range(y.shape[2]):
                y[i, j, k] = RF * dctResult[i, j, k] + (1.0 - RF) * y[i, j, k]


def ComputeSteps(Lambda, x, y, W, s, RF):
    dctData = np.empty(Lambda.shape, dtype=Lambda.dtype)
    for i in tqdm(range(s.shape[0])):
        ComputeDctInput(dctData, x, y, W)
        dctnn(dctData)
        ComputeGammaTransform(dctData, Lambda, s[i])
        idctnn(dctData)
        ComputeDctOutput(dctData, y, RF)

这段代码在我的机器上速度快了 5 倍。您可以使用简单精度而不是双精度进一步加快速度。这使得最终代码比我机器上的原始代码快 7.5 倍。

我也许可以通过基于 GPU 的计算来进一步加快代码速度。困难的部分是在 Python 中找到支持正交归一化的 DCT/IDCT 的 GPU 实现。

Answer 2

可能会很晚，但万一有人像我一样在某个时候来到这里。

我基于

Garcia (2010) 提出的原始 smoothn(x) 代码实现了相同的算法。我对你的代码进行了快速测试，结果我的代码大约快了 4 倍。由于我没有使用 Numba，我想它仍然可以改进很多，正如 Jérôme Richard 的回答 所建议的那样。我还注意到，在我所做的小测试中，向 Scipy 函数添加工人选项会减慢算法的速度（不适用于所有情况）。

我将我的代码上传到这个

存储库，其中包含平滑和修复功能。