diff --git a/docs/output-netcdf-layout.md b/docs/output-netcdf-layout.md index 6f8959c..ec5fc2b 100644 --- a/docs/output-netcdf-layout.md +++ b/docs/output-netcdf-layout.md @@ -16,3 +16,17 @@ Two final products share exactly the same `x`/`y` values, UTM CRS and 30 m affin Scene retrieval auxiliaries use `scene_time`; dT climatology uses `dayofyear`. These temporal dimensions differ from the daily model time axis without changing the spatial grid. Scenario-specific variables receive `_default`/`_calibrated` suffixes. Calibration-disabled runs have no calibrated variable. Four scalar parameters for each scenario and the configuration are stored as attributes; manifests retain input identity and software provenance. Intermediate single-scenario SWB files retain the component API's `rzsm_layer_*` variables under `products/intermediate/`. They support diagnosis and are not additional consolidated final products. Native calibration inputs remain separate under `calibration_inputs/`; they intentionally have a different spatial grid. See [the full native-grid contract](native-grids.md). + +## Streaming large records + +Streamed daily and scene-dependent arrays use temporal chunks of at most +32 entries and spatial chunks of at most 64 pixels. This matches the bounded +copy operations used for final crops, scenario combination and initialization +trimming. Seasonal controls additionally align their y chunks with the row +blocks used to construct the control. + +Older source files with large temporal chunks remain readable. Streaming copies +temporarily increase the current variable's read cache to retain chunk slabs, +capped at 512 MiB per variable, then restore its original cache settings. +These storage choices preserve values, masks, coordinates, units and scientific +provenance; they do not change model or forcing parameters. diff --git a/pysweb/forcing.py b/pysweb/forcing.py index 5998688..18c4202 100644 --- a/pysweb/forcing.py +++ b/pysweb/forcing.py @@ -195,7 +195,9 @@ def write_seasonal_seb(source, destination, *, development_start, development_en from pathlib import Path import json import netCDF4 as nc + from contextlib import ExitStack from pysweb.io.provenance import input_manifest, save_output_manifest + from pysweb.io.netcdf_stream import stream_chunks, sequential_read_cache if block_rows < 1: raise ValueError('Positive control block size required') destination = Path(destination) @@ -210,16 +212,30 @@ def write_seasonal_seb(source, destination, *, development_start, development_en template = ds[['ET','E_t','E_s','ndvi_interp','f_t','etf_interp','reference_et']] # Preserve native coordinates/CRS without first allocating the full cube. skeleton = template.isel(time=slice(0,0)) - skeleton.to_netcdf(temporary,unlimited_dims=['time']) - with nc.Dataset(source) as raw, nc.Dataset(temporary,'a') as output: + encoding = {} + for name in skeleton.data_vars: + if template[name].dims != ('time', 'y', 'x'): + raise ValueError('Seasonal SEB fields must have time/y/x dimensions') + # An explicit encoding avoids xarray dropping chunks because this + # empty skeleton differs from the source's original_shape metadata. + encoding[name] = {key:value for key,value in template[name].encoding.items() + if key in ('dtype', '_FillValue', 'grid_mapping', 'scale_factor', 'add_offset')} + encoding[name].update(zlib=True, complevel=3, + chunksizes=stream_chunks(template[name].dims, template[name].shape, + block_dim='y', block_size=block_rows)) + skeleton.to_netcdf(temporary,unlimited_dims=['time'],encoding=encoding) + with nc.Dataset(source) as raw, nc.Dataset(temporary,'a') as output, ExitStack() as caches: + source_fields = ('etf_interp', 'ndvi_interp', 'reference_et', 'f_t') + for name in source_fields: + caches.enter_context(sequential_read_cache(raw[name], 'y')) output.variables['time'][:] = raw.variables['time'][:] for begin in range(0,ds.sizes['y'],block_rows): rows = slice(begin,min(begin+block_rows,ds.sizes['y'])) - block = template.isel(y=rows).load() - fraction = seasonal_fraction(block.etf_interp.values,dates, + block = {name:np.ma.filled(raw[name][:,rows,:], np.nan) for name in source_fields} + fraction = seasonal_fraction(block['etf_interp'],dates, development_start,development_end,dates,half_window=half_window,minimum_years=minimum_years) - values = dict(etf_interp=fraction,ndvi_interp=block.ndvi_interp.values, - reference_et=block.reference_et.values,f_t=block.f_t.values) + values = dict(etf_interp=fraction,ndvi_interp=block['ndvi_interp'], + reference_et=block['reference_et'],f_t=block['f_t']) values['ET'] = fraction*values['reference_et'] values['E_t'] = values['ET']*values['f_t'] values['E_s'] = values['ET']*(1-values['f_t']) diff --git a/pysweb/io/netcdf_stream.py b/pysweb/io/netcdf_stream.py new file mode 100644 index 0000000..12e98c8 --- /dev/null +++ b/pysweb/io/netcdf_stream.py @@ -0,0 +1,56 @@ +""" +Script: netcdf_stream.py +Objective: Match compressed NetCDF chunks and bounded read caches to streamed array copies. +Author: Yi Yu (with assistance from Codex) +Created: 2026-10-17 +Last updated: 2026-10-17 +Inputs: Variable dimensions, shapes and existing NetCDF chunk metadata. +Outputs: Bounded chunk layouts and temporary per-variable read-cache settings. +Usage: from pysweb.io.netcdf_stream import stream_chunks, sequential_read_cache +Dependencies: contextlib, math, numpy, netCDF4-compatible variables. +""" +from contextlib import contextmanager +from math import ceil, prod +import numpy as np + + +def stream_chunks(dimensions, shape, *, block_dim=None, block_size=32, spatial_size=64): + """Keep compressed chunks within the slabs written by streaming callers.""" + if block_size < 1 or spatial_size < 1: + raise ValueError('Positive streaming chunk sizes required') + if block_dim is None: + block_dim = next((d for d in ('time', 'scene_time') if d in dimensions), None) + if block_dim is None or block_dim not in dimensions: + return None + return tuple(max(1, min(int(size), block_size if dim == block_dim else + 32 if dim in ('time', 'scene_time') else 1 if dim == 'layer' else spatial_size)) + for dim, size in zip(dimensions, shape)) + + +@contextmanager +def sequential_read_cache(variable, block_dim='time', *, max_cache_bytes=512*1024**2): + """Cache one or two source chunk slabs, capped independently of host memory. + + Old files may have multi-year temporal chunks. A small time slice then + touches more spatial chunks than the default cache can retain. Reusing the + slab prevents repeated decompression while callers still read small blocks. + Restore the original cache after each variable, including on failure. + """ + chunks = variable.chunking() + if chunks is None or isinstance(chunks, str) or block_dim not in variable.dimensions: + yield + return + if max_cache_bytes < 1: + raise ValueError('Positive NetCDF cache limit required') + previous = variable.get_var_chunk_cache() + axis = variable.dimensions.index(block_dim) + slab_chunks = prod(ceil(size/chunk) for i, (size, chunk) in + enumerate(zip(variable.shape, chunks)) if i != axis) + slab_bytes = prod(chunks)*np.dtype(variable.dtype).itemsize*slab_chunks + size = min(int(max_cache_bytes), max(int(previous[0]), 2*slab_bytes)) + slots = max(int(previous[1]), 2*slab_chunks+1) + variable.set_var_chunk_cache(size=size, nelems=slots, preemption=1.0) + try: + yield + finally: + variable.set_var_chunk_cache(*previous) diff --git a/pysweb/io/products.py b/pysweb/io/products.py index d1380a6..4028939 100644 --- a/pysweb/io/products.py +++ b/pysweb/io/products.py @@ -15,6 +15,7 @@ import numpy as np from pysweb.io.grids import open_template, spatial from pysweb.io.provenance import input_manifest, save_output_manifest +from pysweb.io.netcdf_stream import stream_chunks, sequential_read_cache def assert_identical_grids(first_path, second_path): @@ -53,16 +54,21 @@ def crop_product(source_path, output_path, template): for name, variable in source.variables.items(): fill = variable.getncattr('_FillValue') if '_FillValue' in variable.ncattrs() else None options = {'zlib': True, 'complevel': 3} if variable.ndim else {} + block_dim = next((d for d in variable.dimensions if d not in slices), None) + shape = tuple(len(output.dimensions[d]) for d in variable.dimensions) + chunks = stream_chunks(variable.dimensions, shape, block_dim=block_dim) + if chunks is not None: + options['chunksizes'] = chunks target = output.createVariable(name, variable.datatype, variable.dimensions, fill_value=fill, **options) target.setncatts({k: variable.getncattr(k) for k in variable.ncattrs() if k != '_FillValue'}) - block_dim = next((d for d in variable.dimensions if d not in slices), None) count = len(source.dimensions[block_dim]) if block_dim else 1 - for begin in range(0, count, 32): - src = tuple(slices.get(d, slice(begin, begin+32) if d == block_dim else slice(None)) - for d in variable.dimensions) - dst = tuple(slice(begin, begin+32) if d == block_dim else slice(None) - for d in variable.dimensions) - target[dst or Ellipsis] = variable[src or Ellipsis] + with sequential_read_cache(variable, block_dim): + for begin in range(0, count, 32): + src = tuple(slices.get(d, slice(begin, begin+32) if d == block_dim else slice(None)) + for d in variable.dimensions) + dst = tuple(slice(begin, begin+32) if d == block_dim else slice(None) + for d in variable.dimensions) + target[dst or Ellipsis] = variable[src or Ellipsis] output['spatial_ref'].GeoTransform = ' '.join(str(v) for v in template.rio.transform().to_gdal()) output.spatial_grid = 'Exact subset of selected Landsat 30 m UTM scene grid' temporary.replace(output_path) diff --git a/pysweb/io/time_reindex.py b/pysweb/io/time_reindex.py index 0ea9ad7..796ef3e 100644 --- a/pysweb/io/time_reindex.py +++ b/pysweb/io/time_reindex.py @@ -15,6 +15,7 @@ import netCDF4 import numpy as np import pandas as pd +from pysweb.io.netcdf_stream import stream_chunks, sequential_read_cache @@ -98,16 +99,18 @@ def reindex_netcdf_time(source, destination, dates, source_indices, *, options = {} if timed: options = dict(zlib=True, complevel=3, - chunksizes=(min(chunk_days,len(dates)),)+tuple(variable.shape[1:])) + chunksizes=stream_chunks(variable.dimensions, + (len(dates),)+variable.shape[1:], block_dim=time_dim, block_size=chunk_days)) target = dst.createVariable(name, variable.dtype, variable.dimensions, fill_value=fill, **options) target.setncatts(attrs) target.set_auto_maskandscale(False) if not timed: target[...] = variable[...] continue - for start in range(0, len(index), chunk_days): - selected = index[start:start+chunk_days] - block = _indexed_block(variable, selected, fill) - target[start:start+len(selected), ...] = block + with sequential_read_cache(variable, time_dim): + for start in range(0, len(index), chunk_days): + selected = index[start:start+chunk_days] + block = _indexed_block(variable, selected, fill) + target[start:start+len(selected), ...] = block temporary.replace(destination) return destination diff --git a/pysweb/ssebop/scene_protocol.py b/pysweb/ssebop/scene_protocol.py index 03c5d0e..eb4b9f4 100644 --- a/pysweb/ssebop/scene_protocol.py +++ b/pysweb/ssebop/scene_protocol.py @@ -254,10 +254,14 @@ def run_scene_workflow(*, date_range, scene_files, gapfill_products, dem, # Supplying encoding replaces xarray's per-variable encoding; preserve the # CF grid-mapping link written by rioxarray for every spatial auxiliary. encoding = {} + from pysweb.io.netcdf_stream import stream_chunks for name, field in result.data_vars.items(): encoding[name] = {key: value for key, value in field.encoding.items() if key in ('grid_mapping', '_FillValue')} encoding[name].update(zlib=True, complevel=3) + chunks = stream_chunks(field.dims, field.shape) + if chunks is not None: + encoding[name]['chunksizes'] = chunks if field.dtype.kind == 'f': encoding[name]['dtype'] = 'float32' else: diff --git a/pysweb/swb/paired.py b/pysweb/swb/paired.py index 9cbfcfa..0d654b5 100644 --- a/pysweb/swb/paired.py +++ b/pysweb/swb/paired.py @@ -18,6 +18,7 @@ from pysweb.swb.run import run_swb_workflow from pysweb.swb.core import SOIL_FILE_STEMS from pysweb.io.products import assert_identical_grids +from pysweb.io.netcdf_stream import stream_chunks, sequential_read_cache def _copy_variable(destination, variable, name=None, time_indices=None): @@ -31,6 +32,10 @@ def _copy_variable(destination, variable, name=None, time_indices=None): raise ValueError(f"Auxiliary {name} has incompatible {dim} dimension") fill = variable.getncattr('_FillValue') if '_FillValue' in variable.ncattrs() else None options = {'zlib': True, 'complevel': 3} if variable.ndim else {} + chunks = stream_chunks(variable.dimensions, + tuple(len(destination.dimensions[d]) for d in variable.dimensions)) + if chunks is not None: + options['chunksizes'] = chunks target = destination.createVariable(name, variable.datatype, variable.dimensions, fill_value=fill, **options) target.setncatts({k:variable.getncattr(k) for k in variable.ncattrs() if k != '_FillValue'}) if 'time' not in variable.dimensions: @@ -38,12 +43,16 @@ def _copy_variable(destination, variable, name=None, time_indices=None): else: axis = variable.dimensions.index('time') count = len(destination.dimensions['time']) - for begin in range(0, count, 32): - source_index = [slice(None)]*variable.ndim - target_index = [slice(None)]*variable.ndim - target_index[axis] = slice(begin, min(begin+32, count)) - source_index[axis] = target_index[axis] if time_indices is None else time_indices[begin:begin+32] - target[tuple(target_index)] = variable[tuple(source_index)] + with sequential_read_cache(variable): + for begin in range(0, count, 32): + source_index = [slice(None)]*variable.ndim + target_index = [slice(None)]*variable.ndim + target_index[axis] = slice(begin, min(begin+32, count)) + selected = target_index[axis] if time_indices is None else time_indices[begin:begin+32] + if not isinstance(selected, slice) and len(selected) and np.all(np.diff(selected) == 1): + selected = slice(int(selected[0]), int(selected[-1])+1) + source_index[axis] = selected + target[tuple(target_index)] = variable[tuple(source_index)] return target diff --git a/tests/io/test_native_grids.py b/tests/io/test_native_grids.py index f189bd7..e7d1b9a 100644 --- a/tests/io/test_native_grids.py +++ b/tests/io/test_native_grids.py @@ -82,6 +82,10 @@ def test_exact_subsets_keep_reference_values_and_final_affine(tmp_path): with xr.open_dataset(cropped,decode_coords='all') as actual: for name in ('ET','elevation','qa'): np.testing.assert_equal(actual[name].values,expected[name].values) + import netCDF4 as nc + with nc.Dataset(cropped) as actual: + assert actual['ET'].chunking()[0] <= 32 + assert actual['qa'].chunking()[0] <= 32 wrong = expected.assign_coords(x=expected.x+15).rio.write_transform(Affine(30,0,500045,0,-30,6099970)) wrong.to_netcdf(tmp_path/'wrong.nc') with pytest.raises(ValueError): diff --git a/tests/io/test_netcdf_stream.py b/tests/io/test_netcdf_stream.py new file mode 100644 index 0000000..9fc416f --- /dev/null +++ b/tests/io/test_netcdf_stream.py @@ -0,0 +1,72 @@ +""" +Script: test_netcdf_stream.py +Objective: Verify bounded compressed layouts, cache restoration and exact streamed auxiliary values. +Author: Yi Yu (with assistance from Codex) +Created: 2026-10-17 +Last updated: 2026-10-17 +Inputs: Small numeric NetCDF fixtures and unpopulated long-calendar variable metadata. +Outputs: Streaming memory/layout and value-preservation assertions. +Usage: python -m pytest tests/io/test_netcdf_stream.py +Dependencies: netCDF4, numpy, pytest, pysweb. +""" +import netCDF4 as nc +import numpy as np +import pytest +from pysweb.io.netcdf_stream import stream_chunks, sequential_read_cache +from pysweb.swb.paired import _copy_variable + + +def test_long_calendar_chunks_follow_temporal_and_row_write_blocks(): + assert stream_chunks(('time', 'y', 'x'), (5844, 167, 167)) == (32, 64, 64) + assert stream_chunks(('scene_time', 'y', 'x'), (577, 202, 221)) == (32, 64, 64) + assert stream_chunks(('time', 'layer', 'y', 'x'), (12414, 5, 167, 167)) == (32, 1, 64, 64) + assert stream_chunks(('time', 'y', 'x'), (5844, 167, 167), + block_dim='y', block_size=8) == (32, 8, 64) + assert stream_chunks(('y', 'x'), (167, 167)) is None + with pytest.raises(ValueError, match='Positive'): + stream_chunks(('time',), (10,), block_size=0) + + +def test_source_cache_is_capped_and_restored_after_failure(tmp_path): + with nc.Dataset(tmp_path/'old_layout.nc', 'w') as ds: + for dim, size in [('time', 6000), ('y', 128), ('x', 128)]: + ds.createDimension(dim, size) + # Metadata only: model the real long-calendar chunk geometry without + # allocating or writing a large fixture in the test suite. + variable = ds.createVariable('ET', 'f4', ('time', 'y', 'x'), + zlib=True, chunksizes=(1500, 32, 32)) + variable.set_var_chunk_cache(size=16*1024**2, nelems=1009, preemption=0.75) + original = variable.get_var_chunk_cache() + with pytest.raises(RuntimeError, match='copy failed'): + with sequential_read_cache(variable, max_cache_bytes=128*1024**2): + assert original[0] < variable.get_var_chunk_cache()[0] <= 128*1024**2 + raise RuntimeError('copy failed') + assert variable.get_var_chunk_cache() == original + + +@pytest.mark.parametrize('format', ['NETCDF3_CLASSIC', 'NETCDF4']) +def test_unchunked_sources_remain_readable(tmp_path, format): + with nc.Dataset(tmp_path/'contiguous.nc', 'w', format=format) as ds: + ds.createDimension('time', 3) + variable = ds.createVariable('ET', 'f4', ('time',)) + variable[:] = [1., 2., 3.] + with sequential_read_cache(variable): + np.testing.assert_array_equal(variable[:], [1., 2., 3.]) + + +@pytest.mark.parametrize('indices', [np.arange(1, 5), np.array([0, 2, 3, 5])]) +def test_auxiliary_copy_preserves_nonleading_time_values_and_metadata(tmp_path, indices): + values = np.arange(36, dtype='float32').reshape(3, 6, 2) + values[1, 3, 0] = np.nan + with nc.Dataset(tmp_path/'source.nc', 'w') as source, nc.Dataset(tmp_path/'copy.nc', 'w') as target: + for dim, size in [('y', 3), ('time', 6), ('x', 2)]: + source.createDimension(dim, size) + target.createDimension(dim, len(indices) if dim == 'time' else size) + variable = source.createVariable('forcing', 'f4', ('y', 'time', 'x'), + fill_value=np.nan, zlib=True) + variable.units = 'mm day-1' + variable[:] = values + copied = _copy_variable(target, variable, time_indices=indices) + np.testing.assert_equal(np.ma.filled(copied[:], np.nan), values[:, indices, :]) + assert copied.units == variable.units + assert copied.chunking()[1] <= 32 diff --git a/tests/test_protocols.py b/tests/test_protocols.py index 701e3a3..2d41b50 100644 --- a/tests/test_protocols.py +++ b/tests/test_protocols.py @@ -399,16 +399,27 @@ def test_seasonal_fraction_support_and_holdout_isolation(tmp_path): result=seasonal_fraction(missing,dates,'2018-01-01','2019-12-31',dates) assert np.isnan(result[:,0,0]).all() grid=raster(np.ones((2,2)),Affine(30,0,500000,0,-30,6100000)) - fraction=grid.expand_dims(time=dates).copy(data=values) + fraction=grid.expand_dims(time=dates).copy(data=missing) ds=xr.Dataset(dict(ET=fraction*5,E_t=fraction*3,E_s=fraction*2, ndvi_interp=xr.ones_like(fraction)*.65,f_t=xr.ones_like(fraction)*.6, etf_interp=fraction,reference_et=xr.ones_like(fraction)*5)) source=tmp_path/'source.nc';destination=tmp_path/'seasonal.nc' - ds.to_netcdf(source) + ds.reference_et[0,0,1]=np.nan + ds.to_netcdf(source,encoding={'etf_interp':{'_FillValue':-9999.,'grid_mapping':'spatial_ref'}, + 'reference_et':{'_FillValue':-9999.,'grid_mapping':'spatial_ref'}}) + with xr.open_dataset(source,decode_coords='all') as encoded: + assert encoded.rio.crs==grid.rio.crs write_seasonal_seb(source,destination,development_start='2018-01-01',development_end='2019-12-31',block_rows=1) with xr.open_dataset(destination,decode_coords='all') as control: assert control.rio.crs==grid.rio.crs - np.testing.assert_allclose(control.ET,1.) - np.testing.assert_allclose(control.E_t,.6) + expected=np.ones_like(values) + expected[:,0,0]=np.nan + expected[0,0,1]=np.nan + np.testing.assert_allclose(control.ET,expected) + np.testing.assert_allclose(control.E_t,expected*.6) np.testing.assert_allclose(control.ET,control.E_t+control.E_s) np.testing.assert_array_equal(control.ndvi_interp,ds.ndvi_interp) + import netCDF4 as nc + with nc.Dataset(destination) as control: + assert control['ET'].chunking()[0]<=32 + assert control['ET'].chunking()[1]==1