From 026768cc8d05ba48d44e8ec3db514207e563b321 Mon Sep 17 00:00:00 2001 From: Feiyang Wu Date: Sun, 12 Jul 2026 22:42:30 -0400 Subject: [PATCH] fix(loader): make zarr shard shape a multiple of chunk shape for short trajectories Trajectories shorter than shard_size collapsed the shard dim to the trajectory length (e.g. T=341), which is generally not a multiple of chunk_size (64), so zarr v3 rejected the array ('chunk_shape not divisible by inner chunk_shape'). Round the leading shard dim down to a multiple of the leading chunk dim. Long trajectories (T>=shard_size) are unaffected. --- iltools/datasets/lafan1/loader.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/iltools/datasets/lafan1/loader.py b/iltools/datasets/lafan1/loader.py index 93c5914..ce3a45a 100644 --- a/iltools/datasets/lafan1/loader.py +++ b/iltools/datasets/lafan1/loader.py @@ -1376,8 +1376,19 @@ def _save_trajectory_data( if array.shape[0] == 0: continue - chunks = [min(chunk_size, array.shape[0])] + list(array.shape[1:]) - shards = [min(shard_size, array.shape[0])] + list(array.shape[1:]) + # zarr v3 sharding requires the shard shape to be an exact multiple of + # the inner chunk shape along every dimension. For trajectories shorter + # than shard_size the shard collapses to the trajectory length, which is + # generally NOT a multiple of chunk_size (e.g. T=341, chunk=64), so round + # the leading shard dim down to a multiple of the leading chunk dim. Long + # trajectories (T >= shard_size) are unaffected (shard_size is a multiple + # of chunk_size by construction). + length = int(array.shape[0]) + chunk_dim0 = max(1, min(chunk_size, length)) + shard_dim0 = max(chunk_dim0, min(shard_size, length)) + shard_dim0 = (shard_dim0 // chunk_dim0) * chunk_dim0 + chunks = [chunk_dim0] + list(array.shape[1:]) + shards = [shard_dim0] + list(array.shape[1:]) ds = traj_group.create_array( key, shape=array.shape,