binsparse 0.1.0__py3-none-any.whl

binsparse/__init__.py

@@ -0,0 +1,66 @@
+"""Python reference implementation of the Binsparse specification."""
+
+from .container import (
+    BINSPARSE_HEADER,
+    BinsparseContainer,
+    HDF5BinsparseContainer,
+    InMemoryBinsparseContainer,
+    NPZBinsparseContainer,
+    ZarrBinsparseContainer,
+)
+from .errors import BinsparseParseError
+from .io import load_binsparse, save_binsparse
+from .tensor import (
+    BinsparseLevel,
+    BinsparseTensor,
+    COOCMatrix,
+    COOMatrix,
+    COORMatrix,
+    CSCMatrix,
+    CSRMatrix,
+    CustomTensor,
+    DCSCMatrix,
+    DCSRMatrix,
+    DMATCMatrix,
+    DMATMatrix,
+    DMATRMatrix,
+    DVECVector,
+    CVECVector,
+    DenseLevel,
+    ElementLevel,
+    IndexableLevel,
+    SparseLevel,
+)
+from .version import BINSPARSE_VERSION
+
+__all__ = [
+    "BINSPARSE_HEADER",
+    "BINSPARSE_VERSION",
+    "BinsparseContainer",
+    "BinsparseLevel",
+    "BinsparseParseError",
+    "BinsparseTensor",
+    "COOCMatrix",
+    "COOMatrix",
+    "COORMatrix",
+    "CSCMatrix",
+    "CSRMatrix",
+    "CVECVector",
+    "CustomTensor",
+    "DCSCMatrix",
+    "DCSRMatrix",
+    "DMATCMatrix",
+    "DMATMatrix",
+    "DMATRMatrix",
+    "DVECVector",
+    "DenseLevel",
+    "ElementLevel",
+    "HDF5BinsparseContainer",
+    "InMemoryBinsparseContainer",
+    "IndexableLevel",
+    "load_binsparse",
+    "NPZBinsparseContainer",
+    "SparseLevel",
+    "save_binsparse",
+    "ZarrBinsparseContainer",
+]

binsparse/container.py

@@ -0,0 +1,328 @@
+"""Container adapters for reading and writing Binsparse data."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import MutableMapping
+import json
+import re
+from typing import Any
+
+import numpy as np
+
+from .errors import BinsparseParseError
+
+
+BINSPARSE_HEADER = "binsparse"
+
+
+def _wrap_header(value: dict[str, Any]) -> dict[str, dict[str, Any]]:
+    return {BINSPARSE_HEADER: value}
+
+
+def _unwrap_header(value: Any) -> dict[str, Any]:
+    if not isinstance(value, dict):
+        raise BinsparseParseError("Binsparse metadata must be a JSON object")
+    try:
+        header = value[BINSPARSE_HEADER]
+    except KeyError as error:
+        raise BinsparseParseError("missing 'binsparse' JSON namespace") from error
+    if not isinstance(header, dict):
+        raise BinsparseParseError("the 'binsparse' JSON namespace must be an object")
+    return header
+
+
+dtype_to_str = {
+    np.dtype("int8"): "int8",
+    np.dtype("int16"): "int16",
+    np.dtype("int32"): "int32",
+    np.dtype("int64"): "int64",
+    np.dtype("uint8"): "uint8",
+    np.dtype("uint16"): "uint16",
+    np.dtype("uint32"): "uint32",
+    np.dtype("uint64"): "uint64",
+    np.dtype("float32"): "float32",
+    np.dtype("float64"): "float64",
+    np.dtype("bool"): "bint8",
+}
+
+str_to_dtype = {value: key for key, value in dtype_to_str.items()}
+
+
+class BinsparseContainer(ABC):
+    """Common interface to a Binsparse binary container or container group."""
+
+    def __init__(self) -> None:
+        self.data_types: dict[str, str] = {}
+
+    def read_header(self) -> dict[str, Any]:
+        """Return the decoded Binsparse JSON descriptor."""
+        header = self._read_header()
+        if not isinstance(header.get("data_types"), dict):
+            raise BinsparseParseError("data_types must be an object")
+        return header
+
+    @abstractmethod
+    def _read_header(self) -> dict[str, Any]:
+        """Read and decode the Binsparse header from the backend."""
+
+    def write_header(self, value: dict[str, Any]) -> None:
+        """Store a Binsparse JSON descriptor."""
+        if not isinstance(value, dict):
+            raise TypeError("the binsparse header must be a dictionary")
+        header = {**value, "data_types": self.data_types.copy()}
+        json.dumps(header)
+        self._write_header(header)
+
+    @abstractmethod
+    def _write_header(self, value: dict[str, Any]) -> None:
+        """Write a validated Binsparse header to the backend."""
+
+    def read_buffer(
+        self,
+        key: str,
+        expected_size: int | None = None,
+        *,
+        copy: bool | None = None,
+    ) -> np.ndarray:
+        """Read and decode a named binary array."""
+        header = self.read_header()
+        try:
+            declared = header["data_types"][key]
+        except KeyError as error:
+            raise BinsparseParseError(f"missing data type for buffer {key!r}") from error
+        if not isinstance(declared, str):
+            raise BinsparseParseError(f"data type for buffer {key!r} must be a string")
+        def decode(data_type: str, data: np.ndarray) -> np.ndarray:
+            if (match := re.fullmatch(r"iso\[(.*)\]", data_type)) is not None:
+                if expected_size is None:
+                    raise BinsparseParseError(
+                        "expected_size is required when reading an ISO buffer"
+                    )
+                decoded = decode(match.group(1), data)
+                if decoded.size != 1:
+                    raise BinsparseParseError("an ISO buffer must contain one value")
+                return np.broadcast_to(decoded.reshape(1), (expected_size,))
+            if (match := re.fullmatch(r"complex\[(.*)\]", data_type)) is not None:
+                decoded = decode(match.group(1), data)
+                if decoded.dtype not in {np.dtype("float32"), np.dtype("float64")}:
+                    raise BinsparseParseError("complex values require float32 or float64")
+                if decoded.ndim != 1 or decoded.size % 2 != 0:
+                    raise BinsparseParseError("invalid complex value buffer")
+                complex_dtype = (
+                    np.complex64 if decoded.dtype == np.float32 else np.complex128
+                )
+                return np.ascontiguousarray(decoded).view(complex_dtype)
+            if re.fullmatch(r"[^\[\]]+", data_type) is None:
+                raise BinsparseParseError(f"unknown Binsparse type wrapper {data_type!r}")
+            try:
+                dtype = str_to_dtype[data_type]
+            except KeyError as error:
+                raise BinsparseParseError(f"unknown Binsparse type {data_type!r}") from error
+            return np.asarray(data, dtype=dtype)  # type: ignore[call-overload]
+        encoded = self._read_buffer(key)
+        data = decode(declared, encoded)
+        if expected_size is not None and data.size != expected_size:
+            raise BinsparseParseError(
+                f"buffer {key!r} has size {data.size}, expected {expected_size}"
+            )
+        shares_memory = np.shares_memory(data, encoded)
+        if copy is False and data.size > 0 and not shares_memory:
+            raise ValueError(f"copy=False cannot decode buffer {key!r} without copying")
+        if copy is True and shares_memory:
+            return data.copy()
+        return data
+
+    @abstractmethod
+    def _read_buffer(self, key: str) -> np.ndarray:
+        """Read a named binary array without applying Binsparse decoding."""
+
+    def write_buffer(
+        self,
+        key: str,
+        value: np.ndarray,
+        *,
+        copy: bool | None = None,
+    ) -> None:
+        """Encode, create or replace a named array and record its data type."""
+        def encode(
+            data: np.ndarray,
+            copy: bool | None,
+            *,
+            detect_iso: bool = True,
+        ) -> tuple[np.ndarray, str]:
+            if (
+                detect_iso
+                and data.size > 0
+                and data.ndim == 1
+                and data.strides == (0,)
+            ):
+                encoded, data_type = encode(data[:1], copy, detect_iso=False)
+                return encoded, f"iso[{data_type}]"
+            if data.dtype == np.bool_:
+                encoded = data.view(np.uint8)
+                return encoded.copy() if copy is True else encoded, "bint8"
+            if np.issubdtype(data.dtype, np.complexfloating):
+                if data.dtype not in {np.dtype("complex64"), np.dtype("complex128")}:
+                    raise TypeError(f"unsupported complex dtype: {data.dtype}")
+                if not data.flags.c_contiguous:
+                    if copy is False:
+                        raise ValueError(
+                            f"copy=False cannot encode buffer {key!r} without copying"
+                        )
+                    data = np.ascontiguousarray(data)
+                dtype = np.dtype("float32" if data.dtype == np.complex64 else "float64")
+                encoded, data_type = encode(data.view(dtype), copy, detect_iso=False)
+                return encoded, f"complex[{data_type}]"
+            try:
+                data_type = dtype_to_str[data.dtype]
+            except KeyError as error:
+                raise TypeError(f"unsupported Binsparse dtype: {data.dtype}") from error
+            return data.copy() if copy is True else data, data_type
+
+        if not isinstance(value, np.ndarray):
+            raise TypeError("buffer value must be a NumPy ndarray")
+        encoded, data_type = encode(value, copy)
+        self._write_buffer(key, encoded)
+        self.data_types[key] = data_type
+
+    @abstractmethod
+    def _write_buffer(self, key: str, value: np.ndarray) -> None:
+        """Write an already encoded named binary array."""
+
+
+class InMemoryBinsparseContainer(BinsparseContainer):
+    """Store a Binsparse descriptor and its arrays directly in memory."""
+
+    def __init__(
+        self,
+        header: dict[str, Any] | None = None,
+        buffers: list[np.ndarray] | None = None,
+    ) -> None:
+        super().__init__()
+        self.header = {} if header is None else header
+        self.buffers = [] if buffers is None else buffers
+        data_types = self.header.get("data_types")
+        if isinstance(data_types, dict):
+            self.data_types.update(data_types)
+
+    def _read_header(self) -> dict[str, Any]:
+        return self.header
+
+    def _write_header(self, value: dict[str, Any]) -> None:
+        self.header.clear()
+        self.header.update(value)
+
+    def _read_buffer(self, key: str) -> np.ndarray:
+        try:
+            index = list(self.header["data_types"]).index(key)
+        except (KeyError, ValueError) as error:
+            raise KeyError(f"in-memory container has no buffer {key!r}") from error
+        try:
+            return self.buffers[index]
+        except IndexError as error:
+            raise KeyError(f"in-memory container has no buffer {key!r}") from error
+
+    def _write_buffer(self, key: str, value: np.ndarray) -> None:
+        if key in self.data_types:
+            self.buffers[list(self.data_types).index(key)] = value
+        else:
+            self.buffers.append(value)
+
+
+class HDF5BinsparseContainer(BinsparseContainer):
+    """Adapt an h5py ``Container`` or ``Group``."""
+
+    def __init__(self, group: Any):
+        super().__init__()
+        self.group = group
+
+    def _read_header(self) -> dict[str, Any]:
+        try:
+            value = self.group.attrs[BINSPARSE_HEADER]
+        except KeyError as error:
+            raise KeyError("HDF5 group has no 'binsparse' attribute") from error
+        return _unwrap_header(json.loads(value))
+
+    def _write_header(self, value: dict[str, Any]) -> None:
+        self.group.attrs[BINSPARSE_HEADER] = json.dumps(
+            _wrap_header(value), indent=2, sort_keys=True, separators=(",", ": ")
+        )
+
+    def _read_buffer(self, key: str) -> np.ndarray:
+        return np.asarray(self.group[key][()])
+
+    def _write_buffer(self, key: str, value: np.ndarray) -> None:
+        if key in self.group:
+            del self.group[key]
+        self.group.create_dataset(key, data=value)
+
+class ZarrBinsparseContainer(BinsparseContainer):
+    """Adapt a Zarr group without requiring Zarr as a dependency."""
+
+    def __init__(self, group: Any):
+        super().__init__()
+        self.group = group
+
+    def _read_header(self) -> dict[str, Any]:
+        try:
+            return _unwrap_header(self.group.attrs[BINSPARSE_HEADER])
+        except KeyError as error:
+            raise KeyError("Zarr group has no 'binsparse' attribute") from error
+
+    def _write_header(self, value: dict[str, Any]) -> None:
+        self.group.attrs[BINSPARSE_HEADER] = _wrap_header(value)
+
+    def _read_buffer(self, key: str) -> np.ndarray:
+        return np.asarray(self.group[key][...])
+
+    def _write_buffer(self, key: str, value: np.ndarray) -> None:
+        if key in self.group:
+            del self.group[key]
+        if hasattr(self.group, "create_array"):
+            self.group.create_array(key, data=value)
+        else:
+            self.group.create_dataset(key, data=value)
+
+class NPZBinsparseContainer(BinsparseContainer):
+    """Adapt a mutable mapping of NPZ entry names to NumPy arrays."""
+
+    def __init__(self, file: MutableMapping[str, np.ndarray]):
+        super().__init__()
+        self.file = file
+
+    def _read_header(self) -> dict[str, Any]:
+        try:
+            value = self.file[BINSPARSE_HEADER]
+        except KeyError as error:
+            raise KeyError("NPZ archive has no 'binsparse' entry") from error
+        return _unwrap_header(json.loads(str(value.item())))
+
+    def _write_header(self, value: dict[str, Any]) -> None:
+        self.file[BINSPARSE_HEADER] = np.asarray(
+            json.dumps(
+                _wrap_header(value),
+                indent=2,
+                sort_keys=True,
+                separators=(",", ": "),
+            )
+        )
+
+    def _read_buffer(self, key: str) -> np.ndarray:
+        if key == BINSPARSE_HEADER:
+            raise KeyError("use read_header() to access binsparse metadata")
+        return self.file[key]
+
+    def _write_buffer(self, key: str, value: np.ndarray) -> None:
+        if key == BINSPARSE_HEADER:
+            raise KeyError("use write_header() to set binsparse metadata")
+        self.file[key] = value
+
+__all__ = [
+    "BINSPARSE_HEADER",
+    "BinsparseContainer",
+    "HDF5BinsparseContainer",
+    "InMemoryBinsparseContainer",
+    "ZarrBinsparseContainer",
+    "NPZBinsparseContainer",
+]

binsparse/conversions/__init__.py

@@ -0,0 +1,17 @@
+"""Optional adapters between Binsparse tensors and third-party array libraries."""
+
+from .numpy import from_numpy, to_numpy
+from .scipy import from_scipy, to_scipy
+from .sparse import from_sparse, to_sparse
+from .torch import from_torch, to_torch
+
+__all__ = [
+    "from_numpy",
+    "from_scipy",
+    "from_sparse",
+    "from_torch",
+    "to_numpy",
+    "to_scipy",
+    "to_sparse",
+    "to_torch",
+]

binsparse/conversions/numpy.py

@@ -0,0 +1,64 @@
+"""Conversions between Binsparse tensors and NumPy arrays."""
+
+from typing import Literal
+
+import numpy as np
+
+from binsparse.tensor import (
+    BinsparseTensor,
+    CustomTensor,
+    DenseLevel,
+    DMATCMatrix,
+    DMATRMatrix,
+    DVECVector,
+    ElementLevel,
+)
+
+
+def from_numpy(value: np.ndarray, *, copy: bool | None = None) -> BinsparseTensor:
+    """Convert a NumPy array to a dense Binsparse tensor."""
+    if not isinstance(value, np.ndarray):
+        raise TypeError("expected a NumPy ndarray")
+    array = np.array(value, copy=copy, order="C")
+    values = array.reshape(-1)
+    common = (tuple(array.shape), int(array.size))
+    if array.ndim == 1:
+        return DVECVector(*common, values=values)
+    if array.ndim == 2:
+        return DMATRMatrix(*common, values=values)
+    level = ElementLevel(values)
+    return CustomTensor(
+        *common,
+        level=level if array.ndim == 0 else DenseLevel(array.ndim, level),
+    )
+
+
+def to_numpy(tensor: BinsparseTensor, *, copy: bool | None = None) -> np.ndarray:
+    """Convert a dense Binsparse tensor to a NumPy array."""
+    order: Literal["C", "F"] = "C"
+    if isinstance(tensor, (DVECVector, DMATRMatrix)):
+        values = tensor.values
+    elif isinstance(tensor, DMATCMatrix):
+        values = tensor.values
+        order = "F"
+    elif (
+        isinstance(tensor, CustomTensor)
+        and tensor.transpose is None
+        and isinstance(tensor.level, ElementLevel)
+        and not tensor.shape
+    ):
+        values = tensor.level.values
+    elif (
+        isinstance(tensor, CustomTensor)
+        and tensor.transpose is None
+        and isinstance(tensor.level, DenseLevel)
+        and tensor.level.rank == len(tensor.shape)
+        and isinstance(tensor.level.level, ElementLevel)
+    ):
+        values = tensor.level.level.values
+    else:
+        raise TypeError(f"cannot convert {type(tensor).__name__} to NumPy")
+    return np.array(values, copy=copy).reshape(tensor.shape, order=order)
+
+
+__all__ = ["from_numpy", "to_numpy"]

binsparse/conversions/scipy.py

@@ -0,0 +1,98 @@
+"""Conversions between Binsparse tensors and SciPy sparse arrays."""
+
+from typing import Any
+
+from binsparse.tensor import BinsparseTensor, COORMatrix, CSCMatrix, CSRMatrix
+
+
+def _scipy_sparse() -> Any:
+    try:
+        import scipy.sparse as scipy_sparse
+    except ImportError as error:
+        raise ImportError(
+            "SciPy conversions require the 'scipy' extra: pip install binsparse[scipy]"
+        ) from error
+    return scipy_sparse
+
+
+def _prepare(value: Any, copy: bool | None) -> Any:
+    if value.has_canonical_format:
+        return value.copy() if copy is True else value
+    if copy is False:
+        raise ValueError("copy=False cannot canonicalize a SciPy sparse array")
+    result = value.copy()
+    result.sum_duplicates()
+    return result
+
+
+def from_scipy(value: Any, *, copy: bool | None = None) -> BinsparseTensor:
+    """Convert a two-dimensional SciPy CSR, CSC, or COO object to Binsparse."""
+    scipy_sparse = _scipy_sparse()
+    if not scipy_sparse.issparse(value) or value.ndim != 2:
+        raise TypeError("expected a two-dimensional SciPy sparse array or matrix")
+    if value.format not in {"coo", "csr", "csc"}:
+        raise TypeError(f"unsupported SciPy sparse format {value.format!r}")
+    value = _prepare(value, copy)
+
+    shape = tuple(value.shape)
+    count = int(value.data.size)
+    values = value.data
+    if value.format == "csr":
+        return CSRMatrix(
+            shape,
+            count,
+            pointers_to_1=value.indptr,
+            indices_1=value.indices,
+            values=values,
+        )
+    if value.format == "csc":
+        return CSCMatrix(
+            shape,
+            count,
+            pointers_to_1=value.indptr,
+            indices_1=value.indices,
+            values=values,
+        )
+    if value.format == "coo":
+        return COORMatrix(
+            shape,
+            count,
+            indices_0=value.row,
+            indices_1=value.col,
+            values=values,
+        )
+    raise TypeError(f"unsupported SciPy sparse format {value.format!r}")
+
+
+def to_scipy(tensor: BinsparseTensor, *, copy: bool | None = None) -> Any:
+    """Convert a Binsparse CSR, CSC, or COO matrix to a SciPy sparse array."""
+    scipy_sparse = _scipy_sparse()
+    if tensor.fill is True and tensor.fill_value != 0:
+        raise ValueError("SciPy conversion requires a zero fill value")
+    if isinstance(tensor, CSCMatrix):
+        result = scipy_sparse.csc_array(
+            (tensor.values, tensor.indices_1, tensor.pointers_to_1),
+            shape=tensor.shape,
+            copy=copy,
+        )
+    elif isinstance(tensor, CSRMatrix):
+        result = scipy_sparse.csr_array(
+            (tensor.values, tensor.indices_1, tensor.pointers_to_1),
+            shape=tensor.shape,
+            copy=copy,
+        )
+    elif isinstance(tensor, COORMatrix):
+        result = scipy_sparse.coo_array(
+            (tensor.values, (tensor.indices_0, tensor.indices_1)),
+            shape=tensor.shape,
+            copy=copy,
+        )
+        # Binsparse COOR requires row-major sorted, duplicate-free coordinates.
+        # SciPy's constructor cannot infer that invariant from buffer inputs.
+        result.has_canonical_format = True
+    else:
+        raise TypeError(f"cannot convert {type(tensor).__name__} to SciPy")
+    return result
+
+
+__all__ = ["from_scipy", "to_scipy"]

binsparse/conversions/sparse.py

@@ -0,0 +1,92 @@
+"""Conversions between Binsparse tensors and PyData/Sparse arrays."""
+
+from typing import Any
+
+import numpy as np
+
+from binsparse.tensor import (
+    BinsparseTensor,
+    COORMatrix,
+    CustomTensor,
+    ElementLevel,
+    SparseLevel,
+)
+
+
+def _sparse() -> Any:
+    try:
+        import sparse
+    except ImportError as error:
+        raise ImportError(
+            "PyData/Sparse conversions require the 'sparse' extra: "
+            "pip install binsparse[sparse]"
+        ) from error
+    return sparse
+
+
+def from_sparse(value: Any, *, copy: bool | None = None) -> BinsparseTensor:
+    """Convert an N-dimensional PyData/Sparse COO array to Binsparse."""
+    sparse = _sparse()
+    if not isinstance(value, sparse.COO) or value.ndim < 1:
+        raise TypeError("expected a non-scalar PyData/Sparse COO array")
+    if copy is True:
+        value = value.copy(deep=True)
+    fill_value = np.asarray(value.fill_value).item()
+    shape = tuple(value.shape)
+    count = int(value.data.size)
+    indices = tuple(value.coords[dimension, :] for dimension in range(value.ndim))
+    values = value.data
+    if value.ndim == 2:
+        return COORMatrix(
+            shape,
+            count,
+            fill=True,
+            fill_value=fill_value,
+            indices_0=indices[0],
+            indices_1=indices[1],
+            values=values,
+        )
+    return CustomTensor(
+        shape,
+        count,
+        fill=True,
+        fill_value=fill_value,
+        level=SparseLevel(value.ndim, ElementLevel(values), indices),
+    )
+
+
+def to_sparse(tensor: BinsparseTensor, *, copy: bool | None = None) -> Any:
+    """Convert a flat Binsparse COO tensor to a PyData/Sparse COO array."""
+    sparse = _sparse()
+    indices: tuple[np.ndarray, ...]
+    if isinstance(tensor, COORMatrix):
+        indices = (tensor.indices_0, tensor.indices_1)
+        values = tensor.values
+    elif (
+        isinstance(tensor, CustomTensor)
+        and tensor.transpose is None
+        and isinstance(tensor.level, SparseLevel)
+        and tensor.level.rank == len(tensor.shape)
+        and tensor.level.pointers_to_next is None
+        and isinstance(tensor.level.level, ElementLevel)
+    ):
+        indices = tensor.level.indices
+        values = tensor.level.level.values
+    else:
+        raise TypeError(f"cannot convert {type(tensor).__name__} to PyData/Sparse")
+    if copy is False:
+        raise ValueError(
+            "copy=False requires Binsparse to support storing COO indices in a "
+            "single coordinate matrix"
+        )
+    coords = np.stack(indices)
+    fill_value = tensor.fill_value if tensor.fill is True else 0
+    return sparse.COO(
+        coords,
+        np.array(values, copy=True) if copy is True else values,
+        shape=tensor.shape,
+        fill_value=fill_value,
+    )
+
+
+__all__ = ["from_sparse", "to_sparse"]