def test_5(self):
w = aobj.expand_diagonals(self.b, (ds.DiagAxis("p", 3),))
B = np.zeros((3, 3, 3))
for i in range(3):
for j in range(3):
B[i, j, j] = (j + 1) * 10 ** i
self.assertAllclose(w, B)
dims = (ds.DiagAxis("f", 3), ds.IndepAxis("p", 3),
ds.DerivAxis("p", 3), )
self.assertEqual(w.dims, dims)
def test_2(self):
w = aobj.expand_diagonals(self.a, (ds.DiagAxis("f", 3),))
self.assertAllclose(w, np.array([[1, 0, 0], [0, 2, 0], [0, 0, 3]]))
dims = (ds.IndepAxis("f", 3), ds.DerivAxis("f", 3),)
self.assertEqual(w.dims, dims)