def test_index_take(self):
a = Index("A", ["a", "b", "c", "d"])
self.assertEqual(a.take([0, 2]).name, "A") # keep name
self.assertTrue(np.array_equal(a.take([0, 2]).values, ["a", "c"]))
self.assertTrue(np.array_equal(a.take([2, 0]).values, ["c", "a"]))
self.assertRaises(ValueError, a.take,
[0, 2, 0]) # duplicate values in Index
def test_stack(self):
c = year_quarter_cube()
d = year_quarter_cube()
country_axis = Index("country", ["GB", "FR"])
e = stack([c, d], country_axis)
self.assertEqual(e.values.shape, (2, 3, 4))
# the merged axis go first
self.assertEqual(tuple(e.dims), ("country", "year", "quarter"))
# axis with the same name already exists
c = year_quarter_cube()
d = year_quarter_cube()
year_axis = Index("year", [2000, 2001])
self.assertRaises(ValueError, stack, [c, d], year_axis)
# different number of cubes and axis length
c = year_quarter_cube()
d = year_quarter_cube()
country_axis = Index("country", ["GB", "FR", "DE"])
self.assertRaises(ValueError, stack, [c, d], country_axis)
# cubes do not have uniform shapes
c = year_quarter_cube()
d = year_quarter_weekday_cube()
country_axis = Index("country", ["GB", "FR"])
self.assertRaises(LookupError, stack, [c, d], country_axis)
# the previous example if O, if automatic broadcasting is allowed
e = stack([c, d], country_axis, broadcast=True)
self.assertEqual(e.ndim, 4)
# broadcast axes go last
self.assertEqual(tuple(e.dims), ("country", "year", "quarter", "weekday"))
def year_quarter_weekday_cube():
"""Creates 3D cube with axes "year", "quarter", "weekday" with shape (3, 4, 7)."""
values = np.arange(3 * 4 * 7).reshape(3, 4, 7)
ax1 = Index("year", [2014, 2015, 2016])
ax2 = Index("quarter", ["Q1", "Q2", "Q3", "Q4"])
ax3 = Index("weekday", ["mon", "tue", "wed", "thu", "fri", "sat", "sun"])
return Cube([ax1, ax2, ax3],values)
def test_concatenate(self):
values = np.arange(12).reshape(3, 4)
ax1 = Index("year", [2014, 2015, 2016])
ax2 = Index("month", ["jan", "feb", "mar", "apr"])
c = Cube([ax1, ax2],values)
values = np.arange(12).reshape(4, 3)
ax3 = Index("year", [2014, 2015, 2016])
ax4 = Index("month", ["may", "jun", "jul", "aug"])
d = Cube([ax4, ax3],values)
e = concatenate([c, d], "month")
self.assertEqual(e.ndim, 2)
self.assertEqual(e.shape, (8, 3)) # the joined axis is always the first
self.assertTrue(is_axis(e.axis("month")))
e = concatenate([c, d], "month", as_index=True)
self.assertEqual(e.ndim, 2)
self.assertEqual(e.shape, (8, 3))
self.assertTrue(is_indexed(e.axis("month")))
# duplicate index values
self.assertRaises(ValueError, concatenate, [c, c], "month", as_index=True)
# broadcasting of an axis
countries = Axis("country", ["DE", "FR"])
f = d.insert_axis(countries)
g = concatenate([c, f], "month", broadcast=True)
self.assertEqual(g.ndim, 3)
self.assertEqual(g.shape, (8, 3, 2))
# if automatic broadcasting is not allowed
self.assertRaises(LookupError, concatenate, [c, f], "month", broadcast=False)
def test_operator_in(self):
a = Index("A", [10, 20, 30])
b = Index("B", ["ab", "bc", "cd", "de"])
self.assertTrue(20 in a)
self.assertFalse(40 in a)
self.assertTrue("bc" in b)
self.assertFalse("ef" in b)
def test_compress(self):
a = Index("A", ["a", "b", "c", "d"])
selector = [True, False, True, False]
b = a.compress(selector)
c = a[np.array(selector)]
self.assertTrue(np.array_equal(b.values, c.values))
self.assertEqual(a.name, b.name) # keep name
self.assertTrue(np.array_equal(b.values, a.values.compress(selector)))
def test_create_index(self):
a = Index("A", [10, 20, 30])
self.assertEqual(a.name, "A")
self.assertEqual(len(a), 3)
a = Index("Dim", ["a", "b", "c", "d"])
self.assertEqual(a.name, "Dim")
self.assertEqual(len(a), 4)
# duplicate values
self.assertRaises(ValueError, Index, "A", ["a", "b", "a"])
self.assertRaises(ValueError, Index, "A", [0, 1, 1])
# invalid Index name
self.assertRaises(TypeError, Index, 1, [1, 2, 3])
def test_squeeze(self):
"""Removes axes which have only one element."""
ax1 = Index("A", [1]) # has only one element, thus can be collapsed
ax2 = Index("B", [1, 2, 3])
c = Cube([ax1, ax2],[[1, 2, 3]] )
self.assertEqual(c.ndim, 2)
d = c.squeeze()
self.assertEqual(d.ndim, 1)
self.assertEqual(d.axis(0).name, "B")
c = Cube([ax2, ax1],[[1], [2], [3]])
self.assertEqual(c.ndim, 2)
d = c.squeeze()
self.assertEqual(d.ndim, 1)
self.assertEqual(d.axis(0).name, "B")
ax3 = Index("C", [1]) # has only one element, thus can be collapsed
c = Cube([ax1, ax3],[[1]])
self.assertEqual(c.ndim, 2)
d = c.squeeze() # will collapse both axes
self.assertEqual(d.ndim, 0)
def test_indexof(self):
a = Index("A", [10, 20, 30])
b = Index("Dim", ["ab", "bc", "cd", "de"])
# a single value
self.assertEqual(a.indexof(10), 0)
self.assertEqual(b.indexof("cd"), 2)
self.assertEqual(b.indexof(["cd"]), 2)
# multiple values
self.assertTrue(np.array_equal(a.indexof([10, 30]), [0, 2]))
self.assertTrue(np.array_equal(b.indexof(["de", "cd"]), [3, 2]))
# non-existent value raises KeyError (similar to dictionary lookup)
self.assertRaises(KeyError, a.indexof, 0)
self.assertRaises(KeyError, b.indexof, "ef")
self.assertRaises(KeyError, b.indexof, None)
self.assertRaises(KeyError, a.indexof, [0, 1])
self.assertRaises(KeyError, b.indexof, ["de", "ef"])
def test_create_cube(self):
ax1 = Index("A", [10, 20, 30])
ax2 = Index("B", ["a", "b", "c", "d"])
ax3 = Index("C", [1.1, 1.2])
# test Cube.zeros()
a = Cube.zeros([ax1, ax3])
self.assertTrue(np.array_equal(a.values, [[0, 0], [0, 0], [0, 0]]))
# test Cube.ones()
a = Cube.ones([ax1, ax3])
self.assertTrue(np.array_equal(a.values, [[1, 1], [1, 1], [1, 1]]))
# test Cube.full()
a = Cube.full([ax1, ax3], np.inf)
self.assertTrue(np.array_equal(a.values, [[np.inf, np.inf], [np.inf, np.inf], [np.inf, np.inf]]))
# test Cube.full with NaNs
# note: be careful because NaN != NaN so np.array_equal does not work
a = Cube.full([ax1, ax3], np.nan)
np.testing.assert_equal(a.values, [[np.nan, np.nan], [np.nan, np.nan], [np.nan, np.nan]])
# create one-dimensional cube
values = np.arange(3)
try:
Cube((ax1,),values)
Cube(ax1,values) # no need to pass axes as collection if there is only one axis
except Exception:
self.fail("raised exception unexpectedly")
# two-dimensional cubes
values = np.arange(12).reshape(3, 4)
try:
Cube((ax1, ax2),values)
Cube([ax1, ax2],values)
except Exception:
self.fail("raised exception unexpectedly")
def test_filter(self):
"""Testing function Cube.filter()"""
c = year_quarter_cube()
d = c.filter("year", [2014, 2018]) # 2018 is ignored
self.assertEqual(d.ndim, 2)
self.assertTrue((d.values == c.values[0]).all())
year_filter = Axis("year", range(2010, 2015))
d = c.filter(year_filter)
self.assertEqual(d.ndim, 2)
self.assertTrue((d.values == c.values[0]).all())
year_filter = Index("year", range(2010, 2015))
d = c.filter(year_filter)
self.assertEqual(d.ndim, 2)
self.assertTrue((d.values == c.values[0]).all())
country_filter = Axis("country", ["DE", "FR"]) # this axis is ignored
# filter by two axis filters
quarter_filter = Index("quarter", ["Q1", "Q3"])
d = c.filter([quarter_filter, country_filter, year_filter])
self.assertEqual(d.ndim, 2)
self.assertTrue((d.values == c.values[[0, 0], [0, 2]]).all())
# cube as a filter
yq_cube_filter = Cube.ones([quarter_filter, year_filter, country_filter])
d = c.filter(yq_cube_filter)
self.assertEqual(d.ndim, 2)
self.assertTrue((d.values == c.values[[0, 0], [0, 2]]).all())
# a collection of cubes as a filter
y_cube_filter = Cube.ones([year_filter, country_filter])
q_cube_filter = Cube.ones([country_filter, quarter_filter])
d = c.filter([y_cube_filter, q_cube_filter])
self.assertEqual(d.ndim, 2)
self.assertTrue((d.values == c.values[[0, 0], [0, 2]]).all())
def test_align_axis(self):
c = year_quarter_cube()
ax1 = Axis("year", [2015, 2015, 2014, 2014])
ax2 = Index("quarter", ["Q1", "Q3"])
d = c.align(ax1)
d = d.align(ax2)
# test identity of the new axis
self.assertTrue(d.axis("year") is ax1)
self.assertTrue(d.axis("quarter") is ax2)
# test aligned values
self.assertTrue(np.array_equal(d.values, [[4, 6], [4, 6], [0, 2], [0, 2]]))
def test_contains(self):
a = Index("A", [10, 20, 30])
b = Index("B", ["ab", "bc", "cd", "de"])
# a single value (in this case, operator 'in' is preferred)
self.assertTrue(a.contains(20))
self.assertFalse(a.contains(40))
self.assertTrue(b.contains("bc"))
self.assertFalse(b.contains("ef"))
# multiple values returns one-dimensional numpy array of logical values
self.assertTrue(
np.array_equal(a.contains([0, 10, 20, 40]),
[False, True, True, False]))
self.assertTrue(np.array_equal(b.contains(["ab"]), [True]))
self.assertTrue(
np.array_equal(b.contains(["ab", "ef", "bc"]),
[True, False, True]))
self.assertTrue(
np.array_equal(b.contains(("ab", "ef", "bc")),
[True, False, True]))
def test_writeable(self):
# once index has been created, its values cannot be changed in order not to break lookup function
a = Index("A", [10, 20, 30])
self.assertRaises(ValueError, a.values.__setitem__, 0, 40)
self.assertRaises(ValueError, a.values.sort)
def year_quarter_cube():
"""Creates a sample 2D cube with axes "year" and "quarter" with shape (3, 4)."""
values = np.arange(12).reshape(3, 4)
ax1 = Index("year", [2014, 2015, 2016])
ax2 = Index("quarter", ["Q1", "Q2", "Q3", "Q4"])
return Cube([ax1, ax2],values)
def test_operations(self):
values = np.arange(12).reshape(3, 4)
axc1 = Index("a", [10, 20, 30])
axc2 = Index("b", ["a", "b", "c", "d"])
c = Cube([axc1, axc2],values)
axd1 = Index("a", [10, 20, 30])
axd2 = Index("b", ["a", "b", "c", "d"])
d = Cube([axd1, axd2],values)
x = c * d
self.assertTrue(np.array_equal(x.values, values * values))
e = Cube([Index("a", [10, 20, 30])],[0, 1, 2])
x2 = c * e
self.assertTrue(np.array_equal(x2.values, values * np.array([[0], [1], [2]])))
c3 = Cube([Index("b", ["a", "b", "c", "d"])],[0, 1, 2, 3])
x3 = c * c3
self.assertTrue(np.array_equal(x3.values, values * np.array([0, 1, 2, 3])))
c3 = Cube([Index("b", ["b", "a", "c", "d"])],[0, 1, 2, 3])
x3 = c * c3
self.assertTrue(np.array_equal(x3.values, values * np.array([1, 0, 2, 3])))
values_d = np.array([0, 1])
d = Cube([Index("d", ["d1", "d2"])],values_d)
x = c * d
self.assertEqual(x.ndim, 3)
self.assertEqual(x.axis(0).name, "a")
self.assertEqual(x.axis(1).name, "b")
self.assertEqual(x.axis(2).name, "d")
self.assertTrue(np.array_equal(x.values, values.reshape(3, 4, 1) * values_d))
# operations with scalar
d = 10
x = c * d
self.assertTrue(np.array_equal(x.values, values * d))
x = d * c
self.assertTrue(np.array_equal(x.values, values * d))
# operations with numpy.ndarray
d = np.arange(4)
x = c * d
self.assertTrue(np.array_equal(x.values, values * d))
x = d * c
self.assertTrue(np.array_equal(x.values, values * d))
d = np.arange(3).reshape(3, 1)
x = c * d
self.assertTrue(np.array_equal(x.values, values * d))
x = d * c
self.assertTrue(np.array_equal(x.values, values * d))
# matching Index and Series
values_d = np.array([0, 1])
d = Cube(Axis("a", [10, 10]),values_d)
x = c * d
self.assertTrue(np.array_equal(x.values, values.take([0, 0], 0) * values_d[:, np.newaxis]))
values_d = np.array([0, 1, 2, 3])
d = Cube(Axis("b", ["d", "d", "c", "a"]),values_d)
x = c * d
self.assertTrue(np.array_equal(x.values, values.take([3, 3, 2, 0], 1) * values_d))
# unary plus and minus
c = year_quarter_cube()
self.assertTrue(np.array_equal((+c).values, c.values))
self.assertTrue(np.array_equal((-c).values, -c.values))
c = year_quarter_cube() + 1 # +1 to prevent division by zero error
import operator as op
ops = [op.add, op.mul, op.floordiv, op.truediv, op.sub, op.pow, op.mod, # arithmetics ops
op.eq, op.ne, op.ge, op.le, op.gt, op.lt, # comparison ops
op.rshift, op.lshift] # bitwise ops
# operations with scalar
d = 2
for op in ops:
self.assertTrue(np.array_equal(op(c, d).values, op(c.values, d)))
self.assertTrue(np.array_equal(op(d, c).values, op(d, c.values)))
# oprations with numpy array
d = (np.arange(12).reshape(3, 4) / 6 + 1).astype(np.int) # +1 to prevent division by zero error
for op in ops:
self.assertTrue(np.array_equal(op(c, d).values, op(c.values, d)))
self.assertTrue(np.array_equal(op(d, c).values, op(d, c.values)))