【Hackathon 8th No.17】FuXi 论文复现 by xiaoyewww · Pull Request #1145 · PaddlePaddle/PaddleScience (original) (raw)

def visualize(save_name, vars=[], titles=[], vmin=None, vmax=None):
import cartopy.crs as ccrs
import matplotlib.pyplot as plt
fig, ax = plt.subplots(
len(vars), 1, figsize=(8, 6), subplot_kw={"projection": ccrs.PlateCarree()}
)
def plot(ax, v, title):
v.plot(
ax=ax,
x="lon",
y="lat",
vmin=vmin,
vmax=vmax,
transform=ccrs.PlateCarree(),
add_colorbar=False,
)
# ax.coastlines()
ax.set_title(title)
gl = ax.gridlines(draw_labels=True, linewidth=0.5)
gl.top_labels = False
gl.right_labels = False
for i, v in enumerate(vars):
if len(vars) == 1:
plot(ax, v, titles[i])
else:
plot(ax[i], v, titles[i])
plt.savefig(
save_name, bbox_inches="tight", pad_inches=0.1, transparent="true", dpi=200
)
plt.close()
def test_visualize(step, data_dir):
src_name = os.path.join(data_dir, f"{step:03d}.nc")
ds = xr.open_dataarray(src_name).isel(time=0)
ds = ds.sel(lon=slice(90, 150), lat=slice(50, 0))
print(ds)
u850 = ds.sel(level="U850", step=step)
v850 = ds.sel(level="V850", step=step)
ws850 = np.sqrt(u850**2 + v850**2)
visualize(
f"ws850/{step:03d}.jpg", [ws850], [f"20230725-18+{step:03d}h"], vmin=0, vmax=30
)