How to use the pykrige.uk.UniversalKriging function in PyKrige

def test_uk_with_external_drift(validation_ref):

    data, _, (uk_test_answer, gridx_ref, gridy_ref) = validation_ref

    dem, demx, demy, cellsize, no_data = kt.read_asc_grid(
        os.path.join(BASE_DIR, "test_data/test3_dem.asc")
    )
    uk = UniversalKriging(
        data[:, 0],
        data[:, 1],
        data[:, 2],
        variogram_model="spherical",
        variogram_parameters=[500.0, 3000.0, 0.0],
        anisotropy_scaling=1.0,
        anisotropy_angle=0.0,
        drift_terms=["external_Z"],
        external_drift=dem,
        external_drift_x=demx,
        external_drift_y=demy,
        verbose=False,
    )
    answer, gridx, gridy, cellsize, no_data = kt.read_asc_grid(
        os.path.join(BASE_DIR, "test_data/test3_answer.asc")
    )

def test_custom_variogram(sample_data_2d):
    data, (gridx, gridy, gridx_2), mask_ref = sample_data_2d

    def func(params, dist):
        return params[0] * np.log10(dist + params[1]) + params[2]

    with pytest.raises(ValueError):
        UniversalKriging(data[:, 0], data[:, 1], data[:, 2], variogram_model="mrow")
    with pytest.raises(ValueError):
        UniversalKriging(data[:, 0], data[:, 1], data[:, 2], variogram_model="custom")
    with pytest.raises(ValueError):
        UniversalKriging(
            data[:, 0],
            data[:, 1],
            data[:, 2],
            variogram_model="custom",
            variogram_function=0,
        )
    with pytest.raises(ValueError):
        UniversalKriging(
            data[:, 0],
            data[:, 1],
            data[:, 2],
            variogram_model="custom",

variogram_model="linear",
            variogram_parameters=[1.0, 0.0],
            drift_terms=["external_Z"],
        )
    with pytest.raises(ValueError):
        UniversalKriging(
            data[:, 0],
            data[:, 1],
            data[:, 2],
            variogram_model="linear",
            variogram_parameters=[1.0, 0.0],
            drift_terms=["external_Z"],
            external_drift=dem,
        )
    with pytest.raises(ValueError):
        UniversalKriging(
            data[:, 0],
            data[:, 1],
            data[:, 2],
            variogram_model="linear",
            variogram_parameters=[1.0, 0.0],
            drift_terms=["external_Z"],
            external_drift=dem,
            external_drift_x=dem_x,
            external_drift_y=np.arange(0.0, 5.0, 1.0),
        )

    uk = UniversalKriging(
        data[:, 0],
        data[:, 1],
        data[:, 2],
        variogram_model="linear",

def test_uk_execute(sample_data_2d):

    data, (gridx, gridy, _), mask_ref = sample_data_2d

    uk = UniversalKriging(
        data[:, 0],
        data[:, 1],
        data[:, 2],
        variogram_model="linear",
        drift_terms=["regional_linear"],
    )

    with pytest.raises(ValueError):
        uk.execute("blurg", gridx, gridy)
    with pytest.raises(ValueError):
        uk.execute("grid", gridx, gridy, backend="mrow")

    z, ss = uk.execute("grid", gridx, gridy, backend="vectorized")
    shape = (gridy.size, gridx.size)
    assert z.shape == shape
    assert ss.shape == shape

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

    uk_spec = UniversalKriging(
        data[:, 0],
        data[:, 1],
        data[:, 2],
        variogram_model="linear",
        drift_terms=["specified"],
        specified_drift=[data[:, 0], data[:, 1]],
    )
    with pytest.raises(ValueError):
        uk_spec.execute("grid", gridx, gridy, specified_drift_arrays=[gridx, gridy])
    with pytest.raises(TypeError):
        uk_spec.execute("grid", gridx, gridy, specified_drift_arrays=gridx)
    with pytest.raises(ValueError):
        uk_spec.execute("grid", gridx, gridy, specified_drift_arrays=[xg])
    z_spec, ss_spec = uk_spec.execute(
        "grid", gridx, gridy, specified_drift_arrays=[xg, yg]
    )

data, (gridx, gridy, _), mask_ref = sample_data_2d

    with pytest.raises(ValueError):
        UniversalKriging(data[:, 0], data[:, 1], data[:, 2], variogram_model="blurg")
    with pytest.raises(ValueError):
        UniversalKriging(data[:, 0], data[:, 1], data[:, 2], drift_terms=["external_Z"])
    with pytest.raises(ValueError):
        UniversalKriging(
            data[:, 0],
            data[:, 1],
            data[:, 2],
            drift_terms=["external_Z"],
            external_drift=np.array([0]),
        )
    with pytest.raises(ValueError):
        UniversalKriging(data[:, 0], data[:, 1], data[:, 2], drift_terms=["point_log"])

    uk = UniversalKriging(data[:, 0], data[:, 1], data[:, 2])
    variogram_model = uk.variogram_model
    variogram_parameters = uk.variogram_model_parameters
    anisotropy_scaling = uk.anisotropy_scaling
    anisotropy_angle = uk.anisotropy_angle

    with pytest.raises(ValueError):
        uk.update_variogram_model("blurg")
    uk.update_variogram_model("power", anisotropy_scaling=3.0, anisotropy_angle=45.0)
    # TODO: check that the new parameters are equal to the expected ones
    assert variogram_model != uk.variogram_model
    assert not np.array_equal(variogram_parameters, uk.variogram_model_parameters)
    assert anisotropy_scaling != uk.anisotropy_scaling
    assert anisotropy_angle != uk.anisotropy_angle

def test_uk_calculate_data_point_zscalars(sample_data_2d):

    data, (gridx, gridy, _), mask_ref = sample_data_2d

    dem = np.arange(0.0, 5.1, 0.1)
    dem = np.repeat(dem[np.newaxis, :], 6, axis=0)
    dem_x = np.arange(0.0, 5.1, 0.1)
    dem_y = np.arange(0.0, 6.0, 1.0)

    with pytest.raises(ValueError):
        UniversalKriging(
            data[:, 0],
            data[:, 1],
            data[:, 2],
            variogram_model="linear",
            variogram_parameters=[1.0, 0.0],
            drift_terms=["external_Z"],
        )
    with pytest.raises(ValueError):
        UniversalKriging(
            data[:, 0],
            data[:, 1],
            data[:, 2],
            variogram_model="linear",
            variogram_parameters=[1.0, 0.0],
            drift_terms=["external_Z"],
            external_drift=dem,

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

    uk_spec = UniversalKriging(
        data[:, 0],
        data[:, 1],
        data[:, 2],
        variogram_model="linear",
        drift_terms=["specified"],
        specified_drift=[data[:, 0], data[:, 1]],
    )

data = np.array(
        [
            [61.0, 139.0, 477.0],
            [63.0, 140.0, 696.0],
            [64.0, 129.0, 227.0],
            [68.0, 128.0, 646.0],
            [71.0, 140.0, 606.0],
            [73.0, 141.0, 791.0],
            [75.0, 128.0, 783.0],
        ]
    )
    point = (65.0, 137.0)
    z_answer = 567.54
    ss_answer = 9.044

    uk = UniversalKriging(
        data[:, 0],
        data[:, 1],
        data[:, 2],
        variogram_model="exponential",
        variogram_parameters=[10.0, 9.99, 0.0],
        drift_terms=["regional_linear"],
    )
    z, ss = uk.execute(
        "points", np.array([point[0]]), np.array([point[1]]), backend="vectorized"
    )
    assert z_answer == approx(z[0], rel=0.1)
    assert ss_answer == approx(ss[0], rel=0.1)

    z, ss = uk.execute(
        "points", np.array([61.0]), np.array([139.0]), backend="vectorized"
    )

def test_ok_uk_produce_same_result(validation_ref):

    data, _, (uk_test_answer, gridx_ref, gridy_ref) = validation_ref

    gridx = np.linspace(1067000.0, 1072000.0, 100)
    gridy = np.linspace(241500.0, 244000.0, 100)
    ok = OrdinaryKriging(
        data[:, 0],
        data[:, 1],
        data[:, 2],
        variogram_model="linear",
        verbose=False,
        enable_plotting=False,
    )
    z_ok, ss_ok = ok.execute("grid", gridx, gridy, backend="vectorized")
    uk = UniversalKriging(
        data[:, 0],
        data[:, 1],
        data[:, 2],
        variogram_model="linear",
        verbose=False,
        enable_plotting=False,
    )
    z_uk, ss_uk = uk.execute("grid", gridx, gridy, backend="vectorized")
    assert_allclose(z_ok, z_uk)
    assert_allclose(ss_ok, ss_uk)

    z_ok, ss_ok = ok.execute("grid", gridx, gridy, backend="loop")
    z_uk, ss_uk = uk.execute("grid", gridx, gridy, backend="loop")
    assert_allclose(z_ok, z_uk)
    assert_allclose(ss_ok, ss_uk)