From 5e9175ec0b4511772dc07ba0777812a59b4a349e Mon Sep 17 00:00:00 2001 From: Tom White Date: Thu, 21 Dec 2023 14:26:02 +0000 Subject: [PATCH] Remove `initial_values` in LazyZarrArray as it is no longer used --- cubed/core/plan.py | 13 +------------ cubed/storage/zarr.py | 16 +--------------- cubed/tests/storage/test_zarr.py | 16 +--------------- 3 files changed, 3 insertions(+), 42 deletions(-) diff --git a/cubed/core/plan.py b/cubed/core/plan.py index af7535012..8c1d301e1 100644 --- a/cubed/core/plan.py +++ b/cubed/core/plan.py @@ -5,7 +5,6 @@ import networkx as nx -from cubed.backend_array_api import backend_array_to_numpy_array from cubed.primitive.blockwise import can_fuse_pipelines, fuse from cubed.runtime.pipeline import visit_nodes from cubed.runtime.types import CubedPipeline @@ -362,17 +361,7 @@ def create_zarr_array(lazy_zarr_array, *, config=None): def create_zarr_arrays(lazy_zarr_arrays, reserved_mem): # projected memory is size of largest initial values, or dtype size if there aren't any projected_mem = ( - max( - [ - # TODO: calculate nbytes from size and dtype itemsize - backend_array_to_numpy_array(lza.initial_values).nbytes - if lza.initial_values is not None - else lza.dtype.itemsize - for lza in lazy_zarr_arrays - ], - default=0, - ) - + reserved_mem + max([lza.dtype.itemsize for lza in lazy_zarr_arrays], default=0) + reserved_mem ) num_tasks = len(lazy_zarr_arrays) diff --git a/cubed/storage/zarr.py b/cubed/storage/zarr.py index 60c33874a..d1f17e056 100644 --- a/cubed/storage/zarr.py +++ b/cubed/storage/zarr.py @@ -1,9 +1,7 @@ -from typing import Any, Optional, Union +from typing import Any, Union import zarr -from numpy import ndarray -from cubed.backend_array_api import backend_array_to_numpy_array from cubed.types import T_DType, T_RegularChunks, T_Shape, T_Store @@ -21,7 +19,6 @@ def __init__( dtype: T_DType, chunks: T_RegularChunks, store: T_Store, - initial_values: Optional[ndarray] = None, fill_value: Any = None, **kwargs, ): @@ -35,7 +32,6 @@ def __init__( self.chunks = template.chunks self.store = store - self.initial_values = initial_values self.fill_value = fill_value self.kwargs = kwargs @@ -60,8 +56,6 @@ def create(self, mode: str = "w-") -> zarr.Array: fill_value=self.fill_value, **self.kwargs, ) - if self.initial_values is not None and self.initial_values.size > 0: - target[...] = backend_array_to_numpy_array(self.initial_values) return target def open(self) -> zarr.Array: @@ -91,14 +85,6 @@ def lazy_empty( return LazyZarrArray(shape, dtype, chunks, store, **kwargs) -def lazy_from_array( - array: ndarray, *, dtype: T_DType, chunks: T_RegularChunks, store: T_Store, **kwargs -) -> LazyZarrArray: - return LazyZarrArray( - array.shape, dtype, chunks, store, initial_values=array, **kwargs - ) - - def lazy_full( shape: T_Shape, fill_value: Any, diff --git a/cubed/tests/storage/test_zarr.py b/cubed/tests/storage/test_zarr.py index c7da964bb..e4a6de44a 100644 --- a/cubed/tests/storage/test_zarr.py +++ b/cubed/tests/storage/test_zarr.py @@ -2,7 +2,7 @@ import pytest from numpy.testing import assert_array_equal -from cubed.storage.zarr import lazy_empty, lazy_from_array, lazy_full +from cubed.storage.zarr import lazy_empty, lazy_full def test_lazy_empty(tmp_path): @@ -18,20 +18,6 @@ def test_lazy_empty(tmp_path): arr.open() -def test_lazy_from_array(tmp_path): - zarr_path = tmp_path / "lazy.zarr" - a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=int) - arr = lazy_from_array(a, dtype=a.dtype, chunks=(2, 2), store=zarr_path) - - assert not zarr_path.exists() - with pytest.raises(ValueError): - arr.open() - - arr.create() - assert zarr_path.exists() - assert_array_equal(arr.open()[:], a) - - def test_lazy_full(tmp_path): zarr_path = tmp_path / "lazy.zarr" arr = lazy_full(