forked from knaughten/roms_tools
-
Notifications
You must be signed in to change notification settings - Fork 0
/
cice_vectorplot.py
180 lines (161 loc) · 6.88 KB
/
cice_vectorplot.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
178
179
from netCDF4 import Dataset
from numpy import *
from matplotlib.pyplot import *
from rotate_vector_cice import *
# For any vector in the CICE output (velocities, stresses, etc.) make a
# circumpolar Antarctic plot of its magnitude, overlaid with the vectors
# themselves.
# Input:
# file_path = path to CICE history file
# tstep = timestep in file_path to plot
# xname, yname = names of the x and y vector components in file_path
# cmax = optional maximum magnitude for colour scale
# save = optional boolean flag indicating that the plot should be saved to a
# file, rather than displayed on the screen
# fig_name = if save=True, filename for figure
def cice_vectorplot (file_path, tstep, xname, yname, cmax=None, save=False, fig_name=None):
# Radius of the Earth in metres
r = 6.371e6
# Degrees to radians conversion factor
deg2rad = pi/180
# Side length of blocks to average vectors over (can't plot vector at
# every single point or the plot will be way too crowded)
block = 15
# Read grid (including rotation angle) and vector components
id = Dataset(file_path, 'r')
lon_tmp = id.variables['ULON'][:-15,:]
lat_tmp = id.variables['ULAT'][:-15,:]
angle_tmp = id.variables['ANGLE'][:-15,:]
u_xy_tmp = id.variables[xname][tstep-1,:-15,:]
v_xy_tmp = id.variables[yname][tstep-1,:-15,:]
id.close()
# Wrap periodic boundary by 1 cell
lon = ma.empty([size(lon_tmp,0), size(lon_tmp,1)+1])
lat = ma.empty([size(lat_tmp,0), size(lat_tmp,1)+1])
angle = ma.empty([size(angle_tmp,0), size(angle_tmp,1)+1])
u_xy = ma.empty([size(u_xy_tmp,0), size(u_xy_tmp,1)+1])
v_xy = ma.empty([size(v_xy_tmp,0), size(v_xy_tmp,1)+1])
lon[:,:-1] = lon_tmp
lon[:,-1] = lon_tmp[:,0]
lat[:,:-1] = lat_tmp
lat[:,-1] = lat_tmp[:,0]
angle[:,:-1] = angle_tmp
angle[:,-1] = angle_tmp[:,0]
u_xy[:,:-1] = u_xy_tmp
u_xy[:,-1] = u_xy_tmp[:,0]
v_xy[:,:-1] = v_xy_tmp
v_xy[:,-1] = v_xy_tmp[:,0]
# Rotate from local x-y space to lon-lat space
u, v = rotate_vector_cice(u_xy, v_xy, angle)
# Calculate magnitude of vector
speed = sqrt(u**2 + v**2)
# Convert vector to polar coordinates, rotate to account for longitude in
# circumpolar projection, and convert back to vector components
theta = arctan2(v, u)
theta_circ = theta - lon*deg2rad
u_circ = speed*cos(theta_circ)
v_circ = speed*sin(theta_circ)
# Calculate x and y coordinates for plotting circumpolar projection
x = -(lat+90)*cos(lon*deg2rad+pi/2)
y = (lat+90)*sin(lon*deg2rad+pi/2)
# Average x, y, u_circ, and v_circ over block x block intervals
# Calculate number of blocks
size0 = int(ceil(size(x,0)/float(block)))
size1 = int(ceil((size(x,1)-1)/float(block)))
# Set up arrays for averaged fields
x_block = ma.empty([size0, size1])
y_block = ma.empty([size0, size1])
u_circ_block = ma.empty([size0, size1])
v_circ_block = ma.empty([size0, size1])
# Set up arrays containing boundary indices
posn0 = range(0, size(x,0), block)
posn0.append(size(x,0))
posn1 = range(0, size(x,1), block)
posn1.append(size(x,1))
# Double loop to average each block (can't find a more efficient way to do
# this)
for j in range(size0):
for i in range(size1):
start0 = posn0[j]
end0 = posn0[j+1]
start1 = posn1[i]
end1 = posn1[i+1]
x_block[j,i] = mean(x[start0:end0, start1:end1])
y_block[j,i] = mean(y[start0:end0, start1:end1])
u_circ_block[j,i] = mean(u_circ[start0:end0, start1:end1])
v_circ_block[j,i] = mean(v_circ[start0:end0, start1:end1])
# Set up colour scale levels
if cmax is None:
lev = linspace(0, amax(speed), num=50)
else:
lev = linspace(0, cmax, num=50)
# Make the plot
fig = figure(figsize=(16,12))
fig.add_subplot(1,1,1, aspect='equal')
# Contour speed values at every point
# Use pastel colour map so overlaid vectors will show up
contourf(x, y, speed, lev, cmap='Paired', extend='both')
cbar = colorbar()
cbar.ax.tick_params(labelsize=20)
# Add vectors for each block
quiver(x_block, y_block, u_circ_block, v_circ_block, color='black')
title(xname + ', ' + yname, fontsize=30)
axis('off')
if save:
fig.savefig(fig_name)
else:
fig.show()
# Command-line interface
if __name__ == "__main__":
file_path = raw_input("Path to CICE history file: ")
tstep = int(raw_input("Timestep number (starting at 1): "))
xname = raw_input("Variable name of x-component: ")
yname = raw_input("Variable name of y-component: ")
clim = raw_input("Set upper bound on colour scale (y/n)? ")
cmax = None
if clim == 'y':
cmax = float(raw_input("Upper bound: "))
action = raw_input("Save figure (s) or display in window (d)? ")
if action == 's':
save = True
fig_name = raw_input("File name for figure: ")
elif action == 'd':
save = False
fig_name = None
# Make the plot
cice_vectorplot(file_path, tstep, xname, yname, cmax, save, fig_name)
# Repeat until the user wants to exit
while True:
repeat = raw_input("Make another plot (y/n)? ")
if repeat == 'y':
while True:
# Ask for changes to the input parameters; repeat until the user is finished
changes = raw_input("Enter a parameter to change: (1) file path, (2) timestep number, (3) variable names, (4) colour scale, (5) save/display; or enter to continue: ")
if len(changes) == 0:
# No more changes to parameters
break
else:
if int(changes) == 1:
# New file path
file_path = raw_input("Path to CICE history file: ")
elif int(changes) == 2:
# New timestep number
tstep = int(raw_input("Timestep number (starting at 1): "))
elif int(changes) == 3:
xname = raw_input("Variable name of x-component: ")
yname = raw_input("Variable name of y-component: ")
elif int(changes) == 4:
cmax = None
clim = raw_input("Set upper bound on colour scale (y/n)? ")
if clim == 'y':
cmax = float(raw_input("Upper bound: "))
elif int(changes) == 5:
# Change from display to save, or vice versa
save = not save
if save:
# Get file name for figure
fig_name = raw_input("File name for figure: ")
# Make the plot
cice_vectorplot(file_path, tstep, xname, yname, cmax, save, fig_name)
else:
break