How to use the pykrige.uk3d.UniversalKriging3D function in PyKrige

def test_uk3d_update_variogram_model(sample_data_3d):

    data, (gridx_ref, gridy_ref, gridz_ref), mask_ref = sample_data_3d

    with pytest.raises(ValueError):
        UniversalKriging3D(
            data[:, 0], data[:, 1], data[:, 2], data[:, 3], variogram_model="blurg"
        )

    uk3d = UniversalKriging3D(data[:, 0], data[:, 1], data[:, 2], data[:, 3])
    variogram_model = uk3d.variogram_model
    variogram_parameters = uk3d.variogram_model_parameters
    anisotropy_scaling_y = uk3d.anisotropy_scaling_y
    anisotropy_scaling_z = uk3d.anisotropy_scaling_z
    anisotropy_angle_x = uk3d.anisotropy_angle_x
    anisotropy_angle_y = uk3d.anisotropy_angle_y
    anisotropy_angle_z = uk3d.anisotropy_angle_z

    with pytest.raises(ValueError):
        uk3d.update_variogram_model("blurg")
    uk3d.update_variogram_model(
        "power",

)
    k, ss = uk3d.execute(
        "points", gridx_ref[0], gridy_ref[0], gridz_ref[0], backend="loop"
    )
    assert k.shape == (1,)
    assert ss.shape == (1,)

    data = np.zeros((125, 4))
    z, y, x = np.meshgrid(
        np.arange(0.0, 5.0, 1.0), np.arange(0.0, 5.0, 1.0), np.arange(0.0, 5.0, 1.0)
    )
    data[:, 0] = np.ravel(x)
    data[:, 1] = np.ravel(y)
    data[:, 2] = np.ravel(z)
    data[:, 3] = np.ravel(z)
    k3d = UniversalKriging3D(
        data[:, 0], data[:, 1], data[:, 2], data[:, 3], variogram_model="linear"
    )
    k, ss = k3d.execute(
        "grid",
        np.arange(2.0, 3.0, 0.1),
        np.arange(2.0, 3.0, 0.1),
        np.arange(0.0, 4.0, 1.0),
        backend="vectorized",
    )
    assert_allclose(k[0, :, :], 0.0, atol=0.01)
    assert_allclose(k[1, :, :], 1.0, rtol=1.0e-2)
    assert_allclose(k[2, :, :], 2.0, rtol=1.0e-2)
    assert_allclose(k[3, :, :], 3.0, rtol=1.0e-2)
    k, ss = k3d.execute(
        "grid",
        np.arange(2.0, 3.0, 0.1),

ok3d = OrdinaryKriging3D(
        data[:, 0],
        data[:, 1],
        np.zeros(data[:, 1].shape),
        data[:, 2],
        variogram_model="exponential",
        variogram_parameters=[500.0, 3000.0, 0.0],
    )
    ok_v, oss_v = ok3d.execute(
        "grid", gridx_ref, gridy_ref, np.array([0.0]), backend="vectorized"
    )
    ok_l, oss_l = ok3d.execute(
        "grid", gridx_ref, gridy_ref, np.array([0.0]), backend="loop"
    )

    uk3d = UniversalKriging3D(
        data[:, 0],
        data[:, 1],
        np.zeros(data[:, 1].shape),
        data[:, 2],
        variogram_model="exponential",
        variogram_parameters=[500.0, 3000.0, 0.0],
    )
    uk_v, uss_v = uk3d.execute(
        "grid", gridx_ref, gridy_ref, np.array([0.0]), backend="vectorized"
    )
    assert_allclose(uk_v, ok_v)
    uk_l, uss_l = uk3d.execute(
        "grid", gridx_ref, gridy_ref, np.array([0.0]), backend="loop"
    )
    assert_allclose(uk_l, ok_l)
    assert_allclose(uk_l, uk_v)

data[:, 3],
            variogram_model="linear",
            drift_terms=["specified"],
        )
    with pytest.raises(TypeError):
        UniversalKriging3D(
            data[:, 0],
            data[:, 1],
            data[:, 2],
            data[:, 3],
            variogram_model="linear",
            drift_terms=["specified"],
            specified_drift=data[:, 0],
        )
    with pytest.raises(ValueError):
        UniversalKriging3D(
            data[:, 0],
            data[:, 1],
            data[:, 2],
            data[:, 3],
            variogram_model="linear",
            drift_terms=["specified"],
            specified_drift=[data[:2, 0]],
        )

    uk_spec = UniversalKriging3D(
        data[:, 0],
        data[:, 1],
        data[:, 2],
        data[:, 3],
        variogram_model="linear",
        drift_terms=["specified"],

def test_uk3d_execute(sample_data_3d):

    data, (gridx_ref, gridy_ref, gridz_ref), mask_ref = sample_data_3d

    uk3d = UniversalKriging3D(data[:, 0], data[:, 1], data[:, 2], data[:, 3])

    with pytest.raises(ValueError):
        uk3d.execute("blurg", gridx_ref, gridy_ref, gridz_ref)

    k, ss = uk3d.execute("grid", gridx_ref, gridy_ref, gridz_ref, backend="vectorized")
    shape = (gridz_ref.size, gridy_ref.size, gridx_ref.size)
    assert k.shape == shape
    assert ss.shape == shape
    assert np.amax(k) != np.amin(k)
    assert np.amax(ss) != np.amin(ss)
    assert not np.ma.is_masked(k)

    k, ss = uk3d.execute("grid", gridx_ref, gridy_ref, gridz_ref, backend="loop")
    shape = (gridz_ref.size, gridy_ref.size, gridx_ref.size)
    assert k.shape == shape
    assert ss.shape == shape

func_x = lambda x, y, z: x  # noqa
    func_y = lambda x, y, z: y  # noqa
    func_z = lambda x, y, z: z  # noqa

    with pytest.raises(ValueError):
        UniversalKriging3D(
            data[:, 0],
            data[:, 1],
            data[:, 2],
            data[:, 3],
            variogram_model="linear",
            drift_terms=["functional"],
        )
    with pytest.raises(TypeError):
        UniversalKriging3D(
            data[:, 0],
            data[:, 1],
            data[:, 2],
            data[:, 3],
            variogram_model="linear",
            drift_terms=["functional"],
            functional_drift=func_x,
        )

    uk_func = UniversalKriging3D(
        data[:, 0],
        data[:, 1],
        data[:, 2],
        data[:, 3],
        variogram_model="linear",
        drift_terms=["functional"],

data[:, 3],
            variogram_model="linear",
            drift_terms=["functional"],
        )
    with pytest.raises(TypeError):
        UniversalKriging3D(
            data[:, 0],
            data[:, 1],
            data[:, 2],
            data[:, 3],
            variogram_model="linear",
            drift_terms=["functional"],
            functional_drift=func_x,
        )

    uk_func = UniversalKriging3D(
        data[:, 0],
        data[:, 1],
        data[:, 2],
        data[:, 3],
        variogram_model="linear",
        drift_terms=["functional"],
        functional_drift=[func_x, func_y, func_z],
    )
    z_func, ss_func = uk_func.execute("grid", gridx, gridy, gridz)
    uk_lin = UniversalKriging3D(
        data[:, 0],
        data[:, 1],
        data[:, 2],
        data[:, 3],
        variogram_model="linear",
        drift_terms=["regional_linear"],

ok3d = OrdinaryKriging3D(
        data[:, 0],
        data[:, 1],
        data[:, 2],
        data[:, 3],
        variogram_model="linear",
        variogram_parameters=[1.0, 0.1],
    )
    ok_v, oss_v = ok3d.execute(
        "grid", np.arange(10.0), np.arange(10.0), np.arange(10.0), backend="vectorized"
    )
    ok_l, oss_l = ok3d.execute(
        "grid", np.arange(10.0), np.arange(10.0), np.arange(10.0), backend="loop"
    )

    uk3d = UniversalKriging3D(
        data[:, 0],
        data[:, 1],
        data[:, 2],
        data[:, 3],
        variogram_model="linear",
        variogram_parameters=[1.0, 0.1],
    )
    uk_v, uss_v = uk3d.execute(
        "grid", np.arange(10.0), np.arange(10.0), np.arange(10.0), backend="vectorized"
    )
    assert_allclose(uk_v, ok_v)
    assert_allclose(uss_v, oss_v)
    uk_l, uss_l = uk3d.execute(
        "grid", np.arange(10.0), np.arange(10.0), np.arange(10.0), backend="loop"
    )
    assert_allclose(uk_l, ok_l)

gridz_ref,
            specified_drift_arrays=[gridx_ref, gridy_ref, gridz_ref],
        )
    with pytest.raises(TypeError):
        uk_spec.execute(
            "grid", gridx_ref, gridy_ref, gridz_ref, specified_drift_arrays=gridx_ref
        )
    with pytest.raises(ValueError):
        uk_spec.execute(
            "grid", gridx_ref, gridy_ref, gridz_ref, specified_drift_arrays=[zg]
        )
    z_spec, ss_spec = uk_spec.execute(
        "grid", gridx_ref, gridy_ref, gridz_ref, specified_drift_arrays=[xg, yg, zg]
    )

    uk_lin = UniversalKriging3D(
        data[:, 0],
        data[:, 1],
        data[:, 2],
        data[:, 3],
        variogram_model="linear",
        drift_terms=["regional_linear"],
    )
    z_lin, ss_lin = uk_lin.execute("grid", gridx_ref, gridy_ref, gridz_ref)

    assert_allclose(z_spec, z_lin)
    assert_allclose(ss_spec, ss_lin)

from pykrige.compat import validate_sklearn

validate_sklearn()
from pykrige.ok import OrdinaryKriging
from pykrige.uk import UniversalKriging
from pykrige.ok3d import OrdinaryKriging3D
from pykrige.uk3d import UniversalKriging3D
from sklearn.base import RegressorMixin, BaseEstimator
from sklearn.svm import SVR
from sklearn.metrics import r2_score

krige_methods = {
    "ordinary": OrdinaryKriging,
    "universal": UniversalKriging,
    "ordinary3d": OrdinaryKriging3D,
    "universal3d": UniversalKriging3D,
}

threed_krige = ("ordinary3d", "universal3d")


def validate_method(method):
    if method not in krige_methods.keys():
        raise ValueError(
            "Kriging method must be one of {}".format(krige_methods.keys())
        )


class Krige(RegressorMixin, BaseEstimator):
    """
    A scikit-learn wrapper class for Ordinary and Universal Kriging.
    This works with both Grid/RandomSearchCv for finding the best