def test_transform_vectors_sph_to_cart():
lon = np.arange(0.0, 360.0, 40.0) # longitude
lat = np.arange(0.0, 181.0, 60.0) # colatitude
lev = [1] # elevation (radius)
u, v = np.meshgrid(lon, lat, indexing="ij")
w = u**2 - v**2
uu, vv, ww = pyvista.transform_vectors_sph_to_cart(lon, lat, lev, u, v, w)
assert np.allclose(
[uu[-1, -1], vv[-1, -1], ww[-1, -1]],
[67.80403533828323, 360.8359915416445, -70000.0],
)
wind_level = [RADIUS * 1.2]
# Sequence of axis indices for transpose()
# (1, 0) for 2D arrays
# (2, 1, 0) for 3D arrays
inv_axes = [*range(u_vec.ndim)[::-1]]
# Transform vectors to cartesian coordinates
vectors = np.stack(
[
i.transpose(inv_axes).swapaxes(-2, -1).ravel("C")
for i in pv.transform_vectors_sph_to_cart(
x,
y_polar,
wind_level,
u_vec.transpose(inv_axes),
-v_vec.transpose(
inv_axes
), # Minus sign because y-vector in polar coords is required
w_vec.transpose(inv_axes),
)
],
axis=1,
)
# Scale vectors to make them visible
vectors *= RADIUS * 0.1
# Create a grid for the vectors
grid_winds = pv.grid_from_sph_coords(x, y_polar, wind_level)
# Add vectors to the grid
