| 1 | import netCDF4 as nc |
|---|
| 2 | import ctypes as ct |
|---|
| 3 | import numpy as np |
|---|
| 4 | import os |
|---|
| 5 | import sys |
|---|
| 6 | import math |
|---|
| 7 | from mpi4py import MPI |
|---|
| 8 | |
|---|
| 9 | import time |
|---|
| 10 | |
|---|
| 11 | remap = ct.cdll.LoadLibrary(os.path.realpath('libmapper.so')) |
|---|
| 12 | |
|---|
| 13 | def from_mpas(filename): |
|---|
| 14 | # construct vortex bounds from Mpas grid structure |
|---|
| 15 | f = nc.Dataset(filename) |
|---|
| 16 | # in this case it is must faster to first read the whole file into memory |
|---|
| 17 | # before converting the data structure |
|---|
| 18 | print "read" |
|---|
| 19 | stime = time.time() |
|---|
| 20 | lon_vert = np.array(f.variables["lonVertex"]) |
|---|
| 21 | lat_vert = np.array(f.variables["latVertex"]) |
|---|
| 22 | vert_cell = np.array(f.variables["verticesOnCell"]) |
|---|
| 23 | nvert_cell = np.array(f.variables["nEdgesOnCell"]) |
|---|
| 24 | ncell, nvert = vert_cell.shape |
|---|
| 25 | assert(max(nvert_cell) <= nvert) |
|---|
| 26 | lon = np.zeros(vert_cell.shape) |
|---|
| 27 | lat = np.zeros(vert_cell.shape) |
|---|
| 28 | etime = time.time() |
|---|
| 29 | print "finished read, now convert", etime-stime |
|---|
| 30 | scal = 180.0/math.pi |
|---|
| 31 | for c in range(ncell): |
|---|
| 32 | lat[c,:] = lat_vert[vert_cell[c,:]-1]*scal |
|---|
| 33 | lon[c,:] = lon_vert[vert_cell[c,:]-1]*scal |
|---|
| 34 | # signal "last vertex" by netCDF convetion |
|---|
| 35 | lon[c,nvert_cell[c]] = lon[c,0] |
|---|
| 36 | lat[c,nvert_cell[c]] = lat[c,0] |
|---|
| 37 | print "convert end", time.time() - etime |
|---|
| 38 | return lon, lat |
|---|
| 39 | |
|---|
| 40 | grid_types = { |
|---|
| 41 | "dynamico:mesh": { |
|---|
| 42 | "lon_name": "bounds_lon_i", |
|---|
| 43 | "lat_name": "bounds_lat_i", |
|---|
| 44 | "pole": [0,0,0] |
|---|
| 45 | }, |
|---|
| 46 | "dynamico:vort": { |
|---|
| 47 | "lon_name": "bounds_lon_v", |
|---|
| 48 | "lat_name": "bounds_lat_v", |
|---|
| 49 | "pole": [0,0,0] |
|---|
| 50 | }, |
|---|
| 51 | "dynamico:restart": { |
|---|
| 52 | "lon_name": "lon_i_vertices", |
|---|
| 53 | "lat_name": "lat_i_vertices", |
|---|
| 54 | "pole": [0,0,0] |
|---|
| 55 | }, |
|---|
| 56 | "test:polygon": { |
|---|
| 57 | "lon_name": "bounds_lon", |
|---|
| 58 | "lat_name": "bounds_lat", |
|---|
| 59 | "pole": [0,0,0] |
|---|
| 60 | }, |
|---|
| 61 | "test:latlon": { |
|---|
| 62 | "lon_name": "bounds_lon", |
|---|
| 63 | "lat_name": "bounds_lat", |
|---|
| 64 | "pole": [0,0,1] |
|---|
| 65 | }, |
|---|
| 66 | "mpas": { |
|---|
| 67 | "reader": from_mpas, |
|---|
| 68 | "pole": [0,0,0] |
|---|
| 69 | } |
|---|
| 70 | } |
|---|
| 71 | |
|---|
| 72 | interp_types = { |
|---|
| 73 | "FV1": 1, |
|---|
| 74 | "FV2": 2 |
|---|
| 75 | } |
|---|
| 76 | |
|---|
| 77 | usage = """ |
|---|
| 78 | Usage: python remap.py interp srctype srcfile dsttype dstfile mode outfile |
|---|
| 79 | |
|---|
| 80 | interp: type of interpolation |
|---|
| 81 | choices: |
|---|
| 82 | FV1: first order conservative Finite Volume |
|---|
| 83 | FV2: second order conservative Finite Volume |
|---|
| 84 | |
|---|
| 85 | srctype, dsttype: grid type of source and destination |
|---|
| 86 | choices: """ + " ".join(grid_types.keys()) + """ |
|---|
| 87 | |
|---|
| 88 | srcfile, dstfile: grid file names, should mostly be netCDF file |
|---|
| 89 | |
|---|
| 90 | mode: modus of operation |
|---|
| 91 | choices: |
|---|
| 92 | weights: computes weight and stores them in outfile |
|---|
| 93 | remap: computes the interpolated values on destination grid and stores them in outfile |
|---|
| 94 | |
|---|
| 95 | outfile: output filename |
|---|
| 96 | |
|---|
| 97 | """ |
|---|
| 98 | |
|---|
| 99 | # parse command line arguments |
|---|
| 100 | if not len(sys.argv) == 8: |
|---|
| 101 | print usage |
|---|
| 102 | sys.exit(2) |
|---|
| 103 | |
|---|
| 104 | interp = sys.argv[1] |
|---|
| 105 | try: |
|---|
| 106 | srctype = grid_types[sys.argv[2]] |
|---|
| 107 | except KeyError: |
|---|
| 108 | print "Error: srctype needs to be one of the following: " + " ".join(grid_types.keys()) + "." |
|---|
| 109 | exit(2) |
|---|
| 110 | srcfile = sys.argv[3] |
|---|
| 111 | try: |
|---|
| 112 | dsttype = grid_types[sys.argv[4]] |
|---|
| 113 | except KeyError: |
|---|
| 114 | print "Error: srctype needs to be one of the following: " + " ".join(grid_types.keys()) + "." |
|---|
| 115 | exit(2) |
|---|
| 116 | dstfile = sys.argv[5] |
|---|
| 117 | mode = sys.argv[6] |
|---|
| 118 | outfile = sys.argv[7] |
|---|
| 119 | |
|---|
| 120 | if not mode in ("weights", "remap"): |
|---|
| 121 | print "Error: mode must be of of the following: weights remap." |
|---|
| 122 | exit(2) |
|---|
| 123 | |
|---|
| 124 | remap.mpi_init() |
|---|
| 125 | rank = remap.mpi_rank() |
|---|
| 126 | size = remap.mpi_size() |
|---|
| 127 | |
|---|
| 128 | print rank, "/", size |
|---|
| 129 | |
|---|
| 130 | print "Reading grids from netCDF files." |
|---|
| 131 | |
|---|
| 132 | if "reader" in srctype: |
|---|
| 133 | src_lon, src_lat = srctype["reader"](srcfile) |
|---|
| 134 | else: |
|---|
| 135 | src = nc.Dataset(srcfile) |
|---|
| 136 | # the following two lines do not perform the actual read |
|---|
| 137 | # the file is read later when assigning to the ctypes array |
|---|
| 138 | # -> no unnecessary array copying in memory |
|---|
| 139 | src_lon = src.variables[srctype["lon_name"]] |
|---|
| 140 | src_lat = src.variables[srctype["lat_name"]] |
|---|
| 141 | |
|---|
| 142 | if "reader" in dsttype: |
|---|
| 143 | dst_lon, dst_lat = dsttype["reader"](dstfile) |
|---|
| 144 | else: |
|---|
| 145 | dst = nc.Dataset(dstfile) |
|---|
| 146 | dst_lon = dst.variables[dsttype["lon_name"]] |
|---|
| 147 | dst_lat = dst.variables[dsttype["lat_name"]] |
|---|
| 148 | |
|---|
| 149 | src_ncell, src_nvert = src_lon.shape |
|---|
| 150 | dst_ncell, dst_nvert = dst_lon.shape |
|---|
| 151 | |
|---|
| 152 | def compute_distribution(ncell): |
|---|
| 153 | "Returns the local number and starting position in global array." |
|---|
| 154 | if rank < ncell % size: |
|---|
| 155 | return ncell//size + 1, \ |
|---|
| 156 | (ncell//size + 1)*rank |
|---|
| 157 | else: |
|---|
| 158 | return ncell//size, \ |
|---|
| 159 | (ncell//size + 1)*(ncell%size) + (ncell//size)*(rank - ncell%size) |
|---|
| 160 | |
|---|
| 161 | src_ncell_loc, src_loc_start = compute_distribution(src_ncell) |
|---|
| 162 | dst_ncell_loc, dst_loc_start = compute_distribution(dst_ncell) |
|---|
| 163 | |
|---|
| 164 | print "src", src_ncell_loc, src_loc_start |
|---|
| 165 | print "dst", dst_ncell_loc, dst_loc_start |
|---|
| 166 | |
|---|
| 167 | c_src_lon = (ct.c_double * (src_ncell_loc*src_nvert))() |
|---|
| 168 | c_src_lat = (ct.c_double * (src_ncell_loc*src_nvert))() |
|---|
| 169 | c_dst_lon = (ct.c_double * (dst_ncell_loc*dst_nvert))() |
|---|
| 170 | c_dst_lat = (ct.c_double * (dst_ncell_loc*dst_nvert))() |
|---|
| 171 | |
|---|
| 172 | c_src_lon[:] = nc.numpy.reshape(src_lon[src_loc_start:src_loc_start+src_ncell_loc,:], (len(c_src_lon),1)) |
|---|
| 173 | c_src_lat[:] = nc.numpy.reshape(src_lat[src_loc_start:src_loc_start+src_ncell_loc,:], (len(c_src_lon),1)) |
|---|
| 174 | c_dst_lon[:] = nc.numpy.reshape(dst_lon[dst_loc_start:dst_loc_start+dst_ncell_loc,:], (len(c_dst_lon),1)) |
|---|
| 175 | c_dst_lat[:] = nc.numpy.reshape(dst_lat[dst_loc_start:dst_loc_start+dst_ncell_loc,:], (len(c_dst_lon),1)) |
|---|
| 176 | |
|---|
| 177 | |
|---|
| 178 | print "Calling remap library to compute weights." |
|---|
| 179 | srcpole = (ct.c_double * (3))() |
|---|
| 180 | dstpole = (ct.c_double * (3))() |
|---|
| 181 | srcpole[:] = srctype["pole"] |
|---|
| 182 | dstpole[:] = dsttype["pole"] |
|---|
| 183 | |
|---|
| 184 | c_src_ncell = ct.c_int(src_ncell_loc) |
|---|
| 185 | c_src_nvert = ct.c_int(src_nvert) |
|---|
| 186 | c_dst_ncell = ct.c_int(dst_ncell_loc) |
|---|
| 187 | c_dst_nvert = ct.c_int(dst_nvert) |
|---|
| 188 | order = ct.c_int(interp_types[interp]) |
|---|
| 189 | |
|---|
| 190 | c_nweight = ct.c_int() |
|---|
| 191 | |
|---|
| 192 | print "src:", src_ncell, src_nvert |
|---|
| 193 | print "dst:", dst_ncell, dst_nvert |
|---|
| 194 | |
|---|
| 195 | remap.remap_get_num_weights(c_src_lon, c_src_lat, c_src_nvert, c_src_ncell, srcpole, |
|---|
| 196 | c_dst_lon, c_dst_lat, c_dst_nvert, c_dst_ncell, dstpole, |
|---|
| 197 | order, ct.byref(c_nweight)) |
|---|
| 198 | |
|---|
| 199 | nwgt = c_nweight.value |
|---|
| 200 | |
|---|
| 201 | c_weights = (ct.c_double * nwgt)() |
|---|
| 202 | c_dst_idx = (ct.c_int * nwgt)() |
|---|
| 203 | c_src_idx = (ct.c_int * nwgt)() |
|---|
| 204 | |
|---|
| 205 | remap.remap_get_weights(c_weights, c_src_idx, c_dst_idx) |
|---|
| 206 | |
|---|
| 207 | wgt_glo = MPI.COMM_WORLD.gather(c_weights[:]) |
|---|
| 208 | src_idx_glo = MPI.COMM_WORLD.gather(c_src_idx[:]) |
|---|
| 209 | dst_idx_glo = MPI.COMM_WORLD.gather(c_dst_idx[:]) |
|---|
| 210 | |
|---|
| 211 | |
|---|
| 212 | if rank == 0 and mode == 'weights': |
|---|
| 213 | nwgt_glo = sum(len(wgt) for wgt in wgt_glo) |
|---|
| 214 | |
|---|
| 215 | print "Writing", nwgt_glo, "weights to netCDF-file '" + outfile + "'." |
|---|
| 216 | f = nc.Dataset(outfile,'w') |
|---|
| 217 | f.createDimension('n_src', src_ncell) |
|---|
| 218 | f.createDimension('n_dst', dst_ncell) |
|---|
| 219 | f.createDimension('n_weight', nwgt_glo) |
|---|
| 220 | |
|---|
| 221 | var = f.createVariable('src_idx', 'i', ('n_weight')) |
|---|
| 222 | var[:] = np.hstack(src_idx_glo) + 1 # make indices start from 1 |
|---|
| 223 | var = f.createVariable('dst_idx', 'i', ('n_weight')) |
|---|
| 224 | var[:] = np.hstack(dst_idx_glo) + 1 # make indices start from 1 |
|---|
| 225 | var = f.createVariable('weight', 'd', ('n_weight')) |
|---|
| 226 | var[:] = np.hstack(wgt_glo) |
|---|
| 227 | f.close() |
|---|
| 228 | |
|---|
| 229 | def test_fun(x, y, z): |
|---|
| 230 | return (1-x**2)*(1-y**2)*z |
|---|
| 231 | |
|---|
| 232 | def test_fun_ll(lat, lon): |
|---|
| 233 | #return np.cos(lat*math.pi/180)*np.cos(lon*math.pi/180) |
|---|
| 234 | return 2.0 + np.cos(lat*math.pi/180.)**2 * np.cos(2*lon*math.pi/180.); |
|---|
| 235 | |
|---|
| 236 | #UNUSED |
|---|
| 237 | #def sphe2cart(lat, lon): |
|---|
| 238 | # phi = math.pi/180*lon[:] |
|---|
| 239 | # theta = math.pi/2 - math.pi/180*lat[:] |
|---|
| 240 | # return np.sin(theta)*np.cos(phi), np.sin(theta)*np.sin(phi), np.cos(theta) |
|---|
| 241 | |
|---|
| 242 | if mode == 'remap': |
|---|
| 243 | c_centre_lon = (ct.c_double * src_ncell_loc)() |
|---|
| 244 | c_centre_lat = (ct.c_double * src_ncell_loc)() |
|---|
| 245 | c_areas = (ct.c_double * src_ncell_loc)() |
|---|
| 246 | remap.remap_get_barycentres_and_areas(c_src_lon, c_src_lat, c_src_nvert, c_src_ncell, srcpole, |
|---|
| 247 | c_centre_lon, c_centre_lat, c_areas) |
|---|
| 248 | src_val_loc = test_fun_ll(np.array(c_centre_lat[:]), np.array(c_centre_lon[:])) |
|---|
| 249 | src_val_glo = MPI.COMM_WORLD.gather(src_val_loc) |
|---|
| 250 | |
|---|
| 251 | c_centre_lon = (ct.c_double * dst_ncell_loc)() |
|---|
| 252 | c_centre_lat = (ct.c_double * dst_ncell_loc)() |
|---|
| 253 | c_areas = (ct.c_double * dst_ncell_loc)() |
|---|
| 254 | remap.remap_get_barycentres_and_areas(c_dst_lon, c_dst_lat, c_dst_nvert, c_dst_ncell, dstpole, |
|---|
| 255 | c_centre_lon, c_centre_lat, c_areas) |
|---|
| 256 | dst_val_loc = test_fun_ll(np.array(c_centre_lat[:]), np.array(c_centre_lon[:])) |
|---|
| 257 | |
|---|
| 258 | dst_val_glo = MPI.COMM_WORLD.gather(dst_val_loc) |
|---|
| 259 | dst_areas_glo = MPI.COMM_WORLD.gather(np.array(c_areas[:])) |
|---|
| 260 | dst_centre_lon_glo = MPI.COMM_WORLD.gather(np.array(c_centre_lon[:])) |
|---|
| 261 | dst_centre_lat_glo = MPI.COMM_WORLD.gather(np.array(c_centre_lat[:])) |
|---|
| 262 | |
|---|
| 263 | |
|---|
| 264 | if rank == 0 and mode == 'remap': |
|---|
| 265 | from scipy import sparse |
|---|
| 266 | A = sparse.csr_matrix(sparse.coo_matrix((np.hstack(wgt_glo),(np.hstack(dst_idx_glo),np.hstack(src_idx_glo))))) |
|---|
| 267 | |
|---|
| 268 | src_val = np.hstack(src_val_glo) |
|---|
| 269 | dst_ref = np.hstack(dst_val_glo) |
|---|
| 270 | dst_areas = np.hstack(dst_areas_glo) |
|---|
| 271 | dst_centre_lon = np.hstack(dst_centre_lon_glo) |
|---|
| 272 | dst_centre_lat = np.hstack(dst_centre_lat_glo) |
|---|
| 273 | |
|---|
| 274 | print "source:", src_val.shape |
|---|
| 275 | print "destin:", dst_ref.shape |
|---|
| 276 | dst_val = A*src_val |
|---|
| 277 | err = dst_val - dst_ref |
|---|
| 278 | print "absolute maximum error, maximum value:", np.max(np.abs(err)), np.max(np.abs(dst_ref)) |
|---|
| 279 | print "relative maximum error, normalized L2 error, average target cell size (edgelength of same-area square):" |
|---|
| 280 | print np.max(np.abs(err))/np.max(np.abs(dst_ref)), np.linalg.norm(err)/np.linalg.norm(dst_ref), np.mean(np.sqrt(dst_areas)) |
|---|
| 281 | |
|---|
| 282 | f = nc.Dataset(outfile,'w') |
|---|
| 283 | f.createDimension('n_vert', dst_nvert) |
|---|
| 284 | f.createDimension('n_cell', dst_ncell) |
|---|
| 285 | |
|---|
| 286 | var = f.createVariable('lat', 'd', ('n_cell')) |
|---|
| 287 | var.setncattr("long_name", "latitude") |
|---|
| 288 | var.setncattr("units", "degrees_north") |
|---|
| 289 | var.setncattr("bounds", "bounds_lat") |
|---|
| 290 | var[:] = dst_centre_lat |
|---|
| 291 | var = f.createVariable('lon', 'd', ('n_cell')) |
|---|
| 292 | var.setncattr("long_name", "longitude") |
|---|
| 293 | var.setncattr("units", "degrees_east") |
|---|
| 294 | var.setncattr("bounds", "bounds_lon") |
|---|
| 295 | var[:] = dst_centre_lon |
|---|
| 296 | |
|---|
| 297 | var = f.createVariable('bounds_lon', 'd', ('n_cell','n_vert')) |
|---|
| 298 | var[:] = dst_lon |
|---|
| 299 | var = f.createVariable('bounds_lat', 'd', ('n_cell','n_vert')) |
|---|
| 300 | var[:] = dst_lat |
|---|
| 301 | var = f.createVariable('val', 'd', ('n_cell')) |
|---|
| 302 | var.setncattr("coordinates", "lon lat") |
|---|
| 303 | var[:] = dst_val |
|---|
| 304 | var = f.createVariable('val_ref', 'd', ('n_cell')) |
|---|
| 305 | var.setncattr("coordinates", "lon lat") |
|---|
| 306 | var[:] = dst_ref |
|---|
| 307 | var = f.createVariable('err', 'd', ('n_cell')) |
|---|
| 308 | var.setncattr("coordinates", "lon lat") |
|---|
| 309 | var[:] = err |
|---|
| 310 | var = f.createVariable('area', 'd', ('n_cell')) |
|---|
| 311 | var.setncattr("coordinates", "lon lat") |
|---|
| 312 | var[:] = dst_areas[:] # dest |
|---|
| 313 | f.close() |
|---|
| 314 | |
|---|
| 315 | if not "reader" in srctype: |
|---|
| 316 | src.close() |
|---|
| 317 | if not "reader" in dsttype: |
|---|
| 318 | dst.close() |
|---|
| 319 | |
|---|
