Graphcast：如何完成任务-平芜编程栈

importcdsapiimportdatetimeimportfunctoolsfromgraphcastimportautoregressive,casting,checkpoint,data_utilsasdu,graphcast,normalization,rolloutimporthaikuashkimportisodateimportjaximportmathimportnumpyasnpimportpandasaspdfrompysolar.radiationimportget_radiation_directfrompysolar.solarimportget_altitudeimportpytzimportscipyfromtypingimportDictimportxarray client=cdsapi.Client()# Making a connection to CDS, to fetch data.# The fields to be fetched from the single-level source.singlelevelfields=['10m_u_component_of_wind','10m_v_component_of_wind','2m_temperature','geopotential','land_sea_mask','mean_sea_level_pressure','toa_incident_solar_radiation','total_precipitation']# The fields to be fetched from the pressure-level source.pressurelevelfields=['u_component_of_wind','v_component_of_wind','geopotential','specific_humidity','temperature','vertical_velocity']# The 13 pressure levels.pressure_levels=[50,100,150,200,250,300,400,500,600,700,850,925,1000]# Initializing other required constants.pi=math.pi gap=6# There is a gap of 6 hours between each graphcast prediction.predictions_steps=4# Predicting for 4 timestamps.watts_to_joules=3600first_prediction=datetime.datetime(2024,1,1,18,0)# Timestamp of the first prediction.lat_range=range(-180,181,1)# Latitude range.lon_range=range(0,360,1)# Longitude range.# A utility function used for ease of coding.# Converting the variable to a datetime object.deftoDatetime(dt)->datetime.datetime:ifisinstance(dt,datetime.date)andisinstance(dt,datetime.datetime):returndtelifisinstance(dt,datetime.date)andnotisinstance(dt,datetime.datetime):returndatetime.datetime.combine(dt,datetime.datetime.min.time())elifisinstance(dt,str):if'T'indt:returnisodate.parse_datetime(dt)else:returndatetime.datetime.combine(isodate.parse_date(dt),datetime.datetime.min.time())

输入

在机器学习中，为了获取一些预测，你必须给机器学习模型一些数据，它会基于这些数据给出预测。例如，在预测一个人是否是蝙蝠侠时，输入数据可能是：

他们每晚睡多少小时？
他们的脸上有晒痕吗？
他们在早晨的会议上会打瞌睡吗？
他们的净资产是多少？

同样，Graphcast 也需要某些输入，这些输入通过CDS获得，使用其 Python 库：cdsapi。目前，数据发布者使用创作共用 4.0 国际许可协议，这意味着任何人都可以复制、分发、传输和修改该作品，只要注明原作者。

然而，在使用 cdsapi 获取数据之前，需要进行身份验证，具体说明见此，CDS 提供了相关说明，且过程相对简单。

假设你现在已经获得 CDS 认证，可以创建输入数据，具体步骤如下：

获取单层值：这些值依赖于坐标和时间。所需的输入字段之一是total_precipitation_6hr。顾名思义，这是指从该特定时间戳开始的过去 6 小时的降水量。因此，我们不仅仅获取两个输入时间戳的值，而是需要获取从2024-01-01 00:00 到 12:00的时间戳数据。
获取压力层值：除了依赖于坐标外，还依赖于压力层。因此，在请求数据时，我们会指定所需的压力层数据。在这种情况下，我们只会获取两个输入时间戳的值。
合并单层和压力值：在上述数据上，基于时间、纬度和经度进行内连接操作。
整合年份和天数进度：除了单层和压力字段外，还需要向输入数据中添加四个字段：year_progress_sin、year_progress_cos、day_progress_sin和day_progress_cos。这可以通过graphcast包中提供的函数来实现。

其他小步骤包括：

从 CDS 获取数据后重命名列，因为 CDS 输出的是天气变量的简化形式。
将geopotential变量重命名为geopotential_at_surface，用于单层数据，因为压力层有相同的字段名。
使用math库中的函数，在从 graphcast 获得progress值后，计算 sin 和 cos 值。
将latitude重命名为lat，将longitude重命名为lon，并引入另一个索引：batch，其值为 0。

创建输入数据的代码如下。

# Getting the single and pressure level values.defgetSingleAndPressureValues():client.retrieve('reanalysis-era5-single-levels',{'product_type':'reanalysis','variable':singlelevelfields,'grid':'1.0/1.0','year':[2024],'month':[1],'day':[1],'time':['00:00','01:00','02:00','03:00','04:00','05:00','06:00','07:00','08:00','09:00','10:00','11:00','12:00'],'format':'netcdf'},'single-level.nc')singlelevel=xarray.open_dataset('single-level.nc',engine=scipy.__name__).to_dataframe()singlelevel=singlelevel.rename(columns={col:singlelevelfields[ind]forind,colinenumerate(singlelevel.columns.values.tolist())})singlelevel=singlelevel.rename(columns={'geopotential':'geopotential_at_surface'})# Calculating the sum of the last 6 hours of rainfall.singlelevel=singlelevel.sort_index()singlelevel['total_precipitation_6hr']=singlelevel.groupby(level=[0,1])['total_precipitation'].rolling(window=6,min_periods=1).sum().reset_index(level=[0,1],drop=True)singlelevel.pop('total_precipitation')client.retrieve('reanalysis-era5-pressure-levels',{'product_type':'reanalysis','variable':pressurelevelfields,'grid':'1.0/1.0','year':[2024],'month':[1],'day':[1],'time':['06:00','12:00'],'pressure_level':pressure_levels,'format':'netcdf'},'pressure-level.nc')pressurelevel=xarray.open_dataset('pressure-level.nc',engine=scipy.__name__).to_dataframe()pressurelevel=pressurelevel.rename(columns={col:pressurelevelfields[ind]forind,colinenumerate(pressurelevel.columns.values.tolist())})returnsinglelevel,pressurelevel# Adding sin and cos of the year progress.defaddYearProgress(secs,data):progress=du.get_year_progress(secs)data['year_progress_sin']=math.sin(2*pi*progress)data['year_progress_cos']=math.cos(2*pi*progress)returndata# Adding sin and cos of the day progress.defaddDayProgress(secs,lon:str,data:pd.DataFrame):lons=data.index.get_level_values(lon).unique()progress:np.ndarray=du.get_day_progress(secs,np.array(lons))prxlon={lon:progforlon,proginlist(zip(list(lons),progress.tolist()))}data['day_progress_sin']=data.index.get_level_values(lon).map(lambdax:math.sin(2*pi*prxlon[x]))data['day_progress_cos']=data.index.get_level_values(lon).map(lambdax:math.cos(2*pi*prxlon[x]))returndata# Adding day and year progress.defintegrateProgress(data:pd.DataFrame):fordtindata.index.get_level_values('time').unique():seconds_since_epoch=toDatetime(dt).timestamp()data=addYearProgress(seconds_since_epoch,data)data=addDayProgress(seconds_since_epoch,'longitude'if'longitude'indata.index.nameselse'lon',data)returndata# Adding batch field and renaming some others.defformatData(data:pd.DataFrame)->pd.DataFrame:data=data.rename_axis(index={'latitude':'lat','longitude':'lon'})if'batch'notindata.index.names:data['batch']=0data=data.set_index('batch',append=True)returndataif__name__=='__main__':values:Dict[str,xarray.Dataset]={}single,pressure=getSingleAndPressureValues()values['inputs']=pd.merge(pressure,single,left_index=True,right_index=True,how='inner')values['inputs']=integrateProgress(values['inputs'])values['inputs']=formatData(values['inputs'])

Targets

有 11 个预测字段：

u_component_of_wind,
v_component_of_wind,
geopotential,
specific_humidity,
temperature,
vertical_velocity,
10m_u_component_of_wind,
10m_v_component_of_wind,
2m_temperature,
mean_sea_level_pressure,
total_precipitation.

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/812a9485ce1659178f38f686944419c2.png

由 Ricardo Arce 拍摄，图片来自 Unsplash

传递的目标本质上是一个空的 xarray，用于所有预测字段：

每个coordinate，
predictiontimestamps和
pressure level。

实现这一功能的代码如下所示。

# Includes the packages imported and constants assigned.# The functions created for the inputs also go here.predictionFields=['u_component_of_wind','v_component_of_wind','geopotential','specific_humidity','temperature','vertical_velocity','10m_u_component_of_wind','10m_v_component_of_wind','2m_temperature','mean_sea_level_pressure','total_precipitation_6hr']# Creating an array full of nan values.defnans(*args)->list:returnnp.full((args),np.nan)# Adding or subtracting time.defdeltaTime(dt,**delta)->datetime.datetime:returndt+datetime.timedelta(**delta)defgetTargets(dt,data:pd.DataFrame):# Creating an array consisting of unique values of each index.lat,lon,levels,batch=sorted(data.index.get_level_values('lat').unique().tolist()),sorted(data.index.get_level_values('lon').unique().tolist()),sorted(data.index.get_level_values('level').unique().tolist()),data.index.get_level_values('batch').unique().tolist()time=[deltaTime(dt,hours=days*gap)fordaysinrange(4)]# Creating an empty dataset using latitude, longitude, the pressure levels and each prediction timestamp.target=xarray.Dataset({field:(['lat','lon','level','time'],nans(len(lat),len(lon),len(levels),len(time)))forfieldinpredictionFields},coords={'lat':lat,'lon':lon,'level':levels,'time':time,'batch':batch})returntarget.to_dataframe()if__name__=='__main__':# The code for creating inputs will be here.values['targets']=getTargets(first_prediction,values['inputs'])

强迫因子

正如targets的情况一样，forcings中也包含每个坐标和预测时间戳的值，但不包括压力层级。forcings中的字段包括：

total_incident_solar_radiation，
year_progress_sin，
year_progress_cos，
day_progress_sin，
day_progress_cos。

需要注意的是，以上值是相对于预测时间戳进行分配的。正如在处理inputs时的情况，year和day progress仅依赖于时间戳，而solar radiation来自单层数据源。然而，由于我们正在进行预测，即获取未来的值，对于forcings，太阳辐射值在 CDS 数据集中是不可用的。为此，我们使用pysolar库来模拟太阳辐射值。

# Includes the packages imported and constants assigned.# The functions created for the inputs and targets also go here.# Adding a timezone to datetime.datetime variables.defaddTimezone(dt,tz=pytz.UTC)->datetime.datetime:dt=toDatetime(dt)ifdt.tzinfo==None:returnpytz.UTC.localize(dt).astimezone(tz)else:returndt.astimezone(tz)# Getting the solar radiation value wrt longitude, latitude and timestamp.defgetSolarRadiation(longitude,latitude,dt):altitude_degrees=get_altitude(latitude,longitude,addTimezone(dt))solar_radiation=get_radiation_direct(dt,altitude_degrees)ifaltitude_degrees>0else0returnsolar_radiation*watts_to_joules# Calculating the solar radiation values for timestamps to be predicted.defintegrateSolarRadiation(data:pd.DataFrame):dates=list(data.index.get_level_values('time').unique())coords=[[lat,lon]forlatinlat_rangeforloninlon_range]values=[]# For each data, getting the solar radiation value at a particular coordinate.fordtindates:values.extend(list(map(lambdacoord:{'time':dt,'lon':coord[1],'lat':coord[0],'toa_incident_solar_radiation':getSolarRadiation(coord[1],coord[0],dt)},coords)))# Setting indices.values=pd.DataFrame(values).set_index(keys=['lat','lon','time'])# The forcings dataset will now contain the solar radiation values.returnpd.merge(data,values,left_index=True,right_index=True,how='inner')defgetForcings(data:pd.DataFrame):# Since forcings data does not contain batch as an index, it is dropped.# So are all the columns, since forcings data only has 5, which will be created.forcingdf=data.reset_index(level='level',drop=True).drop(labels=predictionFields,axis=1)# Keeping only the unique indices.forcingdf=pd.DataFrame(index=forcingdf.index.drop_duplicates(keep='first'))# Adding the sin and cos of day and year progress.# Functions are included in the creation of inputs data section.forcingdf=integrateProgress(forcingdf)# Integrating the solar radiation values.forcingdf=integrateSolarRadiation(forcingdf)returnforcingdfif__name__=='__main__':# The code for creating inputs and targets will be here.values['forcings']=getForcings(values['targets'])

后处理输入、目标和强迫因子

现在三大支柱 Graphcast 已经创建完成，我们进入了最后冲刺阶段。就像 NBA 总决赛中，已经赢得了 3 场比赛，现在我们进入最关键的部分，完成任务。

就像科比·布莱恩特曾经说过的，

工作还没有完成。

https://github.com/OpenDocCN/towardsdatascience-blog-zh-2024/raw/master/docs/img/e70074bb06ef6d6bd94183c390f84b2f.png

由 Mike Von 拍摄，图片来自 Unsplash

提到 xarray 时，数据主要有两种类型：

坐标，索引：lat、lon、time……以及
数据变量，列：land_sea_mask、geopotential等等。

每个数据变量包含的每个值，都有一定的坐标分配给它。坐标是数据变量值所依赖的那些坐标。例如，在我们自己的数据中，

land_sea_mask完全依赖于latitude和longitude，即其坐标。
geopotential的坐标是batch、latitude、longitude、time和pressure level。
与此形成鲜明对比，但却合乎逻辑的是，geopotential_at_surface的坐标是latitude和longitude。

因此，在我们继续进行天气预测之前，我们确保每个数据变量都分配到正确的坐标，相关的代码如下所示。

# Includes the packages imported and constants assigned.# The functions created for the inputs, targets and forcings also go here.# A dictionary created, containing each coordinate a data variable requires.classAssignCoordinates:coordinates={'2m_temperature':['batch','lon','lat','time'],'mean_sea_level_pressure':['batch','lon','lat','time'],'10m_v_component_of_wind':['batch','lon','lat','time'],'10m_u_component_of_wind':['batch','lon','lat','time'],'total_precipitation_6hr':['batch','lon','lat','time'],'temperature':['batch','lon','lat','level','time'],'geopotential':['batch','lon','lat','level','time'],'u_component_of_wind':['batch','lon','lat','level','time'],'v_component_of_wind':['batch','lon','lat','level','time'],'vertical_velocity':['batch','lon','lat','level','time'],'specific_humidity':['batch','lon','lat','level','time'],'toa_incident_solar_radiation':['batch','lon','lat','time'],'year_progress_cos':['batch','time'],'year_progress_sin':['batch','time'],'day_progress_cos':['batch','lon','time'],'day_progress_sin':['batch','lon','time'],'geopotential_at_surface':['lon','lat'],'land_sea_mask':['lon','lat'],}defmodifyCoordinates(data:xarray.Dataset):# Parsing through each data variable and removing unneeded indices.forvarinlist(data.data_vars):varArray:xarray.DataArray=data[var]nonIndices=list(set(list(varArray.coords)).difference(set(AssignCoordinates.coordinates[var])))data[var]=varArray.isel(**{coord:0forcoordinnonIndices})data=data.drop_vars('batch')returndatadefmakeXarray(data:pd.DataFrame)->xarray.Dataset:# Converting to xarray.data=data.to_xarray()data=modifyCoordinates(data)returndataif__name__=='__main__':# The code for creating inputs, targets and forcings will be here.values={value:makeXarray(values[value])forvalueinvalues}

使用 Graphcast 进行预测

计算、处理和组装好inputs、targets和forcings后，接下来就是进行predictions的时候了。

现在我们需要模型权重和归一化统计文件，这些文件是由 Deepmind 提供的。

需要下载的文件包括：

stats/diffs_stddev_by_level.nc，
stats/stddev_by_level.nc，
stats/mean_by_level.nc 和
params/GraphCast_small — ERA5 1979–2015 — 分辨率 1.0 — 压力层级 13 — 网格 2to5 — 降水输入和输出.npz。

下方显示了上述文件相对于预测文件的相对路径。保持结构的重要性在于能够成功导入并读取所需的文件。

.├── prediction.py ├── model ├── params ├── GraphCast_small-ERA51979-2015-resolution1.0-pressure levels13-mesh 2to5-precipitationinputandoutput.npz ├── stats ├── diffs_stddev_by_level.nc ├── mean_by_level.nc ├── stddev_by_level.nc

使用Deepmind 提供的预测代码，上述所有功能最终通过下面的代码片段完成预测。

# Includes the packages imported and constants assigned.# The functions created for the inputs, targets and forcings also go here.withopen(r'model/params/GraphCast_small - ERA5 1979-2015 - resolution 1.0 - pressure levels 13 - mesh 2to5 - precipitation input and output.npz','rb')asmodel:ckpt=checkpoint.load(model,graphcast.CheckPoint)params=ckpt.params state={}model_config=ckpt.model_config task_config=ckpt.task_configwithopen(r'model/stats/diffs_stddev_by_level.nc','rb')asf:diffs_stddev_by_level=xarray.load_dataset(f).compute()withopen(r'model/stats/mean_by_level.nc','rb')asf:mean_by_level=xarray.load_dataset(f).compute()withopen(r'model/stats/stddev_by_level.nc','rb')asf:stddev_by_level=xarray.load_dataset(f).compute()defconstruct_wrapped_graphcast(model_config:graphcast.ModelConfig,task_config:graphcast.TaskConfig):predictor=graphcast.GraphCast(model_config,task_config)predictor=casting.Bfloat16Cast(predictor)predictor=normalization.InputsAndResiduals(predictor,diffs_stddev_by_level=diffs_stddev_by_level,mean_by_level=mean_by_level,stddev_by_level=stddev_by_level)predictor=autoregressive.Predictor(predictor,gradient_checkpointing=True)returnpredictor@hk.transform_with_statedefrun_forward(model_config,task_config,inputs,targets_template,forcings):predictor=construct_wrapped_graphcast(model_config,task_config)returnpredictor(inputs,targets_template=targets_template,forcings=forcings)defwith_configs(fn):returnfunctools.partial(fn,model_config=model_config,task_config=task_config)defwith_params(fn):returnfunctools.partial(fn,params=params,state=state)defdrop_state(fn):returnlambda**kw:fn(**kw)[0]run_forward_jitted=drop_state(with_params(jax.jit(with_configs(run_forward.apply))))classPredictor:@classmethoddefpredict(cls,inputs,targets,forcings)->xarray.Dataset:predictions=rollout.chunked_prediction(run_forward_jitted,rng=jax.random.PRNGKey(0),inputs=inputs,targets_template=targets,forcings=forcings)returnpredictionsif__name__=='__main__':# The code for creating inputs, targets, forcings & processing will be here.predictions=Predictor.predict(values['inputs'],values['targets'],values['forcings'])predictions.to_dataframe().to_csv('predictions.csv',sep=',')