Блог

Существует функция inpaint, написанная в Matlab (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. Использование Цитона

Я попытался набирать все переменные, используемые в этих функциях, но это все еще было медленнее, чем собственная реализация python. И более того, это хлопотно компилировать код Cython на моей машине.

3. Использование векторизации Numpy

Я уже заменил функции дискретного косинусного преобразования и его обратные scipy функции функциями, но, похоже, я не могу придумать способы векторизации внутренних циклов for, чтобы сделать это быстро, и возможно ли это вообще. Я попытался профилировать свой код, и узкое место, похоже, заключается в использовании дискретных косинусных преобразований scipy . Есть и другие узкие места, но для меня это не имеет смысла. Я также прикрепил изображение для профилирования.

Это действительно очень помогло бы, если бы существовали возможные способы ускорить этот код. Я не очень хорошо разбираюсь в Python, но из этого я могу многому научиться, особенно возможности моего вопроса.

Алгоритм работает с довольно большим массивом (не помещается в кэш процессора), частично объясняя, почему он немного медленный. Кроме того, DCT/IDCT, как известно, являются дорогостоящими операциями. Тем не менее, вы можете распараллелить алгоритм, используя JIT Numba и workers=-1 опцию для функций scipy. Кроме того, вы можете избежать создания множества дорогостоящих временных массивов, работая на месте. Вот непроверенный результирующий код:

 # 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 раз быстрее, чем исходный на моем компьютере.

Возможно, мне удастся еще больше ускорить код с помощью вычислений на основе графического процессора. Самое сложное-найти реализацию DCT/IDCT на GPU в Python, поддерживающую ортогональную нормализацию.