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