source: lmdz_wrf/trunk/tools/geometry_tools.py @ 2412

Last change on this file since 2412 was 2412, checked in by lfita, 6 years ago

Adding more functions from AR6-zones script
File size: 14.0 KB
Line 
1# Python tools to manage netCDF files.
2# L. Fita, CIMA. Mrch 2019
3# More information at: http://www.xn--llusfb-5va.cat/python/PyNCplot
4#
5# pyNCplot and its component geometry_tools.py comes with ABSOLUTELY NO WARRANTY.
6# This work is licendes under a Creative Commons
7#   Attribution-ShareAlike 4.0 International License (http://creativecommons.org/licenses/by-sa/4.0)
8#
9## Script for geometry calculations and operations as well as definition of different
10###    standard objects and shapes
11
12import numpy as np
13import matplotlib as mpl
14from mpl_toolkits.mplot3d import Axes3D
15import matplotlib.pyplot as plt
16
17####### Contents:
18# deg_deci: Function to pass from degrees [deg, minute, sec] to decimal angles [rad]
19# multi_rotate_2D: Function to rotate multiple vectors by a certain angle in the plain
20# position_sphere: Function to tranform fom a point in lon, lat deg coordinates to
21#   cartesian coordinates over an sphere
22# rotate_2D: Function to rotate a vector by a certain angle in the plain
23# rotate_line2D: Function to rotate a line given by 2 pairs of x,y coordinates by a
24#   certain angle in the plain
25# rotate_lines2D: Function to rotate multiple lines given by mulitple pars of x,y
26#   coordinates by a certain angle in the plain
27# surface_sphere: Function to provide an sphere as matrix of x,y,z coordinates
28# spheric_line: Function to transform a series of locations in lon, lat coordinates
29#   to x,y,z over an 3D spaceFunction to provide coordinates of a line  on a 3D space
30
31## Shapes/objects
32# ellipse_polar: Function to determine an ellipse from its center and polar coordinates
33
34## Plotting
35# plot_sphere: Function to plot an sphere and determine which standard lines will be
36#   also drawn
37
38def deg_deci(angle):
39    """ Function to pass from degrees [deg, minute, sec] to decimal angles [rad]
40      angle: list of [deg, minute, sec] to pass
41    >>> deg_deci([41., 58., 34.])
42    0.732621346072
43    """
44    fname = 'deg_deci'
45
46    deg = np.abs(angle[0]) + np.abs(angle[1])/60. + np.abs(angle[2])/3600.
47
48    if angle[0] < 0.: deg = -deg*np.pi/180.
49    else: deg = deg*np.pi/180.
50
51    return deg
52
53def position_sphere(radii, alpha, beta):
54    """ Function to tranform fom a point in lon, lat deg coordinates to cartesian 
55          coordinates over an sphere
56      radii: radii of the sphere
57      alpha: longitude of the point
58      beta: latitude of the point
59    >>> position_sphere(10., 30., 45.)
60    (0.81031678432964027, -5.1903473778327376, 8.5090352453411846
61    """
62    fname = 'position_sphere'
63
64    xpt = radii*np.cos(beta)*np.cos(alpha)
65    ypt = radii*np.cos(beta)*np.sin(alpha)
66    zpt = radii*np.sin(beta)
67
68    return xpt, ypt, zpt
69
70def surface_sphere(radii,Npts):
71    """ Function to provide an sphere as matrix of x,y,z coordinates
72      radii: radii of the sphere
73      Npts: number of points to discretisize longitues (half for latitudes)
74    """
75    fname = 'surface_sphere'
76
77    sphereup = np.zeros((3,Npts/2,Npts), dtype=np.float)
78    spheredown = np.zeros((3,Npts/2,Npts), dtype=np.float)
79    for ia in range(Npts):
80        alpha = ia*2*np.pi/(Npts-1)
81        for ib in range(Npts/2):
82            beta = ib*np.pi/(2.*(Npts/2-1))
83            sphereup[:,ib,ia] = position_sphere(radii, alpha, beta)
84        for ib in range(Npts/2):
85            beta = -ib*np.pi/(2.*(Npts/2-1))
86            spheredown[:,ib,ia] = position_sphere(radii, alpha, beta)
87
88    return sphereup, spheredown
89
90def spheric_line(radii,lon,lat):
91    """ Function to transform a series of locations in lon, lat coordinates to x,y,z
92          over an 3D space
93      radii: radius of the sphere
94      lon: array of angles along longitudes
95      lat: array of angles along latitudes
96    """
97    fname = 'spheric_line'
98
99    Lint = lon.shape[0]
100    coords = np.zeros((Lint,3), dtype=np.float)
101
102    for iv in range(Lint):
103        coords[iv,:] = position_sphere(radii, lon[iv], lat[iv])
104
105    return coords
106
107def rotate_2D(vector, angle):
108    """ Function to rotate a vector by a certain angle in the plain
109      vector= vector to rotate [y, x]
110      angle= angle to rotate [rad]
111    >>> rotate_2D(np.array([1.,0.]), np.pi/4.)
112    [ 0.70710678 -0.70710678]
113    """
114    fname = 'rotate_2D'
115
116    rotmat = np.zeros((2,2), dtype=np.float)
117
118    rotmat[0,0] = np.cos(angle)
119    rotmat[0,1] = -np.sin(angle)
120    rotmat[1,0] = np.sin(angle)
121    rotmat[1,1] = np.cos(angle)
122
123    rotvector = np.zeros((2), dtype=np.float)
124
125    vecv = np.zeros((2), dtype=np.float)
126
127    # Unifying vector
128    modvec = vector[0]**2+vector[1]**2
129    if modvec != 0: 
130        vecv[0] = vector[1]/modvec
131        vecv[1] = vector[0]/modvec
132
133        rotvec = np.matmul(rotmat, vecv)
134        rotvec = np.where(np.abs(rotvec) < 1.e-7, 0., rotvec)
135
136        rotvector[0] = rotvec[1]*modvec
137        rotvector[1] = rotvec[0]*modvec
138
139    return rotvector
140
141def multi_rotate_2D(vectors, angle):
142    """ Function to rotate multiple vectors by a certain angle in the plain
143      line= matrix of vectors to rotate
144      angle= angle to rotate [rad]
145    >>> square = np.zeros((4,2), dtype=np.float)
146    >>> square[0,:] = [-0.5,-0.5]
147    >>> square[1,:] = [0.5,-0.5]
148    >>> square[2,:] = [0.5,0.5]
149    >>> square[3,:] = [-0.5,0.5]
150    >>> multi_rotate_2D(square, np.pi/4.)
151    [[-0.70710678  0.        ]
152     [ 0.         -0.70710678]
153     [ 0.70710678  0.        ]
154     [ 0.          0.70710678]]
155    """
156    fname = 'multi_rotate_2D'
157
158    rotvecs = np.zeros(vectors.shape, dtype=np.float)
159
160    Nvecs = vectors.shape[0]
161    for iv in range(Nvecs):
162        rotvecs[iv,:] = rotate_2D(vectors[iv,:], angle)
163
164    return rotvecs
165
166def rotate_line2D(line, angle):
167    """ Function to rotate a line given by 2 pairs of x,y coordinates by a certain
168          angle in the plain
169      line= line to rotate as couple of points [[y0,x0], [y1,x1]]
170      angle= angle to rotate [rad]
171    >>> rotate_line2D(np.array([[0.,0.], [1.,0.]]), np.pi/4.)
172    [[ 0.          0.        ]
173     [0.70710678  -0.70710678]]
174    """
175    fname = 'rotate_2D'
176
177    rotline = np.zeros((2,2), dtype=np.float)
178    rotline[0,:] = rotate_2D(line[0,:], angle)
179    rotline[1,:] = rotate_2D(line[1,:], angle)
180
181    return rotline
182
183def rotate_lines2D(lines, angle):
184    """ Function to rotate multiple lines given by mulitple pars of x,y coordinates 
185          by a certain angle in the plain
186      line= matrix of N couples of points [N, [y0,x0], [y1,x1]]
187      angle= angle to rotate [rad]
188    >>> square = np.zeros((4,2,2), dtype=np.float)
189    >>> square[0,0,:] = [-0.5,-0.5]
190    >>> square[0,1,:] = [0.5,-0.5]
191    >>> square[1,0,:] = [0.5,-0.5]
192    >>> square[1,1,:] = [0.5,0.5]
193    >>> square[2,0,:] = [0.5,0.5]
194    >>> square[2,1,:] = [-0.5,0.5]
195    >>> square[3,0,:] = [-0.5,0.5]
196    >>> square[3,1,:] = [-0.5,-0.5]
197    >>> rotate_lines2D(square, np.pi/4.)
198    [[[-0.70710678  0.        ]
199      [ 0.         -0.70710678]]
200
201     [[ 0.         -0.70710678]
202      [ 0.70710678  0.        ]]
203
204     [[ 0.70710678  0.        ]
205      [ 0.          0.70710678]]
206
207     [[ 0.          0.70710678]
208      [-0.70710678  0.        ]]]
209    """
210    fname = 'rotate_lines2D'
211
212    rotlines = np.zeros(lines.shape, dtype=np.float)
213
214    Nlines = lines.shape[0]
215    for il in range(Nlines):
216        line = np.zeros((2,2), dtype=np.float)
217        line[0,:] = lines[il,0,:]
218        line[1,:] = lines[il,1,:]
219
220        rotlines[il,:,:] = rotate_line2D(line, angle)
221
222    return rotlines
223
224####### ###### ##### #### ### ## #
225# Shapes/objects
226
227def ellipse_polar(c, a, b, Nang=100):
228    """ Function to determine an ellipse from its center and polar coordinates
229        FROM: https://en.wikipedia.org/wiki/Ellipse
230      c= coordinates of the center
231      a= distance major axis
232      b= distance minor axis
233      Nang= number of angles to use
234    """
235    fname = 'ellipse_polar'
236
237    if np.mod(Nang,2) == 0: Nang=Nang+1
238 
239    dtheta = 2*np.pi/(Nang-1)
240
241    ellipse = np.zeros((Nang,2), dtype=np.float)
242    for ia in range(Nang):
243        theta = dtheta*ia
244        rad = a*b/np.sqrt( (b*np.cos(theta))**2 + (a*np.sin(theta))**2 )
245        x = rad*np.cos(theta)
246        y = rad*np.sin(theta)
247        ellipse[ia,:] = [y+c[0],x+c[1]]
248
249    return ellipse
250
251
252####### ####### ##### #### ### ## #
253# Plotting
254
255def plot_sphere(iazm=-60., iele=30., dist=10., Npts=100, radii=10,                   \
256  drwsfc=[True,True], colsfc=['#AAAAAA','#646464'],                                  \
257  drwxline = True, linex=[':','b',2.], drwyline = True, liney=[':','r',2.],          \
258  drwzline = True, linez=['-.','g',2.], drwxcline=[True,True],                       \
259  linexc=[['-','#646400',1.],['--','#646400',1.]],                                   \
260  drwequator=[True,True], lineeq=[['-','#AA00AA',1.],['--','#AA00AA',1.]],           \
261  drwgreeenwhich=[True,True], linegw=[['-','k',1.],['--','k',1.]]):
262    """ Function to plot an sphere and determine which standard lines will be also
263        drawn
264      iazm: azimut of the camera form the sphere
265      iele: elevation of the camera form the sphere
266      dist: distance of the camera form the sphere
267      Npts: Resolution for the sphere
268      radii: radius of the sphere
269      drwsfc: whether 'up' and 'down' portions of the sphere should be drawn
270      colsfc: colors of the surface of the sphere portions ['up', 'down']
271      drwxline: whether x-axis line should be drawn
272      linex: properties of the x-axis line ['type', 'color', 'wdith']
273      drwyline: whether y-axis line should be drawn
274      liney: properties of the y-axis line ['type', 'color', 'wdith']
275      drwzline: whether z-axis line should be drawn
276      linez: properties of the z-axis line ['type', 'color', 'wdith']
277      drwequator: whether 'front' and 'back' portions of the Equator should be drawn
278      lineeq: properties of the lines 'front' and 'back' of the Equator
279      drwgreeenwhich: whether 'front', 'back' portions of Greenqhich should be drawn
280      linegw: properties of the lines 'front' and 'back' Greenwhich
281      drwxcline: whether 'front', 'back' 90 line (lon=90., lon=270.) should be drawn
282      linexc: properties of the lines 'front' and 'back' for the 90 line
283    """
284    fname = 'plot_sphere'
285
286    iazmrad = iazm*np.pi/180.
287    ielerad = iele*np.pi/180.
288
289    # 3D surface Sphere
290    sfcsphereu, sfcsphered = surface_sphere(radii,Npts)
291   
292    # greenwhich
293    if iazmrad > np.pi/2. and iazmrad < 3.*np.pi/2.:
294        ia=np.pi-ielerad
295    else:
296        ia=0.-ielerad
297    ea=ia+np.pi
298    da = (ea-ia)/(Npts-1)
299    beta = np.arange(ia,ea+da,da)[0:Npts]
300    alpha = np.zeros((Npts), dtype=np.float)
301    greenwhichc = spheric_line(radii,alpha,beta)
302    ia=ea+0.
303    ea=ia+np.pi
304    da = (ea-ia)/(Npts-1)
305    beta = np.arange(ia,ea+da,da)[0:Npts]
306    greenwhichd = spheric_line(radii,alpha,beta)
307
308    # Equator
309    ia=np.pi-iazmrad/2.
310    ea=ia+np.pi
311    da = (ea-ia)/(Npts-1)
312    alpha = np.arange(ia,ea+da,da)[0:Npts]
313    beta = np.zeros((Npts), dtype=np.float)
314    equatorc = spheric_line(radii,alpha,beta)
315    ia=ea+0.
316    ea=ia+np.pi
317    da = (ea-ia)/(Npts-1)
318    alpha = np.arange(ia,ea+da,da)[0:Npts]
319    equatord = spheric_line(radii,alpha,beta)
320
321    # 90 line
322    if iazmrad > np.pi and iazmrad < 2.*np.pi:
323        ia=3.*np.pi/2. + ielerad
324    else:
325        ia=np.pi/2. - ielerad
326    if ielerad < 0.:
327        ia = ia + np.pi
328    ea=ia+np.pi
329    da = (ea-ia)/(Npts-1)
330    beta = np.arange(ia,ea+da,da)[0:Npts]
331    alpha = np.ones((Npts), dtype=np.float)*np.pi/2.
332    xclinec = spheric_line(radii,alpha,beta)
333    ia=ea+0.
334    ea=ia+np.pi
335    da = (ea-ia)/(Npts-1)
336    beta = np.arange(ia,ea+da,da)[0:Npts]
337    xclined = spheric_line(radii,alpha,beta)
338
339    # x line
340    xline = np.zeros((2,3), dtype=np.float)
341    xline[0,:] = position_sphere(radii, 0., 0.)
342    xline[1,:] = position_sphere(radii, np.pi, 0.)
343
344    # y line
345    yline = np.zeros((2,3), dtype=np.float)
346    yline[0,:] = position_sphere(radii, np.pi/2., 0.)
347    yline[1,:] = position_sphere(radii, 3*np.pi/2., 0.)
348
349    # z line
350    zline = np.zeros((2,3), dtype=np.float)
351    zline[0,:] = position_sphere(radii, 0., np.pi/2.)
352    zline[1,:] = position_sphere(radii, 0., -np.pi/2.)
353
354    fig = plt.figure()
355    ax = fig.gca(projection='3d')
356
357    # Sphere surface
358    if drwsfc[0]:
359        ax.plot_surface(sfcsphereu[0,:,:], sfcsphereu[1,:,:], sfcsphereu[2,:,:],     \
360          color=colsfc[0])
361    if drwsfc[1]:
362        ax.plot_surface(sfcsphered[0,:,:], sfcsphered[1,:,:], sfcsphered[2,:,:],     \
363          color=colsfc[1])
364
365    # greenwhich
366    linev = linegw[0]
367    if drwgreeenwhich[0]:
368        ax.plot(greenwhichc[:,0], greenwhichc[:,1], greenwhichc[:,2], linev[0],      \
369          color=linev[1], linewidth=linev[2],  label='Greenwich')
370    linev = linegw[1]
371    if drwgreeenwhich[1]:
372        ax.plot(greenwhichd[:,0], greenwhichd[:,1], greenwhichd[:,2], linev[0],      \
373          color=linev[1], linewidth=linev[2])
374
375    # Equator
376    linev = lineeq[0]
377    if drwequator[0]:
378        ax.plot(equatorc[:,0], equatorc[:,1], equatorc[:,2], linev[0],               \
379          color=linev[1], linewidth=linev[2], label='Equator')
380    linev = lineeq[1]
381    if drwequator[1]:
382        ax.plot(equatord[:,0], equatord[:,1], equatord[:,2], linev[0],               \
383          color=linev[1], linewidth=linev[2])
384
385    # 90line
386    linev = linexc[0]
387    if drwxcline[0]:
388        ax.plot(xclinec[:,0], xclinec[:,1], xclinec[:,2], linev[0], color=linev[1],  \
389          linewidth=linev[2], label='90-line')
390    linev = linexc[1]
391    if drwxcline[1]:
392        ax.plot(xclined[:,0], xclined[:,1], xclined[:,2], linev[0], color=linev[1],  \
393          linewidth=linev[2])
394
395    # x line
396    linev = linex
397    if drwxline:
398        ax.plot([xline[0,0],xline[1,0]], [xline[0,1],xline[1,1]],                    \
399          [xline[0,2],xline[1,2]], linev[0], color=linev[1], linewidth=linev[2],  label='xline')
400
401    # y line
402    linev = liney
403    if drwyline:
404        ax.plot([yline[0,0],yline[1,0]], [yline[0,1],yline[1,1]],                    \
405          [yline[0,2],yline[1,2]], linev[0], color=linev[1], linewidth=linev[2],  label='yline')
406
407    # z line
408    linev = linez
409    if drwzline:
410        ax.plot([zline[0,0],zline[1,0]], [zline[0,1],zline[1,1]],                    \
411          [zline[0,2],zline[1,2]], linev[0], color=linev[1], linewidth=linev[2],  label='zline')
412
413    plt.legend()
414
415    return fig, ax
416
Note: See TracBrowser for help on using the repository browser.