import xarray as xr
import datetime as dt
import fsspec
import ujson
from kerchunk.grib2 import scan_grib
from kerchunk.combine import MultiZarrToZarr
fsspec
file systems to read grib2 forecast files from AWS and write reference jsons to an aws bucket in this case
fs_read = fsspec.filesystem('s3', anon=True, skip_instance_cache=True)
fs_write = fsspec.filesystem('s3', anon=False)
#fs_write = fsspec.filesystem('') #uncomment this if you intend to write jsons to a local folder
Get the latest forecast
days_avail = fs_read.glob('s3://noaa-hrrr-bdp-pds/hrrr.*')
files = fs_read.glob(f's3://{days_avail[-1]}/conus/*wrfsfcf01.grib2')
files = sorted(['s3://'+f for f in files])
files
['s3://noaa-hrrr-bdp-pds/hrrr.20220804/conus/hrrr.t00z.wrfsfcf01.grib2', 's3://noaa-hrrr-bdp-pds/hrrr.20220804/conus/hrrr.t01z.wrfsfcf01.grib2', 's3://noaa-hrrr-bdp-pds/hrrr.20220804/conus/hrrr.t02z.wrfsfcf01.grib2', 's3://noaa-hrrr-bdp-pds/hrrr.20220804/conus/hrrr.t03z.wrfsfcf01.grib2', 's3://noaa-hrrr-bdp-pds/hrrr.20220804/conus/hrrr.t04z.wrfsfcf01.grib2', 's3://noaa-hrrr-bdp-pds/hrrr.20220804/conus/hrrr.t05z.wrfsfcf01.grib2', 's3://noaa-hrrr-bdp-pds/hrrr.20220804/conus/hrrr.t06z.wrfsfcf01.grib2', 's3://noaa-hrrr-bdp-pds/hrrr.20220804/conus/hrrr.t07z.wrfsfcf01.grib2', 's3://noaa-hrrr-bdp-pds/hrrr.20220804/conus/hrrr.t08z.wrfsfcf01.grib2', 's3://noaa-hrrr-bdp-pds/hrrr.20220804/conus/hrrr.t09z.wrfsfcf01.grib2']
scan_grib
does not require a filter and will happily create a reference file for each available grib message. However when combining the grib messages using MultiZarrToZarr
it is neccassary for the messages to share a coordinate system. Thus to make our lives easier and ensure all reference outputs from scan_grib
share a coordinate system we pass a filter argument.
afilter={'typeOfLevel': 'heightAboveGround', 'level': [2, 10]}
so = {"anon": True}
json_dir = 's3://esip-qhub/noaa/hrrr/jsons/'
def make_json_name(file_url, message_number): #create a unique name for each reference file
date = file_url.split('/')[3].split('.')[1]
name = file_url.split('/')[5].split('.')[1:3]
return f'{json_dir}{date}_{name[0]}_{name[1]}_message{message_number}.json'
def gen_json(file_url):
out = scan_grib(file_url, storage_options=so, filter=afilter) #create the reference using scan_grib
for i, message in enumerate(out): # scan_grib outputs a list containing one reference per grib message
out_file_name = make_json_name(file_url, i) #get name
with fs_write.open(out_file_name, "w") as f:
f.write(ujson.dumps(message)) #write to file
fs_write.rm(json_dir) # clear json directory of old references
%%time
#this step is best run via a cluster
for f in files:
gen_json(f)
CPU times: user 1min 45s, sys: 11.9 s, total: 1min 57s Wall time: 4min 42s
reference_jsons = fs_write.ls(json_dir) #get list of file names
reference_jsons = sorted(['s3://'+f for f in reference_jsons]) #prepend s3 protocol (not neccessary if writing to local filesystem)
#combine individual references into single consolidated reference
mzz = MultiZarrToZarr(reference_jsons,
concat_dims = ['valid_time'],
identical_dims=['latitude', 'longitude', 'heightAboveGround', 'step'])
%%time
d = mzz.translate()
CPU times: user 1.4 s, sys: 52.3 ms, total: 1.45 s Wall time: 20.4 s
#open dataset as zarr object using fsspec reference file system and xarray
fs = fsspec.filesystem("reference", fo=d, remote_protocol='s3', remote_options={'anon':True})
m = fs.get_mapper("")
ds = xr.open_dataset(m, engine="zarr", backend_kwargs=dict(consolidated=False),
chunks={'valid_time':1})
ds
<xarray.Dataset> Dimensions: (valid_time: 10, x: 1059, y: 1799, heightAboveGround: 1, step: 1, time: 1) Coordinates: * heightAboveGround (heightAboveGround) int64 2 * step (step) timedelta64[ns] 01:00:00 * time (time) datetime64[ns] 2022-08-04 * valid_time (valid_time) datetime64[ns] 2022-08-04T01:00:00 ... 20... Dimensions without coordinates: x, y Data variables: 10si (valid_time, x, y) float64 dask.array<chunksize=(1, 1059, 1799), meta=np.ndarray> 10u (valid_time, x, y) float64 dask.array<chunksize=(1, 1059, 1799), meta=np.ndarray> 10v (valid_time, x, y) float64 dask.array<chunksize=(1, 1059, 1799), meta=np.ndarray> 2d (valid_time, x, y) float64 dask.array<chunksize=(1, 1059, 1799), meta=np.ndarray> 2r (valid_time, x, y) float64 dask.array<chunksize=(1, 1059, 1799), meta=np.ndarray> 2sh (valid_time, x, y) float64 dask.array<chunksize=(1, 1059, 1799), meta=np.ndarray> 2t (valid_time, x, y) float64 dask.array<chunksize=(1, 1059, 1799), meta=np.ndarray> latitude (x, y) float64 dask.array<chunksize=(1059, 1799), meta=np.ndarray> longitude (x, y) float64 dask.array<chunksize=(1059, 1799), meta=np.ndarray> pt (valid_time, x, y) float64 dask.array<chunksize=(1, 1059, 1799), meta=np.ndarray> unknown (valid_time, x, y) float64 dask.array<chunksize=(1, 1059, 1799), meta=np.ndarray> Attributes: centre: kwbc centreDescription: US National Weather Service - NCEP edition: 2 subCentre: 0
ds['2t'][-1].plot()
<matplotlib.collections.QuadMesh at 0x7ffa4fcf2070>
ds['2t'][:,500,500].plot()
[<matplotlib.lines.Line2D at 0x7ffa4dea16a0>]