zarr-developers / VirtualiZarr

Create virtual Zarr stores from archival data files using xarray syntax
https://virtualizarr.readthedocs.io/en/latest/
Apache License 2.0
93 stars 16 forks source link

Use case: UCLA-ROMS fluid HPC model output #217

Open TomNicholas opened 1 month ago

TomNicholas commented 1 month ago

At [C]Worthy in addition to the CESM data we are also working with a oceanographic model called UCLA-ROMS. I want to use VirtualiZarr on the netCDF files written out by the ROMS code to create a single virtual Zarr store pointing to the output of an entire model run.

This problem is a pretty typical example of a fluid HPC code - its spatially parallelized (in latitude and longitude), and writes out one netCDF file per node. It also writes a new set of netCDF files for every timestep (or every fixed number of timesteps). So it's a 3D concatenation problem[^1].

There are a few subtleties of this particular dataset that are relevant for VirtualiZarr.

  1. Staggered grids

This is not a problem for xarray, or for zarr, but was a problem for kerchunk, and one of the motivations for making this package.

  1. Nesting

ROMS (Regional Ocean Modelling System) does local high-resolution modelling of just part of the Earth's surface, but because of the need for boundaries you often have to run multiple simulations with spatial grids that are successively "nested" inside one another (until your finest grid covers the real region of interest at the resolution desired).

So a nested ROMS simulation actually has multiple simulations, each of which we could choose to think of as a different netCDF Zarr group. Then once we have https://github.com/zarr-developers/VirtualiZarr/issues/84 we could supporting opening the entire set of nested results an xarray.DataTree object using just xr.open_datatree (see https://github.com/xarray-contrib/datatree).

  1. Variable-length chunks (maybe)

This is the actually hard one. Currently ROMS decomposes its spatial domain along one dimension (lat and/or lon) into a pattern that has lengths something like [X+2, X, ..., X, X+2]. This ultimately comes from the fact that you require 2 boundary conditions to solve 2nd-order PDEs, and these boundary cells are just tacked on to an integer decomposition across nodes.

The problem is that Zarr's current model only allows for chunk lengths along one dimension to be of the form [X, X, ..., X, Y], where Y <= X. So ROMS doesn't fit into Zarr's data model! (and the deviation is pretty much the smallest one imaginable that still screws it up 😞)

There are two ways to solve this: change Zarr, or change ROMS.

Changing ROMS is conceptually simple - if ROMS instead uses a different default domain decomposition, one which always produces [X, X, ..., X, Y] along both lat and lon, then we're golden. That might happen, but it also might not.

Changing Zarr is possible, but much more involved. It has been suggested, and worked on (the "variable-length chunks ZEP"), but that effort kind of stalled. We might get funding to do this in the next year, but if we don't need to do it get ROMS data virtualized then we should avoid this can of worms.

[^1]: It's actually basically an identical structure to the plasma fluid HPC code I worked with during my PhD - could be useful for you FYI @bendudson @johnomotani!

mdsumner commented 1 month ago

Very interested in this, will follow along with local experts and try out some products 👌

TomNicholas commented 1 month ago

What's your interest @mdsumner ?

mdsumner commented 1 month ago

I help with accessibility to these formats is all. Knowing how far this virtualization scheme can go, and where the complex edges are is key to my work.

I'm also obsessed with redundancy in coordinates (and meshes generally), staggered grids are one of those very interesting cases along the spectrums of compact vs materialized. And also some models are born in actual regular map projections but stored and worked in materialized longlat arrays which is another interesting case in intermediate forms here.