import torch import torch.nn as nn import pandas as pd import numpy as np from functools import partial from datetime import datetime, timedelta from pathlib import Path import pickle import dask import dask.array as da import cartopy import cartopy.crs as ccrs import xarray as xr import xarray.ufuncs as xu import matplotlib.pyplot as plt from model.afnonet import AFNONet # download the model code from https://github.com/HFAiLab/FourCastNet/blob/master/model/afnonet.py DATANAMES = ['10m_u_component_of_wind', '10m_v_component_of_wind', '2m_temperature', 'geopotential@1000', 'geopotential@50', 'geopotential@500', 'geopotential@850', 'mean_sea_level_pressure', 'relative_humidity@500', 'relative_humidity@850', 'surface_pressure', 'temperature@500', 'temperature@850', 'total_column_water_vapour', 'u_component_of_wind@1000', 'u_component_of_wind@500', 'u_component_of_wind@850', 'v_component_of_wind@1000', 'v_component_of_wind@500', 'v_component_of_wind@850', 'total_precipitation'] DATAMAP = { 'geopotential': 'z', 'relative_humidity': 'r', 'temperature': 't', 'u_component_of_wind': 'u', 'v_component_of_wind': 'v' } def load_model(): # input size h, w = 720, 1440 x_c, y_c, p_c = 20, 20, 1 backbone_model = AFNONet(img_size=[h, w], in_chans=x_c, out_chans=y_c, norm_layer=partial(nn.LayerNorm, eps=1e-6)) ckpt = torch.load('./backbone.pt', map_location="cpu") backbone_model.load_state_dict(ckpt['model']) precip_model = AFNONet(img_size=[h, w], in_chans=x_c, out_chans=p_c, norm_layer=partial(nn.LayerNorm, eps=1e-6)) ckpt = torch.load('./precipitation.pt', map_location="cpu") precip_model.load_state_dict(ckpt['model']) def imcol(data, img_path, img_name, **kwargs): fig = plt.figure(figsize=(20, 10)) ax = plt.axes(projection=ccrs.PlateCarree()) I = data.plot(ax=ax, transform=ccrs.PlateCarree(), add_colorbar=False, add_labels=False, rasterized=True, **kwargs) ax.coastlines(resolution='110m') dirname = f'{img_path.absolute()}/{img_name}.jpg' plt.axis('off') plt.savefig(dirname, bbox_inches='tight', pad_inches=0.) plt.close(fig) def plot(real_data, pred_data, save_path): cmap_t = 'RdYlBu_r' wind = xu.sqrt(real_data['u10'] ** 2 + real_data['v10'] ** 2) wmin, wmax = wind.values.min(), wind.values.max() wind = xu.sqrt(pred_data['u10'] ** 2 + pred_data['v10'] ** 2) wmin, wmax = min(wind.values.min(), wmin), max(wind.values.max(), wmax) pmin, pmax = real_data['tp'].values.min(), real_data['tp'].values.max() pmin, pmax = min(pred_data['tp'].values.min(), pmin), max(pred_data['tp'].values.max(), pmax) tmin, tmax = real_data['t2m'].values.min(), real_data['t2m'].values.max() tmin, tmax = min(pred_data['t2m'].values.min(), tmin), max(pred_data['t2m'].values.max(), tmax) for i in range(len(real_data.time)): u = real_data['u10'].isel(time=i) v = real_data['v10'].isel(time=i) wind = xu.sqrt(u ** 2 + v ** 2) precip = real_data['tp'].isel(time=i) temp = real_data['t2m'].isel(time=i) datetime = pd.to_datetime(str(wind['time'].values)) datetime = datetime.strftime('%Y-%m-%d %H:%M:%S') print(f'plot {datetime}') imcol(wind, save_path, img_name=f'wind_{datetime}_real', cmap=cmap_t, vmin=wmin, vmax=wmax), imcol(precip, save_path, img_name=f'precipitation_{datetime}_real', cmap=cmap_t, vmin=pmin, vmax=pmax), imcol(temp, save_path, img_name=f'temperature_{datetime}_real', cmap=cmap_t, vmin=tmin, vmax=tmax) for i in range(len(pred_data.time)): u = pred_data['u10'].isel(time=i) v = pred_data['v10'].isel(time=i) wind = xu.sqrt(u ** 2 + v ** 2) precip = pred_data['tp'].isel(time=i) temp = pred_data['t2m'].isel(time=i) datetime = pd.to_datetime(str(wind['time'].values)) datetime = datetime.strftime('%Y-%m-%d %H:%M:%S') print(f'plot {datetime}') imcol(wind, save_path, img_name=f'wind_{datetime}_pred', cmap=cmap_t, vmin=wmin, vmax=wmax), imcol(precip, save_path, img_name=f'precipitation_{datetime}_pred', cmap=cmap_t, vmin=pmin, vmax=pmax), imcol(temp, save_path, img_name=f'temperature_{datetime}_pred', cmap=cmap_t, vmin=tmin, vmax=tmax) def get_pred(sample, scaler, times=None, latitude=None, longitude=None): backbone_model, precip_model = load_model() sample = torch.from_numpy(sample[0]) sample = sample.float() backbone_model.eval() precip_model.eval() pred = [] x = sample.unsqueeze(0).transpose(3, 2).transpose(2, 1) for i in range(len(times)): print(f"predict {times[i]}") with torch.cuda.amp.autocast(): x = backbone_model(x) tmp = x.transpose(1, 2).transpose(2, 3) p = precip_model(x) tmp = tmp.detach().numpy()[0, :, :, :3] * scaler['std'][:3] + scaler['mean'][:3] p = p.detach().numpy()[0, 0, :, :, np.newaxis] * scaler['std'][-1] + scaler['mean'][-1] tmp = np.concatenate([tmp, p], axis=-1) pred.append(tmp) pred = np.asarray(pred) pred_data = xr.Dataset({ 'u10': (['time', 'latitude', 'longitude'], da.from_array(pred[:, :, :, 0], chunks=(7, 720, 1440))), 'v10': (['time', 'latitude', 'longitude'], da.from_array(pred[:, :, :, 1], chunks=(7, 720, 1440))), 't2m': (['time', 'latitude', 'longitude'], da.from_array(pred[:, :, :, 2], chunks=(7, 720, 1440))), 'tp': (['time', 'latitude', 'longitude'], da.from_array(pred[:, :, :, 3], chunks=(7, 720, 1440))), }, coords={'time': (['time'], times), 'latitude': (['latitude'], latitude), 'longitude': (['longitude'], longitude) } ) return pred_data def get_data(start_time, end_time): times = slice(start_time, end_time) with open(f'./scaler.pkl', "rb") as f: # the mean and std of each atmospheric variables scaler = pickle.load(f) # load weather data datas = [] for file in DATANAMES: tmp = xr.open_mfdataset(f'./ERA5_rawdata/{file}/*.nc', combine='by_coords').sel(time=times) if '@' in file: k, v = file.split('@') tmp = tmp.rename_vars({DATAMAP[k]: f'{DATAMAP[k]}@{v}'}) datas.append(tmp) with dask.config.set(**{'array.slicing.split_large_chunks': False}): raw_data = xr.merge(datas, compat="identical", join="inner") data = [] for name in ['u10', 'v10', 't2m', 'z@1000', 'z@50', 'z@500', 'z@850', 'msl', 'r@500', 'r@850', 'sp', 't@500', 't@850', 'tcwv', 'u@1000', 'u@500', 'u@850', 'v@1000', 'v@500', 'v@850']: raw = raw_data[name].values data.append(raw) data = np.stack(data, axis=-1) data = (data - scaler['mean']) / scaler['std'] data = data[:, 1:, :, :] # 721*1440 -> 720*1440 return raw_data[['u10', 'v10', 't2m', 'tp']].sel(expver=1), data, scaler if __name__ == '__main__': start_time = datetime(2023, 1, 1, 0, 0) end_time = datetime(2023, 1, 5, 18, 0) num = int((end_time - start_time) / timedelta(hours=6)) print(f"start_time: {start_time}, end_time: {end_time}, pred_num: {num}") real_data, sample, scaler = get_data(start_time) print(sample.shape) pred_times = [start_time + timedelta(hours=6) * i for i in range(1, num)] pred = get_pred(sample, scaler=scaler, times=pred_times, latitude=real_data.latitude[1:], longitude=real_data.longitude) save_path = Path(f"./output/") save_path.mkdir(parents=True, exist_ok=True) plot(real_data, pred, save_path)