yt_lin_interp_speeds_3d_4d.py

import numpy as np
import yt.utilities.linear_interpolators as lin
from yt.testing import fake_random_ds
import time
import matplotlib.pyplot as plt

times_3d_eval = []
times_3d_init = []
Nvals_3d = 2 ** np.arange(3, 9)
for n in Nvals_3d:
    random_data = np.random.random((n,)*3)

    Nj = complex(0, n)
    fv = dict(zip("xyz", np.mgrid[0.0:1.0:Nj, 0.0:1.0:Nj, 0.0:1.0:Nj]))

    t0 = time.time()
    bins = np.linspace(0.0, 1.0, n)
    shifts = {ax: (1.0 / n) * np.random.random(n) - (0.5 / n) for ax in "xyz"}
    fv["x"] += shifts["x"][:, np.newaxis, np.newaxis]
    fv["y"] += shifts["y"][:, np.newaxis]
    fv["z"] += shifts["z"]
    tfi = lin.TrilinearFieldInterpolator(
        random_data,
        (bins + shifts["x"], bins + shifts["y"], bins + shifts["z"]),
        "xyz",
        True,
    )
    times_3d_init.append(time.time() - t0)

    t0 = time.time()
    _ = tfi(fv)
    times_3d_eval.append(time.time() - t0)

times_4d_eval = []
times_4d_init = []
Nvals_4d = 2 ** np.arange(3, 7,)
for n in Nvals_4d:
    random_data = np.random.random((n,)*4)

    Nj = complex(0, n)
    fv = dict(zip("xyzw", np.mgrid[0.0:1.0:Nj, 0.0:1.0:Nj, 0.0:1.0:Nj, 0.0:1.0:Nj]))

    t0 = time.time()
    bins = np.linspace(0.0, 1.0, n)
    shifts = {ax: (1.0 / n) * np.random.random(n) - (0.5 / n) for ax in "xyzw"}
    fv["x"] += shifts["x"][:, np.newaxis, np.newaxis, np.newaxis]
    fv["y"] += shifts["y"][:, np.newaxis, np.newaxis]
    fv["z"] += shifts["z"][:, np.newaxis]
    fv["w"] += shifts["w"]
    tfi = lin.QuadrilinearFieldInterpolator(
        random_data,
        (
            bins + shifts["x"],
            bins + shifts["y"],
            bins + shifts["z"],
            bins + shifts["w"],
        ),
        "xyzw",
        True,
    )
    times_4d_init.append(time.time() - t0)

    t0 = time.time()
    _ = tfi(fv)
    times_4d_eval.append(time.time() - t0)


f, axs = plt.subplots(nrows = 1, ncols = 2, figsize=(10,5))
axs[0].loglog(Nvals_3d**3, times_3d_init, marker='.', label='3d')
axs[0].loglog(Nvals_4d**4, times_4d_init,  marker='.', label='4d')
axs[1].loglog(Nvals_3d**3, times_3d_eval,  label='3d')
axs[1].loglog(Nvals_4d**4, times_4d_eval, label='4d')
axs[0].set_ylabel('initialization time [s]')
axs[0].set_xlabel('array size')
axs[1].set_ylabel('evaluation time [s]')
axs[1].set_xlabel('array size')
axs[0].legend()
axs[0].set_title('random-sized bins')
plt.show()