import unittest as ut from os.path import join, basename import numpy as np import matplotlib.pyplot as plt import matplotlib.cm as cm from netCDF4 import Dataset as NetCDF from wrf import get_proj_params from wrf.projection import getproj, RotatedLatLon, PolarStereographic PYNGL = True try: import Ngl except ImportError: PYNGL = False BASEMAP = True try: import mpl_toolkits.basemap except ImportError: BASEMAP = False CARTOPY = True try: from cartopy import crs, feature except ImportError: CARTOPY = False FILE_DIR = "/Users/ladwig/Documents/wrf_files/" WRF_FILES = [join(FILE_DIR, "norway", "geo_em.d01.nc"), join(FILE_DIR, "rotated_pole", "EAS_geo_em.d01.nc"), join(FILE_DIR, "rotated_pole", "EUR_geo_em.d01.nc"), join(FILE_DIR, "wrfout_d01_2016-02-25_18_00_00"), join(FILE_DIR, "wrfout_d01_2008-09-29_23-30-00"), join(FILE_DIR, "wrfout_d01_2010-06-13_21:00:00")] def nz_proj(): lats = np.array([[-47.824014, -47.824014], [-32.669853, -32.669853]]) lons = np.array([[163.839595, -179.693502], [163.839595, -179.693502]]) params = {"MAP_PROJ": 6, "CEN_LAT": -41.814869, "CEN_LON": 179.693502, "TRUELAT1": 0, "TRUELAT2": 0, "MOAD_CEN_LAT": -41.814869, "STAND_LON": 180.0 - 179.693502, "POLE_LAT": 48.185131, "POLE_LON": 0.0} return lats, lons, RotatedLatLon(lats=lats, lons=lons, **params) def argentina_proj(): lats = np.array([[-57.144064, -57.144064], [-21.154470, -21.154470]]) lons = np.array([[-86.893797, -37.089724], [-86.893797, -37.089724]]) params = {"MAP_PROJ": 6, "CEN_LAT": -39.222954, "CEN_LON": -65.980109, "TRUELAT1": 0, "TRUELAT2": 0, "MOAD_CEN_LAT": -39.222954, "STAND_LON": 180.0 - -65.980109, "POLE_LAT": 90 + -39.222954, "POLE_LON": 0.0} return lats, lons, RotatedLatLon(lats=lats, lons=lons, **params) def south_polar_proj(): lats = np.array([[-30.0, -30.0], [-30.0, -30.0]]) lons = np.array([[-120, 60], [-120, 60]]) params = {"MAP_PROJ": 2, "CEN_LAT": -90.0, "CEN_LON": 0, "TRUELAT1": -10.0, "MOAD_CEN_LAT": -90.0, "STAND_LON": 0} return lats, lons, PolarStereographic(lats=lats, lons=lons, **params) def north_polar_proj(): lats = np.array([[30.0, 30.0], [30.0, 30.0]]) lons = np.array([[-45, 140], [-45, 140]]) params = {"MAP_PROJ": 2, "CEN_LAT": 90.0, "CEN_LON": 10, "TRUELAT1": 10.0, "MOAD_CEN_LAT": 90.0, "STAND_LON": 10} return lats, lons, PolarStereographic(lats=lats, lons=lons, **params) def dateline_rot_proj(): lats = np.array([[60.627974, 60.627974], [71.717521, 71.717521]]) lons = np.array([[170.332771, -153.456292], [170.332771, -153.456292]]) params = {"MAP_PROJ": 6, "CEN_LAT": 66.335764, "CEN_LON": -173.143792, "TRUELAT1": 0, "TRUELAT2": 0, "MOAD_CEN_LAT": 66.335764, "STAND_LON": 173.143792, "POLE_LAT": 90.0 - 66.335764, "POLE_LON": 180.0} return lats, lons, RotatedLatLon(lats=lats, lons=lons, **params) class WRFProjTest(ut.TestCase): longMessage = True def make_test(wrf_file=None, fixed_case=None): if wrf_file is not None: ncfile = NetCDF(wrf_file) lats, lons, proj_params = get_proj_params(ncfile) proj = getproj(lats=lats, lons=lons, **proj_params) name_suffix = basename(wrf_file) elif fixed_case is not None: name_suffix = fixed_case if fixed_case == "south_rot": lats, lons, proj = nz_proj() elif fixed_case == "arg_rot": lats, lons, proj = argentina_proj() elif fixed_case == "south_polar": lats, lons, proj = south_polar_proj() elif fixed_case == "north_polar": lats, lons, proj = north_polar_proj() elif fixed_case == "dateline_rot": lats, lons, proj = dateline_rot_proj() print("wrf proj4: {}".format(proj.proj4())) if PYNGL: # PyNGL plotting wks_type = bytes("png") wks = Ngl.open_wks(wks_type, bytes("pyngl_{}".format(name_suffix))) mpres = proj.pyngl() map = Ngl.map(wks, mpres) Ngl.delete_wks(wks) if BASEMAP: # Basemap plotting fig = plt.figure(figsize=(10, 10)) ax = fig.add_axes([0.1, 0.1, 0.8, 0.8]) # Define and plot the meridians and parallels min_lat = np.amin(lats) max_lat = np.amax(lats) min_lon = np.amin(lons) max_lon = np.amax(lons) parallels = np.arange(np.floor(min_lat), np.ceil(max_lat), (max_lat - min_lat)/5.0) meridians = np.arange(np.floor(min_lon), np.ceil(max_lon), (max_lon - min_lon)/5.0) bm = proj.basemap() bm.drawcoastlines(linewidth=.5) # bm.drawparallels(parallels,labels=[1,1,1,1],fontsize=10) # bm.drawmeridians(meridians,labels=[1,1,1,1],fontsize=10) print("basemap proj4: {}".format(bm.proj4string)) plt.savefig("basemap_{}.png".format(name_suffix)) plt.close(fig) if CARTOPY: # Cartopy plotting fig = plt.figure(figsize=(10, 10)) ax = plt.axes([0.1, 0.1, 0.8, 0.8], projection=proj.cartopy()) print("cartopy proj4: {}".format(proj.cartopy().proj4_params)) ax.coastlines('50m', linewidth=0.8) # print proj.x_extents() # print proj.y_extents() ax.set_xlim(proj.cartopy_xlim()) ax.set_ylim(proj.cartopy_ylim()) ax.gridlines() plt.savefig("cartopy_{}.png".format(name_suffix)) plt.close(fig) if __name__ == "__main__": for wrf_file in WRF_FILES: test_func = make_test(wrf_file=wrf_file) setattr(WRFProjTest, "test_proj", test_func) test_func2 = make_test(fixed_case="south_rot") setattr(WRFProjTest, "test_south_rot", test_func2) test_func3 = make_test(fixed_case="arg_rot") setattr(WRFProjTest, "test_arg_rot", test_func3) test_func4 = make_test(fixed_case="south_polar") setattr(WRFProjTest, "test_south_polar", test_func4) test_func5 = make_test(fixed_case="north_polar") setattr(WRFProjTest, "test_north_polar", test_func5) test_func6 = make_test(fixed_case="dateline_rot") setattr(WRFProjTest, "test_dateline_rot", test_func6) ut.main()