tflowline.py - pism - [fork] customized build of PISM, the parallel ice sheet model (tillflux branch)
(HTM) git clone git://src.adamsgaard.dk/pism
(DIR) Log
(DIR) Files
(DIR) Refs
(DIR) LICENSE
---
tflowline.py (7873B)
---
1 #!/usr/bin/env python3
2
3 # @package flowline
4 # \brief This script aids in setting up and visualizing flow-line modeling runs.
5 #
6 # \details This script expands and collapses NetCDF datasets along the 'x' or
7 # 'y' dimension. Run flowline.py --help to see the command-line options.
8 #
9 # To set up a data-set for flow-line modeling, create a NetCDF file with an 'x'
10 # dimension and appropriate variables (usually x(x), acab(x), artm(x),
11 # bheatflx(x), thk(x) and topg(x)), then run
12 # \code
13 # flowline.py -e -o foo_pism.nc foo.nc
14 # \endcode
15 # Then bootstrap from foo_pism.nc, adding -periodicity y to PISM command-line options.
16 #
17 # To collapse a PISM flow-line model output for plotting, etc, run
18 # \code
19 # flowline.py -c -o bar.nc pism_output.nc
20 # \endcode
21
22 from optparse import OptionParser
23
24 try:
25 from netCDF4 import Dataset as CDF
26 except:
27 print("netCDF4 is not installed!")
28 sys.exit(1)
29
30 import numpy as np
31
32 from sys import argv, exit
33 from time import asctime
34
35
36 def get_slice(dimensions, x=None, y=None):
37 """Get an x- or y-slice of a variable."""
38 All = slice(None)
39
40 if not dimensions:
41 return All # so that it does not break processing "mapping"
42
43 index_list = [All] * len(dimensions)
44
45 if x != None:
46 try:
47 index_list[dimensions.index('x')] = x
48 except:
49 pass
50
51 if y != None:
52 try:
53 index_list[dimensions.index('y')] = y
54 except:
55 pass
56
57 return index_list
58
59
60 def permute(variable, output_order=('t', 'z', 'zb', 'y', 'x')):
61 """Permute dimensions of a NetCDF variable to match the output storage order."""
62 input_dimensions = variable.dimensions
63
64 # filter out irrelevant dimensions
65 dimensions = [x for x in output_order if x in input_dimensions]
66
67 # create the mapping
68 mapping = [dimensions.index(x) for x in input_dimensions]
69
70 if mapping:
71 return np.transpose(variable[:], mapping)
72 else:
73 return variable[:] # so that it does not break processing "mapping"
74
75
76 def copy_dim(nc1, nc2, name, direction):
77 """Copy a dimension from nc1 to nc2."""
78 if (name == direction):
79 return
80 dim1 = nc1.dimensions[name]
81 dim2 = nc2.createDimension(name, len(dim1))
82
83
84 def copy_attributes(var1, var2):
85 """Copy attributes of var1 to var2. Skips _FillValue."""
86 for each in var1.ncattrs():
87 if each != "_FillValue":
88 setattr(var2, each, getattr(var1, each))
89
90
91 def process(input, output, direction, collapse):
92 """Process the file 'input', expanding or collapsing data according to
93 'action' and 'direction'. Saves result in 'output'."""
94 try:
95 nc = CDF(input)
96 except:
97 print("ERROR: Can't open %s" % input)
98 exit(1)
99
100 try:
101 out = CDF(output, 'w', format="NETCDF3_CLASSIC")
102 except:
103 print("ERROR: Can't open %s" % output)
104 exit(1)
105
106 copy_attributes(nc, out)
107
108 for name in list(nc.dimensions.keys()):
109 copy_dim(nc, out, name, direction)
110
111 if collapse:
112 for name in list(nc.variables.keys()):
113 if name == direction:
114 continue
115 collapse_var(nc, out, name, direction)
116 message = "Collapsed using flowline.py"
117 else:
118 out.createDimension(direction, 3)
119
120 if direction == 'x':
121 dim = 'y'
122 else:
123 dim = 'x'
124
125 var1 = nc.variables[dim]
126 delta = np.diff(var1[:])[0]
127
128 var2 = out.createVariable(direction, 'f8', (direction,))
129 var2.axis = "%s" % direction.upper()
130 var2.long_name = "%s-coordinate in Cartesian system" % direction.upper()
131 var2.standard_name = "projection_%s_coordinate" % direction
132 var2.units = var1.units
133 var2[:] = [-delta, 0, delta]
134
135 for name in list(nc.variables.keys()):
136 expand_var(nc, out, name, direction)
137
138 message = asctime() + ': ' + ' '.join(argv) + '\n'
139 if 'history' in out.ncattrs():
140 out.history = message + out.history # prepend to history string
141 else:
142 out.history = message
143 out.close()
144
145
146 def collapse_var(nc, out, name, direction):
147 """Saves a collapsed (according to 'direction')
148 copy of a variable 'name' in 'nc' to 'out'."""
149 var1 = nc.variables[name]
150 N = (len(nc.dimensions[direction]) - 1) / 2
151
152 print("Processing %s..." % name)
153 dims = var1.dimensions
154 if len(dims) > 1: # only collapse spatial fields
155 dims = [x for x in dims if x != direction]
156
157 try:
158 fill_value = var1._FillValue
159 var2 = out.createVariable(name, var1.dtype,
160 dimensions=dims, fill_value=fill_value)
161 except:
162 var2 = out.createVariable(name, var1.dtype,
163 dimensions=dims)
164
165 copy_attributes(var1, var2)
166
167 if direction == 'x':
168 var2[:] = var1[get_slice(var1.dimensions, x=N)]
169 elif direction == 'y':
170 var2[:] = var1[get_slice(var1.dimensions, y=N)]
171
172
173 def expand_var(nc, out, name, direction):
174 """Saves an expanded (according to 'direction')
175 copy of a variable 'name' in 'nc' to 'out'."""
176 if name == direction:
177 return
178
179 var1 = nc.variables[name]
180
181 print("Processing %s..." % name)
182
183 # Copy coordinate variables and stop:
184 if name in ['t', 'z', 'y', 'x', 'zb']:
185 var2 = out.createVariable(name, var1.dtype, (name,))
186 var2[:] = var1[:]
187 copy_attributes(var1, var2)
188 return
189
190 dims = var1.dimensions
191 if len(dims) == 1:
192 dims = ('y', 'x')
193 elif len(dims) == 2:
194 dims = ('t', 'y', 'x')
195 elif len(dims) == 3:
196 if name == "litho_temp": # litho_temp is the only variable depending on 'zb'.
197 dims = ('t', 'zb', 'y', 'x')
198 else:
199 dims = ('t', 'z', 'y', 'x')
200
201 var2 = out.createVariable(name, var1.dtype, dims)
202 copy_attributes(var1, var2)
203
204 for j in range(3):
205 if direction == 'x':
206 var2[get_slice(var2.dimensions, x=j)] = permute(var1)
207 elif direction == 'y':
208 var2[get_slice(var2.dimensions, y=j)] = permute(var1)
209
210
211 parser = OptionParser()
212 parser.usage = "usage: %prog -o foo.nc -d {x,y} {--collapse,--expand} file.nc"
213 parser.description = "Collapses or expands a NetCDF file in a specified direction."
214
215 parser.add_option("-d", "--direction", dest="direction",
216 help="direction: one of x,y")
217 parser.add_option("-o", "--output", dest="output_filename",
218 help="output file")
219 parser.add_option("-e", "--expand", dest="expand", action="store_true",
220 help="expand")
221 parser.add_option("-c", "--collapse", dest="collapse", action="store_true",
222 help="collapse")
223
224 (opts, args) = parser.parse_args()
225
226 if len(args) != 1:
227 print("ERROR: File argument is missing.")
228 exit(1)
229
230
231 if (opts.expand and opts.collapse) or ((not opts.expand) and (not opts.collapse)):
232 print("ERROR: exactly one of -e and -c is required.")
233 exit(1)
234
235 if not opts.direction:
236 if opts.collapse or (not opts.expand):
237 opts.direction = 'y'
238 else:
239 nc = CDF(args[0])
240 try:
241 x = nc.variables['x']
242 opts.direction = 'y'
243 except:
244 opts.direction = 'x'
245 nc.close()
246 elif opts.direction not in ['x', 'y']:
247 print("ERROR: Please specify direction using the -d option. (Choose one of x,y.)")
248 exit(1)
249
250 if (not opts.output_filename):
251 print("ERROR: Please specify the output file name using the -o option.")
252 exit(1)
253
254 if opts.collapse:
255 print("Collapsing %s in the %s direction, writing to %s..." % (args[0], opts.direction, opts.output_filename))
256 else:
257 print("Expanding %s in the %s direction, writing to %s..." % (args[0], opts.direction, opts.output_filename))
258
259 process(args[0], opts.output_filename, opts.direction, opts.collapse or (not opts.expand))