-
Notifications
You must be signed in to change notification settings - Fork 359
/
vector_transform.py
177 lines (155 loc) · 6.35 KB
/
vector_transform.py
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
168
169
170
171
172
173
174
175
176
177
# Copyright Crown and Cartopy Contributors
#
# This file is part of Cartopy and is released under the BSD 3-clause license.
# See LICENSE in the root of the repository for full licensing details.
"""
Generic functionality to support Cartopy vector transforms.
"""
import numpy as np
try:
from scipy.interpolate import griddata
except ImportError as e:
raise ImportError("Regridding vectors requires scipy.") from e
def _interpolate_to_grid(nx, ny, x, y, *scalars, **kwargs):
"""
Interpolate two vector components and zero or more scalar fields,
which can be irregular, to a regular grid.
Parameters
----------
nx
Number of points at which to interpolate in x direction.
ny
Number of points at which to interpolate in y direction.
x
Array of source points in x direction.
y
Array of source points in y direction.
Other Parameters
----------------
scalars
Zero or more scalar fields to regrid along with the vector
components.
target_extent
The extent in the target CRS that the grid should occupy, in the
form ``(x-lower, x-upper, y-lower, y-upper)``. Defaults to cover
the full extent of the vector field.
"""
target_extent = kwargs.get('target_extent', None)
if target_extent is None:
target_extent = (x.min(), x.max(), y.min(), y.max())
x0, x1, y0, y1 = target_extent
xr = x1 - x0
yr = y1 - y0
points = np.column_stack([(x.ravel() - x0) / xr, (y.ravel() - y0) / yr])
x_grid, y_grid = np.meshgrid(np.linspace(0, 1, nx),
np.linspace(0, 1, ny))
s_grid_tuple = tuple()
for s in scalars:
s_grid_tuple += (griddata(points, s.ravel(), (x_grid, y_grid),
method='linear'),)
return (x_grid * xr + x0, y_grid * yr + y0) + s_grid_tuple
def vector_scalar_to_grid(src_crs, target_proj, regrid_shape, x, y, u, v,
*scalars, **kwargs):
"""
Transform and interpolate a vector field to a regular grid in the
target projection.
Parameters
----------
src_crs
The :class:`~cartopy.crs.CRS` that represents the coordinate
system the vectors are defined in.
target_proj
The :class:`~cartopy.crs.Projection` that represents the
projection the vectors are to be transformed to.
regrid_shape
The regular grid dimensions. If a single integer then the grid
will have that number of points in the x and y directions. A
2-tuple of integers specify the size of the regular grid in the
x and y directions respectively.
x, y
The x and y coordinates, in the source CRS coordinates,
where the vector components are located.
u, v
The grid eastward and grid northward components of the
vector field respectively. Their shapes must match.
Other Parameters
----------------
scalars
Zero or more scalar fields to regrid along with the vector
components. Each scalar field must have the same shape as the
vector components.
target_extent
The extent in the target CRS that the grid should occupy, in the
form ``(x-lower, x-upper, y-lower, y-upper)``. Defaults to cover
the full extent of the vector field.
Returns
-------
x_grid, y_grid
The x and y coordinates of the regular grid points as
2-dimensional arrays.
u_grid, v_grid
The eastward and northward components of the vector field on
the regular grid.
scalars_grid
The scalar fields on the regular grid. The number of returned
scalar fields is the same as the number that were passed in.
"""
x = np.asanyarray(x)
y = np.asanyarray(y)
u = np.asanyarray(u)
v = np.asanyarray(v)
if u.shape != v.shape:
raise ValueError('u and v must be the same shape')
if x.shape != u.shape:
x, y = np.meshgrid(x, y)
if not (x.shape == y.shape == u.shape):
raise ValueError('x and y coordinates are not compatible '
'with the shape of the vector components')
if scalars:
np_like_scalars = ()
for s in scalars:
s = np.asanyarray(s)
np_like_scalars = np_like_scalars + (s,)
if s.shape != u.shape:
raise ValueError('scalar fields must have the same '
'shape as the vector components')
scalars = np_like_scalars
try:
nx, ny = regrid_shape
except TypeError:
nx = ny = regrid_shape
if target_proj == src_crs:
# Just immediately regrid, interpolate and return
return _interpolate_to_grid(nx, ny, x, y, u, v, *scalars, **kwargs)
# We need to transform the vectors from the source to target frame
# Convert coordinates to the target projection.
proj_xyz = target_proj.transform_points(src_crs, x, y)
targetx, targety = proj_xyz[..., 0], proj_xyz[..., 1]
# Create the grid in the target frame
gridx, gridy = _interpolate_to_grid(nx, ny, targetx, targety, **kwargs)
# Bring the x/y target grid coordinates back into the source frame
src_xyz = src_crs.transform_points(target_proj, gridx, gridy)
# Mask the invalid points that were outside the domain
src_xyz = np.ma.array(src_xyz, mask=~np.isfinite(src_xyz))
sourcex, sourcey = src_xyz[..., 0], src_xyz[..., 1]
# Now interpolate in the source frame
x0, x1 = sourcex.min(), sourcex.max()
y0, y1 = sourcey.min(), sourcey.max()
xr = x1 - x0
yr = y1 - y0
# We also need to transform the original source points to the source
# projection to account for original points outside the wrapped domain
xyz = src_crs.transform_points(src_crs, x, y)
x, y = xyz[..., 0], xyz[..., 1]
points = np.column_stack([(x.ravel() - x0) / xr,
(y.ravel() - y0) / yr])
newx = (sourcex - x0) / xr
newy = (sourcey - y0) / yr
s_grid_tuple = tuple()
for s in (u, v) + scalars:
s_grid_tuple += (griddata(points, s.ravel(), (newx, newy),
method='linear'),)
u, v = s_grid_tuple[0], s_grid_tuple[1]
# Finally, transform the vectors (in the source frame) to the target CRS.
u, v = target_proj.transform_vectors(src_crs, sourcex, sourcey, u, v)
return (gridx, gridy, u, v) + s_grid_tuple[2:]