def assert_same_as_dense(expression, format_out, **tensor_pairs):
tensors_in_format = {
name: Tensor.from_lol(data, format=format)
for name, (data, format) in tensor_pairs.items()
}
tensors_as_dense = {
name: Tensor.from_lol(data)
for name, (data, _) in tensor_pairs.items()
}
actual = evaluate(expression, format_out, **tensors_in_format)
expected = evaluate(expression,
''.join('d' for c in format_out if c in ('d', 's')),
**tensors_as_dense)
assert actual == expected
def test_csr_matrix_vector_product():
A = Tensor.from_aos([[1, 0], [0, 1], [1, 2]], [2.0, -2.0, 4.0],
dimensions=(2, 3),
format='ds')
x = Tensor.from_aos([[0], [1], [2]], [3.0, 2.5, 2.0],
dimensions=(3, ),
format='d')
expected = Tensor.from_aos([[0], [1]], [-5.0, 14.0],
dimensions=(2, ),
format='d')
function = tensor_method('y(i) = A(i,j) * x(j)', dict(A='ds', x='d'), 'd')
actual = function(A, x)
assert actual == expected
actual = evaluate('y(i) = A(i,j) * x(j)', 'd', A=A, x=x)
assert actual == expected
def test_csr_matrix_plus_csr_matrix():
A = Tensor.from_aos([[1, 0], [0, 1], [1, 2]], [2.0, -2.0, 4.0],
dimensions=(2, 3),
format='ds')
B = Tensor.from_aos([[1, 1], [1, 2], [0, 2]], [-3.0, 4.0, 3.5],
dimensions=(2, 3),
format='ds')
expected = Tensor.from_aos([[1, 0], [0, 1], [1, 2], [1, 1], [0, 2]],
[2.0, -2.0, 8.0, -3.0, 3.5],
dimensions=(2, 3),
format='ds')
function = tensor_method('C(i,j) = A(i,j) + B(i,j)', dict(A='ds', B='ds'),
'ds')
actual = function(A, B)
assert actual == expected
actual = evaluate('C(i,j) = A(i,j) * B(i,j)', 'ds', A=A, B=B)
assert actual == expected
def test_copy_2(dense, format_in, format_out):
a = Tensor.from_lol(dense, format=format_in)
actual = evaluate('b(i,j) = a(i,j)', format_out, a=a)
assert actual == a
def run_eval():
# Generate a random expression so that the cache cannot be hit
return evaluate(f'y{randrange(1024)}(i) = A(i,j) * x(j)',
'd',
A=A,
x=x)