[6331] | 1 | #!/usr/bin/env python3 |
---|
| 2 | |
---|
| 3 | import pyoasis |
---|
| 4 | from pyoasis import OASIS |
---|
| 5 | import numpy as np |
---|
| 6 | import math |
---|
| 7 | |
---|
| 8 | comp = pyoasis.Component("writer") |
---|
| 9 | |
---|
| 10 | print(comp) |
---|
| 11 | |
---|
| 12 | comm_rank = comp.localcomm.rank |
---|
| 13 | comm_size = comp.localcomm.size |
---|
| 14 | |
---|
| 15 | nx_loc = 18 |
---|
| 16 | ny_loc = 18 |
---|
| 17 | nx_global = comm_size*nx_loc |
---|
| 18 | ny_global = ny_loc |
---|
| 19 | |
---|
| 20 | dp_conv = math.pi/180. |
---|
| 21 | |
---|
| 22 | partition = pyoasis.BoxPartition(comm_rank*nx_loc, nx_loc, ny_loc, nx_global) |
---|
| 23 | |
---|
| 24 | dx = 360.0/nx_global |
---|
| 25 | dy = 180.0/ny_global |
---|
| 26 | |
---|
| 27 | lon = np.array([-180. + comm_rank*nx_loc*dx + float(i)*dx |
---|
| 28 | + dx/2.0 for i in range(nx_loc)], |
---|
| 29 | dtype=np.float64) |
---|
| 30 | lon = np.tile(lon, (ny_loc, 1)).T |
---|
| 31 | |
---|
| 32 | lat = np.array([-90.0 + float(j)*dy + dy/2.0 for j in range(ny_loc)], |
---|
| 33 | dtype=np.float64) |
---|
| 34 | lat = np.tile(lat, (nx_loc, 1)) |
---|
| 35 | |
---|
| 36 | grid = pyoasis.Grid('pyoa', nx_global, ny_global, lon, lat, partition) |
---|
| 37 | |
---|
| 38 | ncrn = 4 |
---|
| 39 | clo = pyoasis.asarray(np.zeros((nx_loc, ny_loc, ncrn), dtype=np.float64)) |
---|
| 40 | clo[:, :, 0] = lon[:, :] - dx/2.0 |
---|
| 41 | clo[:, :, 1] = lon[:, :] + dx/2.0 |
---|
| 42 | clo[:, :, 2] = clo[:, :, 1] |
---|
| 43 | clo[:, :, 3] = clo[:, :, 0] |
---|
| 44 | cla = pyoasis.asarray(np.zeros((nx_loc, ny_loc, ncrn), dtype=np.float64)) |
---|
| 45 | cla[:, :, 0] = lat[:, :] - dy/2.0 |
---|
| 46 | cla[:, :, 1] = cla[:, :, 0] |
---|
| 47 | cla[:, :, 2] = lat[:, :] + dy/2.0 |
---|
| 48 | cla[:, :, 3] = cla[:, :, 2] |
---|
| 49 | |
---|
| 50 | grid.set_corners(clo, cla) |
---|
| 51 | |
---|
| 52 | msk = np.zeros((nx_loc, ny_loc), dtype=np.int32) |
---|
| 53 | if comm_rank == 0: |
---|
| 54 | msk[4:6, 2:16] = 1 |
---|
| 55 | msk[6:11, (8, 9, 14, 15)] = 1 |
---|
| 56 | msk[11, (8, 9, 10, 13, 14, 15)] = 1 |
---|
| 57 | msk[12, 9:15] = 1 |
---|
| 58 | msk[13, 10:14] = 1 |
---|
| 59 | elif comm_rank == 1: |
---|
| 60 | msk[(3, 14), 14:16] = 1 |
---|
| 61 | msk[(4, 13), 12:16] = 1 |
---|
| 62 | msk[(5, 12), 11:15] = 1 |
---|
| 63 | msk[(6, 11), 10:13] = 1 |
---|
| 64 | msk[(7, 10), 9:12] = 1 |
---|
| 65 | msk[8:10, 2:11] = 1 |
---|
| 66 | elif comm_rank == 2: |
---|
| 67 | msk[(4, 13), 4:14] = 1 |
---|
| 68 | msk[(5, 12), 3:15] = 1 |
---|
| 69 | msk[(6, 11), 2:5] = 1 |
---|
| 70 | msk[(6, 11), 13:16] = 1 |
---|
| 71 | msk[7:11, 2:4] = 1 |
---|
| 72 | msk[7:11, 14:16] = 1 |
---|
| 73 | |
---|
| 74 | grid.set_mask(msk, companion="STFC") |
---|
| 75 | |
---|
| 76 | frc = np.ones((nx_loc, ny_loc), dtype=np.float64) |
---|
| 77 | frc = np.where(msk == 1, 0.0, 1.0) |
---|
| 78 | |
---|
| 79 | grid.set_frac(frc, companion="STFC") |
---|
| 80 | |
---|
| 81 | area = np.zeros((nx_loc, ny_loc), dtype=np.float64) |
---|
| 82 | area[:, :] = dp_conv * \ |
---|
| 83 | np.abs(np.sin(cla[:, :, 2] * dp_conv) - |
---|
| 84 | np.sin(cla[:, :, 0] * dp_conv)) * \ |
---|
| 85 | np.abs(clo[:, :, 1]-clo[:, :, 0]) |
---|
| 86 | |
---|
| 87 | grid.set_area(area) |
---|
| 88 | |
---|
| 89 | angle = np.zeros((nx_loc, ny_loc), dtype=np.float64) |
---|
| 90 | |
---|
| 91 | grid.set_angle(angle) |
---|
| 92 | |
---|
| 93 | if comm_rank == 0: |
---|
| 94 | nx_mono = 180 |
---|
| 95 | ny_mono = 90 |
---|
| 96 | dxm = 360.0 / nx_mono |
---|
| 97 | dym = 180.0 / ny_mono |
---|
| 98 | |
---|
| 99 | lonm = np.array([float(i) * dxm |
---|
| 100 | + dxm/2.0 for i in range(nx_mono)], |
---|
| 101 | dtype=np.float64) |
---|
| 102 | lonm = np.tile(lonm, (ny_mono, 1)).T |
---|
| 103 | |
---|
| 104 | latm = np.array([-90.0 + float(j) * dym |
---|
| 105 | + dym/2.0 for j in range(ny_mono)], |
---|
| 106 | dtype=np.float64) |
---|
| 107 | latm = np.tile(latm, (nx_mono, 1)) |
---|
| 108 | |
---|
| 109 | grid2 = pyoasis.Grid('mono', nx_mono, ny_mono, lonm, latm) |
---|
| 110 | |
---|
| 111 | ncrnm = 4 |
---|
| 112 | clom = pyoasis.asarray(np.zeros((nx_mono, ny_mono, ncrnm), |
---|
| 113 | dtype=np.float64)) |
---|
| 114 | clom[:, :, 0] = lonm[:, :] - dxm/2.0 |
---|
| 115 | clom[:, :, 1] = lonm[:, :] + dxm/2.0 |
---|
| 116 | clom[:, :, 2] = clom[:, :, 1] |
---|
| 117 | clom[:, :, 3] = clom[:, :, 0] |
---|
| 118 | clam = pyoasis.asarray(np.zeros((nx_mono, ny_mono, ncrnm), |
---|
| 119 | dtype=np.float64)) |
---|
| 120 | clam[:, :, 0] = latm[:, :] - dym/2.0 |
---|
| 121 | clam[:, :, 1] = clam[:, :, 0] |
---|
| 122 | clam[:, :, 2] = latm[:, :] + dym/2.0 |
---|
| 123 | clam[:, :, 3] = clam[:, :, 2] |
---|
| 124 | grid2.set_corners(clom, clam) |
---|
| 125 | |
---|
| 126 | mskm = np.zeros((nx_mono, ny_mono), dtype=np.int32) |
---|
| 127 | mskm = np.where(np.power(lonm[:, :]-180., 2) + |
---|
| 128 | np.power(latm[:, :], 2) < 30*30, 1, 0) |
---|
| 129 | grid2.set_mask(mskm) |
---|
| 130 | |
---|
| 131 | frcm = np.ones((nx_mono, ny_mono), dtype=np.float64) |
---|
| 132 | frcm = np.where(mskm == 1, 0.0, 1.0) |
---|
| 133 | |
---|
| 134 | grid2.set_frac(frcm) |
---|
| 135 | |
---|
| 136 | aream = np.zeros((nx_mono, ny_mono), dtype=np.float64) |
---|
| 137 | aream[:, :] = dp_conv * np.abs(np.sin(clam[:, :, 2] * dp_conv) |
---|
| 138 | - np.sin(clam[:, :, 0] * dp_conv))\ |
---|
| 139 | * np.abs(clom[:, :, 1]-clom[:, :, 0]) |
---|
| 140 | |
---|
| 141 | grid2.set_area(aream) |
---|
| 142 | |
---|
| 143 | anglem = np.zeros((nx_mono, ny_mono), dtype=np.float64) |
---|
| 144 | |
---|
| 145 | grid2.set_angle(anglem) |
---|
| 146 | |
---|
| 147 | grid2.write() |
---|
| 148 | |
---|
| 149 | grid.write() |
---|
| 150 | |
---|
| 151 | var_out = pyoasis.Var("FSENDOCN", partition, OASIS.OUT) |
---|
| 152 | var_in = pyoasis.Var("FRECVATM", partition, OASIS.IN) |
---|
| 153 | |
---|
| 154 | comp.enddef() |
---|
| 155 | |
---|
| 156 | date = int(0) |
---|
| 157 | field = pyoasis.asarray(msk) |
---|
| 158 | var_out.put(date, field) |
---|
| 159 | var_in.get(date, field) |
---|
| 160 | |
---|
| 161 | try: |
---|
| 162 | import netCDF4 |
---|
| 163 | import cartopy.crs as ccrs |
---|
| 164 | import matplotlib.pyplot as plt |
---|
| 165 | import matplotlib.collections |
---|
| 166 | from matplotlib.colors import ListedColormap |
---|
| 167 | except ImportError: |
---|
| 168 | if comm_rank == 0: |
---|
| 169 | print("The example completed correctly\n") |
---|
| 170 | print("but no plotting library available.\n") |
---|
| 171 | print("Install netcdf, matplotlib and cartopy\n") |
---|
| 172 | print("for plotting the results\n") |
---|
| 173 | print("In the meantime you can visualize the\n") |
---|
| 174 | print("contents of work/masks.nc with ncview work/masks.nc") |
---|
| 175 | exit(0) |
---|
| 176 | |
---|
| 177 | if comm_rank == 0: |
---|
| 178 | def caramelbleucm(): |
---|
| 179 | rcol = np.hstack((np.linspace(0.1, 0.0, 30 - 0 + 1)[:-1], |
---|
| 180 | np.linspace(0.0, 0.8, 47 - 30 + 1)[:-1], |
---|
| 181 | np.linspace(0.8, 1.0, 52 - 47 + 1)[:-1], |
---|
| 182 | np.linspace(1.0, 1.0, 70 - 52 + 1)[:-1], |
---|
| 183 | np.linspace(1.0, 1.0, 100 - 70 + 1))) |
---|
| 184 | gcol = np.hstack((np.linspace(0.1, 0.9, 30 - 0 + 1)[:-1], |
---|
| 185 | np.linspace(0.9, 1.0, 47 - 30 + 1)[:-1], |
---|
| 186 | np.linspace(1.0, 1.0, 52 - 47 + 1)[:-1], |
---|
| 187 | np.linspace(1.0, 0.9, 70 - 52 + 1)[:-1], |
---|
| 188 | np.linspace(0.9, 0.1, 100 - 70 + 1))) |
---|
| 189 | bcol = np.hstack((np.linspace(1.0, 1.0, 30 - 0 + 1)[:-1], |
---|
| 190 | np.linspace(1.0, 1.0, 47 - 30 + 1)[:-1], |
---|
| 191 | np.linspace(1.0, 0.8, 52 - 47 + 1)[:-1], |
---|
| 192 | np.linspace(0.8, 0.0, 70 - 52 + 1)[:-1], |
---|
| 193 | np.linspace(0.0, 0.1, 100 - 70 + 1))) |
---|
| 194 | alph = np.linspace(1.0, 1.0, 101) |
---|
| 195 | |
---|
| 196 | cm = np.array((np.transpose(rcol), np.transpose(gcol), |
---|
| 197 | np.transpose(bcol), np.transpose(alph))) |
---|
| 198 | cm = np.transpose(cm) |
---|
| 199 | newmap = ListedColormap(cm, name='CaramelBleu') |
---|
| 200 | return newmap |
---|
| 201 | |
---|
| 202 | dgrid = 'pyoa' |
---|
| 203 | gf = netCDF4.Dataset('grids.nc', 'r') |
---|
| 204 | lons = gf.variables[dgrid + '.lon'][:, :].flatten() |
---|
| 205 | lats = gf.variables[dgrid + '.lat'][:, :].flatten() |
---|
| 206 | n_points = lons.size |
---|
| 207 | dgrid_corners = len(gf.dimensions['crn_' + dgrid]) |
---|
| 208 | dlon = gf.variables[dgrid + '.clo'][:].reshape(dgrid_corners, -1) |
---|
| 209 | dlat = gf.variables[dgrid + '.cla'][:].reshape(dgrid_corners, -1) |
---|
| 210 | gf.close() |
---|
| 211 | |
---|
| 212 | mf = netCDF4.Dataset('masks.nc', 'r') |
---|
| 213 | msi = mf.variables[dgrid + '.msk'][:].flatten() |
---|
| 214 | da_msk = msi == 1 |
---|
| 215 | mf.close() |
---|
| 216 | |
---|
| 217 | da_lonlat = np.transpose(np.array([dlon, dlat])) |
---|
| 218 | da_lonlat = np.delete(da_lonlat, np.where(da_msk), axis=0) |
---|
| 219 | |
---|
| 220 | dp_conv = math.pi / 180. |
---|
| 221 | field1 = 2.0 + np.sin(2.0 * lats * dp_conv) ** 4.0 * \ |
---|
| 222 | np.cos(4.0 * lons * dp_conv) |
---|
| 223 | |
---|
| 224 | field1 = np.delete(field1, np.where(da_msk)) |
---|
| 225 | |
---|
| 226 | ti_str = "Grid and mask created by PyOasis" |
---|
| 227 | fig = plt.figure(0, figsize=(8.25, 4.125), frameon=True) |
---|
| 228 | plt.suptitle(ti_str) |
---|
| 229 | cmap = caramelbleucm() |
---|
| 230 | sd_proj = ccrs.PlateCarree() |
---|
| 231 | di_ax = plt.subplot(111, projection=sd_proj) |
---|
| 232 | di_ax.set_global() |
---|
| 233 | # di_ax.coastlines(resolution='110m', linewidth=0.5) |
---|
| 234 | di_pc = matplotlib.collections.PolyCollection(da_lonlat) |
---|
| 235 | di_pc.set_array(field1) |
---|
| 236 | di_pc.set_cmap(cmap) |
---|
| 237 | di_pc.set_edgecolor('black') |
---|
| 238 | di_pc.set_linewidth(0.3) |
---|
| 239 | di_ax.add_collection(di_pc) |
---|
| 240 | di_gl = di_ax.gridlines(crs=ccrs.PlateCarree(), draw_labels=True, |
---|
| 241 | linewidth=0.0, linestyle=':', color='gray') |
---|
| 242 | di_gl.top_labels = False |
---|
| 243 | di_gl.right_labels = False |
---|
| 244 | di_ax.set_title('Precomputed mask') |
---|
| 245 | |
---|
| 246 | plt.subplots_adjust(left=0.10, right=0.90, wspace=0.05, hspace=0.) |
---|
| 247 | plt.show() |
---|
| 248 | |
---|
| 249 | del comp |
---|