1import numpy as np
2
3from yt.testing import assert_allclose_units, fake_random_ds, requires_file
4from yt.units import cm, s
5from yt.utilities.answer_testing.framework import data_dir_load
6from yt.visualization.volume_rendering.off_axis_projection import off_axis_projection
7
8
9def random_unit_vector(prng):
10    v = prng.random_sample(3)
11    while (v == 0).all():
12        v = prng.random_sample(3)
13    return v / np.sqrt((v ** 2).sum())
14
15
16def random_velocity_vector(prng):
17    return 2e5 * prng.random_sample(3) - 1e5
18
19
20def compare_vector_conversions(data_source):
21    prng = np.random.RandomState(8675309)
22    normals = [[1, 0, 0], [0, 1, 0], [0, 0, 1]] + [
23        random_unit_vector(prng) for i in range(2)
24    ]
25    bulk_velocities = [random_velocity_vector(prng) for i in range(2)]
26
27    for bv in bulk_velocities:
28        bulk_velocity = bv * cm / s
29        data_source.set_field_parameter("bulk_velocity", bulk_velocity)
30        data_source.clear_data()
31
32        vmag = data_source[("gas", "velocity_magnitude")]
33        vrad = data_source[("gas", "velocity_spherical_radius")]
34
35        for normal in normals:
36            data_source.set_field_parameter("normal", normal)
37            data_source.clear_data()
38
39            assert_allclose_units(
40                vrad, data_source[("gas", "velocity_spherical_radius")]
41            )
42
43            vmag_new = data_source[("gas", "velocity_magnitude")]
44            assert_allclose_units(vmag, vmag_new)
45
46            vmag_cart = np.sqrt(
47                (data_source[("gas", "velocity_x")] - bulk_velocity[0]) ** 2
48                + (data_source[("gas", "velocity_y")] - bulk_velocity[1]) ** 2
49                + (data_source[("gas", "velocity_z")] - bulk_velocity[2]) ** 2
50            )
51            assert_allclose_units(vmag, vmag_cart)
52
53            vmag_cyl = np.sqrt(
54                data_source[("gas", "velocity_cylindrical_radius")] ** 2
55                + data_source[("gas", "velocity_cylindrical_theta")] ** 2
56                + data_source[("gas", "velocity_cylindrical_z")] ** 2
57            )
58            assert_allclose_units(vmag, vmag_cyl)
59
60            vmag_sph = np.sqrt(
61                data_source[("gas", "velocity_spherical_radius")] ** 2
62                + data_source[("gas", "velocity_spherical_theta")] ** 2
63                + data_source[("gas", "velocity_spherical_phi")] ** 2
64            )
65            assert_allclose_units(vmag, vmag_sph)
66
67            for i, d in enumerate("xyz"):
68                assert_allclose_units(
69                    data_source[("gas", f"velocity_{d}")] - bulk_velocity[i],
70                    data_source[("gas", f"relative_velocity_{d}")],
71                )
72
73        for i, ax in enumerate("xyz"):
74            data_source.set_field_parameter("axis", i)
75            data_source.clear_data()
76            assert_allclose_units(
77                data_source[("gas", "velocity_los")],
78                data_source[("gas", f"relative_velocity_{ax}")],
79            )
80
81        for i, ax in enumerate("xyz"):
82            prj = data_source.ds.proj(
83                ("gas", "velocity_los"), i, weight_field=("gas", "density")
84            )
85            assert_allclose_units(
86                prj[("gas", "velocity_los")], prj[("gas", f"velocity_{ax}")]
87            )
88
89        data_source.clear_data()
90        ax = [0.1, 0.2, -0.3]
91        data_source.set_field_parameter("axis", ax)
92        ax /= np.sqrt(np.dot(ax, ax))
93        vlos = data_source[("gas", "relative_velocity_x")] * ax[0]
94        vlos += data_source[("gas", "relative_velocity_y")] * ax[1]
95        vlos += data_source[("gas", "relative_velocity_z")] * ax[2]
96        assert_allclose_units(data_source[("gas", "velocity_los")], vlos)
97
98        buf_los = off_axis_projection(
99            data_source,
100            data_source.ds.domain_center,
101            ax,
102            0.5,
103            128,
104            ("gas", "velocity_los"),
105            weight=("gas", "density"),
106        )
107
108        buf_x = off_axis_projection(
109            data_source,
110            data_source.ds.domain_center,
111            ax,
112            0.5,
113            128,
114            ("gas", "relative_velocity_x"),
115            weight=("gas", "density"),
116        )
117
118        buf_y = off_axis_projection(
119            data_source,
120            data_source.ds.domain_center,
121            ax,
122            0.5,
123            128,
124            ("gas", "relative_velocity_y"),
125            weight=("gas", "density"),
126        )
127
128        buf_z = off_axis_projection(
129            data_source,
130            data_source.ds.domain_center,
131            ax,
132            0.5,
133            128,
134            ("gas", "relative_velocity_z"),
135            weight=("gas", "density"),
136        )
137
138        vlos = buf_x * ax[0] + buf_y * ax[1] + buf_z * ax[2]
139
140        assert_allclose_units(buf_los, vlos, rtol=1.0e-6)
141
142
143def test_vector_component_conversions_fake():
144    ds = fake_random_ds(16)
145    ad = ds.all_data()
146    compare_vector_conversions(ad)
147
148
149g30 = "IsolatedGalaxy/galaxy0030/galaxy0030"
150
151
152@requires_file(g30)
153def test_vector_component_conversions_real():
154    ds = data_dir_load(g30)
155    sp = ds.sphere(ds.domain_center, (10, "kpc"))
156    compare_vector_conversions(sp)
157