Bound NetCDF chunks for long native-grid records

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.

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'])

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)

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)

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

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:

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,19 +32,27 @@ 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:
         target[...] = variable[...]
     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
 
 

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):

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