1import warnings
2
3import numpy as np
4
5import yt.utilities.linear_interpolators as lin
6from yt.testing import assert_array_almost_equal, assert_array_equal, fake_random_ds
7from yt.utilities.lib.interpolators import ghost_zone_interpolate
8
9
10def setup():
11    pass
12
13
14def test_linear_interpolator_1d():
15    random_data = np.random.random(64)
16    fv = {"x": np.mgrid[0.0:1.0:64j]}
17    # evenly spaced bins
18    ufi = lin.UnilinearFieldInterpolator(random_data, (0.0, 1.0), "x", True)
19    assert_array_equal(ufi(fv), random_data)
20
21    # randomly spaced bins
22    size = 64
23    shift = (1.0 / size) * np.random.random(size) - (0.5 / size)
24    fv["x"] += shift
25    ufi = lin.UnilinearFieldInterpolator(
26        random_data, np.linspace(0.0, 1.0, size) + shift, "x", True
27    )
28    assert_array_almost_equal(ufi(fv), random_data, 15)
29
30
31def test_linear_interpolator_2d():
32    random_data = np.random.random((64, 64))
33    # evenly spaced bins
34    fv = {ax: v for ax, v in zip("xyz", np.mgrid[0.0:1.0:64j, 0.0:1.0:64j])}
35    bfi = lin.BilinearFieldInterpolator(random_data, (0.0, 1.0, 0.0, 1.0), "xy", True)
36    assert_array_equal(bfi(fv), random_data)
37
38    # randomly spaced bins
39    size = 64
40    bins = np.linspace(0.0, 1.0, size)
41    shifts = {ax: (1.0 / size) * np.random.random(size) - (0.5 / size) for ax in "xy"}
42    fv["x"] += shifts["x"][:, np.newaxis]
43    fv["y"] += shifts["y"]
44    bfi = lin.BilinearFieldInterpolator(
45        random_data, (bins + shifts["x"], bins + shifts["y"]), "xy", True
46    )
47    assert_array_almost_equal(bfi(fv), random_data, 15)
48
49
50def test_linear_interpolator_3d():
51    random_data = np.random.random((64, 64, 64))
52    # evenly spaced bins
53    fv = {
54        ax: v for ax, v in zip("xyz", np.mgrid[0.0:1.0:64j, 0.0:1.0:64j, 0.0:1.0:64j])
55    }
56    tfi = lin.TrilinearFieldInterpolator(
57        random_data, (0.0, 1.0, 0.0, 1.0, 0.0, 1.0), "xyz", True
58    )
59    assert_array_almost_equal(tfi(fv), random_data)
60
61    # randomly spaced bins
62    size = 64
63    bins = np.linspace(0.0, 1.0, size)
64    shifts = {ax: (1.0 / size) * np.random.random(size) - (0.5 / size) for ax in "xyz"}
65    fv["x"] += shifts["x"][:, np.newaxis, np.newaxis]
66    fv["y"] += shifts["y"][:, np.newaxis]
67    fv["z"] += shifts["z"]
68    tfi = lin.TrilinearFieldInterpolator(
69        random_data,
70        (bins + shifts["x"], bins + shifts["y"], bins + shifts["z"]),
71        "xyz",
72        True,
73    )
74    assert_array_almost_equal(tfi(fv), random_data, 15)
75
76
77def test_ghost_zone_extrapolation():
78    ds = fake_random_ds(16)
79
80    g = ds.index.grids[0]
81    vec = g.get_vertex_centered_data(
82        [("index", "x"), ("index", "y"), ("index", "z")], no_ghost=True
83    )
84    for i, ax in enumerate("xyz"):
85        xc = g[("index", ax)]
86
87        tf = lin.TrilinearFieldInterpolator(
88            xc,
89            (
90                g.LeftEdge[0] + g.dds[0] / 2.0,
91                g.RightEdge[0] - g.dds[0] / 2.0,
92                g.LeftEdge[1] + g.dds[1] / 2.0,
93                g.RightEdge[1] - g.dds[1] / 2.0,
94                g.LeftEdge[2] + g.dds[2] / 2.0,
95                g.RightEdge[2] - g.dds[2] / 2.0,
96            ),
97            ["x", "y", "z"],
98            truncate=True,
99        )
100
101        lx, ly, lz = np.mgrid[
102            g.LeftEdge[0] : g.RightEdge[0] : (g.ActiveDimensions[0] + 1) * 1j,
103            g.LeftEdge[1] : g.RightEdge[1] : (g.ActiveDimensions[1] + 1) * 1j,
104            g.LeftEdge[2] : g.RightEdge[2] : (g.ActiveDimensions[2] + 1) * 1j,
105        ]
106        xi = tf({"x": lx, "y": ly, "z": lz})
107
108        xz = np.zeros(g.ActiveDimensions + 1)
109        ghost_zone_interpolate(
110            1,
111            xc,
112            np.array([0.5, 0.5, 0.5], dtype="f8"),
113            xz,
114            np.array([0.0, 0.0, 0.0], dtype="f8"),
115        )
116
117        ii = (lx, ly, lz)[i]
118        assert_array_equal(ii, vec[("index", ax)])
119        assert_array_equal(ii, xi)
120        assert_array_equal(ii, xz)
121
122
123def test_get_vertex_centered_data():
124    ds = fake_random_ds(16)
125    g = ds.index.grids[0]
126
127    vec_list = g.get_vertex_centered_data([("gas", "density")], no_ghost=True)
128    with warnings.catch_warnings(record=True) as w:
129        warnings.simplefilter("always")
130        vec_str = g.get_vertex_centered_data(("gas", "density"), no_ghost=True)
131        assert len(w) == 1
132        assert issubclass(w[-1].category, DeprecationWarning)
133        assert "requires list of fields" in str(w[-1].message)
134    with warnings.catch_warnings(record=True) as w:
135        warnings.simplefilter("always")
136        vec_tuple = g.get_vertex_centered_data(("gas", "density"), no_ghost=True)
137        assert len(w) == 1
138        assert issubclass(w[-1].category, DeprecationWarning)
139        assert (
140            "get_vertex_centered_data() requires list of fields, rather than "
141            "a single field as an argument."
142        ) in str(w[-1].message)
143    assert_array_equal(vec_list[("gas", "density")], vec_str)
144    assert_array_equal(vec_list[("gas", "density")], vec_tuple)
145