def test_index(self):
"""Test indexing/slicing into coefficients.
"""
# Index.
sum_coeffs = cu.add(self.x.coefficients(), self.y.coefficients())
key = ku.validate_key((1, 0), self.x.shape)
coeffs = cu.index(sum_coeffs, key)
self.assertItemsEqual(coeffs.keys(), [self.x, self.y])
blocks = coeffs[self.y]
self.assertEqual(len(blocks), 1)
self.assertEqual(blocks[0].size, (1,2))
self.assertEqual(blocks[0][0,0], 0)
# Slice.
sum_coeffs = cu.add(self.A.coefficients(), self.C.coefficients())
key = ku.validate_key((slice(None, None, None), 1), self.A.shape)
coeffs = cu.index(sum_coeffs, key)
self.assertItemsEqual(coeffs.keys(), [self.A, s.CONSTANT])
# Variable.
blocks = coeffs[self.A]
self.assertEqual(len(blocks), 1)
self.assertEqual(blocks[0].size, (2,4))
# Constant.
blocks = coeffs[s.CONSTANT]
self.assertEqual(len(blocks), 1)
self.assertEqual(blocks[0].size, (2,1))
def test_index(self):
"""Test indexing/slicing into coefficients.
"""
# Index.
sum_coeffs = cu.add(self.x.coefficients(), self.y.coefficients())
key = ku.validate_key((1, 0), self.x.shape)
coeffs = cu.index(sum_coeffs, key)
self.assertItemsEqual(coeffs.keys(), [self.x, self.y])
blocks = coeffs[self.y]
self.assertEqual(len(blocks), 1)
self.assertEqual(blocks[0].shape, (1, 2))
self.assertEqual(blocks[0][0, 0], 0)
# Slice.
sum_coeffs = cu.add(self.A.coefficients(), self.C.coefficients())
key = ku.validate_key((slice(None, None, None), 1), self.A.shape)
coeffs = cu.index(sum_coeffs, key)
self.assertItemsEqual(coeffs.keys(), [self.A, s.CONSTANT])
# Variable.
blocks = coeffs[self.A]
self.assertEqual(len(blocks), 1)
self.assertEqual(blocks[0].shape, (2, 4))
# Constant.
blocks = coeffs[s.CONSTANT]
self.assertEqual(len(blocks), 1)
self.assertEqual(blocks[0].shape, (2, 1))
def __init__(self, expr, key, orig_key=None) -> None:
# Format and validate key.
if orig_key is None:
self._orig_key = key
self.key = ku.validate_key(key, expr.shape)
else:
self._orig_key = orig_key
self.key = key
super(index, self).__init__(expr)
def __init__(self, expr, key):
# Format and validate key.
self.key = ku.validate_key(key, expr._dcp_attr.shape)
super(index, self).__init__(expr)
def __init__(self, expr, key):
# Format and validate key.
self.key = ku.validate_key(key, expr.size)
super(index, self).__init__(expr)
