diff --git a/CHANGELOG.md b/CHANGELOG.md index 0ae138eaf..0f1a373fe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,12 @@ Write the date in place of the "Unreleased" in the case a new version is release # Changelog +## Unreleased + +### Added + +- `data_type` and `coord_data_type` properties for sparse arrays in `COOAdapter` and `COOStructure`. + ## v0.1.0b10 (2024-10-11) - Add kwarg to client logout to auto-clear default identity. diff --git a/tiled/_tests/test_protocols.py b/tiled/_tests/test_protocols.py index e54921cd5..a458ae106 100644 --- a/tiled/_tests/test_protocols.py +++ b/tiled/_tests/test_protocols.py @@ -247,9 +247,13 @@ def test_sparseadapter_protocol(mocker: MockFixture) -> None: mock_call4 = mocker.patch.object(CustomSparseAdapter, "specs") mock_call5 = mocker.patch.object(CustomSparseAdapter, "metadata") - structure = COOStructure(shape=(2 * 5,), chunks=((5, 5),)) - array = numpy.random.rand(2, 512, 512) + + structure = COOStructure( + shape=(2 * 5,), + chunks=((5, 5),), + data_type=BuiltinDtype.from_numpy_dtype(array.dtype), + ) blocks: Dict[Tuple[int, ...], Tuple[NDArray[Any], Any]] = {(1,): (array, (1,))} metadata: JSON = {"foo": "bar"} anyslice = (1, 1, 1) diff --git a/tiled/_tests/test_writing.py b/tiled/_tests/test_writing.py index 571aaad59..320dba4e3 100644 --- a/tiled/_tests/test_writing.py +++ b/tiled/_tests/test_writing.py @@ -26,6 +26,7 @@ from ..mimetypes import PARQUET_MIMETYPE from ..queries import Key from ..server.app import build_app +from ..structures.array import BuiltinDtype from ..structures.core import Spec, StructureFamily from ..structures.data_source import DataSource from ..structures.sparse import COOStructure @@ -245,7 +246,13 @@ def test_write_sparse_chunked(tree): "sparse", [ DataSource( - structure=COOStructure(shape=(2 * N,), chunks=((N, N),)), + structure=COOStructure( + shape=(2 * N,), + chunks=((N, N),), + data_type=BuiltinDtype.from_numpy_dtype( + numpy.dtype("float64") + ), + ), structure_family="sparse", ) ], diff --git a/tiled/adapters/sparse.py b/tiled/adapters/sparse.py index 60b122d5c..0bb62172c 100644 --- a/tiled/adapters/sparse.py +++ b/tiled/adapters/sparse.py @@ -6,6 +6,7 @@ import sparse from numpy._typing import NDArray +from ..structures.array import BuiltinDtype from ..structures.core import Spec, StructureFamily from ..structures.sparse import COOStructure from .array import slice_and_shape_from_block_and_chunks @@ -49,6 +50,8 @@ def from_arrays( dims=dims, shape=shape, chunks=tuple((dim,) for dim in shape), + data_type=BuiltinDtype.from_numpy_dtype(data.dtype), + coord_data_type=BuiltinDtype.from_numpy_dtype(coords.dtype), resizable=False, ) return cls( @@ -133,6 +136,8 @@ def from_global_ref( dims=dims, shape=shape, chunks=chunks, + data_type=BuiltinDtype.from_numpy_dtype(data.dtype), + coord_data_type=BuiltinDtype.from_numpy_dtype(coords.dtype), resizable=False, ) return cls( diff --git a/tiled/client/container.py b/tiled/client/container.py index 7a87ce507..bb20ec2f0 100644 --- a/tiled/client/container.py +++ b/tiled/client/container.py @@ -886,6 +886,7 @@ def write_sparse( >>> x.write_block(coords=[[2, 4]], data=[3.1, 2.8], block=(0,)) >>> x.write_block(coords=[[0, 1]], data=[6.7, 1.2], block=(1,)) """ + from ..structures.array import BuiltinDtype from ..structures.sparse import COOStructure structure = COOStructure( @@ -893,6 +894,8 @@ def write_sparse( # This method only supports single-chunk COO arrays. chunks=tuple((dim,) for dim in shape), dims=dims, + data_type=BuiltinDtype.from_numpy_dtype(data.dtype), + coord_data_type=BuiltinDtype.from_numpy_dtype(coords.dtype), ) client = self.new( StructureFamily.sparse, diff --git a/tiled/server/pydantic_sparse.py b/tiled/server/pydantic_sparse.py index 145f272d0..7d6113e2d 100644 --- a/tiled/server/pydantic_sparse.py +++ b/tiled/server/pydantic_sparse.py @@ -2,12 +2,17 @@ import pydantic +from ..structures.array import BuiltinDtype, Endianness, Kind, StructDtype from ..structures.sparse import SparseLayout class COOStructure(pydantic.BaseModel): shape: Tuple[int, ...] # tuple of ints like (3, 3) chunks: Tuple[Tuple[int, ...], ...] # tuple-of-tuples-of-ints like ((3,), (3,)) + data_type: Optional[Union[BuiltinDtype, StructDtype]] = None + coord_data_type: Optional[BuiltinDtype] = BuiltinDtype( + Endianness("little"), Kind("u"), 8 + ) # numpy 'uint' dtype dims: Optional[Tuple[str, ...]] = None # None or tuple of names like ("x", "y") resizable: Union[bool, Tuple[bool, ...]] = False layout: SparseLayout = SparseLayout.COO diff --git a/tiled/structures/sparse.py b/tiled/structures/sparse.py index 354d150d7..a84e67499 100644 --- a/tiled/structures/sparse.py +++ b/tiled/structures/sparse.py @@ -2,6 +2,8 @@ from dataclasses import dataclass from typing import Optional, Tuple, Union +from .array import BuiltinDtype, Endianness, Kind, StructDtype + class SparseLayout(str, enum.Enum): # Only COO is currently supported, but this lays a path @@ -13,6 +15,10 @@ class SparseLayout(str, enum.Enum): class COOStructure: chunks: Tuple[Tuple[int, ...], ...] # tuple-of-tuples-of-ints like ((3,), (3,)) shape: Tuple[int, ...] # tuple of ints like (3, 3) + data_type: Optional[Union[BuiltinDtype, StructDtype]] = None + coord_data_type: Optional[BuiltinDtype] = BuiltinDtype( + Endianness("little"), Kind("u"), 8 + ) # numpy 'uint' dtype dims: Optional[Tuple[str, ...]] = None # None or tuple of names like ("x", "y") resizable: Union[bool, Tuple[bool, ...]] = False layout: SparseLayout = SparseLayout.COO @@ -20,7 +26,17 @@ class COOStructure: @classmethod def from_json(cls, structure): + data_type = structure.get("data_type", None) + if data_type is not None and "fields" in data_type: + data_type = StructDtype.from_json(data_type) + else: + data_type = BuiltinDtype.from_json(data_type) + coord_data_type = structure.get( + "coord_data_type", {"endianness": "little", "kind": "u", "itemsize": 8} + ) return cls( + data_type=data_type, + coord_data_type=BuiltinDtype.from_json(coord_data_type), chunks=tuple(map(tuple, structure["chunks"])), shape=tuple(structure["shape"]), dims=structure["dims"],