Skip to content
Snippets Groups Projects
Commit e4587446 authored by GUEZ Lionel's avatar GUEZ Lionel
Browse files

Allow the module to be imported

The motivation is to be able to combine a plot of velocity with a plot
of eddy contours.
parent 44af9dd6
No related branches found
No related tags found
No related merge requests found
......@@ -15,112 +15,116 @@ import netCDF4
import wind_cartopy
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
parser.add_argument(
"-s",
"--scale",
default=20,
type=float,
help="scale of arrows for the velocity field",
)
parser.add_argument(
"-w",
"--window",
help="choose a limited plot window",
type=float,
nargs=4,
metavar=("llcrnrlon", "llcrnrlat", "urcrnrlon", "urcrnrlat"),
)
parser.add_argument(
"--save", metavar="format", help="Save file to specified format"
)
parser.add_argument(
"-u",
"--undefined",
action="store_true",
help="plot points where velocity is not defined",
)
parser.add_argument("input_file", help="NetCDF file containing velocity")
args = parser.parse_args()
with netCDF4.Dataset(args.input_file) as f:
if "lon" in f.variables:
lon = "lon"
lat = "lat"
else:
lon = "longitude"
lat = "latitude"
longitude = f[lon][:]
latitude = f[lat][:]
if "time" in f["ugos"].dimensions:
ugos = f["ugos"][0]
vgos = f["vgos"][0]
else:
ugos = f["ugos"][:]
vgos = f["vgos"][:]
if args.window is None:
lon_mask = np.ones(len(longitude), dtype=bool)
lat_mask = np.ones(len(latitude), dtype=bool)
else:
llcrnrlon, llcrnrlat, urcrnrlon, urcrnrlat = args.window
if urcrnrlon - llcrnrlon > 360:
sys.exit("bad values of urcrnrlon and llcrnrlon")
longitude += np.ceil((llcrnrlon - longitude) / 360) * 360
# (in [llcrnrlon, llcrnrlon + 2 pi[)
lon_mask = longitude <= urcrnrlon
lat_mask = np.logical_and(latitude >= llcrnrlat, latitude <= urcrnrlat)
longitude = longitude[lon_mask]
latitude = latitude[lat_mask]
src_crs = ccrs.PlateCarree()
# Use a conformal projection for quiver:
projection = ccrs.Stereographic(
central_latitude=latitude.mean(), central_longitude=longitude.mean()
)
##projection = ccrs.NorthPolarStereo()
fig = plt.figure()
ax = plt.axes(projection=projection)
if args.undefined:
undef_velocity = np.logical_or(ugos.mask, vgos.mask)
lon_2d, lat_2d = np.meshgrid(longitude, latitude)
ax.plot(
lon_2d[undef_velocity].reshape(-1),
lat_2d[undef_velocity].reshape(-1),
transform=src_crs,
marker="*",
color="violet",
linestyle="None",
def plot_velocity():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
else:
quiver_return = wind_cartopy.plot(
ax,
longitude,
latitude,
ugos[lat_mask][:, lon_mask],
vgos[lat_mask][:, lon_mask],
scale=args.scale,
scale_units="width",
parser.add_argument(
"-s",
"--scale",
default=20,
type=float,
help="scale of arrows for the velocity field",
)
ax.quiverkey(
quiver_return, 0.9, 0.9, 1, r"1 m s$^{-1}$", coordinates="figure"
parser.add_argument(
"-w",
"--window",
help="choose a limited plot window",
type=float,
nargs=4,
metavar=("llcrnrlon", "llcrnrlat", "urcrnrlon", "urcrnrlat"),
)
parser.add_argument(
"--save", metavar="format", help="Save file to specified format"
)
parser.add_argument(
"-u",
"--undefined",
action="store_true",
help="plot points where velocity is not defined",
)
parser.add_argument("input_file", help="NetCDF file containing velocity")
args = parser.parse_args()
with netCDF4.Dataset(args.input_file) as f:
if "lon" in f.variables:
lon = "lon"
lat = "lat"
else:
lon = "longitude"
lat = "latitude"
longitude = f[lon][:]
latitude = f[lat][:]
if "time" in f["ugos"].dimensions:
ugos = f["ugos"][0]
vgos = f["vgos"][0]
else:
ugos = f["ugos"][:]
vgos = f["vgos"][:]
if args.window is None:
lon_mask = np.ones(len(longitude), dtype=bool)
lat_mask = np.ones(len(latitude), dtype=bool)
else:
llcrnrlon, llcrnrlat, urcrnrlon, urcrnrlat = args.window
ax.gridlines(draw_labels=True)
ax.coastlines()
if urcrnrlon - llcrnrlon > 360:
sys.exit("bad values of urcrnrlon and llcrnrlon")
longitude += np.ceil((llcrnrlon - longitude) / 360) * 360
# (in [llcrnrlon, llcrnrlon + 2 pi[)
lon_mask = longitude <= urcrnrlon
lat_mask = np.logical_and(latitude >= llcrnrlat, latitude <= urcrnrlat)
longitude = longitude[lon_mask]
latitude = latitude[lat_mask]
src_crs = ccrs.PlateCarree()
# Use a conformal projection for quiver:
projection = ccrs.Stereographic(
central_latitude=latitude.mean(), central_longitude=longitude.mean()
)
##projection = ccrs.NorthPolarStereo()
fig = plt.figure()
ax = plt.axes(projection=projection)
if args.undefined:
undef_velocity = np.logical_or(ugos.mask, vgos.mask)
lon_2d, lat_2d = np.meshgrid(longitude, latitude)
ax.plot(
lon_2d[undef_velocity].reshape(-1),
lat_2d[undef_velocity].reshape(-1),
transform=src_crs,
marker="*",
color="violet",
linestyle="None",
)
else:
quiver_return = wind_cartopy.plot(
ax,
longitude,
latitude,
ugos[lat_mask][:, lon_mask],
vgos[lat_mask][:, lon_mask],
scale=args.scale,
scale_units="width",
)
ax.quiverkey(
quiver_return, 0.9, 0.9, 1, r"1 m s$^{-1}$", coordinates="figure"
)
ax.gridlines(draw_labels=True)
ax.coastlines()
if args.save:
fig.savefig(f"plot_velocity.{args.save}")
print(f'Created "plot_velocity.{args.save}".')
else:
plt.show()
if args.save:
fig.savefig(f"plot_velocity.{args.save}")
print(f'Created "plot_velocity.{args.save}".')
else:
plt.show()
if __name__ == "__main__":
plot_velocity()
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment