1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
|
"""
Plot data in spherical coordinates
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Generate and visualize meshes from data in longitude-latitude coordinates.
"""
from __future__ import annotations
import numpy as np
import pyvista as pv
def _cell_bounds(points, bound_position=0.5):
"""
Calculate coordinate cell boundaries.
Parameters
----------
points: numpy.ndarray
One-dimensional array of uniformly spaced values of shape (M,).
bound_position: bool, optional
The desired position of the bounds relative to the position
of the points.
Returns
-------
bounds: numpy.ndarray
Array of shape (M+1,)
Examples
--------
>>> a = np.arange(-1, 2.5, 0.5)
>>> a
array([-1. , -0.5, 0. , 0.5, 1. , 1.5, 2. ])
>>> cell_bounds(a)
array([-1.25, -0.75, -0.25, 0.25, 0.75, 1.25, 1.75, 2.25])
"""
if points.ndim != 1:
raise ValueError("Only 1D points are allowed.")
diffs = np.diff(points)
delta = diffs[0] * bound_position
return np.concatenate([[points[0] - delta], points + delta])
# Seed random number generator for reproducible plots
rng = np.random.default_rng(seed=0)
# First, create some dummy data
# Approximate radius of the Earth
RADIUS = 6371.0
# Longitudes and latitudes
x = np.arange(0, 360, 5)
y = np.arange(-90, 91, 10)
y_polar = 90.0 - y # grid_from_sph_coords() expects polar angle
xx, yy = np.meshgrid(x, y)
# x- and y-components of the wind vector
u_vec = np.cos(np.radians(xx)) # zonal
v_vec = np.sin(np.radians(yy)) # meridional
# Scalar data
scalar = u_vec**2 + v_vec**2
# Create arrays of grid cell boundaries, which have shape of (x.shape[0] + 1)
xx_bounds = _cell_bounds(x)
yy_bounds = _cell_bounds(y_polar)
# Vertical levels
# in this case a single level slightly above the surface of a sphere
levels = [RADIUS * 1.01]
# %%
# Create a structured grid
grid_scalar = pv.grid_from_sph_coords(xx_bounds, yy_bounds, levels)
# And fill its cell arrays with the scalar data
grid_scalar.cell_data["example"] = np.array(scalar).swapaxes(-2, -1).ravel("C")
# Make a plot
p = pv.Plotter()
p.add_mesh(pv.Sphere(radius=RADIUS))
p.add_mesh(grid_scalar, clim=[0.1, 2.0], opacity=0.5, cmap="plasma")
p.show()
# %%
# Visualize vectors in spherical coordinates
# Vertical wind
w_vec = rng.random(u_vec.shape)
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
grid_winds.point_data["example"] = vectors
# Show the result
p = pv.Plotter()
p.add_mesh(pv.Sphere(radius=RADIUS))
p.add_mesh(grid_winds.glyph(orient="example", scale="example", tolerance=0.005))
p.show()
# %%
# Isurfaces of 3D data in spherical coordinates
# Number of vertical levels
nlev = 10
# Dummy 3D scalar data
scalar_3d = (
scalar.repeat(nlev).reshape((*scalar.shape, nlev)) * np.arange(nlev)[np.newaxis, np.newaxis, :]
).transpose(2, 0, 1)
z_scale = 10
z_offset = RADIUS * 1.1
# Now it's not a single level but an array of levels
levels = z_scale * (np.arange(scalar_3d.shape[0] + 1)) ** 2 + z_offset
# Create a structured grid by transforming coordinates
grid_scalar_3d = pv.grid_from_sph_coords(xx_bounds, yy_bounds, levels)
# Add data to the grid
grid_scalar_3d.cell_data["example"] = np.array(scalar_3d).swapaxes(-2, -1).ravel("C")
# Create a set of isosurfaces
surfaces = grid_scalar_3d.cell_data_to_point_data().contour(isosurfaces=[1, 5, 10, 15])
# Show the result
p = pv.Plotter()
p.add_mesh(pv.Sphere(radius=RADIUS))
p.add_mesh(surfaces)
p.show()
|
