import inspect

import itertools

import pickle

import sys

import types

import weakref

import numpy as np

import pytest

from numpy.testing import assert_array_equal

import graphblas as gb

from graphblas import agg, backend, binary, dtypes, indexunary, monoid, select, semiring, unary

from graphblas.core import _supports_udfs as supports_udfs

from graphblas.exceptions import (

DimensionMismatch,

DomainMismatch,

EmptyObject,

IndexOutOfBound,

InvalidObject,

InvalidValue,

OutputNotEmpty,

UdfParseError,

)

from .conftest import autocompute, compute, pypy

from graphblas import Matrix, Scalar, Vector # isort:skip (for dask-graphblas)

suitesparse = backend == "suitesparse"

if suitesparse:

ss_version_major = gb.core.ss.version_major

@pytest.fixture

def A():

data = [

[3, 0, 3, 5, 6, 0, 6, 1, 6, 2, 4, 1],

[0, 1, 2, 2, 2, 3, 3, 4, 4, 5, 5, 6],

[3, 2, 3, 1, 5, 3, 7, 8, 3, 1, 7, 4],

]

return Matrix.from_coo(*data)

@pytest.fixture

def v():

data = [[1, 3, 4, 6], [1, 1, 2, 0]]

return Vector.from_coo(*data)

def test_new():

u = Vector(dtypes.INT8, 17)

assert u.dtype == "INT8"

assert u.nvals == 0

assert u.size == 17

def test_large_vector():

u = Vector.from_coo([0, 2**59], [0, 1])

assert u.size == 2**59 + 1

assert u[2**59].new() == 1

with pytest.raises(InvalidValue):

Vector.from_coo([0, 2**64 - 2], [0, 1])

with pytest.raises(OverflowError):

Vector.from_coo([0, 2**64], [0, 1])

def test_dup(v):

u = v.dup()

assert u is not v

assert u.dtype == v.dtype

assert u.nvals == v.nvals

assert u.size == v.size

# Ensure they are not the same backend object

v[0] = 1000

assert u[0].new() != 1000

# extended functionality

w = Vector.from_coo([0, 1], [0, 2.5], dtype=dtypes.FP64)

x = w.dup(dtype=dtypes.INT64)

assert x.isequal(Vector.from_coo([0, 1], [0, 2], dtype=dtypes.INT64), check_dtype=True)

x = w.dup(mask=w.V)

assert x.isequal(Vector.from_coo([1], [2.5], dtype=dtypes.FP64), check_dtype=True)

x = w.dup(dtype=dtypes.INT64, mask=w.V)

assert x.isequal(Vector.from_coo([1], [2], dtype=dtypes.INT64), check_dtype=True)

def test_dup_clear(v):

w = v.dup(clear=True)

assert w.dtype == v.dtype

assert w.nvals == 0

assert w.size == v.size

x = v.dup(dtypes.INT8, clear=True)

assert x.dtype == dtypes.INT8

assert x.nvals == 0

assert x.size == v.size

def test_from_coo():

u = Vector.from_coo([0, 1, 3], [True, False, True])

assert u.size == 4

assert u.nvals == 3

assert u.dtype == bool

u2 = Vector.from_coo([0, 1, 3], [12.3, 12.4, 12.5], size=17)

assert u2.size == 17

assert u2.nvals == 3

assert u2.dtype == float

u3 = Vector.from_coo([0, 1, 1], [1, 2, 3], size=10, dup_op=binary.times)

assert u3.size == 10

assert u3.nvals == 2 # duplicates were combined

assert u3.dtype == int

assert u3[1].new() == 6 # 2*3

with pytest.raises(ValueError, match="Duplicate indices found"):

# Duplicate indices requires a dup_op

Vector.from_coo([0, 1, 1], [True, True, True])

with pytest.raises(ValueError, match="No indices provided. Unable to infer size."):

Vector.from_coo([], [])

# Changed: Assume empty value is float64 (like numpy)

# with pytest.raises(ValueError, match="No values provided. Unable to determine type"):

w = Vector.from_coo([], [], size=10)

assert w.size == 10

assert w.nvals == 0

assert w.dtype == dtypes.FP64

with pytest.raises(ValueError, match="No indices provided. Unable to infer size"):

Vector.from_coo([], [], dtype=dtypes.INT64)

u4 = Vector.from_coo([], [], size=10, dtype=dtypes.INT64)

u5 = Vector(dtypes.INT64, size=10)

assert u4.isequal(u5, check_dtype=True)

# we check index dtype if given numpy array

with pytest.raises(ValueError, match="indices must be integers, not float64"):

Vector.from_coo(np.array([1.2, 3.4]), [1, 2])

# but coerce index if given Python lists (we defer to numpy casting)

u6 = Vector.from_coo([1.2, 3.4], [1, 2])

assert u6.isequal(Vector.from_coo([1, 3], [1, 2]))

# mis-matched sizes

with pytest.raises(ValueError, match="`indices` and `values` lengths must match"):

Vector.from_coo([0], [1, 2])

def test_from_coo_scalar():

u = Vector.from_coo([0, 1, 3], 7)

assert u.size == 4

assert u.nvals == 3

assert u.dtype == dtypes.INT64

if suitesparse:

assert u.ss.is_iso

assert u.ss.iso_value == 7

assert u.reduce(monoid.any).new() == 7

# ignore duplicate indices; iso trumps duplicates!

u = Vector.from_coo([0, 1, 1, 3], 7)

assert u.size == 4

assert u.nvals == 3

if suitesparse:

assert u.ss.is_iso

assert u.ss.iso_value == 7

assert u.reduce(monoid.any).new() == 7

with pytest.raises(ValueError, match="dup_op must be None"):

Vector.from_coo([0, 1, 1, 3], 7, dup_op=binary.plus)

u[0] = 0

if suitesparse:

with pytest.raises(ValueError, match="not iso"):

u.ss.iso_value

def test_from_coo_default_values():

u = Vector.from_coo([1, 3])

v = Vector.from_coo([1, 3], [1.0, 1.0])

assert v.isequal(u, check_dtype=True)

def test_clear(v):

v.clear()

assert v.nvals == 0

assert v.size == 7

def test_resize(v):

assert v.size == 7

assert v.nvals == 4

v.resize(20)

assert v.size == 20

assert v.nvals == 4

assert compute(v[19].new().value) is None

v.resize(4)

assert v.size == 4

assert v.nvals == 2

def test_size(v):

assert v.size == 7

def test_nvals(v):

assert v.nvals == 4

def test_build(v):

assert v.nvals == 4

v.clear()

v.build([0, 6], [1, 2])

assert v.nvals == 2

with pytest.raises(OutputNotEmpty):

v.build([1, 5], [3, 4])

assert v.nvals == 2 # should be unchanged

# We can clear though

v.build([1, 2, 5], [2, 3, 4], clear=True)

assert v.nvals == 3

v.clear()

with pytest.raises(IndexOutOfBound):

v.build([0, 11], [1, 1])

w = Vector(int, size=3)

w.build([0, 11], [1, 1], size=12)

assert w.isequal(Vector.from_coo([0, 11], [1, 1]))

@pytest.mark.skipif("not suitesparse")

def test_ss_build_scalar(v):

with pytest.raises(OutputNotEmpty):

v.ss.build_scalar([1, 5], 3)

v.clear()

v.ss.build_scalar([1, 5], 3)

assert v.nvals == 2

assert v.ss.is_iso

def test_extract_values(v):

idx, vals = v.to_coo()

np.testing.assert_array_equal(idx, (1, 3, 4, 6))

np.testing.assert_array_equal(vals, (1, 1, 2, 0))

assert idx.dtype == np.uint64

assert vals.dtype == np.int64

idx, vals = v.to_coo(dtype=int)

np.testing.assert_array_equal(idx, (1, 3, 4, 6))

np.testing.assert_array_equal(vals, (1, 1, 2, 0))

assert idx.dtype == np.uint64

assert vals.dtype == np.int64

idx, vals = v.to_coo(dtype=float)

np.testing.assert_array_equal(idx, (1, 3, 4, 6))

np.testing.assert_array_equal(vals, (1, 1, 2, 0))

assert idx.dtype == np.uint64

assert vals.dtype == np.float64

def test_extract_input_mask():

v = Vector.from_coo([0, 1, 2], [0, 1, 2])

m = Vector.from_coo([0, 2], [0, 2])

result = v[[0, 1]].new(input_mask=m.S)

expected = Vector.from_coo([0], [0], size=2)

assert result.isequal(expected)

# again

result.clear()

result(input_mask=m.S) << v[[0, 1]]

assert result.isequal(expected)

with pytest.raises(ValueError, match="Size of `input_mask` does not match size of input"):

v[[0, 2]].new(input_mask=expected.S)

with pytest.raises(TypeError, match="`input_mask` argument may only be used for extract"):

v(input_mask=m.S) << 1

def test_extract_element(v):

assert v[1].new() == 1

with pytest.raises(TypeError, match="enable automatic"):

v[1].value

assert v[6].new() == 0

with pytest.raises(TypeError, match="Invalid type for index"):

v[object()]

with pytest.raises(IndexError):

v[100]

s = Scalar(int)

s << v[1]

assert s == 1

def test_set_element(v):

assert compute(v[0].new().value) is None

assert v[1].new() == 1

v[0] = 12

v[1] << 9

assert v[0].new() == 12

assert v[1].new() == 9

def test_remove_element(v):

assert v[1].new() == 1

del v[1]

assert compute(v[1].new().value) is None

assert v[4].new() == 2

# with pytest.raises(TypeError, match="Remove Element only supports"):

del v[1:3] # Now okay

def test_vxm(v, A):

w = v.vxm(A, semiring.plus_times).new()

result = Vector.from_coo([0, 2, 3, 4, 5, 6], [3, 3, 0, 8, 14, 4])

assert w.isequal(result)

def test_vxm_transpose(v, A):

w = v.vxm(A.T, semiring.plus_times).new()

result = Vector.from_coo([0, 1, 6], [5, 16, 13])

assert w.isequal(result)

def test_vxm_nonsquare(v):

A = Matrix.from_coo([0, 3], [0, 1], [10, 20], nrows=7, ncols=2)

u = Vector(v.dtype, size=2)

u().update(v.vxm(A, semiring.min_plus))

result = Vector.from_coo([1], [21])

assert u.isequal(result)

w1 = v.vxm(A, semiring.min_plus).new()

assert w1.isequal(u)

# Test the transpose case

v2 = Vector.from_coo([0, 1], [1, 2])

w2 = v2.vxm(A.T, semiring.min_plus).new()

assert w2.size == 7

def test_vxm_mask(v, A):

val_mask = Vector.from_coo([0, 1, 2, 3, 4], [True, False, False, True, True], size=7)

struct_mask = Vector.from_coo([0, 3, 4], [False, False, False], size=7)

u = v.dup()

u(struct_mask.S) << v.vxm(A, semiring.plus_times)

result = Vector.from_coo([0, 1, 3, 4, 6], [3, 1, 0, 8, 0], size=7)

assert u.isequal(result)

u = v.dup()

u(~~struct_mask.S) << v.vxm(A, semiring.plus_times)

assert u.isequal(result)

u = v.dup()

u(~struct_mask.S) << v.vxm(A, semiring.plus_times)

result2 = Vector.from_coo([2, 3, 4, 5, 6], [3, 1, 2, 14, 4], size=7)

assert u.isequal(result2)

u = v.dup()

u(replace=True, mask=val_mask.V) << v.vxm(A, semiring.plus_times)

result3 = Vector.from_coo([0, 3, 4], [3, 0, 8], size=7)

assert u.isequal(result3)

u = v.dup()

u(replace=True, mask=~~val_mask.V) << v.vxm(A, semiring.plus_times)

assert u.isequal(result3)

w = v.vxm(A, semiring.plus_times).new(mask=val_mask.V)

assert w.isequal(result3)

def test_vxm_accum(v, A):

w1 = v.dup()

w1(binary.plus) << v.vxm(A, semiring.plus_times)

result = Vector.from_coo([0, 1, 2, 3, 4, 5, 6], [3, 1, 3, 1, 10, 14, 4], size=7)

assert w1.isequal(result)

# Allow monoids

w2 = v.dup()

w2(monoid.plus) << v.vxm(A, semiring.plus_times)

assert w2.isequal(result)

w3 = v.dup()

w3(accum=monoid.plus) << v.vxm(A, semiring.plus_times)

assert w3.isequal(result)

# allow strings for accum

w4 = v.dup()

w4("+") << v.vxm(A, semiring.plus_times)

assert w4.isequal(result)

w5 = v.dup()

w5(accum="plus") << v.vxm(A, semiring.plus_times)

assert w5.isequal(result)

def test_ewise_mult(v):

# Binary, Monoid, and Semiring

v2 = Vector.from_coo([0, 3, 5, 6], [2, 3, 2, 1])

result = Vector.from_coo([3, 6], [3, 0])

w = v.ewise_mult(v2, binary.times).new()

assert w.isequal(result)

w << v.ewise_mult(v2, monoid.times)

assert w.isequal(result)

with pytest.raises(TypeError, match="Expected type: BinaryOp, Monoid"):

v.ewise_mult(v2, semiring.plus_times)

def test_ewise_mult_change_dtype(v):

# We want to divide by 2, converting ints to floats

v2 = Vector.from_coo([1, 3, 4, 6], [2, 2, 2, 2])

assert v.dtype == dtypes.INT64

assert v2.dtype == dtypes.INT64

result = Vector.from_coo([1, 3, 4, 6], [0.5, 0.5, 1.0, 0], dtype=dtypes.FP64)

w = v.ewise_mult(v2, binary.cdiv[dtypes.FP64]).new()

assert w.isequal(result), w

# Here is the potentially surprising way to do things

# Division is still done with ints, but results are then stored as floats

result2 = Vector.from_coo([1, 3, 4, 6], [0.0, 0.0, 1.0, 0.0], dtype=dtypes.FP64)

w2 = v.ewise_mult(v2, binary.cdiv).new(dtype=dtypes.FP64)

assert w2.isequal(result2), w2

# Try with boolean dtype via auto-conversion

result3 = Vector.from_coo([1, 3, 4, 6], [True, True, False, True])

w3 = v.ewise_mult(v2, binary.lt).new()

assert w3.isequal(result3), w3

def test_ewise_add(v):

# Binary, Monoid, and Semiring

v2 = Vector.from_coo([0, 3, 5, 6], [2, 3, 2, 1])

result = Vector.from_coo([0, 1, 3, 4, 5, 6], [2, 1, 3, 2, 2, 1])

v.ewise_add(v2, binary.minus) # okay now

w = v.ewise_add(v2, binary.max).new() # ok if the binaryop is part of a monoid

assert w.isequal(result)

w = v.ewise_add(v2, binary.max).new()

assert w.isequal(result)

w.update(v.ewise_add(v2, monoid.max))

assert w.isequal(result)

with pytest.raises(TypeError, match="Expected type: BinaryOp, Monoid"):

v.ewise_add(v2, semiring.max_times)

# default is plus

w = v.ewise_add(v2).new()

result = v.ewise_add(v2, monoid.plus).new()

assert w.isequal(result)

# what about default for bool?

b1 = Vector.from_coo([0, 1, 2, 3], [True, False, True, False])

b2 = Vector.from_coo([0, 1, 2, 3], [True, True, False, False])

with pytest.raises(KeyError, match="plus does not work"):

b1.ewise_add(b2).new()

# with pytest.raises(TypeError, match="for BOOL datatype"):

binary.plus(b1 | b2) # now okay (btw, `binary.plus[bool].monoid is None`)

def test_extract(v):

w = Vector(v.dtype, 3)

result = Vector.from_coo([0, 1], [1, 1], size=3)

w << v[[1, 3, 5]]

assert w.isequal(result)

w() << v[1::2]

assert w.isequal(result)

w2 = v[1::2].new()

assert w2.isequal(w)

def test_extract_array(v):

w = Vector(v.dtype, 3)

result = Vector.from_coo(np.array([0, 1]), np.array([1, 1]), size=3)

w << v[np.array([1, 3, 5])]

assert w.isequal(result)

def test_extract_with_vector(v):

with pytest.raises(TypeError, match="Invalid type for index"):

v[v].new()

with pytest.raises(TypeError, match="Invalid type for index"):

v[v.S].new()

def test_extract_fancy_scalars(v):

assert v.dtype == dtypes.INT64

s = v[1].new()

assert s == 1

assert s.dtype == dtypes.INT64

assert v.dtype == dtypes.INT64

s = v[1].new(dtype=float)

assert s == 1.0

assert s.dtype == dtypes.FP64

t = Scalar(float)

# with pytest.raises(TypeError, match="is not supported"):

t(accum=binary.plus) << s # Now okay

assert t == 1.0

# with pytest.raises(TypeError, match="is not supported"):

t(accum=binary.plus) << 1 # Now okay

assert t == 2.0

with pytest.raises(TypeError, match="Mask not allowed for Scalars"):

t(mask=t) << s

s << v[1]

assert s.value == 1

t = Scalar(float)

t << v[1]

assert t.value == 1.0

t = Scalar(float)

t() << v[1]

assert t.value == 1.0

# with pytest.raises(TypeError, match="Scalar accumulation with extract element"):

with pytest.raises(TypeError, match="autocompute"):

t(accum=binary.plus) << v[1] # Now okay, but autocomputes

with gb.config.set(autocompute=True):

t(accum=binary.plus) << v[1] # Now okay, but autocomputes

assert t == 2.0

def test_extract_negative_indices(v):

assert v[-1].new() == 0

assert compute(v[-v.size].new().value) is None

assert v[[-v.size]].new().nvals == 0

assert v[Scalar.from_value(-4)].new() == 1

w = v[[-1, -3]].new()

assert w.isequal(Vector.from_coo([0, 1], [0, 2]))

with pytest.raises(IndexError):

v[-v.size - 1]

with pytest.raises(IndexError):

v[Scalar.from_value(-v.size - 1)]

with pytest.raises(IndexError):

v[[-v.size - 1]]

def test_assign(v):

u = Vector.from_coo([0, 2], [9, 8])

result = Vector.from_coo([0, 1, 3, 4, 6], [9, 1, 1, 8, 0])

w = v.dup()

w[[0, 2, 4]] = u

assert w.isequal(result)

w = v.dup()

w[:5:2] << u

assert w.isequal(result)

with pytest.raises(TypeError):

w[:] << u()

with pytest.raises(TypeError, match="Invalid type for index: Vector."):

w[w] = 1

def test_assign_scalar(v):

result = Vector.from_coo([1, 3, 4, 5, 6], [9, 9, 2, 9, 0])

w = v.dup()

w[[1, 3, 5]] = 9

assert w.isequal(result)

w = v.dup()

w[1::2] = 9

assert w.isequal(result)

w = Vector.from_coo([0, 1, 2], [1, 1, 1])

s = Scalar.from_value(9)

w[0] = s

assert w.isequal(Vector.from_coo([0, 1, 2], [9, 1, 1]))

w[:] = s

assert w.isequal(Vector.from_coo([0, 1, 2], [9, 9, 9]))

with pytest.raises(TypeError, match="Bad type for arg"):

w[:] = object()

with pytest.raises(TypeError, match="Bad type for arg"):

w[1] = object()

w << 2

assert w.isequal(Vector.from_coo([0, 1, 2], [2, 2, 2]))

w[0] = Scalar(int)

assert w.isequal(Vector.from_coo([1, 2], [2, 2]))

def test_assign_scalar_mask(v):

mask = Vector.from_coo([1, 2, 5, 6], [0, 0, 1, 0])

result = Vector.from_coo([1, 3, 4, 5, 6], [1, 1, 2, 5, 0])

w = v.dup()

w[:](mask.V) << 5

assert w.isequal(result)

w = v.dup()

w(mask.V) << 5

assert w.isequal(result)

w = v.dup()

w(mask.V)[:] << 5

assert w.isequal(result)

result2 = Vector.from_coo([0, 1, 2, 3, 4, 6], [5, 5, 5, 5, 5, 5])

w = v.dup()

w[:](~mask.V) << 5

assert w.isequal(result2)

w = v.dup()

w(~mask.V) << 5

assert w.isequal(result2)

w = v.dup()

w(~mask.V)[:] << 5

assert w.isequal(result2)

result3 = Vector.from_coo([1, 2, 3, 4, 5, 6], [5, 5, 1, 2, 5, 5])

w = v.dup()

w[:](mask.S) << 5

assert w.isequal(result3)

w = v.dup()

w(mask.S) << 5

assert w.isequal(result3)

w = v.dup()

w(mask.S)[:] << 5

assert w.isequal(result3)

result4 = Vector.from_coo([0, 1, 3, 4, 6], [5, 1, 5, 5, 0])

w = v.dup()

w[:](~mask.S) << 5

assert w.isequal(result4)

w = v.dup()

w(~mask.S) << Scalar.from_value(5)

assert w.isequal(result4)

w = v.dup()

w(~mask.S)[:] << 5

assert w.isequal(result4)

def test_subassign(A):

v = Vector.from_coo([0, 1, 2], [0, 1, 2])

w = Vector.from_coo([0, 1], [10, 20])

m = Vector.from_coo([1], [True])

v[[0, 1]](m.S) << w

result1 = Vector.from_coo([0, 1, 2], [0, 20, 2])

assert v.isequal(result1)

with pytest.raises(DimensionMismatch):

v[[0, 1]](v.S) << w

with pytest.raises(DimensionMismatch):

v[[0, 1]](m.S) << v

v[[0, 1]](m.S) << 100

result2 = Vector.from_coo([0, 1, 2], [0, 100, 2])

assert v.isequal(result2)

with pytest.raises(DimensionMismatch):

v[[0, 1]](v.S) << 99

with pytest.raises(TypeError, match="Mask object must be type Vector"):

v[[0, 1]](A.S) << 88

with pytest.raises(TypeError, match="Mask object must be type Vector"):

v[[0, 1]](A.S) << w

# It may be nice for these to also raise

v[[0, 1]](A.S)

v[0](m.S)

v[0](v.S, replace=True)

def test_assign_scalar_with_mask():

v = Vector.from_coo([0, 1, 2], [1, 2, 3])

m = Vector.from_coo([0, 2], [False, True])

w1 = Vector.from_coo([0], [50])

w3 = Vector.from_coo([0, 1, 2], [10, 20, 30])

v(m.V)[:] << w3

result = Vector.from_coo([0, 1, 2], [1, 2, 30])

assert v.isequal(result)

v(m.V)[:] << 100

result = Vector.from_coo([0, 1, 2], [1, 2, 100])

assert v.isequal(result)

v(m.V, accum=binary.plus)[2] << 1000

result = Vector.from_coo([0, 1, 2], [1, 2, 1100])

assert v.isequal(result)

with pytest.raises(TypeError, match="Single element assign does not accept a submask"):

v[2](w1.S) << w1

with pytest.raises(TypeError, match="Single element assign does not accept a submask"):

v[2](w1.S) << 7

v[[2]](w1.S) << 7

result = Vector.from_coo([0, 1, 2], [1, 2, 7])

assert v.isequal(result)

def test_assign_list():

v = Vector(int, 4)

v[[0, 1]] = [2, 3]

expected = Vector.from_coo([0, 1], [2, 3], size=4)

assert v.isequal(expected)

v[::2] = np.arange(2)

expected = Vector.from_coo([0, 1, 2], [0, 3, 1], size=4)

assert v.isequal(expected)

with pytest.raises(TypeError):

v[0] = [1]

with pytest.raises(TypeError):

v()[0] = [1]

with pytest.raises(ValueError, match="shape mismatch"):

v[[0, 1]] = [1, 2, 3]

with pytest.raises(ValueError, match="shape mismatch"):

v[[0, 1]] = [[1, 2]]

with pytest.raises(ValueError, match="shape mismatch"):

v[[0, 1]] = [[1], [2]]

with pytest.raises(TypeError):

v[[0, 1]] = [2, object()]

def test_apply(v):

result = Vector.from_coo([1, 3, 4, 6], [-1, -1, -2, 0])

w = v.apply(unary.ainv).new()

assert w.isequal(result)

with pytest.raises(TypeError, match="apply only accepts UnaryOp with no scalars or BinaryOp"):

v.apply(semiring.min_plus)

def test_apply_binary(v):

result_right = Vector.from_coo([1, 3, 4, 6], [False, False, True, False])

w_right = v.apply(binary.gt, right=1).new()

w_right2 = v.apply(binary.gt, right=Scalar.from_value(1)).new()

assert w_right.isequal(result_right)

assert w_right2.isequal(result_right)

result_left = Vector.from_coo([1, 3, 4, 6], [1, 1, 0, 2])

w_left = v.apply(binary.minus, left=2).new()

w_left2 = v.apply(binary.minus, left=Scalar.from_value(2)).new()

assert w_left.isequal(result_left)

assert w_left2.isequal(result_left)

with pytest.raises(TypeError):

v.apply(binary.plus, left=v)

with pytest.raises(TypeError):

v.apply(binary.plus, right=v)

with pytest.raises(TypeError, match="Cannot provide both"):

v.apply(binary.plus, left=1, right=1)

# accept monoids

w1 = v.apply(binary.plus, left=1).new()

w2 = v.apply(monoid.plus, left=1).new()

w3 = v.apply(monoid.plus, right=1).new()

assert w1.isequal(w2)

assert w1.isequal(w3)

def test_apply_empty(v):

s = Scalar(int, is_cscalar=True)

with pytest.raises(EmptyObject):

v.apply(binary.plus, s).new()

def test_apply_indexunary(v):

idx = [1, 3, 4, 6]

vr = Vector.from_coo(idx, idx)

r1 = v.apply(indexunary.index).new()

r2 = v.apply("index").new()

r3 = indexunary.index(v).new()

assert r1.isequal(vr)

assert r2.isequal(vr)

assert r3.isequal(vr)

v2 = Vector.from_coo(idx, [False, False, True, False])

s2 = Scalar(dtypes.INT64)

s2.value = 2

w1 = v.apply("==", s2).new()

w2 = v.apply(indexunary.valueeq, s2).new()

w3 = v.apply(select.valueeq, s2).new()

w4 = indexunary.valueeq(v, s2).new()

assert w1.isequal(v2)

assert w2.isequal(v2)

assert w3.isequal(v2)

assert w4.isequal(v2)

with pytest.raises(TypeError, match="left"):

v.apply(indexunary.valueeq, left=s2)

def test_select(v):

result = Vector.from_coo([1, 3], [1, 1], size=7)

w1 = v.select(select.valueeq, 1).new()

w2 = v.select("==", 1).new()

w3 = select.rowle(v, 3).new()

w4 = v.select("index<=", 3).new()

w5 = select.value(v == 1).new()

w6 = select.index(v < 4).new()

assert w1.isequal(result)

assert w2.isequal(result)

assert w3.isequal(result)

assert w4.isequal(result)

assert w5.isequal(result)

assert w6.isequal(result)

with pytest.raises(TypeError, match="thunk"):

v.select(select.valueeq, object())

w7 = v.select("==").new()

expected = Vector.from_coo([6], 0)

assert w7.isequal(expected)

w8 = select.index(v >= 4).new()

expected = Vector.from_coo([4, 6], [2, 0], size=7)

assert w8.isequal(expected)

@autocompute

def test_select_bools_and_masks(v):

# Select with boolean and masks

result = Vector.from_coo([1, 3], [1, 1], size=7)

w7 = v.select((v == 1).new()).new()

assert w7.isequal(result)

w7b = v.select(v == 1).new() # we rewrite!

assert w7b.isequal(result)

w7c = v.select(~(v != 1)).new()

assert w7c.isequal(result)

w8 = v.select(w7.S).new()

assert w8.isequal(result)

w7[4] = 0

w9 = v.select(w7.V).new()

assert w9.isequal(result)

with pytest.raises(TypeError, match="thunk"):

v.select(w7.V, 777)

with pytest.raises(TypeError, match="thunk"):

v.select(v == 1, 777)

with pytest.raises(TypeError):

v.select(v.outer(v).new().S)

# Make sure we use `replace`

w9 << 1

w9 << v.select(w7.V)

assert w9.isequal(result)

# and with masks

w9 << 1

w9[1] = 0

w9(w9.V) << v.select(w7.V)

result2 = Vector.from_coo([1, 3], [0, 1], size=7)

assert w9.isequal(result2)

w9 << 1

w9[1] = 0

w9(w9.V, replace=True) << v.select(w7.V)

del result2[1]

assert w9.isequal(result2)

w9 << 1

w9[1] = 0

w9(w9.V, binary.plus) << v.select(w7.V)

w8 << 1

w8[1] = 0

w8(w8.V, binary.plus) << v.dup(mask=w7.V)

assert w8.isequal(w9)

w9 << 1

w9(binary.plus) << v.select(w7.V)

w8 << 1

w8(binary.plus) << v.dup(mask=w7.V)

assert w8.isequal(w9)

@pytest.mark.skipif("not supports_udfs")

@pytest.mark.slow

def test_indexunary_udf(v):

def twox_minusthunk(x, row, col, thunk): # pragma: no cover (numba)

return 2 * x - thunk

indexunary.register_new("twox_minusthunk", twox_minusthunk, lazy=True)

assert hasattr(indexunary, "twox_minusthunk")

assert not hasattr(select, "twox_minusthunk")

with pytest.raises(ValueError, match="SelectOp must have BOOL return type"):

select.register_anonymous(twox_minusthunk)

with pytest.raises(TypeError, match="must be a function"):

select.register_anonymous(object())

expected = Vector.from_coo([1, 3, 4, 6], [-2, -2, 0, -4], size=7)

result = indexunary.twox_minusthunk(v, 4).new()

assert result.isequal(expected)

assert pickle.loads(pickle.dumps(indexunary.triu)) is indexunary.triu

assert indexunary.twox_minusthunk[int]._numba_func(1, 2, 3, 4) == twox_minusthunk(1, 2, 3, 4)

assert indexunary.twox_minusthunk[int].orig_func is twox_minusthunk

delattr(indexunary, "twox_minusthunk")

def ii(x, idx, _, thunk): # pragma: no cover (numba)

return idx // 2 >= thunk

def iin(n):

def inner(x, idx, _, thunk): # pragma: no cover (numba)

return idx // n >= thunk

return inner

select.register_new("ii", ii, lazy=True)

select.register_new("iin", iin, parameterized=True)

assert "ii" in dir(select)

assert "ii" in dir(indexunary)

assert hasattr(select, "ii")

assert hasattr(indexunary, "ii")

ii_apply = indexunary.register_anonymous(ii)

expected = Vector.from_coo([1, 3, 4, 6], [False, False, True, True], size=7)

result = ii_apply(v, 2).new()

assert result.isequal(expected)

result = v.apply(indexunary.iin(2), 2).new()

assert result.isequal(expected)

result = v.apply(indexunary.register_anonymous(iin, parameterized=True)(2), 2).new()

assert result.isequal(expected)

ii_select = select.register_anonymous(ii)

expected = Vector.from_coo([4, 6], [2, 0], size=7)

result = ii_select(v, 2).new()

assert result.isequal(expected)

result = v.select(select.iin(2), 2).new()

assert result.isequal(expected)

result = v.select(select.register_anonymous(iin, parameterized=True)(2), 2).new()

assert result.isequal(expected)

delattr(indexunary, "ii")

delattr(select, "ii")

delattr(indexunary, "iin")

delattr(select, "iin")

with pytest.raises(UdfParseError, match="Unable to parse function using Numba"):

indexunary.register_new("bad", lambda x, row, col, thunk: result) # pragma: no branch

def test_reduce(v):

s = v.reduce(monoid.plus).new()

assert s == 4

assert s.dtype == dtypes.INT64

assert v.reduce(binary.plus).new() == 4

with pytest.raises(TypeError, match="Expected type: Monoid"):

v.reduce(binary.minus)

with pytest.raises(TypeError, match="to get the Monoid"):

v.reduce(semiring.plus_times)

with pytest.raises(TypeError, match="part of a Monoid for BOOL datatype"):

v.reduce(binary.plus[bool])

# Test accum

s(accum=binary.times) << v.reduce(monoid.plus)

assert s == 16

# Test default for non-bool

assert v.reduce().new() == 4

# Test default for bool

b1 = Vector.from_coo([0, 1], [True, False])

with pytest.raises(KeyError, match="plus does not work"):

# KeyError here is kind of weird

b1.reduce()

# v is not empty, so it shouldn't matter how we reduce

for allow_empty, is_cscalar in itertools.product([True, False], [True, False]):

t = v.reduce(allow_empty=allow_empty).new(is_cscalar=is_cscalar)

assert t == 4

def test_reduce_empty():

w = Vector(int, 5)

s = Scalar.from_value(16)

s(accum=binary.times) << w.reduce(monoid.plus, allow_empty=True)

assert s == 16

s(accum=binary.times) << w.reduce(monoid.plus, allow_empty=False)

assert s == 0

assert w.reduce(monoid.plus, allow_empty=True).new().is_empty

assert w.reduce(monoid.plus, allow_empty=False).new() == 0

assert w.reduce(agg.sum, allow_empty=True).new().is_empty

assert w.reduce(agg.sum, allow_empty=False).new() == 0

assert w.reduce(agg.mean, allow_empty=True).new().is_empty

with pytest.raises(ValueError, match="allow_empty=False not allowed when using Aggregators"):

w.reduce(agg.mean, allow_empty=False)

def test_reduce_agg(v):

s = v.reduce(agg.sum).new()

assert s.dtype == "INT64"

assert s == 4

s = v.reduce(agg.sum[float]).new()

assert s.dtype == "FP64"

assert s == 4

assert v.reduce(agg.prod).new() == 0

assert v.reduce(agg.count).new() == 4

assert v.reduce(agg.count_nonzero).new() == 3

assert v.reduce(agg.count_zero).new() == 1

assert v.reduce(agg.sum_of_squares).new() == 6

assert v.reduce(agg.hypot).new().isclose(6**0.5)

assert v.reduce(agg.logaddexp).new().isclose(np.log(1 + 2 * np.e + np.e**2))

assert v.reduce(agg.logaddexp2).new().isclose(np.log2(9))

assert v.reduce(agg.mean).new() == 1

assert v.reduce(agg.peak_to_peak).new() == 2

assert v.reduce(agg.varp).new().isclose(0.5)

assert v.reduce(agg.vars).new().isclose(2 / 3)

assert v.reduce(agg.stdp).new().isclose(0.5**0.5)

assert v.reduce(agg.stds).new().isclose((2 / 3) ** 0.5)

assert v.reduce(agg.L0norm).new() == 3

assert v.reduce(agg.L1norm).new() == 4

assert v.reduce(agg.L2norm).new().isclose(6**0.5)

assert v.reduce(agg.Linfnorm).new() == 2

assert v.reduce(agg.exists).new() == 1

w = binary.plus(v, 1).new()

assert w.reduce(agg.geometric_mean).new().isclose(12**0.25)

assert w.reduce(agg.harmonic_mean).new().isclose(12 / 7)

silly = agg.Aggregator(

"silly",

composite=[agg.varp, agg.stdp],

finalize=lambda x, y, opts: binary.times(x & y),

types=[agg.varp],

)

s = v.reduce(silly).new()

assert s.isclose(0.5**1.5)

s = Vector(int, size=5).reduce(silly).new()

assert s.is_empty

def test_reduce_agg_count_is_int64(v):

"""Aggregators that count should default to INT64 return dtype."""

assert v.dtype == dtypes.INT64

res = v.reduce(agg.count).new()

assert res.dtype == dtypes.INT64

assert res == 4

res = v.dup(dtypes.INT8).reduce(agg.count).new()

assert res.dtype == dtypes.INT64

assert res == 4

# Allow return dtype to be specified

res = v.dup(dtypes.INT8).reduce(agg.count[dtypes.INT16]).new()

assert res.dtype == dtypes.INT16

assert res == 4

@pytest.mark.skipif("not suitesparse")

def test_reduce_agg_argminmax(v):

assert v.reduce(agg.ss.argmin).new() == 6

assert v.reduce(agg.ss.argmax).new() == 4

silly = agg.Aggregator(

"silly",

composite=[agg.ss.argmin, agg.ss.argmax],

finalize=lambda x, y, opts: binary.plus(x & y),

types=[agg.ss.argmin],

)

s = v.reduce(silly).new()

assert s == 10

@pytest.mark.skipif("not suitesparse")

def test_reduce_agg_firstlast(v):

empty = Vector(int, size=4)

assert empty.reduce(agg.ss.first).new().is_empty

assert empty.reduce(agg.ss.last).new().is_empty

assert v.reduce(agg.ss.first).new() == 1

assert v.reduce(agg.ss.last).new() == 0

silly = agg.Aggregator(

"silly",

composite=[agg.ss.first, agg.ss.last],

finalize=lambda x, y, opts: binary.plus(x & y),

types=[agg.ss.first],

)

s = v.reduce(silly).new()

assert s == 1