Can this inpaint interpolation be made faster in Python?

There is an inpaint function written in Matlab (Inpaintn) using discrete cosine transforms for filling missing values in multidimensional datasets according to this paper Garcia et. al. (2012). I tried porting this code (inpaintn.m) into Python as follows,

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

The code works fine, but I have been trying to find ways to make this code run faster especially since my datasets are large. The advantage of using this type of interpolation is that I can feed an entire 3D dataset (with time and grid coordinates) to fill the missing values, instead of doing it for each time coordinate.

Here is an example dataset using the 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

I have tried some ways but they weren't successful,

1. Using Numba

The jit decorator made it slower, even with options like parallel/fastmath/vectorize,nopython=True .

2. Using Cython

I tried typesetting all variables used in these functions but it was still slower than native python implementation. And moreso, it is hassle compiling Cython code on my machine.

3. Using Numpy Vectorization

I have already replaced the discrete cosine transform functions and its inverse with scipy functions, but I cannot seem to think of ways to vectorize the inner for loops to make it fast, and whether it is even possible. I have tried profiling my code, and the bottleneck seems to be in the discrete cosine transforms using scipy. There are other bottlenecks but it doesn't make sense to me. I have attached an image for the profiling as well.

It would really help a lot if there are feasible ways to speed up this code. I am not really advanced in Python, but there is a lot that I can learn from this, especially the feasibility of my question.