python-graphblas/graphblas/core/expr.py at main · python-graphblas/python-graphblas

560 lines (468 loc) · 20.7 KB
import numpy as np
from .. import backend
from ..dtypes import _INDEX
from . import lib, utils
from .descriptor import lookup as descriptor_lookup
from .utils import _CArray, output_type
class _AllIndices:
    __slots__ = "_carg", "name", "_expr_name"
    def __init__(self):
        self._carg = lib.GrB_ALL
        self.name = "GrB_ALL"
        self._expr_name = ":"
    def __reduce__(self):
        return "_ALL_INDICES"
_ALL_INDICES = _AllIndices()
class AxisIndex:
    __slots__ = "size", "index", "cscalar", "dimsize"
    def __init__(self, size, index, cscalar, dimsize):
        self.size = size
        self.index = index
        self.cscalar = cscalar
        self.dimsize = dimsize
    @property
    def _carg(self):
        return self.index._carg
    @property
    def name(self):
        return self.index.name
    @property
    def _expr_name(self):
        if self.size is None:
            return f"{self.index.value}"
        idx = self._py_index()
        if type(idx) is slice:
            rv = f"{'' if idx.start is None else idx.start}:{'' if idx.stop is None else idx.stop}"
            if idx.step is not None:
                return f"{rv}:{idx.step}"
            return rv
        if idx.size < 6:
            return f"[{', '.join(map(str, idx))}]"
        return f"[{', '.join(map(str, idx[:3]))}, ...]"
    def _py_index(self):
        """Convert resolved index back into a valid Python index."""
        if self.size is None:
            return self.index.value
        if self.index is _ALL_INDICES:
            return slice(None)
        if backend != "suitesparse":
            return self.index.array
        from .slice import _gxb_backwards, _gxb_range, _gxb_stride
        if self.cscalar is _gxb_backwards:
            start, stop, step = self.index.array.tolist()
            size = self.dimsize
            stop -= size + 1
            step = -step
        elif self.cscalar is _gxb_range:
            start, stop = self.index.array.tolist()
            step = None
            stop += 1
        elif self.cscalar is _gxb_stride:
            start, stop, step = self.index.array.tolist()
            stop += 1
        else:
            return self.index.array
        if (
            start == 0
            and (step is None or step > 0)
            or start == self.dimsize - 1
            and step is not None
            and step < 0
            start = None
        if (
            stop == self.dimsize
            and (step is None or step > 0)
            or stop == -self.dimsize - 1
            and step is not None
            and step < 0
            stop = None
        return slice(start, stop, step)
class IndexerResolver:
    __slots__ = "obj", "indices", "shape", "__weakref__"
    def __init__(self, obj, indices):
        self.obj = obj
        if indices is Ellipsis:
            from .scalar import _as_scalar
            from .vector import Vector
            if type(obj) is Vector:
                self.indices = [
                    AxisIndex(
                        obj._size,
                        _ALL_INDICES,
                        _as_scalar(obj._size, _INDEX, is_cscalar=True),
                        obj._size,
                self.shape = (obj._size,)
            else:
                self.indices = [
                    AxisIndex(
                        obj._nrows,
                        _ALL_INDICES,
                        _as_scalar(obj._nrows, _INDEX, is_cscalar=True),
                        obj._nrows,
                    AxisIndex(
                        obj._ncols,
                        _ALL_INDICES,
                        _as_scalar(obj._ncols, _INDEX, is_cscalar=True),
                        obj._ncols,
                self.shape = (obj._nrows, obj._ncols)
        else:
            self.indices = self.parse_indices(indices, obj.shape)
            self.shape = tuple(index.size for index in self.indices if index.size is not None)
    @property
    def is_single_element(self):
        return not self.shape
    @property
    def py_indices(self):
        if self.obj.ndim > 1:
            return tuple(index._py_index() for index in self.indices)
        return self.indices[0]._py_index()
    def parse_indices(self, indices, shape):
        """Returns
        -------
            [(rows, rowsize), (cols, colsize)] for Matrix
            [(idx, idx_size)] for Vector
        Within each tuple, if the index is of type int, the size will be None
        """
        if len(shape) == 1:
            if type(indices) is tuple:
                raise TypeError(f"Index for {type(self.obj).__name__} cannot be a tuple")
            # Convert to tuple for consistent processing
            indices = (indices,)
        elif type(indices) is not tuple or len(indices) != 2:
            raise TypeError(f"Index for {type(self.obj).__name__} must be a 2-tuple")
        out = []
        for i, idx in enumerate(indices):
            typ = output_type(idx)
            if typ is tuple:
                raise TypeError(
                    f"Index in position {i} cannot be a tuple; must use slice or list or int"
            out.append(self.parse_index(idx, typ, shape[i]))
        return out
    def parse_index(self, index, typ, size):
        from .scalar import _as_scalar
        if np.issubdtype(typ, np.integer):
            if index >= size:
                raise IndexError(f"Index out of range: index={index}, size={size}")
            if np.issubdtype(typ, np.signedinteger) and index < 0:
                index = index + size
                if index < 0:
                    raise IndexError(f"Index out of range: index={index - size}, size={size}")
            return AxisIndex(None, _as_scalar(int(index), _INDEX, is_cscalar=True), None, size)
        if typ is list:
            pass
        elif typ is slice:
            from .slice import slice_to_index
            return slice_to_index(index, size)
        elif typ is np.ndarray:
            if len(index.shape) != 1:
                raise TypeError(f"Invalid number of dimensions for index: {len(index.shape)}")
            if not np.issubdtype(index.dtype, np.integer):
                raise TypeError(f"Invalid dtype for index: {index.dtype}")
            if np.issubdtype(index.dtype, np.signedinteger):
                is_negative = index < 0
                if is_negative.any():
                    index = np.where(is_negative, index + size, index)
                    if (index < 0).any():
                        bad_index = index[index < 0][0] - size
                        raise IndexError(f"Index out of range: index={bad_index}, size={size}")
            return AxisIndex(
                len(index), _CArray(index), _as_scalar(len(index), _INDEX, is_cscalar=True), size
        else:
            from .scalar import Scalar
            if typ is Scalar:
                if not np.issubdtype(index.dtype.np_type, np.integer):
                    raise TypeError(f"An integer is required for indexing.  Got: {index.dtype}")
                value = int(index)
                if np.issubdtype(index.dtype.np_type, np.signedinteger) and value < 0:
                    value = value + size
                    if value < 0:
                        raise IndexError(f"Index out of range: index={value - size}, size={size}")
                return AxisIndex(None, _as_scalar(value, _INDEX, is_cscalar=True), None, size)
            from .matrix import Matrix, TransposedMatrix
            from .vector import Vector
            if typ is Vector or typ is Matrix:
                raise TypeError(
                    f"Invalid type for index: {typ.__name__}.\n"
                    "If you want to apply a mask, perhaps do something like "
                    f"`x.dup(mask={index.name}.S)`.\n"
                    "If you want to assign with a mask, perhaps do something like "
                    f"`x(mask={index.name}.S) << value`."
            if typ is TransposedMatrix:
                raise TypeError(f"Invalid type for index: {typ.__name__}.")
            try:
                index = list(index)
            except Exception:
                from .mask import Mask
                if isinstance(index, Mask):
                    raise TypeError(
                        f"Invalid type for index: {typ.__name__}.\n"
                        "If you want to apply a mask, perhaps do something like "
                        f"`x.dup(mask={index.name})`.\n"
                        "If you want to assign with a mask, perhaps do something like "
                        f"`x(mask={index.name}) << value`."
                    ) from None
                raise TypeError(
                    f"Invalid type for index: {typ}; unable to convert to list"
                ) from None
        return self.parse_index(np.array(index), np.ndarray, size)
    def get_index(self, dim):
        """Return a new IndexerResolver with index for the selected dimension."""
        rv = object.__new__(IndexerResolver)
        rv.obj = self.obj
        rv.indices = (self.indices[dim],)
        return rv
class Assigner:
    __slots__ = "updater", "resolved_indexes", "is_submask", "__weakref__"
    def __init__(self, updater, resolved_indexes, *, is_submask):
        # We could check here whether mask dimensions match index dimensions.
        # We could also check for valid `updater.kwargs` if `resolved_indexes.is_single_element`.
        self.updater = updater
        self.resolved_indexes = resolved_indexes
        self.is_submask = is_submask
        if updater.kwargs.get("input_mask") is not None:
            raise TypeError("`input_mask` argument may only be used for extract")
    def update(self, expr):
        # Occurs when user calls `C[index](...).update(expr)` or `C(...)[index].update(expr)`
        self.updater._setitem(self.resolved_indexes, expr, is_submask=self.is_submask)
    def __lshift__(self, expr):
        # Occurs when user calls `C[index](...) << expr` or `C(...)[index] << expr`
        self.updater._setitem(self.resolved_indexes, expr, is_submask=self.is_submask)
    def __eq__(self, other):
        raise TypeError(f"__eq__ not defined for objects of type {type(self)}.")
    __hash__ = None
    def __bool__(self):
        raise TypeError(f"__bool__ not defined for objects of type {type(self)}.")
class AmbiguousAssignOrExtract:
    __slots__ = "parent", "resolved_indexes", "_value", "__weakref__"
    _is_scalar = False
    def __init__(self, parent, resolved_indexes):
        self.parent = parent
        self.resolved_indexes = resolved_indexes
        self._value = None
    def __call__(self, *args, **kwargs):
        # Occurs when user calls `C[index](params)`
        # Reverse the call order so we can parse the call args and kwargs
        updater = self.parent(*args, **kwargs)
        is_submask = updater.kwargs.get("mask") is not None
        return Assigner(updater, self.resolved_indexes, is_submask=is_submask)
    def __lshift__(self, expr, **opts):
        # Occurs when user calls `C[index] << expr`
        self.update(expr, **opts)
    def update(self, expr, **opts):
        # Occurs when user calls `C[index].update(expr)`
        if getattr(self.parent, "_is_transposed", False):
            raise TypeError("'TransposedMatrix' object does not support item assignment")
        Updater(self.parent, opts=opts)._setitem(self.resolved_indexes, expr, is_submask=False)
    def new(self, dtype=None, *, mask=None, input_mask=None, name=None, **opts):
        """Force extraction of the indexes into a new object.
        dtype and mask are the only controllable parameters.
        """
        if input_mask is not None:
            if mask is not None:
                raise TypeError("mask and input_mask arguments cannot both be given")
            from .base import _check_mask
            input_mask = _check_mask(input_mask, self.parent)
            mask = self._input_mask_to_mask(input_mask, **opts)
        delayed_extractor = self.parent._prep_for_extract(self.resolved_indexes)
        return delayed_extractor.new(dtype, mask=mask, name=name, **opts)
    def _extract_delayed(self):
        """Return an Expression object, treating this as an extract call."""
        return self.parent._prep_for_extract(self.resolved_indexes)
    def _input_mask_to_mask(self, input_mask, **opts):
        from .vector import Vector
        if type(input_mask.parent) is Vector and type(self.parent) is not Vector:
            rowidx, colidx = self.resolved_indexes.indices
            if rowidx.size is None:
                if self.parent._ncols != input_mask.parent._size:
                    raise ValueError(
                        "Size of `input_mask` Vector does not match ncols of Matrix:\n"
                        f"{self.parent.name}.ncols != {input_mask.parent.name}.size  -->  "
                        f"{self.parent._ncols} != {input_mask.parent._size}"
                mask_expr = input_mask.parent._prep_for_extract(self.resolved_indexes.get_index(1))
            elif colidx.size is None:
                if self.parent._nrows != input_mask.parent._size:
                    raise ValueError(
                        "Size of `input_mask` Vector does not match nrows of Matrix:\n"
                        f"{self.parent.name}.nrows != {input_mask.parent.name}.size  -->  "
                        f"{self.parent._nrows} != {input_mask.parent._size}"
                mask_expr = input_mask.parent._prep_for_extract(self.resolved_indexes.get_index(0))
            else:
                raise TypeError(
                    "Got Vector `input_mask` when extracting a submatrix from a Matrix.  "
                    "Vector `input_mask` with a Matrix (or TransposedMatrix) input is "
                    "only valid when extracting from a single row or column."
        elif input_mask.parent.shape != self.parent.shape:
            if type(self.parent) is Vector:
                attr = "size"
                shape1 = self.parent._size
                shape2 = input_mask.parent._size
            else:
                attr = "shape"
                shape1 = self.parent.shape
                shape2 = input_mask.parent.shape
            raise ValueError(
                f"{attr.capitalize()} of `input_mask` does not match {attr} of input:\n"
                f"{self.parent.name}.{attr} != {input_mask.parent.name}.{attr}  -->  "
                f"{shape1} != {shape2}"
        else:
            mask_expr = input_mask.parent._prep_for_extract(self.resolved_indexes)
        mask_value = mask_expr.new(name="mask_temp", **opts)
        return type(input_mask)(mask_value)
    def _repr_html_(self):
        from . import formatting
        return formatting.format_index_expression_html(self)
    def __repr__(self):
        from . import formatting
        return formatting.format_index_expression(self)
    def _format_expr(self):
        indices = ", ".join(index._expr_name for index in self.resolved_indexes.indices)
        return f"{self.parent.name}[{indices}]"
    def _format_expr_html(self):
        indices = ", ".join(index._expr_name for index in self.resolved_indexes.indices)
        return f"{self.parent._name_html}[{indices}]"
    @property
    def dtype(self):
        return self.parent.dtype
class Updater:
    __slots__ = "parent", "kwargs", "opts", "__weakref__"
    def __init__(self, parent, *, opts, **kwargs):
        self.parent = parent
        self.kwargs = kwargs
        self.opts = opts
    def __getitem__(self, keys):
        # Occurs when user calls `C(params)[index]`
        # Need something prepared to receive `<<` or `.update()`
        if self.parent._is_scalar:
            raise TypeError("Indexing not supported for Scalars")
        resolved_indexes = IndexerResolver(self.parent, keys)
        return Assigner(self, resolved_indexes, is_submask=False)
    def _setitem(self, resolved_indexes, obj, *, is_submask):
        # Occurs when user calls C(params)[index] = expr
        if resolved_indexes.is_single_element and not self.kwargs:
            # Fast path using assignElement
            if self.opts:
                # Ignore opts for now
                desc = descriptor_lookup(**self.opts)  # noqa: F841 (keep desc in scope for context)
            self.parent._assign_element(resolved_indexes, obj)
        else:
            mask = self.kwargs.get("mask")
            replace = self.kwargs.get("replace", False)
            expr, alt_mask = self.parent._prep_for_assign(
                resolved_indexes,
                mask,
                is_submask,
                replace,
                self.opts,
            if alt_mask is mask:
                self.update(expr)
            else:
                kwargs = dict(self.kwargs, mask=alt_mask)
                Updater(self.parent, opts=self.opts, **kwargs).update(expr)
    def __setitem__(self, keys, obj):
        if self.parent._is_scalar:
            raise TypeError("Indexing not supported for Scalars")
        resolved_indexes = IndexerResolver(self.parent, keys)
        self._setitem(resolved_indexes, obj, is_submask=False)
    def __delitem__(self, keys):
        # Occurs when user calls `del C(params)[index]`
        if self.parent._is_scalar:
            raise TypeError("Indexing not supported for Scalars")
        resolved_indexes = IndexerResolver(self.parent, keys)
        if resolved_indexes.is_single_element:
            self.parent._delete_element(resolved_indexes)
        else:
            # Delete selection by assigning an empty scalar
            from .scalar import Scalar
            scalar = Scalar(
                self.parent.dtype, is_cscalar=False, name="s_empty"  # pragma: is_grbscalar
            self._setitem(resolved_indexes, scalar, is_submask=False)
    def __lshift__(self, expr):
        # Occurs when user calls `C(params) << expr`
        self.parent._update(expr, opts=self.opts, **self.kwargs)
    def update(self, expr):
        # Occurs when user calls `C(params).update(expr)`
        self.parent._update(expr, opts=self.opts, **self.kwargs)
    def __eq__(self, other):
        raise TypeError(f"__eq__ not defined for objects of type {type(self)}.")
    __hash__ = None
    def __bool__(self):
        raise TypeError(f"__bool__ not defined for objects of type {type(self)}.")
class InfixExprBase:
    __slots__ = "left", "right", "_expr", "__weakref__"
    _is_scalar = False
    def __init__(self, left, right):
        self.left = left
        self.right = right
        self._expr = None
    def new(self, dtype=None, *, mask=None, name=None, **opts):
        if (
            mask is None
            and self._expr is not None
            and self._expr._value is not None
            and (dtype is None or self._expr._value.dtype == dtype)
            rv = self._expr._value
            if name is not None:
                rv.name = name
            self._expr._value = None
            return rv
        expr = self._to_expr()
        return expr.new(dtype, mask=mask, name=name, **opts)
    def _to_expr(self):
        if self._expr is None:
            # Rely on the default operator for `x @ y`
            self._expr = getattr(self.left, self.method_name)(self.right)
        return self._expr
    def _get_value(self, attr=None, default=None):
        expr = self._to_expr()
        return expr._get_value(attr=attr, default=default)
    def _format_expr(self):
        return f"{self.left.name} {self._infix} {self.right.name}"
    def _format_expr_html(self):
        return f"{self.left._name_html} {self._infix} {self.right._name_html}"
    def _repr_html_(self):
        from . import formatting
        if self.output_type.__name__ == "VectorExpression":
            return formatting.format_vector_infix_expression_html(self)
        if self.output_type.__name__ == "MatrixExpression":
            return formatting.format_matrix_infix_expression_html(self)
        return formatting.format_scalar_infix_expression_html(self)
    def __repr__(self):
        from . import formatting
        if self.output_type.__name__ == "VectorExpression":
            return formatting.format_vector_infix_expression(self)
        if self.output_type.__name__ == "MatrixExpression":
            return formatting.format_matrix_infix_expression(self)
        return formatting.format_scalar_infix_expression(self)
    @property
    def dtype(self):
        return self._to_expr().dtype
    @property
    def _value(self):
        if self._expr is None:
            return None
        return self._expr._value
    @_value.setter
    def _value(self, val):
        self._to_expr()._value = val
utils._output_types[AmbiguousAssignOrExtract] = AmbiguousAssignOrExtract
utils._output_types[Assigner] = Assigner
utils._output_types[Updater] = Updater
Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

FilesExpand file tree

expr.py

Latest commit

History

expr.py

File metadata and controls
