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))