def test_subclasspreservation(self):
# Checks that masked_array(...,subok=True) preserves the class.
x = np.arange(5)
m = [0, 0, 1, 0, 0]
xinfo = [(i, j) for (i, j) in zip(x, m)]
xsub = MSubArray(x, mask=m, info={'xsub':xinfo})
#
mxsub = masked_array(xsub, subok=False)
self.assertTrue(not isinstance(mxsub, MSubArray))
self.assertTrue(isinstance(mxsub, MaskedArray))
assert_equal(mxsub._mask, m)
#
mxsub = asarray(xsub)
self.assertTrue(not isinstance(mxsub, MSubArray))
self.assertTrue(isinstance(mxsub, MaskedArray))
assert_equal(mxsub._mask, m)
#
mxsub = masked_array(xsub, subok=True)
self.assertTrue(isinstance(mxsub, MSubArray))
assert_equal(mxsub.info, xsub.info)
assert_equal(mxsub._mask, xsub._mask)
#
mxsub = asanyarray(xsub)
self.assertTrue(isinstance(mxsub, MSubArray))
assert_equal(mxsub.info, xsub.info)
assert_equal(mxsub._mask, m)
def test_attributepropagation(self):
x = array(arange(5), mask=[0]+[1]*4)
my = masked_array(subarray(x))
ym = msubarray(x)
#
z = (my+1)
self.assertTrue(isinstance(z, MaskedArray))
self.assertTrue(not isinstance(z, MSubArray))
self.assertTrue(isinstance(z._data, SubArray))
assert_equal(z._data.info, {})
#
z = (ym+1)
self.assertTrue(isinstance(z, MaskedArray))
self.assertTrue(isinstance(z, MSubArray))
self.assertTrue(isinstance(z._data, SubArray))
self.assertTrue(z._data.info['added'] > 0)
# Test that inplace methods from data get used (gh-4617)
ym += 1
self.assertTrue(isinstance(ym, MaskedArray))
self.assertTrue(isinstance(ym, MSubArray))
self.assertTrue(isinstance(ym._data, SubArray))
self.assertTrue(ym._data.info['iadded'] > 0)
#
ym._set_mask([1, 0, 0, 0, 1])
assert_equal(ym._mask, [1, 0, 0, 0, 1])
ym._series._set_mask([0, 0, 0, 0, 1])
assert_equal(ym._mask, [0, 0, 0, 0, 1])
#
xsub = subarray(x, info={'name':'x'})
mxsub = masked_array(xsub)
self.assertTrue(hasattr(mxsub, 'info'))
assert_equal(mxsub.info, xsub.info)
def test_out(self):
x = masked_array(np.arange(30).reshape(10, 3))
x[:3] = x[-3:] = masked
out = masked_array(np.ones(10))
r = median(x, axis=1, out=out)
assert_equal(r, out)
assert_(type(r) == MaskedArray)
def test_subclass_repr(self):
"""test that repr uses the name of the subclass
and 'array' for np.ndarray"""
x = np.arange(5)
mx = masked_array(x, mask=[True, False, True, False, False])
self.assertTrue(repr(mx).startswith("masked_array"))
xsub = SubArray(x)
mxsub = masked_array(xsub, mask=[True, False, True, False, False])
self.assertTrue(repr(mxsub).startswith("masked_{0}(data = [-- 1 -- 3 4]".format(SubArray.__name__)))
def test_subclass_repr(self):
"""test that repr uses the name of the subclass
and 'array' for np.ndarray"""
x = np.arange(5)
mx = masked_array(x, mask=[True, False, True, False, False])
assert_startswith(repr(mx), 'masked_array')
xsub = SubArray(x)
mxsub = masked_array(xsub, mask=[True, False, True, False, False])
assert_startswith(repr(mxsub),
'masked_{0}(data=[--, 1, --, 3, 4]'.format(SubArray.__name__))
def test_subclass_str(self):
"""test str with subclass that has overridden str, setitem"""
# first without override
x = np.arange(5)
xsub = SubArray(x)
mxsub = masked_array(xsub, mask=[True, False, True, False, False])
self.assertTrue(str(mxsub) == "[-- 1 -- 3 4]")
xcsub = ComplicatedSubArray(x)
assert_raises(ValueError, xcsub.__setitem__, 0, np.ma.core.masked_print_option)
mxcsub = masked_array(xcsub, mask=[True, False, True, False, False])
self.assertTrue(str(mxcsub) == "myprefix [-- 1 -- 3 4] mypostfix")
def test_unique_onmaskedarray(self):
# Test unique on masked data w/use_mask=True
data = masked_array([1, 1, 1, 2, 2, 3], mask=[0, 0, 1, 0, 1, 0])
test = unique(data, return_index=True, return_inverse=True)
assert_equal(test[0], masked_array([1, 2, 3, -1], mask=[0, 0, 0, 1]))
assert_equal(test[1], [0, 3, 5, 2])
assert_equal(test[2], [0, 0, 3, 1, 3, 2])
#
data.fill_value = 3
data = masked_array(data=[1, 1, 1, 2, 2, 3], mask=[0, 0, 1, 0, 1, 0], fill_value=3)
test = unique(data, return_index=True, return_inverse=True)
assert_equal(test[0], masked_array([1, 2, 3, -1], mask=[0, 0, 0, 1]))
assert_equal(test[1], [0, 3, 5, 2])
assert_equal(test[2], [0, 0, 3, 1, 3, 2])
def test_unique_allmasked(self):
# Test all masked
data = masked_array([1, 1, 1], mask=True)
test = unique(data, return_index=True, return_inverse=True)
assert_equal(test[0], masked_array([1, ], mask=[True]))
assert_equal(test[1], [0])
assert_equal(test[2], [0, 0, 0])
#
# Test masked
data = masked
test = unique(data, return_index=True, return_inverse=True)
assert_equal(test[0], masked_array(masked))
assert_equal(test[1], [0])
assert_equal(test[2], [0])
def test_2d(self):
# Tests median w/ 2D
(n, p) = (101, 30)
x = masked_array(np.linspace(-1., 1., n),)
x[:10] = x[-10:] = masked
z = masked_array(np.empty((n, p), dtype=float))
z[:, 0] = x[:]
idx = np.arange(len(x))
for i in range(1, p):
np.random.shuffle(idx)
z[:, i] = x[idx]
assert_equal(median(z[:, 0]), 0)
assert_equal(median(z), 0)
assert_equal(median(z, axis=0), np.zeros(p))
assert_equal(median(z.T, axis=1), np.zeros(p))
def test_masked_binary_operations2(self):
# Tests domained_masked_binary_operation
(x, mx) = self.data
xmx = masked_array(mx.data.__array__(), mask=mx.mask)
self.assertTrue(isinstance(divide(mx, mx), mmatrix))
self.assertTrue(isinstance(divide(mx, x), mmatrix))
assert_equal(divide(mx, mx), divide(xmx, xmx))
def test_allany_onmatrices(self):
x = np.array([[0.13, 0.26, 0.90],
[0.28, 0.33, 0.63],
[0.31, 0.87, 0.70]])
X = np.matrix(x)
m = np.array([[True, False, False],
[False, False, False],
[True, True, False]], dtype=np.bool_)
mX = masked_array(X, mask=m)
mXbig = (mX > 0.5)
mXsmall = (mX < 0.5)
assert_(not mXbig.all())
assert_(mXbig.any())
assert_equal(mXbig.all(0), np.matrix([False, False, True]))
assert_equal(mXbig.all(1), np.matrix([False, False, True]).T)
assert_equal(mXbig.any(0), np.matrix([False, False, True]))
assert_equal(mXbig.any(1), np.matrix([True, True, True]).T)
assert_(not mXsmall.all())
assert_(mXsmall.any())
assert_equal(mXsmall.all(0), np.matrix([True, True, False]))
assert_equal(mXsmall.all(1), np.matrix([False, False, False]).T)
assert_equal(mXsmall.any(0), np.matrix([True, True, False]))
assert_equal(mXsmall.any(1), np.matrix([True, True, False]).T)
def test_edges(self):
# Tests unmasked_edges
data = masked_array(np.arange(25).reshape(5, 5),
mask=[[0, 0, 1, 0, 0],
[0, 0, 0, 1, 1],
[1, 1, 0, 0, 0],
[0, 0, 0, 0, 0],
[1, 1, 1, 0, 0]],)
test = notmasked_edges(data, None)
assert_equal(test, [0, 24])
test = notmasked_edges(data, 0)
assert_equal(test[0], [(0, 0, 1, 0, 0), (0, 1, 2, 3, 4)])
assert_equal(test[1], [(3, 3, 3, 4, 4), (0, 1, 2, 3, 4)])
test = notmasked_edges(data, 1)
assert_equal(test[0], [(0, 1, 2, 3, 4), (0, 0, 2, 0, 3)])
assert_equal(test[1], [(0, 1, 2, 3, 4), (4, 2, 4, 4, 4)])
#
test = notmasked_edges(data.data, None)
assert_equal(test, [0, 24])
test = notmasked_edges(data.data, 0)
assert_equal(test[0], [(0, 0, 0, 0, 0), (0, 1, 2, 3, 4)])
assert_equal(test[1], [(4, 4, 4, 4, 4), (0, 1, 2, 3, 4)])
test = notmasked_edges(data.data, -1)
assert_equal(test[0], [(0, 1, 2, 3, 4), (0, 0, 0, 0, 0)])
assert_equal(test[1], [(0, 1, 2, 3, 4), (4, 4, 4, 4, 4)])
#
data[-2] = masked
test = notmasked_edges(data, 0)
assert_equal(test[0], [(0, 0, 1, 0, 0), (0, 1, 2, 3, 4)])
assert_equal(test[1], [(1, 1, 2, 4, 4), (0, 1, 2, 3, 4)])
test = notmasked_edges(data, -1)
assert_equal(test[0], [(0, 1, 2, 4), (0, 0, 2, 3)])
assert_equal(test[1], [(0, 1, 2, 4), (4, 2, 4, 4)])
def test_clump_unmasked(self):
# Test clump_unmasked
a = masked_array(np.arange(10))
a[[0, 1, 2, 6, 8, 9]] = masked
test = clump_unmasked(a)
control = [slice(3, 6), slice(7, 8), ]
assert_equal(test, control)
def test_subclass_items(self):
"""test that getter and setter go via baseclass"""
x = np.arange(5)
xcsub = ComplicatedSubArray(x)
mxcsub = masked_array(xcsub, mask=[True, False, True, False, False])
# getter should return a ComplicatedSubArray, even for single item
# first check we wrote ComplicatedSubArray correctly
self.assertTrue(isinstance(xcsub[1], ComplicatedSubArray))
self.assertTrue(isinstance(xcsub[1:4], ComplicatedSubArray))
# now that it propagates inside the MaskedArray
self.assertTrue(isinstance(mxcsub[1], ComplicatedSubArray))
self.assertTrue(mxcsub[0] is masked)
self.assertTrue(isinstance(mxcsub[1:4].data, ComplicatedSubArray))
# also for flattened version (which goes via MaskedIterator)
self.assertTrue(isinstance(mxcsub.flat[1].data, ComplicatedSubArray))
self.assertTrue(mxcsub[0] is masked)
self.assertTrue(isinstance(mxcsub.flat[1:4].base, ComplicatedSubArray))
# setter should only work with ComplicatedSubArray input
# first check we wrote ComplicatedSubArray correctly
assert_raises(ValueError, xcsub.__setitem__, 1, x[4])
# now that it propagates inside the MaskedArray
assert_raises(ValueError, mxcsub.__setitem__, 1, x[4])
assert_raises(ValueError, mxcsub.__setitem__, slice(1, 4), x[1:4])
mxcsub[1] = xcsub[4]
mxcsub[1:4] = xcsub[1:4]
# also for flattened version (which goes via MaskedIterator)
assert_raises(ValueError, mxcsub.flat.__setitem__, 1, x[4])
assert_raises(ValueError, mxcsub.flat.__setitem__, slice(1, 4), x[1:4])
mxcsub.flat[1] = xcsub[4]
mxcsub.flat[1:4] = xcsub[1:4]
def test_pickling_subbaseclass(self):
# Test pickling w/ a subclass of ndarray
a = masked_array(np.matrix(list(range(10))), mask=[1, 0, 1, 0, 0] * 2)
a_pickled = pickle.loads(a.dumps())
assert_equal(a_pickled._mask, a._mask)
assert_equal(a_pickled, a)
assert_(isinstance(a_pickled._data, np.matrix))
def test_ediff1d(self):
# Tests mediff1d
x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1])
control = array([1, 1, 1, 4], mask=[1, 0, 0, 1])
test = ediff1d(x)
assert_equal(test, control)
assert_equal(test.data, control.data)
assert_equal(test.mask, control.mask)
def test_clump_masked(self):
# Test clump_masked
a = masked_array(np.arange(10))
a[[0, 1, 2, 6, 8, 9]] = masked
#
test = clump_masked(a)
control = [slice(0, 3), slice(6, 7), slice(8, 10)]
assert_equal(test, control)
def test_2d_waxis(self):
# Tests median w/ 2D arrays and different axis.
x = masked_array(np.arange(30).reshape(10, 3))
x[:3] = x[-3:] = masked
assert_equal(median(x), 14.5)
assert_equal(median(x, axis=0), [13.5, 14.5, 15.5])
assert_equal(median(x, axis=1), [0, 0, 0, 10, 13, 16, 19, 0, 0, 0])
assert_equal(median(x, axis=1).mask, [1, 1, 1, 0, 0, 0, 0, 1, 1, 1])
def test_2d_waxis(self):
# Tests median w/ 2D arrays and different axis.
x = masked_array(np.arange(30).reshape(10, 3))
x[:3] = x[-3:] = masked
assert_equal(median(x), 14.5)
assert_equal(median(x, axis=0), [13.5, 14.5, 15.5])
assert_equal(median(x, axis=1), [0, 0, 0, 10, 13, 16, 19, 0, 0, 0])
assert_equal(median(x, axis=1).mask, [1, 1, 1, 0, 0, 0, 0, 1, 1, 1])
def test_compressed(self):
a = masked_array(np.matrix([1, 2, 3, 4]), mask=[0, 0, 0, 0])
b = a.compressed()
assert_equal(b, a)
assert_(isinstance(b, np.matrix))
a[0, 0] = masked
b = a.compressed()
assert_equal(b, [[2, 3, 4]])
def test_unique_onlist(self):
# Test unique on list
data = [1, 1, 1, 2, 2, 3]
test = unique(data, return_index=True, return_inverse=True)
self.assertTrue(isinstance(test[0], MaskedArray))
assert_equal(test[0], masked_array([1, 2, 3], mask=[0, 0, 0]))
assert_equal(test[1], [0, 3, 5])
assert_equal(test[2], [0, 0, 0, 1, 1, 2])
def test_ediff1d(self):
"Tests mediff1d"
x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1])
control = array([1, 1, 1, 4], mask=[1, 0, 0, 1])
test = ediff1d(x)
assert_equal(test, control)
assert_equal(test.data, control.data)
assert_equal(test.mask, control.mask)
def test_clump_masked(self):
# Test clump_masked
a = masked_array(np.arange(10))
a[[0, 1, 2, 6, 8, 9]] = masked
#
test = clump_masked(a)
control = [slice(0, 3), slice(6, 7), slice(8, 10)]
assert_equal(test, control)
def test_atleast2d(self):
# Test atleast_2d
a = masked_array([0, 1, 2], mask=[0, 1, 0])
b = atleast_2d(a)
assert_equal(b.shape, (1, 3))
assert_equal(b.mask.shape, b.data.shape)
assert_equal(a.shape, (3, ))
assert_equal(a.mask.shape, a.data.shape)
def test_pickling_subbaseclass(self):
# Test pickling w/ a subclass of ndarray
a = masked_array(np.matrix(list(range(10))), mask=[1, 0, 1, 0, 0] * 2)
for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
a_pickled = pickle.loads(pickle.dumps(a, protocol=proto))
assert_equal(a_pickled._mask, a._mask)
assert_equal(a_pickled, a)
assert_(isinstance(a_pickled._data, np.matrix))
def test_atleast2d(self):
# Test atleast_2d
a = masked_array([0, 1, 2], mask=[0, 1, 0])
b = atleast_2d(a)
assert_equal(b.shape, (1, 3))
assert_equal(b.mask.shape, b.data.shape)
assert_equal(a.shape, (3,))
assert_equal(a.mask.shape, a.data.shape)
def test_ediff1d(self):
# Tests mediff1d
x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1])
control = array([1, 1, 1, 4], mask=[1, 0, 0, 1])
test = ediff1d(x)
assert_equal(test, control)
assert_equal(test.filled(0), control.filled(0))
assert_equal(test.mask, control.mask)
def test_unique_onlist(self):
# Test unique on list
data = [1, 1, 1, 2, 2, 3]
test = unique(data, return_index=True, return_inverse=True)
self.assertTrue(isinstance(test[0], MaskedArray))
assert_equal(test[0], masked_array([1, 2, 3], mask=[0, 0, 0]))
assert_equal(test[1], [0, 3, 5])
assert_equal(test[2], [0, 0, 0, 1, 1, 2])
def test_pickling_subbaseclass(self):
# Test pickling w/ a subclass of ndarray
a = masked_array(np.matrix(list(range(10))), mask=[1, 0, 1, 0, 0] * 2)
for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
a_pickled = pickle.loads(pickle.dumps(a, protocol=proto))
assert_equal(a_pickled._mask, a._mask)
assert_equal(a_pickled, a)
assert_(isinstance(a_pickled._data, np.matrix))
def test_compressed(self):
a = masked_array(np.matrix([1, 2, 3, 4]), mask=[0, 0, 0, 0])
b = a.compressed()
assert_equal(b, a)
assert_(isinstance(b, np.matrix))
a[0, 0] = masked
b = a.compressed()
assert_equal(b, [[2, 3, 4]])
def test_2d(self):
"Tests mr_ on 2D arrays."
a_1 = rand(5, 5)
a_2 = rand(5, 5)
m_1 = np.round_(rand(5, 5), 0)
m_2 = np.round_(rand(5, 5), 0)
b_1 = masked_array(a_1, mask=m_1)
b_2 = masked_array(a_2, mask=m_2)
d = mr_['1', b_1, b_2] # append columns
self.assertTrue(d.shape == (5, 10))
assert_array_equal(d[:, :5], b_1)
assert_array_equal(d[:, 5:], b_2)
assert_array_equal(d.mask, np.r_['1', m_1, m_2])
d = mr_[b_1, b_2]
self.assertTrue(d.shape == (10, 5))
assert_array_equal(d[:5,:], b_1)
assert_array_equal(d[5:,:], b_2)
assert_array_equal(d.mask, np.r_[m_1, m_2])
def test_clump_unmasked(self):
# Test clump_unmasked
a = masked_array(np.arange(10))
a[[0, 1, 2, 6, 8, 9]] = masked
test = clump_unmasked(a)
control = [slice(3, 6), slice(7, 8), ]
assert_equal(test, control)
self.check_clump(clump_unmasked)
def test_data_subclassing(self):
# Tests whether the subclass is kept.
x = np.arange(5)
m = [0, 0, 1, 0, 0]
xsub = SubArray(x)
xmsub = masked_array(xsub, mask=m)
self.assertTrue(isinstance(xmsub, MaskedArray))
assert_equal(xmsub._data, xsub)
self.assertTrue(isinstance(xmsub._data, SubArray))
def test_2d(self):
"Tests mr_ on 2D arrays."
a_1 = rand(5, 5)
a_2 = rand(5, 5)
m_1 = np.round_(rand(5, 5), 0)
m_2 = np.round_(rand(5, 5), 0)
b_1 = masked_array(a_1, mask=m_1)
b_2 = masked_array(a_2, mask=m_2)
d = mr_['1', b_1, b_2] # append columns
self.assertTrue(d.shape == (5, 10))
assert_array_equal(d[:, :5], b_1)
assert_array_equal(d[:, 5:], b_2)
assert_array_equal(d.mask, np.r_['1', m_1, m_2])
d = mr_[b_1, b_2]
self.assertTrue(d.shape == (10, 5))
assert_array_equal(d[:5, :], b_1)
assert_array_equal(d[5:, :], b_2)
assert_array_equal(d.mask, np.r_[m_1, m_2])
def test_count_mean_with_matrix(self):
m = masked_array(np.matrix([[1, 2], [3, 4]]), mask=np.zeros((2, 2)))
assert_equal(m.count(axis=0).shape, (1, 2))
assert_equal(m.count(axis=1).shape, (2, 1))
# Make sure broadcasting inside mean and var work
assert_equal(m.mean(axis=0), [[2., 3.]])
assert_equal(m.mean(axis=1), [[1.5], [3.5]])
def test_data_subclassing(self):
# Tests whether the subclass is kept.
x = np.arange(5)
m = [0, 0, 1, 0, 0]
xsub = SubArray(x)
xmsub = masked_array(xsub, mask=m)
assert_(isinstance(xmsub, MaskedArray))
assert_equal(xmsub._data, xsub)
assert_(isinstance(xmsub._data, SubArray))
def test_count_mean_with_matrix(self):
m = masked_array(np.matrix([[1, 2], [3, 4]]), mask=np.zeros((2, 2)))
assert_equal(m.count(axis=0).shape, (1, 2))
assert_equal(m.count(axis=1).shape, (2, 1))
# Make sure broadcasting inside mean and var work
assert_equal(m.mean(axis=0), [[2., 3.]])
assert_equal(m.mean(axis=1), [[1.5], [3.5]])
def test_1d(self):
# Tests mr_ on 1D arrays.
assert_array_equal(mr_[1, 2, 3, 4, 5, 6], array([1, 2, 3, 4, 5, 6]))
b = ones(5)
m = [1, 0, 0, 0, 0]
d = masked_array(b, mask=m)
c = mr_[d, 0, 0, d]
self.assertTrue(isinstance(c, MaskedArray))
assert_array_equal(c, [1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1])
assert_array_equal(c.mask, mr_[m, 0, 0, m])
def test_view(self):
# Test view w/ flexible dtype
iterator = list(zip(np.arange(10), np.random.rand(10)))
data = np.array(iterator)
a = masked_array(iterator, dtype=[('a', float), ('b', float)])
a.mask[0] = (1, 0)
test = a.view((float, 2), np.matrix)
assert_equal(test, data)
assert_(isinstance(test, np.matrix))
assert_(not isinstance(test, MaskedArray))
def test_subclass_nomask_items(self):
x = np.arange(5)
xcsub = ComplicatedSubArray(x)
mxcsub_nomask = masked_array(xcsub)
assert_(isinstance(mxcsub_nomask[1, ...].data, ComplicatedSubArray))
assert_(isinstance(mxcsub_nomask[0, ...].data, ComplicatedSubArray))
assert_(isinstance(mxcsub_nomask[1], ComplicatedSubArray))
assert_(isinstance(mxcsub_nomask[0], ComplicatedSubArray))
def test_subclass_nomask_items(self):
x = np.arange(5)
xcsub = ComplicatedSubArray(x)
mxcsub_nomask = masked_array(xcsub)
self.assertTrue(isinstance(mxcsub_nomask[1,...].data, ComplicatedSubArray))
self.assertTrue(isinstance(mxcsub_nomask[0,...].data, ComplicatedSubArray))
self.assertTrue(isinstance(mxcsub_nomask[1], ComplicatedSubArray))
self.assertTrue(isinstance(mxcsub_nomask[0], ComplicatedSubArray))
def test_matrix_indexing(self):
# Tests conversions and indexing
x1 = np.matrix([[1, 2, 3], [4, 3, 2]])
x2 = masked_array(x1, mask=[[1, 0, 0], [0, 1, 0]])
x3 = masked_array(x1, mask=[[0, 1, 0], [1, 0, 0]])
x4 = masked_array(x1)
# test conversion to strings
str(x2) # raises?
repr(x2) # raises?
# tests of indexing
assert_(type(x2[1, 0]) is type(x1[1, 0]))
assert_(x1[1, 0] == x2[1, 0])
assert_(x2[1, 1] is masked)
assert_equal(x1[0, 2], x2[0, 2])
assert_equal(x1[0, 1:], x2[0, 1:])
assert_equal(x1[:, 2], x2[:, 2])
assert_equal(x1[:], x2[:])
assert_equal(x1[1:], x3[1:])
x1[0, 2] = 9
x2[0, 2] = 9
assert_equal(x1, x2)
x1[0, 1:] = 99
x2[0, 1:] = 99
assert_equal(x1, x2)
x2[0, 1] = masked
assert_equal(x1, x2)
x2[0, 1:] = masked
assert_equal(x1, x2)
x2[0, :] = x1[0, :]
x2[0, 1] = masked
assert_(allequal(getmask(x2), np.array([[0, 1, 0], [0, 1, 0]])))
x3[1, :] = masked_array([1, 2, 3], [1, 1, 0])
assert_(allequal(getmask(x3)[1], masked_array([1, 1, 0])))
assert_(allequal(getmask(x3[1]), masked_array([1, 1, 0])))
x4[1, :] = masked_array([1, 2, 3], [1, 1, 0])
assert_(allequal(getmask(x4[1]), masked_array([1, 1, 0])))
assert_(allequal(x4[1], masked_array([1, 2, 3])))
x1 = np.matrix(np.arange(5) * 1.0)
x2 = masked_values(x1, 3.0)
assert_equal(x1, x2)
assert_(allequal(masked_array([0, 0, 0, 1, 0], dtype=MaskType),
x2.mask))
assert_equal(3.0, x2.fill_value)
def test_1d(self):
# Tests mr_ on 1D arrays.
assert_array_equal(mr_[1, 2, 3, 4, 5, 6], array([1, 2, 3, 4, 5, 6]))
b = ones(5)
m = [1, 0, 0, 0, 0]
d = masked_array(b, mask=m)
c = mr_[d, 0, 0, d]
self.assertTrue(isinstance(c, MaskedArray))
assert_array_equal(c, [1, 1, 1, 1, 1, 0, 0, 1, 1, 1, 1, 1])
assert_array_equal(c.mask, mr_[m, 0, 0, m])
def test_view(self):
# Test view w/ flexible dtype
iterator = list(zip(np.arange(10), np.random.rand(10)))
data = np.array(iterator)
a = masked_array(iterator, dtype=[('a', float), ('b', float)])
a.mask[0] = (1, 0)
test = a.view((float, 2), np.matrix)
assert_equal(test, data)
assert_(isinstance(test, np.matrix))
assert_(not isinstance(test, MaskedArray))
def test_matrix_indexing(self):
# Tests conversions and indexing
x1 = np.matrix([[1, 2, 3], [4, 3, 2]])
x2 = masked_array(x1, mask=[[1, 0, 0], [0, 1, 0]])
x3 = masked_array(x1, mask=[[0, 1, 0], [1, 0, 0]])
x4 = masked_array(x1)
# test conversion to strings
str(x2) # raises?
repr(x2) # raises?
# tests of indexing
assert_(type(x2[1, 0]) is type(x1[1, 0]))
assert_(x1[1, 0] == x2[1, 0])
assert_(x2[1, 1] is masked)
assert_equal(x1[0, 2], x2[0, 2])
assert_equal(x1[0, 1:], x2[0, 1:])
assert_equal(x1[:, 2], x2[:, 2])
assert_equal(x1[:], x2[:])
assert_equal(x1[1:], x3[1:])
x1[0, 2] = 9
x2[0, 2] = 9
assert_equal(x1, x2)
x1[0, 1:] = 99
x2[0, 1:] = 99
assert_equal(x1, x2)
x2[0, 1] = masked
assert_equal(x1, x2)
x2[0, 1:] = masked
assert_equal(x1, x2)
x2[0, :] = x1[0, :]
x2[0, 1] = masked
assert_(allequal(getmask(x2), np.array([[0, 1, 0], [0, 1, 0]])))
x3[1, :] = masked_array([1, 2, 3], [1, 1, 0])
assert_(allequal(getmask(x3)[1], masked_array([1, 1, 0])))
assert_(allequal(getmask(x3[1]), masked_array([1, 1, 0])))
x4[1, :] = masked_array([1, 2, 3], [1, 1, 0])
assert_(allequal(getmask(x4[1]), masked_array([1, 1, 0])))
assert_(allequal(x4[1], masked_array([1, 2, 3])))
x1 = np.matrix(np.arange(5) * 1.0)
x2 = masked_values(x1, 3.0)
assert_equal(x1, x2)
assert_(allequal(masked_array([0, 0, 0, 1, 0], dtype=MaskType),
x2.mask))
assert_equal(3.0, x2.fill_value)
def test_testAverage2(self):
# More tests of average.
w1 = [0, 1, 1, 1, 1, 0]
w2 = [[0, 1, 1, 1, 1, 0], [1, 0, 0, 0, 0, 1]]
x = arange(6, dtype=np.float_)
assert_equal(average(x, axis=0), 2.5)
assert_equal(average(x, axis=0, weights=w1), 2.5)
y = array([arange(6, dtype=np.float_), 2.0 * arange(6)])
assert_equal(average(y, None), np.add.reduce(np.arange(6)) * 3. / 12.)
assert_equal(average(y, axis=0), np.arange(6) * 3. / 2.)
assert_equal(
average(y, axis=1),
[average(x, axis=0), average(x, axis=0) * 2.0])
assert_equal(average(y, None, weights=w2), 20. / 6.)
assert_equal(average(y, axis=0, weights=w2), [0., 1., 2., 3., 4., 10.])
assert_equal(
average(y, axis=1),
[average(x, axis=0), average(x, axis=0) * 2.0])
m1 = zeros(6)
m2 = [0, 0, 1, 1, 0, 0]
m3 = [[0, 0, 1, 1, 0, 0], [0, 1, 1, 1, 1, 0]]
m4 = ones(6)
m5 = [0, 1, 1, 1, 1, 1]
assert_equal(average(masked_array(x, m1), axis=0), 2.5)
assert_equal(average(masked_array(x, m2), axis=0), 2.5)
assert_equal(average(masked_array(x, m4), axis=0).mask, [True])
assert_equal(average(masked_array(x, m5), axis=0), 0.0)
assert_equal(count(average(masked_array(x, m4), axis=0)), 0)
z = masked_array(y, m3)
assert_equal(average(z, None), 20. / 6.)
assert_equal(average(z, axis=0), [0., 1., 99., 99., 4.0, 7.5])
assert_equal(average(z, axis=1), [2.5, 5.0])
assert_equal(average(z, axis=0, weights=w2),
[0., 1., 99., 99., 4.0, 10.0])
def test_ediff1d_toend(self):
# Test ediff1d w/ to_end
x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1])
test = ediff1d(x, to_end=masked)
control = array([1, 1, 1, 4, 0], mask=[1, 0, 0, 1, 1])
assert_equal(test, control)
assert_equal(test.data, control.data)
assert_equal(test.mask, control.mask)
#
test = ediff1d(x, to_end=[1, 2, 3])
control = array([1, 1, 1, 4, 1, 2, 3], mask=[1, 0, 0, 1, 0, 0, 0])
assert_equal(test, control)
assert_equal(test.data, control.data)
assert_equal(test.mask, control.mask)
def test_ediff1d_tobegin(self):
# Test ediff1d w/ to_begin
x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1])
test = ediff1d(x, to_begin=masked)
control = array([0, 1, 1, 1, 4], mask=[1, 1, 0, 0, 1])
assert_equal(test, control)
assert_equal(test.data, control.data)
assert_equal(test.mask, control.mask)
#
test = ediff1d(x, to_begin=[1, 2, 3])
control = array([1, 2, 3, 1, 1, 1, 4], mask=[0, 0, 0, 1, 0, 0, 1])
assert_equal(test, control)
assert_equal(test.data, control.data)
assert_equal(test.mask, control.mask)
def check_clump(self, f):
for i in range(1, 7):
for j in range(2**i):
k = np.arange(i, dtype=int)
ja = np.full(i, j, dtype=int)
a = masked_array(2**k)
a.mask = (ja & (2**k)) != 0
s = 0
for sl in f(a):
s += a.data[sl].sum()
if f == clump_unmasked:
assert_equal(a.compressed().sum(), s)
else:
a.mask = ~a.mask
assert_equal(a.compressed().sum(), s)
def test_ediff1d_tobegin_toend(self):
# Test ediff1d w/ to_begin and to_end
x = masked_array(np.arange(5), mask=[1, 0, 0, 0, 1])
test = ediff1d(x, to_end=masked, to_begin=masked)
control = array([0, 1, 1, 1, 4, 0], mask=[1, 1, 0, 0, 1, 1])
assert_equal(test, control)
assert_equal(test.filled(0), control.filled(0))
assert_equal(test.mask, control.mask)
#
test = ediff1d(x, to_end=[1, 2, 3], to_begin=masked)
control = array([0, 1, 1, 1, 4, 1, 2, 3],
mask=[1, 1, 0, 0, 1, 0, 0, 0])
assert_equal(test, control)
assert_equal(test.filled(0), control.filled(0))
assert_equal(test.mask, control.mask)
def test_complex(self):
# Test with complex data.
# (Regression test for https://github.com/numpy/numpy/issues/2684)
mask = np.array([[0, 0, 0, 1, 0],
[0, 1, 0, 0, 0]], dtype=bool)
a = masked_array([[0, 1+2j, 3+4j, 5+6j, 7+8j],
[9j, 0+1j, 2+3j, 4+5j, 7+7j]],
mask=mask)
av = average(a)
expected = np.average(a.compressed())
assert_almost_equal(av.real, expected.real)
assert_almost_equal(av.imag, expected.imag)
av0 = average(a, axis=0)
expected0 = average(a.real, axis=0) + average(a.imag, axis=0)*1j
assert_almost_equal(av0.real, expected0.real)
assert_almost_equal(av0.imag, expected0.imag)
av1 = average(a, axis=1)
expected1 = average(a.real, axis=1) + average(a.imag, axis=1)*1j
assert_almost_equal(av1.real, expected1.real)
assert_almost_equal(av1.imag, expected1.imag)
# Test with the 'weights' argument.
wts = np.array([[0.5, 1.0, 2.0, 1.0, 0.5],
[1.0, 1.0, 1.0, 1.0, 1.0]])
wav = average(a, weights=wts)
expected = np.average(a.compressed(), weights=wts[~mask])
assert_almost_equal(wav.real, expected.real)
assert_almost_equal(wav.imag, expected.imag)
wav0 = average(a, weights=wts, axis=0)
expected0 = (average(a.real, weights=wts, axis=0) +
average(a.imag, weights=wts, axis=0)*1j)
assert_almost_equal(wav0.real, expected0.real)
assert_almost_equal(wav0.imag, expected0.imag)
wav1 = average(a, weights=wts, axis=1)
expected1 = (average(a.real, weights=wts, axis=1) +
average(a.imag, weights=wts, axis=1)*1j)
assert_almost_equal(wav1.real, expected1.real)
assert_almost_equal(wav1.imag, expected1.imag)
def test_contiguous(self):
# Tests notmasked_contiguous
a = masked_array(np.arange(24).reshape(3, 8),
mask=[[0, 0, 0, 0, 1, 1, 1, 1],
[1, 1, 1, 1, 1, 1, 1, 1],
[0, 0, 0, 0, 0, 0, 1, 0], ])
tmp = notmasked_contiguous(a, None)
assert_equal(tmp[-1], slice(23, 24, None))
assert_equal(tmp[-2], slice(16, 22, None))
assert_equal(tmp[-3], slice(0, 4, None))
#
tmp = notmasked_contiguous(a, 0)
self.assertTrue(len(tmp[-1]) == 1)
self.assertTrue(tmp[-2] is None)
assert_equal(tmp[-3], tmp[-1])
self.assertTrue(len(tmp[0]) == 2)
#
tmp = notmasked_contiguous(a, 1)
assert_equal(tmp[0][-1], slice(0, 4, None))
self.assertTrue(tmp[1] is None)
assert_equal(tmp[2][-1], slice(7, 8, None))
assert_equal(tmp[2][-2], slice(0, 6, None))
def test_testAverage3(self):
# Yet more tests of average!
a = arange(6)
b = arange(6) * 3
r1, w1 = average([[a, b], [b, a]], axis=1, returned=1)
assert_equal(shape(r1), shape(w1))
assert_equal(r1.shape, w1.shape)
r2, w2 = average(ones((2, 2, 3)), axis=0, weights=[3, 1], returned=1)
assert_equal(shape(w2), shape(r2))
r2, w2 = average(ones((2, 2, 3)), returned=1)
assert_equal(shape(w2), shape(r2))
r2, w2 = average(ones((2, 2, 3)), weights=ones((2, 2, 3)), returned=1)
assert_equal(shape(w2), shape(r2))
a2d = array([[1, 2], [0, 4]], float)
a2dm = masked_array(a2d, [[False, False], [True, False]])
a2da = average(a2d, axis=0)
assert_equal(a2da, [0.5, 3.0])
a2dma = average(a2dm, axis=0)
assert_equal(a2dma, [1.0, 3.0])
a2dma = average(a2dm, axis=None)
assert_equal(a2dma, 7. / 3.)
a2dma = average(a2dm, axis=1)
assert_equal(a2dma, [1.5, 4.0])
def test_subclass_items(self):
"""test that getter and setter go via baseclass"""
x = np.arange(5)
xcsub = ComplicatedSubArray(x)
mxcsub = masked_array(xcsub, mask=[True, False, True, False, False])
# getter should return a ComplicatedSubArray, even for single item
# first check we wrote ComplicatedSubArray correctly
assert_(isinstance(xcsub[1], ComplicatedSubArray))
assert_(isinstance(xcsub[1, ...], ComplicatedSubArray))
assert_(isinstance(xcsub[1:4], ComplicatedSubArray))
# now that it propagates inside the MaskedArray
assert_(isinstance(mxcsub[1], ComplicatedSubArray))
assert_(isinstance(mxcsub[1, ...].data, ComplicatedSubArray))
assert_(mxcsub[0] is masked)
assert_(isinstance(mxcsub[0, ...].data, ComplicatedSubArray))
assert_(isinstance(mxcsub[1:4].data, ComplicatedSubArray))
# also for flattened version (which goes via MaskedIterator)
assert_(isinstance(mxcsub.flat[1].data, ComplicatedSubArray))
assert_(mxcsub.flat[0] is masked)
assert_(isinstance(mxcsub.flat[1:4].base, ComplicatedSubArray))
# setter should only work with ComplicatedSubArray input
# first check we wrote ComplicatedSubArray correctly
assert_raises(ValueError, xcsub.__setitem__, 1, x[4])
# now that it propagates inside the MaskedArray
assert_raises(ValueError, mxcsub.__setitem__, 1, x[4])
assert_raises(ValueError, mxcsub.__setitem__, slice(1, 4), x[1:4])
mxcsub[1] = xcsub[4]
mxcsub[1:4] = xcsub[1:4]
# also for flattened version (which goes via MaskedIterator)
assert_raises(ValueError, mxcsub.flat.__setitem__, 1, x[4])
assert_raises(ValueError, mxcsub.flat.__setitem__, slice(1, 4), x[1:4])
mxcsub.flat[1] = xcsub[4]
mxcsub.flat[1:4] = xcsub[1:4]
def test_allany_onmatrices(self):
x = np.array([[0.13, 0.26, 0.90], [0.28, 0.33, 0.63],
[0.31, 0.87, 0.70]])
X = npm.matrix(x)
m = np.array(
[[True, False, False], [False, False, False], [True, True, False]],
dtype=np.bool_)
mX = masked_array(X, mask=m)
mXbig = (mX > 0.5)
mXsmall = (mX < 0.5)
assert_(not mXbig.all())
assert_(mXbig.any())
assert_equal(mXbig.all(0), npm.matrix([False, False, True]))
assert_equal(mXbig.all(1), npm.matrix([False, False, True]).T)
assert_equal(mXbig.any(0), npm.matrix([False, False, True]))
assert_equal(mXbig.any(1), npm.matrix([True, True, True]).T)
assert_(not mXsmall.all())
assert_(mXsmall.any())
assert_equal(mXsmall.all(0), npm.matrix([True, True, False]))
assert_equal(mXsmall.all(1), npm.matrix([False, False, False]).T)
assert_equal(mXsmall.any(0), npm.matrix([True, True, False]))
assert_equal(mXsmall.any(1), npm.matrix([True, True, False]).T)
def test_flat(self):
# Test that flat can return items even for matrices [#4585, #4615]
# test simple access
test = masked_array(np.matrix([[1, 2, 3]]), mask=[0, 0, 1])
assert_equal(test.flat[1], 2)
assert_equal(test.flat[2], masked)
assert_(np.all(test.flat[0:2] == test[0, 0:2]))
# Test flat on masked_matrices
test = masked_array(np.matrix([[1, 2, 3]]), mask=[0, 0, 1])
test.flat = masked_array([3, 2, 1], mask=[1, 0, 0])
control = masked_array(np.matrix([[3, 2, 1]]), mask=[1, 0, 0])
assert_equal(test, control)
# Test setting
test = masked_array(np.matrix([[1, 2, 3]]), mask=[0, 0, 1])
testflat = test.flat
testflat[:] = testflat[[2, 1, 0]]
assert_equal(test, control)
testflat[0] = 9
# test that matrices keep the correct shape (#4615)
a = masked_array(np.matrix(np.eye(2)), mask=0)
b = a.flat
b01 = b[:2]
assert_equal(b01.data, np.array([[1., 0.]]))
assert_equal(b01.mask, np.array([[False, False]]))
def test_ravel(self):
a = masked_array(np.matrix([1, 2, 3, 4, 5]), mask=[[0, 1, 0, 0, 0]])
aravel = a.ravel()
assert_equal(aravel.shape, (1, 5))
assert_equal(aravel._mask.shape, a.shape)
def test_dot(self):
# Tests dot product
n = np.arange(1, 7)
#
m = [1, 0, 0, 0, 0, 0]
a = masked_array(n, mask=m).reshape(2, 3)
b = masked_array(n, mask=m).reshape(3, 2)
c = dot(a, b, strict=True)
assert_equal(c.mask, [[1, 1], [1, 0]])
c = dot(b, a, strict=True)
assert_equal(c.mask, [[1, 1, 1], [1, 0, 0], [1, 0, 0]])
c = dot(a, b, strict=False)
assert_equal(c, np.dot(a.filled(0), b.filled(0)))
c = dot(b, a, strict=False)
assert_equal(c, np.dot(b.filled(0), a.filled(0)))
#
m = [0, 0, 0, 0, 0, 1]
a = masked_array(n, mask=m).reshape(2, 3)
b = masked_array(n, mask=m).reshape(3, 2)
c = dot(a, b, strict=True)
assert_equal(c.mask, [[0, 1], [1, 1]])
c = dot(b, a, strict=True)
assert_equal(c.mask, [[0, 0, 1], [0, 0, 1], [1, 1, 1]])
c = dot(a, b, strict=False)
assert_equal(c, np.dot(a.filled(0), b.filled(0)))
assert_equal(c, dot(a, b))
c = dot(b, a, strict=False)
assert_equal(c, np.dot(b.filled(0), a.filled(0)))
#
m = [0, 0, 0, 0, 0, 0]
a = masked_array(n, mask=m).reshape(2, 3)
b = masked_array(n, mask=m).reshape(3, 2)
c = dot(a, b)
assert_equal(c.mask, nomask)
c = dot(b, a)
assert_equal(c.mask, nomask)
#
a = masked_array(n, mask=[1, 0, 0, 0, 0, 0]).reshape(2, 3)
b = masked_array(n, mask=[0, 0, 0, 0, 0, 0]).reshape(3, 2)
c = dot(a, b, strict=True)
assert_equal(c.mask, [[1, 1], [0, 0]])
c = dot(a, b, strict=False)
assert_equal(c, np.dot(a.filled(0), b.filled(0)))
c = dot(b, a, strict=True)
assert_equal(c.mask, [[1, 0, 0], [1, 0, 0], [1, 0, 0]])
c = dot(b, a, strict=False)
assert_equal(c, np.dot(b.filled(0), a.filled(0)))
#
a = masked_array(n, mask=[0, 0, 0, 0, 0, 1]).reshape(2, 3)
b = masked_array(n, mask=[0, 0, 0, 0, 0, 0]).reshape(3, 2)
c = dot(a, b, strict=True)
assert_equal(c.mask, [[0, 0], [1, 1]])
c = dot(a, b)
assert_equal(c, np.dot(a.filled(0), b.filled(0)))
c = dot(b, a, strict=True)
assert_equal(c.mask, [[0, 0, 1], [0, 0, 1], [0, 0, 1]])
c = dot(b, a, strict=False)
assert_equal(c, np.dot(b.filled(0), a.filled(0)))
#
a = masked_array(n, mask=[0, 0, 0, 0, 0, 1]).reshape(2, 3)
b = masked_array(n, mask=[0, 0, 1, 0, 0, 0]).reshape(3, 2)
c = dot(a, b, strict=True)
assert_equal(c.mask, [[1, 0], [1, 1]])
c = dot(a, b, strict=False)
assert_equal(c, np.dot(a.filled(0), b.filled(0)))
c = dot(b, a, strict=True)
assert_equal(c.mask, [[0, 0, 1], [1, 1, 1], [0, 0, 1]])
c = dot(b, a, strict=False)
assert_equal(c, np.dot(b.filled(0), a.filled(0)))
def test_neg_axis(self):
x = masked_array(np.arange(30).reshape(10, 3))
x[:3] = x[-3:] = masked
assert_equal(median(x, axis=-1), median(x, axis=1))
def plot_matrix(matrix, roi_names, networks, threshold=None, **kwargs):
"""
This function is used to plot connections in square matrix form.
Parameters
----------
matrix : numpy array (n x n) float
The values of connectivity between each of n ROI
roi_names : list of n string
The names of each of the n ROI
networks : list of p string
List of names representing the networks subdivision
threshold : int
Indicates the value of the most important connections
ticks_type : {'networks', 'roi'}, optional
Indicates if the tick names should be ROI or networks
ticks_color : list of colors, optional
The list in matplotlib formats of colors used to
color the ticks names, this should be in line with
the ticks_type choice: p colors if we choose 'networks'
facecolor : string, optional
As in matplotlib it indicates the background color of
the plot
Returns
-------
f : matplotlib figure
The figure just composed.
"""
_plot_cfg = {
'ticks_type': 'networks',
'ticks_color': sns.color_palette(
"Paired",
n_colors=matrix.shape[0]), #get_atlas_info('findlab')[1],
'facecolor': 'k',
'zscore': True
}
_plot_cfg.update(kwargs)
facecolor_ = _plot_cfg['facecolor']
if facecolor_ == 'k':
ax_color = 'white'
else:
ax_color = 'k'
facecolor_ = 'white'
f = plt.figure(figsize=(16., 12.), facecolor=facecolor_, dpi=300)
#f = plt.figure()
a = f.add_subplot(111)
max_value = np.max(np.abs(matrix))
# plot gray values
if _plot_cfg['zscore'] == True:
z_matrix = (matrix - matrix.mean()) / matrix.std()
else:
z_matrix = matrix
max_value = np.max(np.abs(z_matrix))
ax = a.imshow(
z_matrix,
interpolation='nearest',
cmap=plt.cm.RdBu_r,
#cmap='Greys',
alpha=0.5,
vmax=max_value,
vmin=max_value * -1)
if threshold is not None:
thresh_matrix = masked_array(z_matrix, (np.abs(z_matrix) < threshold))
ax = a.imshow(
thresh_matrix,
interpolation='nearest',
cmap=plt.cm.bwr,
#cmap='gray',
vmax=max_value,
vmin=max_value * -1,
)
min_ = -0.5
max_ = len(networks) + 0.5
### Draw networks separation lines ###
network_ticks = []
network_name, indices = np.unique(networks, return_index=True)
counter = -0.5
colors_ = []
for net in np.unique(networks):
items_idx = np.nonzero(networks == net)
items = items_idx[0].shape[0]
ix = np.nonzero(networks == net)
if _plot_cfg['ticks_type'] == 'networks':
tick_ = items_idx[0].mean()
colors_.append(_plot_cfg['ticks_color'])
else:
tick_ = items_idx[0]
colors_.append(_plot_cfg['ticks_color'][tick_])
counter = counter + items
network_ticks.append(tick_)
a.axvline(x=counter, ymin=min_, ymax=max_)
a.axhline(y=counter, xmin=min_, xmax=max_)
if _plot_cfg['ticks_type'] == 'networks':
ticks_labels = np.unique(networks)
else:
ticks_labels = roi_names
network_ticks = np.hstack(network_ticks)
colors_ = np.hstack(colors_)
a.set_yticks(network_ticks)
a.set_yticklabels(ticks_labels, fontsize=15)
a.set_xticks(network_ticks)
a.set_xticklabels(ticks_labels, fontsize=15, rotation=80)
#colors_[colors_ == facecolor_] = ax_color
#colors_[colors_ == 'beige'] = 'darkgray'
[t.set_color(colors_[i]) for i, t in enumerate(a.xaxis.get_ticklabels())]
[t.set_color(colors_[i]) for i, t in enumerate(a.yaxis.get_ticklabels())]
cbar = f.colorbar(ax)
[
t.set_color(ax_color)
for i, t in enumerate(cbar.ax.yaxis.get_ticklabels())
]
return f