Stack And Split Variables Across Files

Stack And Split Variables Across Files#

Many solvers store a vector field as one file per component (ux-0001.bin, uy-0001.bin, uz-0001.bin) and a set of scalar fractions as numbered variables (phi1-0001.bin, phi2-0001.bin). In memory, you would rather work with one array u carrying a component dimension i, and one array phi carrying a fraction dimension n.

xarray-binfile describes that encoding once, with a VariableStack, and uses it in both directions:

split on write: u with i = ["x", "y", "z"] becomes the variables ux, uy and uz, each written to its own files
stack on read: ux, uy and uz found on disk are concatenated back into u

Nothing about u, x, y, z or phi is hard-coded in the library: the dimension name, the name template and the allowed values all come from the declaration.

import pathlib
import tempfile

import numpy as np
import xarray as xr

import xarray_binfile  # noqa: F401  (registers the .binary_engine accessors)
from xarray_binfile.conventions import Layout, StepIndexedFiles, VariableStack

directory_holder = tempfile.TemporaryDirectory()
directory = pathlib.Path(directory_holder.name)

Declare the stacks#

A VariableStack has three parts:

dim: the dimension that lives in the variable names rather than in the files
template: how a variable name is built from the base {name} and the value of dim
values: the values the dimension may take, in the order the stacked coordinate should follow

The template is parsed with FilenamePattern, so integer values can carry zero-fill ("{name}{n:02d}") and are converted back to int on read. values is required because templates without a separator would otherwise match almost any name: pp fits "{name}{i}" as p + p.

velocity = VariableStack("i", "{name}{i}", values=("x", "y", "z"))
scalars = VariableStack("n", "{name}{n:d}", values=range(1, 10), names=("phi",))
velocity, scalars

(VariableStack(dim='i', template='{name}{i}', values=('x', 'y', 'z'), names=None, attrs=None, min_components=2),
 VariableStack(dim='n', template='{name}{n:d}', values=(1, 2, 3, 4, 5, 6, 7, 8, 9), names=frozenset({'phi'}), attrs=None, min_components=1))

The optional names restricts stacking to given base names. Without it, a vorticity component written as w3 would be picked up by scalars as w with n = [3], which is harmless but probably not what you want.

Attach them to a convention#

Every shipped convention accepts stacks. On write, arrays are split in the order the stacks are listed, before the per-file split along time; on read, open rebuilds the stacked arrays by default.

layout = Layout(
    {
        "x": np.linspace(0.0, 1.0, 8, dtype=np.float32),
        "y": np.linspace(0.0, 2.0, 6, dtype=np.float32),
        "z": np.linspace(0.0, 0.5, 4, dtype=np.float32),
    },
    dtype=np.float32,
)
convention = StepIndexedFiles(
    layout,
    pattern="{name}-{step:03d}.bin",
    time_step=0.25,
    stacks=[velocity, scalars],
)
convention

StepIndexedFiles(layout=Layout(coords={'x': array([0.        , 0.14285715, 0.2857143 , 0.42857143, 0.5714286 ,
       0.71428573, 0.85714287, 1.        ], dtype=float32), 'y': array([0. , 0.4, 0.8, 1.2, 1.6, 2. ], dtype=float32), 'z': array([0.        , 0.16666667, 0.33333334, 0.5       ], dtype=float32)}, dtype=<class 'numpy.float32'>, order='C', coord_attrs=None), pattern=FilenamePattern(template='{name}-{step:03d}.bin', exact_width=False), check_coords=True, stacks=(VariableStack(dim='i', template='{name}{i}', values=('x', 'y', 'z'), names=None, attrs=None, min_components=2), VariableStack(dim='n', template='{name}{n:d}', values=(1, 2, 3, 4, 5, 6, 7, 8, 9), names=frozenset({'phi'}), attrs=None, min_components=1)), name_of=None, names=None, time_dim='time', time_dtype=None, time_step=0.25)

Split on write#

Writing the dataset through the convention produces one file per component, per fraction and per step. u never appears on disk; pp is written as is.

dataset.binary_engine.to_file(convention.writer, directory)
sorted(path.name for path in directory.glob("*.bin"))

['phi1-000.bin',
 'phi1-001.bin',
 'phi1-002.bin',
 'phi2-000.bin',
 'phi2-001.bin',
 'phi2-002.bin',
 'pp-000.bin',
 'pp-001.bin',
 'pp-002.bin',
 'ux-000.bin',
 'ux-001.bin',
 'ux-002.bin',
 'uy-000.bin',
 'uy-001.bin',
 'uy-002.bin',
 'uz-000.bin',
 'uz-001.bin',
 'uz-002.bin']

The split happens before the file split, so the per-file slices are plain (x, y, z) arrays. You can inspect the plan without writing anything by iterating the writer:

[(spec.filename, spec.sub_array.dims) for spec in convention.writer(dataset["u"])][:4]

[('ux-000.bin', ('x', 'y', 'z')),
 ('ux-001.bin', ('x', 'y', 'z')),
 ('ux-002.bin', ('x', 'y', 'z')),
 ('uy-000.bin', ('x', 'y', 'z'))]

Stack on read#

open lists the files that follow the pattern, opens them lazily with xarray.open_mfdataset, and applies the stacks. The result has the same variables as the original dataset, with i and n as proper coordinates.

opened = convention.open(directory, chunks={"time": 1})
opened

<xarray.Dataset> Size: 14kB
Dimensions:  (x: 8, y: 6, z: 4, time: 3, n: 2, i: 3)
Coordinates:
  * x        (x) float32 32B 0.0 0.1429 0.2857 0.4286 0.5714 0.7143 0.8571 1.0
  * y        (y) float32 24B 0.0 0.4 0.8 1.2 1.6 2.0
  * z        (z) float32 16B 0.0 0.1667 0.3333 0.5
  * time     (time) float64 24B 0.0 0.25 0.5
  * n        (n) int64 16B 1 2
  * i        (i) <U1 12B 'x' 'y' 'z'
Data variables:
    pp       (x, y, z, time) float32 2kB dask.array<chunksize=(8, 6, 4, 1), meta=np.ndarray>
    phi      (n, x, y, z, time) float32 5kB dask.array<chunksize=(1, 8, 6, 4, 1), meta=np.ndarray>
    u        (i, x, y, z, time) float32 7kB dask.array<chunksize=(1, 8, 6, 4, 1), meta=np.ndarray>

for name in dataset.data_vars:
    xr.testing.assert_allclose(
        opened[name].transpose(*dataset[name].dims).load(), dataset[name]
    )
print("round trip matches")

round trip matches

Stacking is lazy, so a Dask-backed dataset stays Dask-backed, and i and n can be used like any other coordinate:

Opt out, or stack by hand#

Pass stack=False to see the variables exactly as they are on disk. The same declarations are available as functions for datasets opened another way: VariableStack.stack for one dimension, stack_variables for a sequence, and split_variables for the inverse.

raw = convention.open(directory, stack=False)
sorted(raw.data_vars)

['phi1', 'phi2', 'pp', 'ux', 'uy', 'uz']

from xarray_binfile.conventions import split_variables, stack_variables

restacked = stack_variables(raw, [velocity, scalars])
pieces = [piece.name for piece in split_variables(dataset["u"], [velocity])]
sorted(restacked.data_vars), pieces

(['phi', 'pp', 'u'], ['ux', 'uy', 'uz'])

Guard rails#

A value that is not declared cannot be written, because the reader would have no way to stack it back:

bad = dataset["u"].assign_coords(i=["x", "y", "w"])
try:
    list(convention.writer(bad))
except ValueError as error:
    print(error)

Coordinate value(s) ['w'] on 'i' are not listed in the VariableStack values ['x', 'y', 'z']; files written for them could not be stacked back.

directory_holder.cleanup()