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