SABER
sampling_grids.py
Go to the documentation of this file.
1 #!/usr/bin/env python3
2 
3 import argparse
4 from netCDF4 import Dataset
5 import matplotlib
6 import matplotlib.cm as cm
7 import matplotlib.pyplot as plt
8 import matplotlib.tri as tri
9 import numpy as np
10 import numpy.ma as ma
11 import os
12 
13 def sampling_grids(testdata, test, mpi, omp, suffix, testfig):
14  # c3 plot
15  ic3_plot = [0,2,4,8]
16 
17  # c2 plot
18  ic2_plot = [0,2,4,8]
19 
20  # Loop over files
21  first = True
22  for impi in range(0, int(mpi)):
23  # Open file
24  f = Dataset(testdata + "/" + test + "/test_" + mpi + "-" + omp + "_" + suffix + "_" + mpi.zfill(6) + "-" + format(impi+1, '06d') + ".nc", "r", format="NETCDF4")
25 
26  # Get _FillValue
27  _FillValue = f.__dict__["_FillValue"]
28 
29  for group in f.groups:
30  # Check what is in the file
31  if first:
32  sc3 = ("lon" in f.groups[group].variables)
33  local = (group == "hdiag")
34  vbal = (group == "vbal")
35 
36  # Get data
37  nc0a = f["lon_c0a"].shape[0]
38  if sc3:
39  nc1a = f.groups[group]["lon"].shape[2]
40  if local:
41  nc2a = f.groups[group]["lon_local"].shape[1]
42  if vbal:
43  nc2b = f.groups[group]["lon_vbal"].shape[1]
44  if first:
45  nl0 = f["gmask_c0a"].shape[0]
46  lon_c0 = np.zeros((0))
47  lat_c0 = np.zeros((0))
48  gmask_c0 = np.zeros((0,nl0))
49  if sc3:
50  nc3 = f.groups[group]["lon"].shape[1]
51  lon = np.zeros((0, nl0, nc3))
52  lat = np.zeros((0, nl0, nc3))
53  if local:
54  nc1max = f.groups[group]["lon_local"].shape[2]
55  lon_local = np.zeros((0, nl0, nc1max))
56  lat_local = np.zeros((0, nl0, nc1max))
57  if vbal:
58  nc1max = f.groups[group]["lon_vbal"].shape[2]
59  lon_vbal = np.zeros((0, nl0, nc1max))
60  lat_vbal = np.zeros((0, nl0, nc1max))
61  first = False
62  for ic0a in range(0, nc0a):
63  lon_c0 = np.append(lon_c0, f["lon_c0a"][ic0a].data)
64  lat_c0 = np.append(lat_c0, f["lat_c0a"][ic0a].data)
65  gmask_c0 = np.append(gmask_c0, [f["gmask_c0a"][:,ic0a].data], axis=0)
66  if sc3:
67  for ic1a in range(0, nc1a):
68  lon = np.append(lon, [f.groups[group]["lon"][:,:,ic1a].data], axis=0)
69  lat = np.append(lat, [f.groups[group]["lat"][:,:,ic1a].data], axis=0)
70  if local:
71  for ic2a in range(0, nc2a):
72  lon_local = np.append(lon_local, [f.groups[group]["lon_local"][:,ic2a,:].data], axis=0)
73  lat_local = np.append(lat_local, [f.groups[group]["lat_local"][:,ic2a,:].data], axis=0)
74  if vbal:
75  for ic2b in range(0, nc2b):
76  lon_vbal = np.append(lon_vbal, [f.groups[group]["lon_vbal"][:,ic2b,:].data], axis=0)
77  lat_vbal = np.append(lat_vbal, [f.groups[group]["lat_vbal"][:,ic2b,:].data], axis=0)
78 
79  # Apply mask
80  if sc3:
81  lon = ma.masked_where(lon==_FillValue, lon, False)
82  lat = ma.masked_where(lat==_FillValue, lat, False)
83  if local:
84  lon_local = ma.masked_where(lon_local==_FillValue, lon_local, False)
85  lat_local = ma.masked_where(lat_local==_FillValue, lat_local, False)
86  if vbal:
87  lon_vbal = ma.masked_where(lon_vbal==_FillValue, lon_vbal, False)
88  lat_vbal = ma.masked_where(lat_vbal==_FillValue, lat_vbal, False)
89 
90  # Total sizes
91  nc0 = lon_c0.shape[0]
92  if sc3:
93  nc1 = lon.shape[0]
94  if local:
95  nc2 = lon_local.shape[0]
96  if vbal:
97  nc2btot = lon_vbal.shape[0]
98 
99  # Set geographical mask
100  for ic0 in range(0, nc0):
101  gmask_c0[ic0,:] = gmask_c0[ic0,:]*(ic0%2+1)
102  gmask_c0 = ma.masked_where(gmask_c0<0.5, gmask_c0, False)
103 
104  # Get levels
105  levs = f.__dict__["nam_levs"].split(":")
106 
107  # Set masked values and levels
108  gmask_c0 = ma.masked_invalid(gmask_c0)
109  vmax = np.max(gmask_c0)
110  gmask_c0 = gmask_c0.filled(fill_value=-1.0e38)
111  levels = [-1.0, 0.0, 1.0, 2.0]
112 
113  if sc3:
114  for ic3 in ic3_plot:
115  # Plots
116  fig, ax = plt.subplots(nrows=nl0)
117  fig.subplots_adjust(hspace=0.4)
118  for il0 in range(0, nl0):
119  ax[il0].set_title("@ " + str(il0))
120  ax[il0].tricontourf(lon_c0, lat_c0, gmask_c0[:,il0], levels=levels, cmap="gray")
121  for ic1 in range(0, nc1):
122  if (not lon.mask[ic1,il0,ic3]) and (not lat.mask[ic1,il0,ic3]):
123  if (np.abs(lon[ic1,il0,0]-lon[ic1,il0,ic3]) < 180.0):
124  ax[il0].plot([lon[ic1,il0,0], lon[ic1,il0,ic3]], [lat[ic1,il0,0], lat[ic1,il0,ic3]], "k-")
125  ax[il0].plot([lon[ic1,il0,0], lon[ic1,il0,ic3]], [lat[ic1,il0,0], lat[ic1,il0,ic3]], marker="o", color="black", ms=0.2)
126 
127  # Save and close figure
128  plt.savefig(testfig + "/test_" + mpi + "-" + omp + "_" + suffix + "_" + group + "_c3_" + str(ic3) + ".jpg", format="jpg", dpi=300)
129  plt.close()
130 
131  if sc3 and (local or vbal):
132  # Local sampling
133 
134  # Plots
135  fig, ax = plt.subplots(nrows=nl0)
136  fig.subplots_adjust(hspace=0.4)
137  for il0 in range(0, nl0):
138  ax[il0].set_title("@ " + str(il0))
139  ax[il0].tricontourf(lon_c0, lat_c0, gmask_c0[:,il0], levels=levels, cmap="gray")
140  for ic1 in range(0, nc1):
141  if (not lon.mask[ic1,il0,0]) and (not lat.mask[ic1,il0,0]):
142  ax[il0].plot(lon[ic1,il0,0], lat[ic1,il0,0], marker="o", color="black", ms=0.2)
143 
144  if local:
145  for ic2 in range(0, nc2):
146  if (not lon_local.mask[ic2,il0,0]) and (not lat_local.mask[ic2,il0,0]):
147  ax[il0].plot(lon_local[ic2,il0,0], lat_local[ic2,il0,0], marker="o", color="red", ms=0.2)
148  if vbal:
149  for ic2 in range(0, nc2btot):
150  if (not lon_vbal.mask[ic2,il0,0]) and (not lat_vbal.mask[ic2,il0,0]):
151  ax[il0].plot(lon_vbal[ic2,il0,0], lat_vbal[ic2,il0,0], marker="o", color="red", ms=0.2)
152 
153  # Save and close figure
154  plt.savefig(testfig + "/test_" + mpi + "-" + omp + "_" + suffix + "_" + group + "_c2_in_c1.jpg", format="jpg", dpi=300)
155  plt.close()
156 
157  if local:
158  # Local mask
159  for ic2 in ic2_plot:
160  # Plots
161  fig, ax = plt.subplots(nrows=nl0)
162  fig.subplots_adjust(hspace=0.4)
163  for il0 in range(0, nl0):
164  ax[il0].set_title("@ " + str(il0))
165  ax[il0].tricontourf(lon_c0, lat_c0, gmask_c0[:,il0], levels=levels, cmap="gray")
166  for ic1 in range(0, nc1max):
167  if (not lon_local.mask[ic2,il0,ic1]) and (not lat_local.mask[ic2,il0,ic1]):
168  if (np.abs(lon_local[ic2,il0,0]-lon_local[ic2,il0,ic1]) < 180.0):
169  ax[il0].plot([lon_local[ic2,il0,0], lon_local[ic2,il0,ic1]], [lat_local[ic2,il0,0], lat_local[ic2,il0,ic1]], "k-")
170  ax[il0].plot([lon_local[ic2,il0,0], lon_local[ic2,il0,ic1]], [lat_local[ic2,il0,0], lat_local[ic2,il0,ic1]], marker="o", color="black", ms=0.2)
171 
172  # Save and close figure
173  plt.savefig(testfig + "/test_" + mpi + "-" + omp + "_" + suffix + "_" + group + "_c2_" + str(ic2) + ".jpg", format="jpg", dpi=300)
174  plt.close()
175 
176  if vbal:
177  # Vertical balance mask
178  for ic2 in ic2_plot:
179  # Plots
180  fig, ax = plt.subplots(nrows=nl0)
181  fig.subplots_adjust(hspace=0.4)
182  for il0 in range(0, nl0):
183  ax[il0].set_title("@ " + str(il0))
184  ax[il0].tricontourf(lon_c0, lat_c0, gmask_c0[:,il0], levels=levels, cmap="gray")
185  for ic1 in range(0, nc1max):
186  if (not lon_vbal.mask[ic2,il0,ic1]) and (not lat_vbal.mask[ic2,il0,ic1]):
187  if (np.abs(lon_vbal[ic2,il0,0]-lon_vbal[ic2,il0,ic1]) < 180.0):
188  ax[il0].plot([lon_vbal[ic2,il0,0], lon_vbal[ic2,il0,ic1]], [lat_vbal[ic2,il0,0], lat_vbal[ic2,il0,ic1]], "k-")
189  ax[il0].plot([lon_vbal[ic2,il0,0], lon_vbal[ic2,il0,ic1]], [lat_vbal[ic2,il0,0], lat_vbal[ic2,il0,ic1]], marker="o", color="black", ms=0.2)
190 
191  # Save and close figure
192  plt.savefig(testfig + "/test_" + mpi + "-" + omp + "_" + suffix + "_" + group + "_c2_" + str(ic2) + ".jpg", format="jpg", dpi=300)
193  plt.close()
sampling_grids.sampling_grids
def sampling_grids(testdata, test, mpi, omp, suffix, testfig)
Definition: sampling_grids.py:13