Dataset Read Guidelines

Dataset Read Guidelines#

GLORYS#

The default approach shown for reading and accessing previously-saved GLORYS output is to access the files directly using xr.open_mfdataset() pointing to a list of files. This is ok, but it can generally be slow. If you would like to speed up your access to previously-saved GLORYS files, you can try subchunking (Example Medium article we follow which uses kerchunk). Note that this works because the files are uncompressed; if you have compressed files you would like to try subchunking with, check out this paper: A new sub-chunking strategy for fast netCDF-4 access in local, remote and cloud infrastructures, chunkindex V1.1.0.

Code is sketched out here that you can copy-paste; it isn’t in the ROMS-Tools codebase because it requires several additional libraries and is not required to be used.

This code changes the files from individual and unchunked into a representation that knows where to find the variables so it can preload the necessary index information. Additionally, it “subchunks” by depth which makes these unchunked files smaller and easier to fit into memory. The speed up on the HPC we have been running on is very large — one year of generating boundary forcing files using this subchunked representation takes about 30 minutes and when pointing to the files directly takes 3 hours.

[4]:

"""
Build kerchunk Parquet references for local GLORYS NetCDF4 files.
Adapted from github.com/rsignell/hycom-kerchunk (0_hycom_generate_refs.ipynb).

Assumes GLORYS files are contiguous + uncompressed.

Set depth_factor and lat_factor to select the subchunk variation (see
the GLORYS case in __main__).
"""
import glob
import fsspec
from pathlib import Path
import xarray as xr
from fsspec.implementations.reference import LazyReferenceMapper
from kerchunk.hdf import SingleHdf5ToZarr
from kerchunk.combine import MultiZarrToZarr
from kerchunk.utils import subchunk
from time import time

# for running in a notebook
import nest_asyncio
nest_asyncio.apply()

def dict_to_parquet(refs_dict, output_path, record_size=100_000):
    """Write a kerchunk refs dict to parquet.

    Uses a two-pass write to work around a LazyReferenceMapper bug where
    .zarray metadata written in the same translate() call isn't visible
    via zmetadata until flushed to disk.
    """
    refs = refs_dict.get("refs", {})
    out = LazyReferenceMapper.create(output_path, record_size=record_size)

    for k, v in refs.items():         # pass 1: .zarray / .zattrs / top-level
        if ".z" in k or "/" not in k:
            out[k] = v
    out.flush()                        # flush so zmetadata can find .zarray

    for k, v in refs.items():         # pass 2: chunk references
        if ".z" not in k and "/" in k:
            out[k] = v
    out.flush()


def subchunk_var(refs, var, depth_factor, lat_factor):
    if depth_factor and depth_factor > 1:
        refs = subchunk(refs, var, depth_factor)   # splits outermost non-unit axis = depth
    if lat_factor and lat_factor > 1:
        refs = subchunk(refs, var, lat_factor)     # depth now == 1, so this hits latitude
    return refs


def build_one(path, data_vars_4d, depth_factor, lat_factor, inline_threshold):
    refs = SingleHdf5ToZarr(path, inline_threshold=inline_threshold).translate()
    for v in data_vars_4d:
        refs = subchunk_var(refs, v, depth_factor, lat_factor)
    return refs


def run(
    files,
    out,
    output_format=".parquet",
    data_vars_4d=("thetao", "so", "uo", "vo"),
    concat_dim="time",
    identical_dims=("latitude", "longitude", "depth"),
    inline_threshold=5000,
    depth_factor=None,
    lat_factor=None,
):
    """Build combined kerchunk references for a set of NetCDF4 files.

    Parameters
    ----------
    files : sequence of str
        Paths to the source NetCDF4 files (already sorted as desired).
    out : str
        Output path. For ".parquet"/".json" the appropriate extension
        should already be appended (see __main__). Ignored for "in-memory".
    output_format : {".parquet", ".json", "in-memory"}
        How to emit the combined references.
    data_vars_4d : sequence of str
        4D data variables to subchunk.
    concat_dim : str
        Dimension along which the per-file references are concatenated.
    identical_dims : sequence of str
        Dimensions identical across files (not concatenated).
    inline_threshold : int
        Bytes below which chunks are inlined rather than referenced.
    depth_factor, lat_factor : int or None
        subchunk factors; see subchunk_var.

    Returns
    -------
    dict
        The combined kerchunk references.
    """
    # 1) per-file refs — parallelise with ProcessPoolExecutor for a year / 17 yr
    #    (avoid the deepcopy-and-patch shortcut from HYCOM: it assumes identical
    #     HDF5 byte offsets across files, which is unsafe for NetCDF4)
    single_refs = [
        build_one(p, data_vars_4d, depth_factor, lat_factor, inline_threshold)
        for p in files
    ]

    # 2) combine across files
    combined = MultiZarrToZarr(
        single_refs,
        remote_protocol="file",
        concat_dims=[concat_dim],
        identical_dims=list(identical_dims),
    ).translate()

    # 3) emit
    if output_format == ".json":
        import ujson

        with open(out, "w") as f:
            ujson.dump(combined, f)
        print(f"References written to {out}")

    elif output_format == ".parquet":
        dict_to_parquet(combined, out)
        print(f"References written to {out}")

    return combined


def build_subchunk_refs(
    input_files,
    out,
    output_format=".parquet",
    data_vars_4d=("thetao", "so", "uo", "vo"),
    concat_dim="time",
    identical_dims=("latitude", "longitude", "depth"),
    inline_threshold=5000,
    depth_factor=None,
    lat_factor=1,
    overwrite=False,
):
    """Build (and optionally write) subchunked kerchunk references.

    Wraps :func:`run` with file globbing, output-path handling, an
    auto-detected depth factor, a skip-if-exists guard, and timing.

    Parameters
    ----------
    input_files : str or sequence of str
        A glob pattern (e.g. ``".../GLOBAL_2024????.nc"``) or an explicit,
        already-sorted list of NetCDF4 file paths. A glob is expanded and
        sorted automatically.
    out : str
        Output path *without* extension. The extension is appended from
        ``output_format`` (ignored for ``"in-memory"``).
    output_format : {".parquet", ".json", "in-memory"}
        How to emit the combined references.
    data_vars_4d : sequence of str
        4D data variables to subchunk.
    concat_dim : str
        Dimension along which per-file references are concatenated (time).
    identical_dims : sequence of str
        Dimensions identical across files (not concatenated).
    inline_threshold : int
        Bytes below which chunks are inlined rather than referenced.
    depth_factor : int or None
        subchunk factor along depth. ``None`` (default) auto-detects the
        depth size from the first file, collapsing depth to 1 chunk.
    lat_factor : int
        subchunk factor along latitude (1 = no latitude subchunking).
    overwrite : bool
        If False (default) and the output already exists, skip and return
        its path without rebuilding.

    Returns
    -------
    str or dict
        The output path for ``".parquet"``/``".json"``, or the combined
        references dict for ``"in-memory"``.
    """
    start_time = time()

    # 1) resolve input files
    if isinstance(input_files, str):
        files = sorted(glob.glob(input_files))
    else:
        files = list(input_files)
    if not files:
        raise FileNotFoundError(f"No input files matched: {input_files!r}")

    # 2) resolve output path / skip-if-exists
    if output_format != "in-memory":
        out = out + output_format
        if Path(out).exists() and not overwrite:
            print(f"{out} already exists, skipping (pass overwrite=True to rebuild)")
            return out

    # 3) auto-detect depth factor (collapse depth to 1 chunk) if not given
    if depth_factor is None:
        with xr.open_dataset(files[0]) as ds:
            depth_factor = ds.sizes["depth"]
        print(f"Using depth_factor={depth_factor} (depth size of {files[0]})")

    # 4) build
    result = run(
        files=files,
        out=out,
        output_format=output_format,
        data_vars_4d=list(data_vars_4d),
        concat_dim=concat_dim,
        identical_dims=list(identical_dims),
        inline_threshold=inline_threshold,
        depth_factor=depth_factor,
        lat_factor=lat_factor,
    )

    print(f"Making subchunked file representation took {time() - start_time:.1f} seconds")
    return out if output_format != "in-memory" else result

Example call:

[5]:

out = build_subchunk_refs(
    input_files=(
        "/anvil/scratch/x-kthyng/GLORYS_GLOBAL/"
        "cmems_mod_glo_phy_my_0.083deg_P1D-m_GLOBAL_2024????.nc"
    ),
    out="/anvil/scratch/x-kthyng/GLORYS_GLOBAL/test-2024-subchunk",
    output_format=".parquet",
    data_vars_4d=["thetao", "so", "uo", "vo"],
    # depth_factor=None  -> auto-detect (collapse depth); set lat_factor=13 for 157-row strips
)

Using depth_factor=50 (depth size of /anvil/scratch/x-kthyng/GLORYS_GLOBAL/cmems_mod_glo_phy_my_0.083deg_P1D-m_GLOBAL_20240501.nc)
References written to /anvil/scratch/x-kthyng/GLORYS_GLOBAL/test-2024-subchunk.parquet
Making subchunked file representation took 2.1 seconds

[7]:

loc = "/anvil/scratch/x-kthyng/GLORYS_GLOBAL/test-2024-subchunk.parquet"
ds = xr.open_dataset(loc, chunks={})
ds

Note that now the variables are chunked in depth which allows for faster processing.

[8]:

ds.thetao.chunks

[8]:

((1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1),
 (1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1,
  1),
 (2041,),
 (4320,))

Dataset Read Guidelines

Contents

Dataset Read Guidelines#

GLORYS#