WRF domain 绘制改进

from netCDF4 import Dataset
import numpy as np
from wrf import getvar, get_cartopy
import matplotlib.pyplot as plt
from cartopy.mpl.gridliner import LONGITUDE_FORMATTER, LATITUDE_FORMATTER
import matplotlib.transforms as mtransforms
import cartopy.crs as ccrs
from cartopy.io.shapereader import BasicReader
import cmaps
from matplotlib.path import Path
from cartopy.mpl.patch import geos_to_path
from mpl_toolkits.axes_grid1.inset_locator import TransformedBbox, BboxPatch, BboxConnector
import shapely.geometry as sgeom
from copy import copy
import concurrent.futures

'''
参考Andrew Dawson 提供的解决方法: https://gist.github.com/ajdawson/dd536f786741e987ae4e
'''
# 给出一个假定为矩形的LineString，返回对应于矩形给定边的直线。
def find_side(ls, side):
    """
    Given a shapely LineString which is assumed to be rectangular, return the
    line corresponding to a given side of the rectangle.
    """
    minx, miny, maxx, maxy = ls.bounds
    points = {'left': [(minx, miny), (minx, maxy)],
              'right': [(maxx, miny), (maxx, maxy)],
              'bottom': [(minx, miny), (maxx, miny)],
              'top': [(minx, maxy), (maxx, maxy)],}
    return sgeom.LineString(points[side])

# 在兰伯特投影的底部X轴上绘制刻度线
def lambert_xticks(ax, ticks):
    """Draw ticks on the bottom x-axis of a Lambert Conformal projection."""
    te = lambda xy: xy[0]
    lc = lambda t, n, b: np.vstack((np.zeros(n) + t, np.linspace(b[2], b[3], n))).T
    xticks, xticklabels = _lambert_ticks(ax, ticks, 'bottom', lc, te)
    ax.xaxis.tick_bottom()
    ax.set_xticks(xticks)
    ax.set_xticklabels([ax.xaxis.get_major_formatter()(xtick) for xtick in xticklabels])

# 在兰伯特投影的左侧y轴上绘制刻度线
def lambert_yticks(ax, ticks):
    """Draw ricks on the left y-axis of a Lamber Conformal projection."""
    te = lambda xy: xy[1]
    lc = lambda t, n, b: np.vstack((np.linspace(b[0], b[1], n), np.zeros(n) + t)).T
    yticks, yticklabels = _lambert_ticks(ax, ticks, 'left', lc, te)
    ax.yaxis.tick_left()
    ax.set_yticks(yticks)
    ax.set_yticklabels([ax.yaxis.get_major_formatter()(ytick) for ytick in yticklabels])

# 获取兰伯特投影中底部X轴或左侧y轴的刻度线位置和标签
def _lambert_ticks(ax, ticks, tick_location, line_constructor, tick_extractor):
    """Get the tick locations and labels for an axis of a Lambert Conformal projection."""
    outline_patch = sgeom.LineString(ax.outline_patch.get_path().vertices.tolist())
    axis = find_side(outline_patch, tick_location)
    n_steps = 30
    extent = ax.get_extent(ccrs.PlateCarree())
    _ticks = []
    for t in ticks:
        xy = line_constructor(t, n_steps, extent)
        proj_xyz = ax.projection.transform_points(ccrs.Geodetic(), xy[:, 0], xy[:, 1])
        xyt = proj_xyz[..., :2]
        ls = sgeom.LineString(xyt.tolist())
        locs = axis.intersection(ls)
        if not locs:
            tick = [None]
        else:
            tick = tick_extractor(locs.xy)
        _ticks.append(tick[0])
    # Remove ticks that aren't visible: 
    ticklabels = copy(ticks)
    while True:
        try:
            index = _ticks.index(None)
        except ValueError:
            break
        _ticks.pop(index)
        ticklabels.pop(index)
    return _ticks, ticklabels

## 子图连线函数
def mark_inset(parent_axes, inset_axes, loc1a=1, loc1b=1, loc2a=2, loc2b=2, **kwargs):
    rect = TransformedBbox(inset_axes.viewLim, parent_axes.transData)

    pp = BboxPatch(rect, fill=False, **kwargs)
    parent_axes.add_patch(pp)

    p1 = BboxConnector(inset_axes.bbox, rect, loc1=loc1a, loc2=loc1b, **kwargs)
    inset_axes.add_patch(p1)
    p1.set_clip_on(False)
    p2 = BboxConnector(inset_axes.bbox, rect, loc1=loc2a, loc2=loc2b, **kwargs)
    inset_axes.add_patch(p2)
    p2.set_clip_on(False)

    return pp, p1, p2


## 等高子图设置
def add_equal_axes(ax, loc, pad, width, projection):
    '''
    在原有的Axes旁新添一个等高或等宽的Axes并返回该对象.

    Parameters
    ----------
    ax : Axes or array_like of Axes
        原有的Axes,也可以为一组Axes构成的数组.

    loc : {'left', 'right', 'bottom', 'top'}
        新Axes相对于旧Axes的位置.

    pad : float
        新Axes与旧Axes的间距.

    width: float
        当loc='left'或'right'时,width表示新Axes的宽度.
        当loc='bottom'或'top'时,width表示新Axes的高度.

    Returns
    -------
    ax_new : Axes
        新Axes对象.
    '''
    # 无论ax是单个还是一组Axes,获取ax的大小位置.
    axes = np.atleast_1d(ax).ravel()
    bbox = mtransforms.Bbox.union([ax.get_position() for ax in axes])

    # 决定新Axes的大小位置.
    if loc == 'left':
        x0_new = bbox.x0 - pad - width
        x1_new = x0_new + width
        y0_new = bbox.y0
        y1_new = bbox.y1
    elif loc == 'right':
        x0_new = bbox.x1 + pad
        x1_new = x0_new + width
        y0_new = bbox.y0
        y1_new = bbox.y1
    elif loc == 'bottom':
        x0_new = bbox.x0
        x1_new = bbox.x1
        y0_new = bbox.y0 - pad - width
        y1_new = y0_new + width
    elif loc == 'top':
        x0_new = bbox.x0
        x1_new = bbox.x1
        y0_new = bbox.y1 + pad
        y1_new = y0_new + width
    # 创建新Axes.
    fig = axes[0].get_figure()
    bbox_new = mtransforms.Bbox.from_extents(x0_new, y0_new, x1_new, y1_new)
    ax_new = fig.add_axes(bbox_new, projection=projection)

    return ax_new

def process_data():
    with concurrent.futures.ThreadPoolExecutor() as executor:
        future_d01 = executor.submit(get_data, '/home/mw/input/GEO3900/geo_em.d01.nc')
        future_d02 = executor.submit(get_data, '/home/mw/input/GEO3900/geo_em.d02.nc')
        future_d03 = executor.submit(get_data, '/home/mw/input/GEO3900/geo_em.d03.nc')
        lon, lat, proj, hgt = future_d01.result()
        lon_1, lat_1,proj, hgt_1 = future_d02.result()
        lon_2, lat_2,proj, hgt_2 = future_d03.result()
 
    return lon, lat, proj, hgt, lon_1, lat_1, hgt_1 ,lon_2, lat_2,hgt_2

def get_data(file):
    with Dataset(file) as f:
        lon = getvar(f, 'lon', meta=False)
        lat = getvar(f, 'lat', meta=False)
        proj = get_cartopy(wrfin=f)

        hgt = getvar(f, 'HGT_M', meta=False)
        hgt[hgt < 0] = 0

    return lon, lat, proj, hgt

lon, lat, proj, hgt, lon_1, lat_1, hgt_1 ,lon_2, lat_2,hgt_2= process_data()

provinces = BasicReader('/home/mw/input/data5246/中国地图/China_provinces/China_provinces.shp')
countries = BasicReader('/home/mw/input/data5246/中国地图/world_countries/world_countries.shp')
cmap = cmaps.MPL_terrain

#绘图部分
fig = plt.figure(figsize=(20, 12),dpi=200)
plt.tight_layout()
# 在figure上添加子图
ax = fig.add_axes([0.09, 0.15, 0.4, 0.7], projection=proj)
contour = ax.contourf(lon, lat, hgt, transform=ccrs.PlateCarree(), cmap=cmap, levels=np.arange(0, 5000 + 100, 100))
ax.add_geometries(provinces.geometries(), linewidth=.5, edgecolor='black', crs=ccrs.PlateCarree(),
                  facecolor='none') 
ax.add_geometries(countries.geometries(), linewidth=1., edgecolor='black', crs=ccrs.PlateCarree(),
facecolor='none')
# 添加经纬度网格和刻度
# *must* call draw in order to get the axis boundary used to add ticks:
fig.canvas.draw()
# Define gridline locations and draw the lines using cartopy's built-in gridliner:
xticks=[80,85,90,95,100,105,110,115,120]
yticks=[10,15,20,25,30,35,40,45,50,55]
ax.gridlines(xlocs=xticks, ylocs=yticks)
font3={'family':'SimHei','size':14,'color':'k'}
ax.gridlines(xlocs=xticks, ylocs=yticks, draw_labels=False, linewidth=0.8, color='k', alpha=0.6, linestyle='--')
# Label the end-points of the gridlines using the custom tick makers:
ax.xaxis.set_major_formatter(LONGITUDE_FORMATTER)
ax.yaxis.set_major_formatter(LATITUDE_FORMATTER)
lambert_xticks(ax, xticks)
lambert_yticks(ax, yticks)
#g1.rotate_labels = False

# 在d01的模拟区域上框出d02的模拟区域范围
ax.plot([lon_1[0, 0], lon_1[-1, 0]], [lat_1[0, 0], lat_1[-1, 0]], color='k', lw=1.75, transform=ccrs.PlateCarree())
ax.plot([lon_1[-1, 0], lon_1[-1, -1]], [lat_1[-1, 0], lat_1[-1, -1]], color='k', lw=1.75, transform=ccrs.PlateCarree())
ax.plot([lon_1[-1, -1], lon_1[0, -1]], [lat_1[-1, -1], lat_1[0, -1]], color='k', lw=1.75, transform=ccrs.PlateCarree())
ax.plot([lon_1[0, -1], lon_1[0, 0]], [lat_1[0, -1], lat_1[0, 0]], color='k', lw=1.75, transform=ccrs.PlateCarree())



#添加标题和颜色条
plt.title('D01 DOMAIN')
cbar = plt.colorbar(contour, ax=ax, orientation='horizontal', pad=0.05, shrink=0.7)
cbar.set_label('height (m)')

#在右上角加一个小图示例
ax_inset = add_equal_axes(ax, loc='right', pad=0.1, width=0.4, projection=proj)
contour_inset = ax_inset.contourf(lon_1, lat_1, hgt_1, transform=ccrs.PlateCarree(), cmap=cmap,
levels=np.arange(0, 5000 + 100, 100))
ax_inset.add_geometries(provinces.geometries(), linewidth=.5, edgecolor='black', crs=ccrs.PlateCarree(),
facecolor='none')
ax_inset.add_geometries(countries.geometries(), linewidth=1., edgecolor='black', crs=ccrs.PlateCarree(),
facecolor='none')
# 添加经纬度网格和刻度
fig.canvas.draw()
# Define gridline locations and draw the lines using cartopy's built-in gridliner:
ax_inset.gridlines(xlocs=xticks, ylocs=yticks)
ax_inset.gridlines(xlocs=xticks, ylocs=yticks, draw_labels=False, linewidth=0.8, color='k', alpha=0.6, linestyle='--')
# Label the end-points of the gridlines using the custom tick makers:
ax_inset.xaxis.set_major_formatter(LONGITUDE_FORMATTER)
ax_inset.yaxis.set_major_formatter(LATITUDE_FORMATTER)
lambert_xticks(ax_inset, xticks)
lambert_yticks(ax_inset, yticks)

#添加子图连线
mark_inset(ax, ax_inset, loc1a=2, loc1b=1, loc2a=3, loc2b=4, fc="none", ec='k', lw=0.75)
plt.show()

provinces = BasicReader('/home/mw/input/data5246/中国地图/China_provinces/China_provinces.shp')
countries = BasicReader('/home/mw/input/data5246/中国地图/world_countries/world_countries.shp')
cmap = cmaps.MPL_terrain

#绘图部分
fig = plt.figure(figsize=(20, 12),dpi=200)
plt.tight_layout()
# 在figure上添加子图
ax = fig.add_axes([0.09, 0.15, 0.4, 0.7], projection=proj)
contour = ax.contourf(lon, lat, hgt, transform=ccrs.PlateCarree(), cmap=cmap, levels=np.arange(0, 5000 + 100, 100))
ax.add_geometries(provinces.geometries(), linewidth=.5, edgecolor='black', crs=ccrs.PlateCarree(),
                  facecolor='none') 
ax.add_geometries(countries.geometries(), linewidth=1., edgecolor='black', crs=ccrs.PlateCarree(),
facecolor='none')
# 添加经纬度网格和刻度
# *must* call draw in order to get the axis boundary used to add ticks:
fig.canvas.draw()
# Define gridline locations and draw the lines using cartopy's built-in gridliner:
xticks=[80,85,90,95,100,105,110,115,120]
yticks=[10,15,20,25,30,35,40,45,50,55]
ax.gridlines(xlocs=xticks, ylocs=yticks)
font3={'family':'SimHei','size':14,'color':'k'}
ax.gridlines(xlocs=xticks, ylocs=yticks, draw_labels=False, linewidth=0.8, color='k', alpha=0.6, linestyle='--')
# Label the end-points of the gridlines using the custom tick makers:
ax.xaxis.set_major_formatter(LONGITUDE_FORMATTER)
ax.yaxis.set_major_formatter(LATITUDE_FORMATTER)
lambert_xticks(ax, xticks)
lambert_yticks(ax, yticks)
#g1.rotate_labels = False

# 在d01的模拟区域上框出d02的模拟区域范围
ax.plot([lon_1[0, 0], lon_1[-1, 0]], [lat_1[0, 0], lat_1[-1, 0]], color='k', lw=1.75, transform=ccrs.PlateCarree())
ax.plot([lon_1[-1, 0], lon_1[-1, -1]], [lat_1[-1, 0], lat_1[-1, -1]], color='k', lw=1.75, transform=ccrs.PlateCarree())
ax.plot([lon_1[-1, -1], lon_1[0, -1]], [lat_1[-1, -1], lat_1[0, -1]], color='k', lw=1.75, transform=ccrs.PlateCarree())
ax.plot([lon_1[0, -1], lon_1[0, 0]], [lat_1[0, -1], lat_1[0, 0]], color='k', lw=1.75, transform=ccrs.PlateCarree())

# 在d02的模拟区域上框出d03的模拟区域范围
ax.plot([lon_2[0, 0], lon_2[-1, 0]], [lat_2[0, 0], lat_2[-1, 0]], color='k', lw=1.75, transform=ccrs.PlateCarree())
ax.plot([lon_2[-1, 0], lon_2[-1, -1]], [lat_2[-1, 0], lat_2[-1, -1]], color='k', lw=1.75, transform=ccrs.PlateCarree())
ax.plot([lon_2[-1, -1], lon_2[0, -1]], [lat_2[-1, -1], lat_2[0, -1]], color='k', lw=1.75, transform=ccrs.PlateCarree())
ax.plot([lon_2[0, -1], lon_2[0, 0]], [lat_2[0, -1], lat_2[0, 0]], color='k', lw=1.75, transform=ccrs.PlateCarree())

#添加标题和颜色条
plt.title('D01 DOMAIN')
cbar = plt.colorbar(contour, ax=ax, orientation='horizontal', pad=0.05, shrink=0.7)
cbar.set_label('height (m)')

#在右上角加一个小图示例
ax_inset = add_equal_axes(ax, loc='right', pad=0.1, width=0.4, projection=proj)
contour_inset = ax_inset.contourf(lon_2, lat_2, hgt_2, transform=ccrs.PlateCarree(), cmap=cmap,
levels=np.arange(0, 5000 + 100, 100))
ax_inset.add_geometries(provinces.geometries(), linewidth=.5, edgecolor='black', crs=ccrs.PlateCarree(),
facecolor='none')
ax_inset.add_geometries(countries.geometries(), linewidth=1., edgecolor='black', crs=ccrs.PlateCarree(),
facecolor='none')
# 添加经纬度网格和刻度
fig.canvas.draw()
# Define gridline locations and draw the lines using cartopy's built-in gridliner:
ax_inset.gridlines(xlocs=xticks, ylocs=yticks)
ax_inset.gridlines(xlocs=xticks, ylocs=yticks, draw_labels=False, linewidth=0.8, color='k', alpha=0.6, linestyle='--')
# Label the end-points of the gridlines using the custom tick makers:
ax_inset.xaxis.set_major_formatter(LONGITUDE_FORMATTER)
ax_inset.yaxis.set_major_formatter(LATITUDE_FORMATTER)
lambert_xticks(ax_inset, xticks)
lambert_yticks(ax_inset, yticks)

#添加子图连线
mark_inset(ax, ax_inset, loc1a=2, loc1b=1, loc2a=3, loc2b=4, fc="none", ec='k', lw=0.75)
plt.show()