-
Notifications
You must be signed in to change notification settings - Fork 3
/
Copy pathplot_layers.py
143 lines (123 loc) · 5.34 KB
/
plot_layers.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
from netCDF4 import Dataset
from numpy import *
from matplotlib.pyplot import *
from calc_z import *
from interp_lon_roms import *
# Plot the terrain-following vertical levels through a given longitude, in a
# latitude vs. depth slice. This is a good way to test out different choices
# for vertical stretching parameters.
# Input:
# grid_path = path to ROMS grid file
# lon0 = longitude to interpolate to (-180 to 180)
# depth_min = deepest depth to plot (negative, in metres)
# Vstretching, theta_s, theta_b, hc, N = vertical stretching parameters (see
# *.in configuration file if unsure)
def plot_layers (grid_path, lon0, depth_min, Vstretching, theta_s, theta_b, hc, N):
# Read grid
id = Dataset(grid_path, 'r')
lon_2d = id.variables['lon_rho'][:,:]
lat_2d = id.variables['lat_rho'][:,:]
h = id.variables['h'][:,:]
zice = id.variables['zice'][:,:]
mask = id.variables['mask_rho'][:,:]
id.close()
# Calculate a 3D field of depth values
z_3d, sc_r, Cs_r = calc_z(h, zice, theta_s, theta_b, hc, N, None, Vstretching)
# Figure out how to write longitude in the title
if lon0 < 0:
lon_string = str(int(round(-lon0))) + r'$^{\circ}$W'
else:
lon_string = str(int(round(lon0))) + r'$^{\circ}$E'
# Get longitude in the range (0,360) as per ROMS convention
if lon0 < 0:
lon0 += 360
# Get a 3D land mask
mask_3d = tile(mask, (N,1,1))
# Interpolate land mask, depth, longitude to the given longitude
mask_interp, z, lat = interp_lon_roms(mask_3d, z_3d, lat_2d, lon_2d, lon0)
# Simple array containing layer numbers, same size as z
layer = zeros(shape(z))
for k in range(N):
layer[k,:] = k+1
# Mask out land
layer = ma.masked_where(mask_interp==0, layer)
# Choose latitude bounds based on land mask
mask_sum = sum(mask_interp, axis=0)
# Find southernmost and northernmost unmasked j-indices
edges = ma.flatnotmasked_edges(mask_sum)
j_min = edges[0]
j_max = edges[1]
if j_min == 0:
# There are ocean points right to the southern boundary
# Don't do anything special
lat_min = min(lat[:,j_min])
else:
# There is land everywhere at the southern boundary
# Show the last 2 degrees of this land mask
lat_min = min(lat[:,j_min]) - 2
if j_max == size(mask_sum) - 1:
# There are ocean points right to the northern boundary
# Don't do anything special
lat_max = max(lat[:,j_max])
else:
# There is land everywhere at the northern boundary
# Show the first 2 degrees of this land mask
lat_max = max(lat[:,j_max]) + 2
#lat_min = -85
#lat_max = -75
# Contour levels
lev = range(1, N)
# Plot
fig = figure(figsize=(18,6)) #(18,12))
contour(lat, z, layer, lev, colors='k')
title(lon_string) #(r"ROMS terrain-following levels through the Ross Ice Shelf cavity (180$^{\circ}$E)", fontsize=24)
xlabel('Latitude')
ylabel('Depth (m)')
xlim([lat_min, lat_max])
ylim([depth_min, 0])
#text(-82, -200, "Ice shelf", fontsize=24)
#text(-82, -650, "Sea floor", fontsize=24)
#fig.savefig('ross_levels.png')
fig.show()
# Put longitude back to original range in case the user wants to go again
if lon0 > 180:
lon0 -= 360
# Command-line interface
if __name__ == "__main__":
grid_path = raw_input("Path to grid file: ")
lon0 = float(raw_input("Enter longitude (-180 to 180): "))
depth_min = -1*float(raw_input("Deepest depth to plot (positive, metres): "))
Vstretching = int(raw_input("Vstretching (2 or 4): "))
theta_s = float(raw_input("theta_s: "))
theta_b = float(raw_input("theta_b: "))
hc = float(raw_input("hc: "))
N = int(raw_input("N: "))
plot_layers (grid_path, lon0, depth_min, Vstretching, theta_s, theta_b, hc, N)
# Keep repeating until the user wants to exit
while True:
repeat = raw_input("Make another plot (y/n)? ")
if repeat == 'y':
while True:
changes = raw_input("Enter a parameter to change: (1) grid file, (2) longitude, (3) deepest depth, (4) Vstretching, (5) theta_s, (6) theta_b, (7) hc, (8) N; or enter to continue: ")
if len(changes) == 0:
break
else:
if int(changes) == 1:
grid_path = raw_input("Path to grid file: ")
elif int(changes) == 2:
lon0 = float(raw_input("Enter longitude (-180 to 180): "))
elif int(changes) == 3:
depth_min = -1*float(raw_input("Deepest depth to plot (positive, metres): "))
elif int(changes) == 4:
Vstretching = int(raw_input("Vstretching (2 or 4): "))
elif int(changes) == 5:
theta_s = float(raw_input("theta_s: "))
elif int(changes) == 6:
theta_b = float(raw_input("theta_b: "))
elif int(changes) == 7:
hc = float(raw_input("hc: "))
elif int(changes) == 8:
N = int(raw_input("N: "))
plot_layers (grid_path, lon0, depth_min, Vstretching, theta_s, theta_b, hc, N)
else:
break