Fixing chunk sizes in zarr output

acoat · Dec 22, 2023 · 0f4c31d · 0f4c31d
1 parent ab4dde1
commit 0f4c31d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/parcels/particlefile.py b/parcels/particlefile.py
@@ -183,7 +183,7 @@ def _extend_zarr_dims(self, Z, store, dtype, axis):
             if len(obs) == Z.shape[1]:
                 obs.append(np.arange(self.chunks[1])+obs[-1]+1)
         else:
-            extra_trajs = max(self.maxids - Z.shape[0], self.chunks[0])
+            extra_trajs = self.maxids - Z.shape[0]
             if len(Z.shape) == 2:
                 a = np.full((extra_trajs, Z.shape[1]), self.fill_value_map[dtype], dtype=dtype)
             else:
@@ -233,7 +233,7 @@ def write(self, pset, time, indices=None):
                 if (self.maxids > len(ids)) or (self.maxids > self.chunks[0]):
                     arrsize = (self.maxids, self.chunks[1])
                 else:
-                    arrsize = (len(ids), 1)
+                    arrsize = (len(ids), self.chunks[1])
                 ds = xr.Dataset(attrs=self.metadata, coords={"trajectory": ("trajectory", pids),
                                                              "obs": ("obs", np.arange(arrsize[1], dtype=np.int32))})
                 attrs = self._create_variables_attribute_dict()