Notebook

In [1]:

import os
import fsspec
from datetime import datetime, timedelta

import ujson
from kerchunk.hdf import SingleHdf5ToZarr 
from kerchunk.combine import MultiZarrToZarr
import xarray as xr
import dask

import numpy as np

from pathlib import Path

Model run¶

In [2]:

hourlyrun='00'

In [3]:

fs = fsspec.filesystem('s3', anon='True', skip_instance_cache=True)

Latest data available¶

In [4]:

latest = sorted(fs.glob('s3://noaa-nwm-pds/nwm*'))[-1]

This creates a dictionary containing a list of the first 9 hours of data for each ensemble¶

In [5]:

ens_files={}
for i in range(1,8):
    ens_files[i] = sorted(fs.glob(f's3://{latest}/medium_range_mem{i}/nwm.t{hourlyrun}z.medium_range.channel_rt_{i}.f00*.nc'))

In [6]:

json_dir = '/home/peterm790/shared/users/petermarsh/ensemble_NWM/jsons/'

Create json reference files from these files¶

In [7]:

def gen_json(u, json_dir):
    fstem = Path(u).stem 
    outf = f'{json_dir}{fstem}.json'
    with fs.open(u, **so) as infile:
        h5chunks = SingleHdf5ToZarr(infile, u, inline_threshold=300)
        with fs2.open(outf, 'wb') as f:
            f.write(ujson.dumps(h5chunks.translate()).encode());

In [3]:

fs2 = fsspec.filesystem("") 

In [4]:

so = dict(mode='rb', anon=True, default_fill_cache=False, default_cache_type='first')

In [10]:

%%time 
for ens in ens_files:
    for file in ens_files[ens]:
        gen_json(file, json_dir)

CPU times: user 12.4 s, sys: 2.64 s, total: 15 s
Wall time: 1min 57s

In [11]:

ens_jsons = {}
for ens in ens_files:
    flist = fs2.glob(f'{json_dir}nwm.t00z.medium_range.channel_rt_{i}*')
    ens_jsons[ens] = flist#[f's3:/{f}' for f in flist]

Create a single kerchunk json for each ensemble member¶

Post process appends the ensemble name to each variable

In [14]:

for ens in ens_jsons:
    def post_process(d):
            keys = list(d.keys())
            for key in keys:
                var = key.split('/')[0]
                try:
                    end = '/'.join(key.split('/')[1:])
                except:
                    end = None
                if var not in ['.zgroup', '.zattrs', 'time', 'feature_id']:
                    d[f'{var}_{ens}/{end}'] = d.pop(key)
                else:
                    pass
            return d
    mzz = MultiZarrToZarr(ens_jsons[ens], 
        remote_protocol='s3',
        remote_options={'anon':True},   #JSON files  
        concat_dims = ['time'],
        identical_dims = ['feature_id', 'reference_time'],
        postprocess = post_process
    )
    d = mzz.translate()
    with fs2.open(f'/home/peterm790/shared/users/petermarsh/ensemble_NWM/{ens}.json', 'wb') as f:
        f.write(ujson.dumps(d).encode());

Create a single virtual dataset containing all 7 ensemble members¶

In [5]:

flist = fs2.glob('/home/peterm790/shared/users/petermarsh/ensemble_NWM/*.json')

In [6]:

mzz = MultiZarrToZarr(flist, 
    remote_protocol='s3',
    remote_options={'anon':True},
    concat_dims = ['time'],
    identical_dims = ['feature_id', 'reference_time']                      
)
d = mzz.translate()

In [7]:

%%time
fs = fsspec.filesystem("reference", fo=d, ref_storage_args={'skip_instance_cache':True},
                       remote_protocol='s3', remote_options={'anon':False})
m = fs.get_mapper("")
ds_kerchunk = xr.open_dataset(m, engine="zarr", backend_kwargs={'consolidated':False}, chunks={})

CPU times: user 273 ms, sys: 68.6 ms, total: 342 ms
Wall time: 2.3 s

In [8]:

ds_kerchunk

Out[8]:

<xarray.Dataset>
Dimensions:           (time: 9, feature_id: 2776738, reference_time: 1)
Coordinates:
  * feature_id        (feature_id) float64 101.0 179.0 ... 1.18e+09 1.18e+09
  * time              (time) datetime64[ns] 2022-06-23T01:00:00 ... 2022-06-2...
Dimensions without coordinates: reference_time
Data variables: (12/56)
    crs_1             (time) object dask.array<chunksize=(1,), meta=np.ndarray>
    crs_2             (time) object dask.array<chunksize=(1,), meta=np.ndarray>
    crs_3             (time) object dask.array<chunksize=(1,), meta=np.ndarray>
    crs_4             (time) object dask.array<chunksize=(1,), meta=np.ndarray>
    crs_5             (time) object dask.array<chunksize=(1,), meta=np.ndarray>
    crs_6             (time) object dask.array<chunksize=(1,), meta=np.ndarray>
    ...                ...
    velocity_2        (time, feature_id) float64 dask.array<chunksize=(1, 925580), meta=np.ndarray>
    velocity_3        (time, feature_id) float64 dask.array<chunksize=(1, 925580), meta=np.ndarray>
    velocity_4        (time, feature_id) float64 dask.array<chunksize=(1, 925580), meta=np.ndarray>
    velocity_5        (time, feature_id) float64 dask.array<chunksize=(1, 925580), meta=np.ndarray>
    velocity_6        (time, feature_id) float64 dask.array<chunksize=(1, 925580), meta=np.ndarray>
    velocity_7        (time, feature_id) float64 dask.array<chunksize=(1, 925580), meta=np.ndarray>
Attributes: (12/20)
    Conventions:                CF-1.6
    NWM_version_number:         v2.1
    TITLE:                      OUTPUT FROM NWM v2.1
    cdm_datatype:               Station
    code_version:               v5.2.0-beta2
    dev:                        dev_ prefix indicates development/internal me...
    ...                         ...
    model_output_type:          channel_rt
    model_output_valid_time:    2022-06-23_01:00:00
    model_total_valid_times:    204
    proj4:                      +proj=lcc +units=m +a=6370000.0 +b=6370000.0 ...
    station_dimension:          feature_id
    stream_order_output:        1

In [ ]: