Эффективное сшивание наборов данных

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

Думайте об измерениях как о нескольких картах высот одного объекта, которые я хочу объединить в одну карту высот. Мои измерения не идеальны и могут иметь некоторое смещение наклона и высоты. Давайте предположим (на данный момент), что мы точно знаем координаты x-y. Вот пример:

import numpy as np
import matplotlib.pyplot as plt

def height_profile(x, y):
    radius = 100
    return np.sqrt(radius**2-x**2-y**2)-radius

np.random.seed(123)

datasets = {}

# DATASET 1
x = np.arange(-8, 2.01, 0.1)
y = np.arange(-3, 7.01, 0.1)

xx, yy = np.meshgrid(x, y)
# height is the actual profile + noise
zz = height_profile(xx, yy) + np.random.randn(*xx.shape)*0.001

datasets[1] = [xx, yy, zz]

plt.figure()
plt.pcolormesh(*datasets[1])
plt.colorbar()

# DATASET 2
x = np.arange(-2, 8.01, 0.1)
y = np.arange(-3, 7.01, 0.1)

xx, yy = np.meshgrid(x, y)
# height is the actual profile + noise + random offset + random tilt
zz = height_profile(xx, yy) + np.random.randn(*xx.shape)*0.001 + np.random.rand() + np.random.rand()*xx*0.1 + np.random.rand()*yy*0.1

datasets[2] = [xx, yy, zz]

plt.figure()
plt.pcolormesh(*datasets[2])
plt.colorbar()

# DATASET 3
x = np.arange(-5, 5.01, 0.1)
y = np.arange(-7, 3.01, 0.1)

xx, yy = np.meshgrid(x, y)
# height is the actual profile + noise + random offset + random tilt
zz = height_profile(xx, yy) + np.random.randn(*xx.shape)*0.001 + np.random.rand() + np.random.rand()*xx*0.1 + np.random.rand()*yy*0.1

datasets[3] = [xx, yy, zz]

plt.figure()
plt.pcolormesh(*datasets[3])
plt.colorbar()

Чтобы объединить три (или более) набора данных, у меня есть следующая стратегия: найти перекрытие между наборами данных, вычислить суммарную разницу высот между наборами данных в областях перекрытия (residual_overlap) и попытаться минимизировать разницу высот (остаточную), используя lmfit. Для применения преобразований к набору данных (наклон, смещение и т. д.) у меня есть специальная функция.

from lmfit import minimize, Parameters
from copy import deepcopy
from itertools import combinations
from scipy.interpolate import griddata

def data_transformation(dataset, idx, params):
    dataset = deepcopy(dataset)
    
    if 'x_offset_{}'.format(idx) in params:
        x_offset = params['x_offset_{}'.format(idx)].value
    else:
        x_offset = 0

    if 'y_offset_{}'.format(idx) in params:
        y_offset = params['y_offset_{}'.format(idx)].value
    else:
        y_offset = 0
    
    if 'tilt_x_{}'.format(idx) in params:
        x_tilt = params['tilt_x_{}'.format(idx)].value
    else:
        x_tilt = 0

    if 'tilt_y_{}'.format(idx) in params:
        y_tilt = params['tilt_y_{}'.format(idx)].value
    else:
        y_tilt = 0

    if 'piston_{}'.format(idx) in params:
        piston = params['piston_{}'.format(idx)].value
    else:
        piston = 0

    _x = dataset[0] - np.mean(dataset[0])
    _y = dataset[1] - np.mean(dataset[1])

    dataset[0] = dataset[0] + x_offset
    dataset[1] = dataset[1] + y_offset
    dataset[2] = dataset[2] + 2 * (x_tilt * _x + y_tilt * _y) + piston

    return dataset

def residual_overlap(dataset_0, dataset_1):
    xy_0 = np.stack((dataset_0[0].flatten(), dataset_0[1].flatten()), axis=1)
    xy_1 = np.stack((dataset_1[0].flatten(), dataset_1[1].flatten()), axis=1)
    difference = griddata(xy_0, dataset_0[2].flatten(), xy_1) - \
                 dataset_1[2].flatten()

    return difference

def residual(params, datasets):
    datasets = deepcopy(datasets)

    for idx in datasets:
        datasets[idx] = data_transformation(
            datasets[idx], idx, params)

    residuals = []

    for combination in combinations(list(datasets), 2):
        residuals.append(residual_overlap(
            datasets[combination[0]], datasets[combination[1]]))

    residuals = np.concatenate(residuals)
    residuals[np.isnan(residuals)] = 0

    return residuals

def minimize_datasets(params, datasets, **minimizer_kw):
    minimize_fnc = lambda *args, **kwargs: residual(*args, **kwargs)

    datasets = deepcopy(datasets)

    min_result = minimize(minimize_fnc, params,
                          args=(datasets, ), **minimizer_kw)

    return min_result

Вышиваю так:

params = Parameters()
params.add('tilt_x_2', 0)
params.add('tilt_y_2', 0)
params.add('piston_2', 0)
params.add('tilt_x_3', 0)
params.add('tilt_y_3', 0)
params.add('piston_3', 0)

fit_result = minimize_datasets(params, datasets)

plt.figure()
plt.pcolormesh(*data_transformation(datasets[1], 1, fit_result.params), alpha=0.3, vmin=-0.5, vmax=0)
plt.pcolormesh(*data_transformation(datasets[2], 2, fit_result.params), alpha=0.3, vmin=-0.5, vmax=0)
plt.pcolormesh(*data_transformation(datasets[3], 3, fit_result.params), alpha=0.3, vmin=-0.5, vmax=0)
plt.colorbar()

Как видите, это работает, но сшивание этих небольших наборов данных на моем компьютере занимает около минуты. На самом деле у меня больше и больше наборов данных.

Видите ли вы способ улучшить производительность сшивания?

Редактировать: Как и было предложено, я запустил профилировщик, и он показывает, что 99,5% времени тратится на функцию griddata. Этот используется для интерполяции точек данных из набора данных_0 в местоположения набора данных_1. Если я переключаю метод на ближайший, время выполнения падает примерно до секунды, но интерполяции не происходит. Есть ли шанс улучшить скорость интерполяции?