binsparse 0.1.0__py3-none-any.whl

reference_cli/converters.py

@@ -0,0 +1,178 @@
+"""Command-line adapters for the Binsparse compliance tests."""
+
+from __future__ import annotations
+
+import json
+import sys
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+from binsparse.io import load_binsparse, save_binsparse
+from .reformat import binsparse_to_coo, coo_to_binsparse, reformat
+from binsparse.tensor import (
+    _PREDEFINED_LEVELS,
+    BinsparseLevel,
+    BinsparseTensor,
+    CustomTensor,
+    DenseLevel,
+    ElementLevel,
+    SparseLevel,
+)
+
+
+def npy_to_binsparse_main(argv: Sequence[str] | None = None) -> int:
+    """Implement ``npy_to_binsparse``."""
+    args = sys.argv[1:] if argv is None else list(argv)
+    if len(args) != 5:
+        print(
+            "usage: npy_to_binsparse <tensor_in> <pattern_in> <fill_value_in> "
+            "<header_in> <tensor_out>",
+            file=sys.stderr,
+        )
+        return 2
+    tensor_in, pattern_in, fill_value_in, header_in, tensor_out = args
+
+    dense = np.load(tensor_in, allow_pickle=False)
+    pattern = np.asarray(np.load(pattern_in, allow_pickle=False), dtype=bool)
+    fill_value = _load_fill_value(fill_value_in)
+    header = _load_header(header_in)
+
+    tensor = _tensor_from_npy(dense, pattern, fill_value, header)
+    save_binsparse(tensor, tensor_out, header=header)
+    return 0
+
+
+def binsparse_to_npy_main(argv: Sequence[str] | None = None) -> int:
+    """Implement ``binsparse_to_npy``."""
+    args = sys.argv[1:] if argv is None else list(argv)
+    if len(args) != 4:
+        print(
+            "usage: binsparse_to_npy <tensor_in> <tensor_out> "
+            "<pattern_out> <fill_value_out>",
+            file=sys.stderr,
+        )
+        return 2
+    tensor_in, tensor_out, pattern_out, fill_value_out = args
+
+    tensor = load_binsparse(tensor_in, alias=False)
+    assert isinstance(tensor, CustomTensor)
+    dense, pattern, fill_value = _npy_from_tensor(tensor)
+    np.save(tensor_out, dense)
+    np.save(pattern_out, pattern)
+    np.save(fill_value_out, np.asarray(fill_value))
+    return 0
+
+
+def binsparse_to_binsparse_main(argv: Sequence[str] | None = None) -> int:
+    """Implement ``binsparse_to_binsparse``."""
+    args = sys.argv[1:] if argv is None else list(argv)
+    if len(args) != 2:
+        print(
+            "usage: binsparse_to_binsparse <tensor_in> <tensor_out>",
+            file=sys.stderr,
+        )
+        return 2
+    tensor_in, tensor_out = args
+
+    tensor = load_binsparse(tensor_in)
+    save_binsparse(tensor, tensor_out)
+    return 0
+
+
+def _load_header(path: str) -> dict[str, Any]:
+    with Path(path).open() as file:
+        header = json.load(file)
+    if not isinstance(header, dict):
+        raise TypeError("header_in must contain a JSON object")
+    return header
+
+
+def _load_fill_value(path: str) -> Any:
+    fill_value = np.load(path, allow_pickle=False)
+    if fill_value.size != 1:
+        raise ValueError("fill_value_in must contain exactly one value")
+    return fill_value.reshape(-1)[0]
+
+
+def _tensor_from_npy(
+    dense: np.ndarray,
+    pattern: np.ndarray,
+    fill_value: Any,
+    header: dict[str, Any],
+) -> BinsparseTensor:
+    if dense.shape != pattern.shape:
+        raise ValueError("tensor_in and pattern_in must have the same shape")
+    format, transpose = _format_from_header(header)
+    target_transpose = transpose or tuple(range(dense.ndim))
+    stored_shape = tuple(dense.shape[dimension] for dimension in target_transpose)
+
+    entries = [
+        (
+            tuple(int(coord[dimension]) for dimension in target_transpose),
+            dense[tuple(coord)],
+        )
+        for coord in np.argwhere(pattern)
+    ]
+    custom = coo_to_binsparse(format, entries, shape=stored_shape)
+    custom.shape = tuple(dense.shape)
+    custom.transpose = None if transpose is None else tuple(transpose)
+    custom.fill = True
+    custom.fill_value = fill_value
+    return reformat(custom, header)
+
+
+def _format_from_header(
+    header: dict[str, Any],
+) -> tuple[dict[str, Any], tuple[int, ...] | None]:
+    format_name = header["format"]
+
+    if format_name == "custom":
+        custom = header.get("custom")
+        if not isinstance(custom, dict) or "level" not in custom:
+            raise ValueError("custom format requires custom.level")
+        transpose = custom.get("transpose")
+        return custom["level"], None if transpose is None else tuple(transpose)
+
+    try:
+        format, transpose = _PREDEFINED_LEVELS[format_name]
+    except KeyError as error:
+        raise ValueError(f"unknown Binsparse format {format_name!r}") from error
+    return format, transpose
+
+
+def _npy_from_tensor(tensor: CustomTensor) -> tuple[np.ndarray, np.ndarray, Any]:
+    fill_value = tensor.fill_value if tensor.fill is True else 0
+    values = _element_values(tensor.level)
+    dtype = (
+        np.result_type(values.dtype, np.asarray(fill_value).dtype)
+        if values.size > 0
+        else np.asarray(fill_value).dtype
+    )
+    dense = np.full(tensor.shape, fill_value, dtype=dtype)
+    pattern = np.zeros(tensor.shape, dtype=bool)
+    transpose = tensor.transpose or tuple(range(len(tensor.shape)))
+    for stored_coord, value in binsparse_to_coo(tensor):
+        coord = tuple(stored_coord[transpose.index(d)] for d in range(len(transpose)))
+        dense[coord] = value
+        pattern[coord] = True
+    return dense, pattern, fill_value
+
+
+def _element_values(level: BinsparseLevel | None) -> np.ndarray:
+    match level:
+        case ElementLevel(values):
+            return values
+        case DenseLevel(_, child) | SparseLevel(_, child, _, _):
+            return _element_values(child)
+        case _:
+            raise TypeError(f"unsupported level type {type(level).__name__}")
+
+
+__all__ = [
+    "binsparse_to_binsparse_main",
+    "binsparse_to_npy_main",
+    "npy_to_binsparse_main",
+]

reference_cli/framework.py

@@ -0,0 +1,81 @@
+"""Compliance-test executables which round-trip through optional frameworks."""
+
+from __future__ import annotations
+
+import sys
+from collections.abc import Callable, Sequence
+from typing import Any
+
+import numpy as np
+
+from .converters import (
+    _load_fill_value,
+    _load_header,
+    _npy_from_tensor,
+    _tensor_from_npy,
+)
+from binsparse.conversions import (
+    from_numpy,
+    from_scipy,
+    from_sparse,
+    to_numpy,
+    to_scipy,
+    to_sparse,
+)
+from binsparse.io import load_binsparse, save_binsparse
+from .reformat import alias_to_custom
+from binsparse.tensor import BinsparseTensor
+
+Converter = tuple[Callable[[BinsparseTensor], Any], Callable[[Any], BinsparseTensor]]
+
+
+def _converter(name: str) -> Converter:
+    converters: dict[str, Converter] = {
+        "numpy": (to_numpy, from_numpy),
+        "scipy": (to_scipy, from_scipy),
+        "sparse": (to_sparse, from_sparse),
+    }
+    try:
+        return converters[name]
+    except KeyError as error:
+        raise ValueError(f"unknown framework {name!r}") from error
+
+
+def _roundtrip(tensor: BinsparseTensor, framework: str) -> BinsparseTensor:
+    to_framework, from_framework = _converter(framework)
+    return from_framework(to_framework(tensor))
+
+
+def main(argv: Sequence[str] | None = None) -> int:
+    args = list(sys.argv[1:] if argv is None else argv)
+    if len(args) < 2:
+        return 2
+    framework, operation, *paths = args
+    if operation == "npy_to_binsparse" and len(paths) == 5:
+        tensor_in, pattern_in, fill_in, header_in, tensor_out = paths
+        header = _load_header(header_in)
+        tensor = _tensor_from_npy(
+            np.load(tensor_in, allow_pickle=False),
+            np.asarray(np.load(pattern_in, allow_pickle=False), dtype=bool),
+            _load_fill_value(fill_in),
+            header,
+        )
+        save_binsparse(_roundtrip(tensor, framework), tensor_out, header=header)
+        return 0
+    if operation == "binsparse_to_npy" and len(paths) == 4:
+        tensor_in, tensor_out, pattern_out, fill_out = paths
+        tensor = alias_to_custom(_roundtrip(load_binsparse(tensor_in), framework))
+        dense, pattern, fill = _npy_from_tensor(tensor)
+        np.save(tensor_out, dense)
+        np.save(pattern_out, pattern)
+        np.save(fill_out, np.asarray(fill))
+        return 0
+    if operation == "binsparse_to_binsparse" and len(paths) == 2:
+        tensor_in, tensor_out = paths
+        save_binsparse(_roundtrip(load_binsparse(tensor_in), framework), tensor_out)
+        return 0
+    return 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

reference_cli/reformat.py

@@ -0,0 +1,290 @@
+"""Reformat Binsparse tensors through coordinate lists."""
+
+from collections.abc import Iterable, Iterator
+from itertools import groupby, product
+from typing import Any, cast
+
+import numpy as np
+
+from binsparse.container import InMemoryBinsparseContainer, str_to_dtype
+from binsparse.tensor import (
+    _PREDEFINED_LEVELS,
+    BinsparseLevel,
+    BinsparseTensor,
+    CustomTensor,
+    DenseLevel,
+    ElementLevel,
+    SparseLevel,
+)
+
+_FORMAT_CANONICAL = {
+    "DMAT": "DMATR",
+    "COO": "COOR",
+}
+
+
+def _canonical_format(format_name: str) -> str:
+    return _FORMAT_CANONICAL.get(format_name, format_name)
+
+
+def alias_to_custom(tensor: BinsparseTensor) -> CustomTensor:
+    """Return the custom-level representation of any Binsparse tensor."""
+    container = InMemoryBinsparseContainer()
+    tensor.serialize(container, alias=False)
+    custom = CustomTensor.parse(container)
+    assert isinstance(custom, CustomTensor)
+    return custom
+
+
+def custom_to_alias(
+    tensor: BinsparseTensor, format_name: str | None = None
+) -> BinsparseTensor:
+    """Return the predefined alias for a custom tensor, when one exists."""
+    container = InMemoryBinsparseContainer()
+    tensor.serialize(container, alias=True)
+    alias = BinsparseTensor.parse(container)
+    if (
+        format_name is not None
+        and _canonical_format(alias.format) != _canonical_format(format_name)
+    ):
+        raise ValueError(f"custom layout is not compatible with {format_name!r}")
+    return alias
+
+
+def binsparse_to_coo(
+    tensor: CustomTensor,
+) -> Iterator[tuple[tuple[Any, ...], Any]]:
+    """Iterate over a tensor's explicitly stored coordinate/value pairs."""
+    shape = (
+        tensor.shape
+        if tensor.transpose is None
+        else tuple(tensor.shape[dimension] for dimension in tensor.transpose)
+    )
+    root = tensor.level
+    assert root is not None
+
+    def level_to_coo(
+        level: BinsparseLevel, coords: list[tuple[int, ...]], shape: tuple[int, ...]
+    ) -> Iterable[tuple[tuple[Any, ...], Any]]:
+        match level:
+            case ElementLevel(values):
+                return zip(coords, values, strict=True)
+            case DenseLevel(rank, child):
+                coords = [
+                    (*coord, *suffix)
+                    for coord in coords
+                    for suffix in product(*(range(size) for size in shape[:rank]))
+                ]
+                return level_to_coo(child, coords, shape[rank:])
+            case SparseLevel(rank, child, indices, pointers):
+                pointers = (
+                    np.array([0, indices[0].size])
+                    if pointers is None
+                    else pointers
+                )
+                coords = [
+                    (*coord, *(index[q] for index in indices))
+                    for p, coord in enumerate(coords)
+                    for q in range(pointers[p], pointers[p + 1])
+                ]
+                return level_to_coo(child, coords, shape[rank:])
+            case _:
+                raise TypeError(f"unsupported level type {type(level).__name__}")
+
+    yield from level_to_coo(root, [()], shape)
+
+
+def coo_to_binsparse(
+    format: dict[str, Any],
+    entries: Iterable[tuple[tuple[Any, ...], Any]],
+    *,
+    shape: tuple[int, ...],
+) -> CustomTensor:
+    """Build a custom tensor from a nested format descriptor and COO entries."""
+    entries = sorted((tuple(coord), value) for coord, value in entries)
+    for coord, _ in entries:
+        if len(coord) != len(shape):
+            raise ValueError("coordinate rank does not match tensor rank")
+        if any(i < 0 or i >= n for i, n in zip(coord, shape, strict=True)):
+            raise ValueError("coordinate is out of bounds")
+
+    def coo_to_level(
+        format: dict[str, Any],
+        entries: list[tuple[tuple[Any, ...], Any]],
+        ptr: np.ndarray,
+        shape: tuple[int, ...],
+        root: bool = False,
+    ) -> BinsparseLevel:
+        match format["level_desc"]:
+            case "element":
+                if np.any(np.diff(ptr) != 1):
+                    raise ValueError("each element position must have one value")
+                return ElementLevel(np.asarray([value for _, value in entries]))
+            case "dense":
+                rank = format["rank"]
+                groups = [
+                    [
+                        (coord[rank:], value)
+                        for coord, value in entries[ptr[p] : ptr[p + 1]]
+                        if coord[:rank] == suffix
+                    ]
+                    for p in range(ptr.size - 1)
+                    for suffix in product(*(range(size) for size in shape[:rank]))
+                ]
+                entries = [entry for group in groups for entry in group]
+                ptr = np.cumsum([0, *(len(group) for group in groups)])
+                return DenseLevel(
+                    rank,
+                    coo_to_level(format["level"], entries, ptr, shape[rank:]),
+                )
+            case "sparse":
+                rank = format["rank"]
+                grouped = [
+                    [
+                        (coord, list(group))
+                        for coord, group in groupby(
+                            entries[ptr[p] : ptr[p + 1]],
+                            lambda entry: entry[0][:rank],
+                        )
+                    ]
+                    for p in range(ptr.size - 1)
+                ]
+                groups = [
+                    [(coord[rank:], value) for coord, value in group]
+                    for fiber in grouped
+                    for _, group in fiber
+                ]
+                entries = [entry for group in groups for entry in group]
+                child_ptr = np.cumsum([0, *(len(group) for group in groups)])
+                new_indices = tuple(
+                    np.asarray([coord[d] for fiber in grouped for coord, _ in fiber])
+                    for d in range(rank)
+                )
+                new_pointers = np.cumsum(
+                    [0, *(len(fiber) for fiber in grouped)], dtype=np.uint64
+                )
+                return SparseLevel(
+                    rank,
+                    coo_to_level(
+                        format["level"], entries, child_ptr, shape[rank:]
+                    ),
+                    new_indices,
+                    None if root else new_pointers,
+                )
+            case _:
+                raise ValueError(f"unknown level descriptor {format['level_desc']!r}")
+
+    return CustomTensor(
+        shape,
+        len(entries),
+        level=coo_to_level(
+            format,
+            entries,
+            np.array([0, len(entries)], dtype=np.uint64),
+            shape,
+            root=True,
+        ),
+    )
+
+def reformat(tensor: BinsparseTensor, header: dict[str, Any]) -> BinsparseTensor:
+    """Reformat *tensor* according to a Binsparse format descriptor."""
+    if isinstance(tensor, CustomTensor):
+        source = tensor
+    else:
+        container = InMemoryBinsparseContainer()
+        tensor.serialize(container, alias=False)
+        parsed = CustomTensor.parse(container)
+        assert isinstance(parsed, CustomTensor)
+        source = parsed
+
+    format_name = header["format"]
+    if format_name == "custom":
+        custom = header["custom"]
+        format = custom["level"]
+        transpose = custom.get("transpose")
+    else:
+        try:
+            format, transpose = _PREDEFINED_LEVELS[format_name]
+        except KeyError as error:
+            raise ValueError(f"unknown Binsparse format {format_name!r}") from error
+
+    source_transpose = source.transpose or tuple(range(len(source.shape)))
+    target_transpose = transpose or tuple(range(len(source.shape)))
+    entries = [
+        (
+            tuple(
+                coord[source_transpose.index(d)]
+                for d in target_transpose
+            ),
+            value,
+        )
+        for coord, value in binsparse_to_coo(source)
+    ]
+    stored_shape = tuple(source.shape[d] for d in target_transpose)
+    result = coo_to_binsparse(format, entries, shape=stored_shape)
+    result.shape = source.shape
+    result.transpose = None if transpose is None else tuple(transpose)
+    result.fill = source.fill
+    result.fill_value = source.fill_value
+
+    container = InMemoryBinsparseContainer()
+    result.serialize(container, alias=format_name != "custom")
+    _recast_buffers(container, header.get("data_types"))
+    alias = BinsparseTensor.parse(container)
+    if _canonical_format(alias.format) != _canonical_format(format_name):
+        raise ValueError(f"custom layout is not compatible with {format_name!r}")
+    return alias
+
+
+def _recast_buffers(container: InMemoryBinsparseContainer, data_types: Any) -> None:
+    if not isinstance(data_types, dict):
+        return
+    for key, data_type in data_types.items():
+        if key not in container.data_types:
+            continue
+        value = container.read_buffer(key, _expected_size(container, key))
+        container.write_buffer(key, _cast_buffer(value, data_type))
+    container.write_header(container.read_header())
+
+
+def _expected_size(
+    container: InMemoryBinsparseContainer,
+    key: str,
+) -> int | None:
+    if key == "fill_value":
+        return 1
+    if key == "values":
+        return int(container.read_header()["number_of_stored_values"])
+    return None
+
+
+def _cast_buffer(value: np.ndarray, data_type: Any) -> np.ndarray:
+    if not isinstance(data_type, str):
+        raise TypeError("Binsparse data type must be a string")
+    if data_type.startswith("iso[") and data_type.endswith("]"):
+        result = _cast_buffer(value, data_type[4:-1])
+        if result.size == 0:
+            return result
+        scalar = result.reshape(-1)[0]
+        if not np.all(result == scalar):
+            raise ValueError("ISO data type requires identical stored values")
+        return np.broadcast_to(np.asarray([scalar], dtype=result.dtype), result.shape)
+    if data_type.startswith("complex[") and data_type.endswith("]"):
+        inner = data_type[8:-1]
+        if inner == "float32":
+            return np.asarray(value, dtype=np.dtype("complex64"))
+        if inner == "float64":
+            return np.asarray(value, dtype=np.dtype("complex128"))
+    try:
+        return np.asarray(value, dtype=cast(np.dtype[Any], str_to_dtype[data_type]))
+    except KeyError as error:
+        raise ValueError(f"unknown Binsparse data type {data_type!r}") from error
+
+
+__all__ = [
+    "alias_to_custom",
+    "binsparse_to_coo",
+    "coo_to_binsparse",
+    "custom_to_alias",
+    "reformat",
+]