From e29d74245d98e2c1ed3f27018fbd70bedb6376f8 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Tue, 4 Aug 2026 02:41:08 -0500 Subject: [PATCH 1/3] Read whole variables at once in mpas_to_xdmf Each index of an extra dimension (e.g. each vertical level) becomes its own field in the XDMF output, and the converter was reading those fields one at a time. A single vertical level is strided across the whole variable on disk, so this meant reading the entire variable once per level. On a 4M-cell, 80-layer mesh, one level of `temperature` took 1.6 s while all 80 levels read together took 3.9 s: an ~80x read amplification that dominated the run time. Instead, read as many indices as possible in a single pass over each variable and slice them apart in memory. The amount read at a time is capped by `max_read_bytes` (2 GB by default) so memory use stays bounded; variables larger than that are read in as few passes as the cap allows. Unwrapping the extra dimensions moves out of `_process_extra_dims` and into the writer, which is what makes the grouped read possible. Writing one 3D variable went from 199 s to 6.5 s (30x) on the mesh above. A cell-blocked variant using HDF5 hyperslabs was also tried and was slower (19.8 s), because the scattered writes cost far more than they saved. Also: * add a `float32` option that writes floating-point fields in single precision, halving the size of the HDF5 files and the time ParaView needs to load them. The mesh geometry stays in double precision so that cell shapes are unaffected, and integer fields are left alone. * apply the vertex-to-kite map after reading rather than as a lazy indexed read from netCDF, which is much slower. * write the connectivity as 32-bit integers when the indices fit. The contents of the HDF5 files are otherwise unchanged; only the order in which attributes are listed in the XDMF differs. Co-Authored-By: Claude Opus 5 --- .../mpas_tools/viz/mpas_to_xdmf/io.py | 368 +++++++++++++++--- .../viz/mpas_to_xdmf/mpas_to_xdmf.py | 44 ++- 2 files changed, 350 insertions(+), 62 deletions(-) diff --git a/conda_package/mpas_tools/viz/mpas_to_xdmf/io.py b/conda_package/mpas_tools/viz/mpas_to_xdmf/io.py index 2a65ce1e4..a334b00c0 100644 --- a/conda_package/mpas_tools/viz/mpas_to_xdmf/io.py +++ b/conda_package/mpas_tools/viz/mpas_to_xdmf/io.py @@ -1,8 +1,10 @@ import glob import importlib.resources +import itertools import os import h5py +import numpy as np import xarray as xr from jinja2 import Template from tqdm import tqdm @@ -15,6 +17,14 @@ from mpas_tools.viz.mpas_to_xdmf.mesh import _get_ds_mesh from mpas_tools.viz.mpas_to_xdmf.time import _set_time +# Fields with extra dimensions (e.g. nVertLevels) are unwrapped into one +# 2D field per index. Reading each of those fields individually is very +# slow, because a single vertical level is strided across the whole variable +# on disk, so the entire variable has to be read once per level. Instead we +# read as many indices at a time as we can, capped by this many bytes so that +# memory use stays bounded on very large meshes. +_DEFAULT_MAX_READ_BYTES = 2 * 1024**3 + def _load_dataset(mesh_filename, time_series_filenames, variables, xtime_var): """ @@ -102,7 +112,15 @@ def _load_dataset(mesh_filename, time_series_filenames, variables, xtime_var): return ds_mesh, ds -def _convert_to_xdmf(ds, ds_mesh, out_dir, quiet=False): +def _convert_to_xdmf( + ds, + ds_mesh, + out_dir, + extra_dims=None, + quiet=False, + float32=False, + max_read_bytes=_DEFAULT_MAX_READ_BYTES, +): """ Convert an xarray Dataset to XDMF + HDF5 format. @@ -114,53 +132,89 @@ def _convert_to_xdmf(ds, ds_mesh, out_dir, quiet=False): The mesh dataset. out_dir : str Directory where XDMF and HDF5 files will be saved. + extra_dims : dict, optional + Dictionary mapping extra dimensions to the indices to write. quiet : bool, optional If True, suppress progress output. Default is False. + float32 : bool, optional + If True, write floating-point fields in single precision. + max_read_bytes : int, optional + Approximate limit on the number of bytes read from ``ds`` at a time. """ os.makedirs(out_dir, exist_ok=True) + kwargs = dict( + out_dir=out_dir, + extra_dims=extra_dims, + quiet=quiet, + float32=float32, + max_read_bytes=max_read_bytes, + ) + if 'nCells' in ds.dims: - _convert_cells_to_xdmf(ds, ds_mesh, out_dir, quiet) + _convert_cells_to_xdmf(ds, ds_mesh, **kwargs) if 'nEdges' in ds.dims: - _convert_edges_to_xdmf(ds, ds_mesh, out_dir, quiet) + _convert_edges_to_xdmf(ds, ds_mesh, **kwargs) if 'nVertices' in ds.dims: - _convert_vertices_to_xdmf(ds, ds_mesh, out_dir, quiet) + _convert_vertices_to_xdmf(ds, ds_mesh, **kwargs) -def _convert_cells_to_xdmf(ds, ds_mesh, out_dir, quiet): +def _convert_cells_to_xdmf(ds, ds_mesh, **kwargs): """ Convert cell-centered data to XDMF + HDF5 format. """ ds_cell_geom = _build_cell_geometry(ds_mesh) cell_vars = [var for var in ds.data_vars if 'nCells' in ds[var].dims] ds_cells = ds[cell_vars] - _write_xdmf(ds_cell_geom, ds_cells, out_dir, suffix='Cells', quiet=quiet) + _write_xdmf( + ds_cell_geom, ds_cells, suffix='Cells', topo_dim='nCells', **kwargs + ) -def _convert_edges_to_xdmf(ds, ds_mesh, out_dir, quiet): +def _convert_edges_to_xdmf(ds, ds_mesh, **kwargs): """ Convert edge-centered data to XDMF + HDF5 format. """ ds_edge_geom = _build_edge_geometry(ds_mesh) edge_vars = [var for var in ds.data_vars if 'nEdges' in ds[var].dims] ds_edges = ds[edge_vars] - _write_xdmf(ds_edge_geom, ds_edges, out_dir, suffix='Edges', quiet=quiet) + _write_xdmf( + ds_edge_geom, ds_edges, suffix='Edges', topo_dim='nEdges', **kwargs + ) -def _convert_vertices_to_xdmf(ds, ds_mesh, out_dir, quiet): +def _convert_vertices_to_xdmf(ds, ds_mesh, **kwargs): """ Convert vertex-centered data to XDMF + HDF5 format. """ ds_vertex_geom = _build_vertex_geometry(ds_mesh) vertex_vars = [var for var in ds.data_vars if 'nVertices' in ds[var].dims] - vert_to_kite_map = ds_vertex_geom['vert_to_kite_map'] - ds_vertices = ds[vertex_vars].isel(nVertices=vert_to_kite_map) + ds_vertices = ds[vertex_vars] + # each vertex is repeated once per kite; the map is applied to each field + # after it has been read rather than as a (much slower) indexed read + vert_to_kite_map = ds_vertex_geom['vert_to_kite_map'].values _write_xdmf( - ds_vertex_geom, ds_vertices, out_dir, suffix='Vertices', quiet=quiet + ds_vertex_geom, + ds_vertices, + suffix='Vertices', + topo_dim='nVertices', + topo_map=vert_to_kite_map, + **kwargs, ) -def _write_xdmf(ds_geom, ds_data, out_dir, suffix, quiet=False): +def _write_xdmf( + ds_geom, + ds_data, + out_dir, + suffix, + topo_dim, + extra_dims=None, + topo_map=None, + quiet=False, + float32=False, + max_read_bytes=_DEFAULT_MAX_READ_BYTES, +): """ Write data to HDF5 and metadata to XDMF format. @@ -174,55 +228,60 @@ def _write_xdmf(ds_geom, ds_data, out_dir, suffix, quiet=False): Directory where XDMF and HDF5 files will be saved. suffix : str Suffix to append to output filenames (e.g., 'Cells', 'Edges'). + topo_dim : str + The dimension the fields are defined on (e.g. ``'nCells'``). + extra_dims : dict, optional + Dictionary mapping extra dimensions to the indices to write. Each + variable is unwrapped into one field per combination of indices. + topo_map : numpy.ndarray, optional + Indices along ``topo_dim`` to map each field onto the geometry. quiet : bool, optional If True, suppress progress output. Default is False. + float32 : bool, optional + If True, write floating-point fields in single precision. + max_read_bytes : int, optional + Approximate limit on the number of bytes read from ``ds_data`` at a + time. """ h5_basename = f'fieldsOn{suffix}.h5' h5_filename = os.path.join(out_dir, h5_basename) xdmf_filename = os.path.join(out_dir, f'fieldsOn{suffix}.xdmf') + variables_metadata = _field_metadata(ds_data, extra_dims) + # Write HDF5 file with h5py.File(h5_filename, 'w') as h5_file: # Write geometry h5_file.create_dataset('Points', data=ds_geom['points'].values) - h5_file.create_dataset('Cells', data=ds_geom['cells'].values) + h5_file.create_dataset('Cells', data=_as_index_type(ds_geom['cells'])) # Calculate total progress steps total_steps = sum( - ds_data.sizes['Time'] if 'Time' in ds_data[var].dims else 1 - for var in ds_data.data_vars + ds_data.sizes['Time'] if field['has_time'] else 1 + for field in variables_metadata ) # Write time-varying and static data with progress bar if quiet: iterator = None else: - iterator = tqdm( - ds_data.data_vars, total=total_steps, desc=f'Writing {suffix}' - ) + iterator = tqdm(total=total_steps, desc=f'Writing {suffix}') for var_name in ds_data.data_vars: if iterator is not None: iterator.set_description(f'Processing {var_name}') - if 'Time' in ds_data[var_name].dims: - for t_idx in range(ds_data.sizes['Time']): - dataset_name = f'{var_name}_t{t_idx}' - da = ds_data[var_name].isel(Time=t_idx) - h5_file.create_dataset(dataset_name, data=da.values) - if iterator is not None: - iterator.update(1) - else: - h5_file.create_dataset(var_name, data=ds_data[var_name].values) - if iterator is not None: - iterator.update(1) - - # Preprocess variable metadata for the template - variables_metadata = [ - { - 'name': var_name, - 'has_time': 'Time' in ds_data[var_name].dims, - } - for var_name in ds_data.data_vars - ] + _write_variable( + h5_file=h5_file, + da=ds_data[var_name], + var_name=var_name, + extra_dims=extra_dims, + topo_dim=topo_dim, + topo_map=topo_map, + float32=float32, + max_read_bytes=max_read_bytes, + iterator=iterator, + ) + if iterator is not None: + iterator.close() # Load XDMF template from external file package = 'mpas_tools.viz.mpas_to_xdmf.templates' @@ -231,7 +290,7 @@ def _write_xdmf(ds_geom, ds_data, out_dir, suffix, quiet=False): xdmf_template = Template(template_file.read()) # Render XDMF file - cells = ds_geom['cells'].values + cells = ds_geom['cells'] times = ds_data['Time'].values if 'Time' in ds_data.dims else [] xdmf_content = xdmf_template.render( @@ -248,6 +307,196 @@ def _write_xdmf(ds_geom, ds_data, out_dir, suffix, quiet=False): xdmf_file.write(xdmf_content) +def _as_index_type(da): + """ + Return the connectivity array as 32-bit integers if the indices fit, + halving the size of the largest static array in the HDF5 file. + """ + values = da.values + if ( + values.dtype.itemsize > 4 + and values.size > 0 + and values.max() < np.iinfo(np.int32).max + ): + values = values.astype(np.int32) + return values + + +def _expand_extra_dims(var_name, da, extra_dims): + """ + Determine the fields a variable is unwrapped into, one per combination of + extra-dimension indices. + + Parameters + ---------- + var_name : str + The name of the variable. + da : xarray.DataArray + The variable to unwrap. + extra_dims : dict or None + Dictionary mapping extra dimensions to the indices to write. + + Returns + ------- + dims : list of str + The extra dimensions present on ``da``, in the order given by + ``extra_dims``. + fields : list of tuple + A ``(name, indices)`` pair for each field, where ``indices`` gives the + index along each dimension in ``dims``. + """ + dims = [dim for dim in (extra_dims or {}) if dim in da.dims] + index_lists = [extra_dims[dim] for dim in dims] + fields = [ + (var_name + ''.join(f'_{index}' for index in indices), indices) + for indices in itertools.product(*index_lists) + ] + return dims, fields + + +def _field_metadata(ds_data, extra_dims): + """ + Build the list of fields to write, in the order they are written, for use + in the XDMF template. + """ + metadata = [] + for var_name in ds_data.data_vars: + da = ds_data[var_name] + _, fields = _expand_extra_dims(var_name, da, extra_dims) + has_time = 'Time' in da.dims + metadata.extend( + {'name': name, 'has_time': has_time} for name, _ in fields + ) + return metadata + + +def _read_group_size(da, dims, extra_dims, topo_dim, max_read_bytes): + """ + Determine how many indices along the first extra dimension can be read at + once without exceeding ``max_read_bytes``. At least one index is always + read. + """ + per_index = da.sizes[topo_dim] * da.dtype.itemsize + for dim in dims[1:]: + indices = extra_dims[dim] + per_index *= max(indices) - min(indices) + 1 + return max(1, int(max_read_bytes // max(per_index, 1))) + + +def _write_variable( + h5_file, + da, + var_name, + extra_dims, + topo_dim, + topo_map, + float32, + max_read_bytes, + iterator, +): + """ + Unwrap a variable into one HDF5 dataset per combination of extra-dimension + indices (and per time index), reading as many indices at a time as the + memory budget allows. + """ + dims, _ = _expand_extra_dims(var_name, da, extra_dims) + has_time = 'Time' in da.dims + time_indices = range(da.sizes['Time']) if has_time else [None] + + if dims: + group_size = _read_group_size( + da, dims, extra_dims, topo_dim, max_read_bytes + ) + first_indices = extra_dims[dims[0]] + groups = [ + first_indices[start : start + group_size] + for start in range(0, len(first_indices), group_size) + ] + else: + groups = [None] + + for t_index in time_indices: + da_t = da.isel(Time=t_index) if has_time else da + time_suffix = f'_t{t_index}' if has_time else '' + for group in groups: + block, offsets, index_lists = _read_block( + da_t, dims, extra_dims, group + ) + _write_block( + h5_file=h5_file, + block=block, + dims=dims, + offsets=offsets, + index_lists=index_lists, + var_name=var_name, + time_suffix=time_suffix, + topo_map=topo_map, + float32=float32, + iterator=iterator, + ) + # release the block before the next one is read, so that only one + # is ever in memory at a time + del block + + +def _read_block(da, dims, extra_dims, group): + """ + Read the smallest contiguous span of ``da`` that covers ``group`` (a set of + indices along ``dims[0]``) together with all the requested indices of the + remaining extra dimensions, in a single pass over the variable on disk. + + Returns the block along with the offset of each extra dimension within it + and the indices that the block was read for. + """ + if group is None: + return da.compute(), {}, [] + + offsets = {dims[0]: min(group)} + selection = {dims[0]: slice(min(group), max(group) + 1)} + for dim in dims[1:]: + indices = extra_dims[dim] + offsets[dim] = min(indices) + selection[dim] = slice(min(indices), max(indices) + 1) + index_lists = [group] + [extra_dims[dim] for dim in dims[1:]] + return da.isel(selection).compute(), offsets, index_lists + + +def _write_block( + h5_file, + block, + dims, + offsets, + index_lists, + var_name, + time_suffix, + topo_map, + float32, + iterator, +): + """ + Write one HDF5 dataset per combination of extra-dimension indices in an + in-memory block. + """ + for indices in itertools.product(*index_lists): + name = ( + var_name + ''.join(f'_{index}' for index in indices) + time_suffix + ) + field = block.isel( + { + dim: index - offsets[dim] + for dim, index in zip(dims, indices, strict=True) + } + ) + values = field.values + if topo_map is not None: + values = values[topo_map] + if float32 and values.dtype.kind == 'f': + values = values.astype(np.float32) + h5_file.create_dataset(name, data=values) + if iterator is not None: + iterator.update(1) + + def _parse_extra_dims(dimension_list, ds): """ Parse and prompt for indices of extra dimensions. @@ -342,9 +591,14 @@ def _parse_indices(index_string, dim_size): def _process_extra_dims(ds, extra_dims): """ - Process extra dimensions in the dataset by ensuring all are covered, - unwrapping variables with extra dimensions into multiple variables with - basic dimensions, and applying slicing or dropping variables as needed. + Process extra dimensions in the dataset by ensuring all are covered and + dropping variables with extra dimensions for which no indices were + selected. + + Variables that are kept retain their extra dimensions here. They are + unwrapped into one field per combination of indices (with the suffix + ``_`` for each extra dimension) as they are written, so that each + variable only needs to be read from disk once. Parameters ---------- @@ -372,23 +626,15 @@ def _process_extra_dims(ds, extra_dims): ) if extra_dims: - for dim, indices in extra_dims.items(): - if not indices: - # Drop variables with the given dimension if the list is empty - ds = ds.drop_vars( - [var for var in ds.data_vars if dim in ds[var].dims] - ) - else: - # Unwrap variables with the extra dimension - vars_to_unwrap = [ - var for var in ds.data_vars if dim in ds[var].dims - ] - for var in vars_to_unwrap: - for index in indices: - # Create a new variable with the suffix `_index` - new_var_name = f'{var}_{index}' - ds[new_var_name] = ds[var].isel({dim: index}) - # Drop the original variable with the extra dimension - ds = ds.drop_vars(var) + # Drop variables with a dimension for which no indices were selected + vars_to_drop = { + var: None + for dim, indices in extra_dims.items() + if not indices + for var in ds.data_vars + if dim in ds[var].dims + } + if vars_to_drop: + ds = ds.drop_vars(list(vars_to_drop)) return ds diff --git a/conda_package/mpas_tools/viz/mpas_to_xdmf/mpas_to_xdmf.py b/conda_package/mpas_tools/viz/mpas_to_xdmf/mpas_to_xdmf.py index 3499f74d0..800fcd9a0 100644 --- a/conda_package/mpas_tools/viz/mpas_to_xdmf/mpas_to_xdmf.py +++ b/conda_package/mpas_tools/viz/mpas_to_xdmf/mpas_to_xdmf.py @@ -33,6 +33,7 @@ """ # noqa: E501 from mpas_tools.viz.mpas_to_xdmf.io import ( + _DEFAULT_MAX_READ_BYTES, _convert_to_xdmf, _load_dataset, _parse_extra_dims, @@ -123,7 +124,14 @@ def load( xtime_var=xtime_var, ) - def convert_to_xdmf(self, out_dir, extra_dims=None, quiet=False): + def convert_to_xdmf( + self, + out_dir, + extra_dims=None, + quiet=False, + float32=False, + max_read_bytes=_DEFAULT_MAX_READ_BYTES, + ): """ Convert the loaded xarray Dataset to XDMF + HDF5 format. @@ -137,6 +145,15 @@ def convert_to_xdmf(self, out_dir, extra_dims=None, quiet=False): included. quiet : bool, optional If True, suppress progress output. + float32 : bool, optional + If True, write floating-point fields in single precision, halving + the size of the HDF5 files (the mesh geometry stays in double + precision). + max_read_bytes : int, optional + Approximate limit on the number of bytes read from the input + dataset at a time. Larger values allow more indices of an extra + dimension (e.g. ``nVertLevels``) to be read in a single pass over + a variable, which is faster but uses more memory. Output ------ @@ -154,7 +171,10 @@ def convert_to_xdmf(self, out_dir, extra_dims=None, quiet=False): ds=self.ds, ds_mesh=self.ds_mesh, out_dir=out_dir, + extra_dims=extra_dims, quiet=quiet, + float32=float32, + max_read_bytes=max_read_bytes, ) @@ -225,6 +245,26 @@ def main(): action='store_true', help='Suppress progress output.', ) + parser.add_argument( + '-f', + '--float32', + action='store_true', + help=( + 'Write floating-point fields in single precision, halving the ' + 'size of the HDF5 files. The mesh geometry stays in double ' + 'precision.' + ), + ) + parser.add_argument( + '--max-read-gb', + type=float, + default=_DEFAULT_MAX_READ_BYTES / 1024**3, + help=( + 'Approximate limit in GB on the amount of data read from the ' + 'input files at a time. Larger values are faster but use more ' + 'memory.' + ), + ) args = parser.parse_args() @@ -243,4 +283,6 @@ def main(): out_dir=args.output_dir, extra_dims=extra_dims, quiet=args.quiet, + float32=args.float32, + max_read_bytes=int(args.max_read_gb * 1024**3), ) From 5cc21cc322eeef9bceb3e171f1df9b488a6a7c7f Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Tue, 4 Aug 2026 02:41:20 -0500 Subject: [PATCH 2/3] Add tests for extra-dim unwrapping in mpas_to_xdmf Cover the behavior that has to be preserved now that extra dimensions are unwrapped as fields are written rather than up front: * unwrapped field names and values match the corresponding slice of the source variable, for one and for two extra dimensions, and indices that were not requested are not written * vertex-centered fields are repeated once per kite of the dual mesh * `max_read_bytes` changes only how many indices are read at a time, so forcing one read per index gives byte-identical output * `float32` casts floating-point fields but not integer fields or the geometry * `_process_extra_dims` keeps dimensions that have selected indices and drops variables with dimensions that do not Co-Authored-By: Claude Opus 5 --- conda_package/tests/test_viz_xdmf.py | 207 +++++++++++++++++++++++++++ 1 file changed, 207 insertions(+) diff --git a/conda_package/tests/test_viz_xdmf.py b/conda_package/tests/test_viz_xdmf.py index 7632565ed..6af142aaf 100644 --- a/conda_package/tests/test_viz_xdmf.py +++ b/conda_package/tests/test_viz_xdmf.py @@ -1,6 +1,7 @@ import os import sys +import h5py import numpy as np import pytest import xarray as xr @@ -18,6 +19,61 @@ TEST_MESH = get_test_data_file('mesh.QU.1920km.151026.nc') +NTIME = 3 +NVERT_LEVELS = 5 +NTRACERS = 2 + + +def _make_test_dataset(ds_mesh): + """ + Build a dataset of fields on cells, edges and vertices with an extra + ``nVertLevels`` dimension (and an ``nTracers`` dimension on one field) for + testing how extra dimensions are unwrapped. + """ + ncells = ds_mesh.sizes['nCells'] + nedges = ds_mesh.sizes['nEdges'] + nvertices = ds_mesh.sizes['nVertices'] + rng = np.random.default_rng(seed=42) + + ds = xr.Dataset() + ds['temperature'] = ( + ('Time', 'nCells', 'nVertLevels'), + rng.random((NTIME, ncells, NVERT_LEVELS)), + ) + ds['tracers'] = ( + ('Time', 'nTracers', 'nCells', 'nVertLevels'), + rng.random((NTIME, NTRACERS, ncells, NVERT_LEVELS)), + ) + ds['bottomDepth'] = ('nCells', rng.random(ncells)) + ds['maxLevelCell'] = ( + 'nCells', + rng.integers(1, NVERT_LEVELS, ncells).astype(np.int32), + ) + ds['normalVelocity'] = ( + ('Time', 'nEdges', 'nVertLevels'), + rng.random((NTIME, nedges, NVERT_LEVELS)), + ) + ds['vorticity'] = ( + ('Time', 'nVertices', 'nVertLevels'), + rng.random((NTIME, nvertices, NVERT_LEVELS)), + ) + return ds + + +def _convert_test_dataset(out_dir, **kwargs): + """ + Convert the test dataset, keeping a non-contiguous subset of the vertical + levels so that the read grouping has to handle gaps. + """ + ds_mesh = xr.open_dataset(TEST_MESH) + ds = _make_test_dataset(ds_mesh) + extra_dims = {'nVertLevels': [0, 2, 4], 'nTracers': [0, 1]} + converter = MpasToXdmf(ds=ds, ds_mesh=ds_mesh) + converter.convert_to_xdmf( + out_dir=str(out_dir), extra_dims=extra_dims, quiet=True, **kwargs + ) + return ds + @pytest.mark.skipif( not os.path.exists(TEST_MESH), reason='Test mesh not available' @@ -126,6 +182,157 @@ def test_process_extra_dims_drop(tmp_path): assert dim not in ds.dims, f'Dimension {dim} should be dropped' +@pytest.mark.skipif( + not os.path.exists(TEST_MESH), reason='Test mesh not available' +) +def test_extra_dims_unwrapped_values(tmp_path): + """ + Each combination of extra-dimension indices becomes its own field, named + with one ``_`` suffix per extra dimension, and holds the values of + the corresponding slice of the source variable. + """ + out_dir = tmp_path / 'out_unwrap' + ds = _convert_test_dataset(out_dir) + + with h5py.File(out_dir / 'fieldsOnCells.h5', 'r') as h5_file: + names = set(h5_file.keys()) + for t_index in range(NTIME): + for level in [0, 2, 4]: + name = f'temperature_{level}_t{t_index}' + expected = ds.temperature.isel( + Time=t_index, nVertLevels=level + ).values + assert np.array_equal(np.asarray(h5_file[name]), expected) + for tracer in range(NTRACERS): + name = f'tracers_{level}_{tracer}_t{t_index}' + expected = ds.tracers.isel( + Time=t_index, nVertLevels=level, nTracers=tracer + ).values + assert np.array_equal(np.asarray(h5_file[name]), expected) + # fields without extra dimensions keep their original name + assert np.array_equal( + np.asarray(h5_file['bottomDepth']), ds.bottomDepth.values + ) + + # levels that were not requested are not written + assert 'temperature_1_t0' not in names + assert 'temperature_3_t0' not in names + + +@pytest.mark.skipif( + not os.path.exists(TEST_MESH), reason='Test mesh not available' +) +def test_vertex_fields_mapped_to_kites(tmp_path): + """ + Vertex-centered fields are repeated once per kite of the dual mesh, so + they are longer than the ``nVertices`` dimension of the source data. + """ + out_dir = tmp_path / 'out_vertices' + ds = _convert_test_dataset(out_dir) + + with h5py.File(out_dir / 'fieldsOnVertices.h5', 'r') as h5_file: + cells = np.asarray(h5_file['Cells']) + field = np.asarray(h5_file['vorticity_2_t1']) + assert field.shape == (cells.shape[0],) + # every written value comes from the source field + source = ds.vorticity.isel(Time=1, nVertLevels=2).values + assert np.all(np.isin(field, source)) + + +@pytest.mark.skipif( + not os.path.exists(TEST_MESH), reason='Test mesh not available' +) +def test_max_read_bytes_does_not_change_results(tmp_path): + """ + ``max_read_bytes`` only controls how many indices are read at a time, so a + limit small enough to force one read per index must give the same output as + a limit large enough to read every index at once. + """ + one_pass = tmp_path / 'out_one_pass' + many_passes = tmp_path / 'out_many_passes' + _convert_test_dataset(one_pass, max_read_bytes=1 << 30) + # 1 byte forces the minimum of one index of the first extra dimension per + # read + _convert_test_dataset(many_passes, max_read_bytes=1) + + for basename in ['fieldsOnCells.h5', 'fieldsOnEdges.h5']: + with ( + h5py.File(one_pass / basename, 'r') as expected, + h5py.File(many_passes / basename, 'r') as actual, + ): + assert set(expected.keys()) == set(actual.keys()) + for name in expected: + assert np.array_equal( + np.asarray(expected[name]), np.asarray(actual[name]) + ), f'{basename}:/{name} differs' + + for basename in ['fieldsOnCells.xdmf', 'fieldsOnEdges.xdmf']: + assert (one_pass / basename).read_text() == ( + many_passes / basename + ).read_text() + + +@pytest.mark.skipif( + not os.path.exists(TEST_MESH), reason='Test mesh not available' +) +def test_float32(tmp_path): + """ + ``float32=True`` writes floating-point fields in single precision but + leaves integer fields and the mesh geometry alone. + """ + out_dir = tmp_path / 'out_float32' + ds = _convert_test_dataset(out_dir, float32=True) + + with h5py.File(out_dir / 'fieldsOnCells.h5', 'r') as h5_file: + assert h5_file['temperature_0_t0'].dtype == np.float32 + assert h5_file['bottomDepth'].dtype == np.float32 + assert h5_file['maxLevelCell'].dtype == np.int32 + # the geometry stays in double precision so cell shapes are unaffected + assert h5_file['Points'].dtype == np.float64 + expected = ds.temperature.isel(Time=0, nVertLevels=0).values + assert np.array_equal( + np.asarray(h5_file['temperature_0_t0']), + expected.astype(np.float32), + ) + + +@pytest.mark.skipif( + not os.path.exists(TEST_MESH), reason='Test mesh not available' +) +def test_process_extra_dims_keeps_selected_dims(tmp_path): + """ + Variables with selected indices keep their extra dimensions, which are + unwrapped as the variables are written rather than up front. + """ + converter = MpasToXdmf() + converter.load(mesh_filename=TEST_MESH) + + extra_dims = { + 'maxEdges': [0, 1], + 'maxEdges2': [], + 'TWO': [], + 'vertexDegree': [], + } + ds = _process_extra_dims(converter.ds, extra_dims=extra_dims) + + assert 'maxEdges' in ds.dims + assert 'verticesOnCell' in ds.data_vars + assert 'maxEdges' in ds.verticesOnCell.dims + for dim in ['maxEdges2', 'TWO', 'vertexDegree']: + assert dim not in ds.dims, f'Dimension {dim} should be dropped' + + +@pytest.mark.skipif( + not os.path.exists(TEST_MESH), reason='Test mesh not available' +) +def test_process_extra_dims_uncovered_dim(): + """An extra dimension that is not listed in ``extra_dims`` is an error.""" + converter = MpasToXdmf() + converter.load(mesh_filename=TEST_MESH) + with pytest.raises(ValueError, match='is not covered'): + _process_extra_dims(converter.ds, extra_dims={'maxEdges': [0]}) + + @pytest.mark.skipif( not os.path.exists(TEST_MESH), reason='Test mesh not available' ) From 5c485e1e2708004ef7b9b53ad23b558c440b9087 Mon Sep 17 00:00:00 2001 From: Xylar Asay-Davis Date: Tue, 4 Aug 2026 02:41:30 -0500 Subject: [PATCH 3/3] Document mpas_to_xdmf performance on large meshes Add a section explaining why reading a vertical level at a time was slow, what `max_read_bytes` trades off, and the two ways to cut the volume of data written (`float32` and requesting only the levels of interest). Also list the `-q`, `-f` and `--max-read-gb` command-line arguments. Co-Authored-By: Claude Opus 5 --- conda_package/docs/mpas_to_xdmf.rst | 33 +++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/conda_package/docs/mpas_to_xdmf.rst b/conda_package/docs/mpas_to_xdmf.rst index 82dd47ee2..8536a2567 100644 --- a/conda_package/docs/mpas_to_xdmf.rst +++ b/conda_package/docs/mpas_to_xdmf.rst @@ -50,6 +50,11 @@ The following arguments are supported: - ``-x, --xtime``: Name of the variable containing time information (optional). - ``-d, --dim-list``: List of dimensions and indices to slice (e.g., ``nVertLevels=0:10:2``). +- ``-q, --quiet``: Suppress progress output. +- ``-f, --float32``: Write floating-point fields in single precision, halving + the size of the HDF5 files. The mesh geometry stays in double precision. +- ``--max-read-gb``: Approximate limit in GB on the amount of data read from + the input files at a time (default 2). See :ref:`xdmf_performance`. Examples -------- @@ -143,6 +148,34 @@ The MPAS to XDMF Converter includes several basic features: - **Selective Variable Conversion**: Users can choose specific variables or groups of variables (e.g., ``allOnCells``) for conversion. +.. _xdmf_performance: + +Performance on Large Meshes +=========================== +Each index of an extra dimension becomes its own field in the output, so a +3D field such as ``temperature(Time, nCells, nVertLevels)`` is written as one +2D field per vertical level. A single vertical level is strided across the +whole variable on disk, so reading the levels one at a time means reading the +entire variable once per level. On a 4-million-cell, 80-layer mesh that read +amplification dominated the run time completely. + +Instead, the converter reads as many indices of an extra dimension as it can +in a single pass over each variable, and slices them apart in memory. The +amount read at a time is capped by ``max_read_bytes`` (``--max-read-gb`` on +the command line, 2 GB by default) so that memory use stays bounded; peak +memory is roughly two to three times that limit, because masking the fill +values makes a copy. Raising the limit lets larger variables be read in a +single pass; lowering it trades speed for memory. + +Two further options help with the sheer volume of data: + +- ``float32`` (``-f``/``--float32``) writes floating-point fields in single + precision. This halves both the size of the HDF5 files and the time + ParaView needs to load them. The mesh geometry stays in double precision so + that the cell shapes are unaffected. +- Passing only the vertical levels you actually want to look at, via + ``extra_dims``/``-d``, avoids writing the rest. + Opening Files in ParaView ========================= Once the conversion is complete, you can open the generated XDMF files in