def test_0d_array(self, cls=np.ndarray):
def sum_to_0d(x):
""" Sum x, returning a 0d array of the same class """
assert_equal(x.ndim, 1)
return np.squeeze(np.sum(x, keepdims=True))
a = np.ones((6, 3)).view(cls)
res = apply_along_axis(sum_to_0d, 0, a)
assert_(isinstance(res, cls))
assert_array_equal(res, np.array([6, 6, 6]).view(cls))
res = apply_along_axis(sum_to_0d, 1, a)
assert_(isinstance(res, cls))
assert_array_equal(res, np.array([3, 3, 3, 3, 3, 3]).view(cls))
def test_preserve_subclass(self):
# this test is particularly malicious because matrix
# refuses to become 1d
def double(row):
return row * 2
m = np.matrix([[0, 1], [2, 3]])
expected = np.matrix([[0, 2], [4, 6]])
result = apply_along_axis(double, 0, m)
assert_(isinstance(result, np.matrix))
assert_array_equal(result, expected)
result = apply_along_axis(double, 1, m)
assert_(isinstance(result, np.matrix))
assert_array_equal(result, expected)
def transform(self, x, use_spln=False, **kwargs):
"""
Apply transform to x
Parameters
----------
x : float-array-convertible
Data to be transformed.
Should support conversion to an array of floats.
use_spln: bool
True - transform using the spline specified in self.slpn.
If self.spln is None, set the spline.
False - transform using self.tfun
kwargs:
Keyword arguments to be passed to self.set_spline.
Only used if use_spln=True & self.spln=None.
Returns
-------
Array of transformed values.
"""
x = asarray(x, dtype=float)
if use_spln:
if self.spln is None:
self.set_spline(x.min(), x.max(), **kwargs)
return apply_along_axis(self.spln, 0, x)
else:
return self.tfun(x, *self.args, **self.kwargs)
def test_scalar_array(self):
class MinimalSubclass(np.ndarray):
pass
a = np.ones((6, 3)).view(MinimalSubclass)
res = apply_along_axis(np.sum, 0, a)
assert isinstance(res, MinimalSubclass)
assert_array_equal(res, np.array([6, 6, 6]).view(MinimalSubclass))
def test_preserve_subclass(self):
def double(row):
return row * 2
class MyNDArray(np.ndarray):
pass
m = np.array([[0, 1], [2, 3]]).view(MyNDArray)
expected = np.array([[0, 2], [4, 6]]).view(MyNDArray)
result = apply_along_axis(double, 0, m)
assert_(isinstance(result, MyNDArray))
assert_array_equal(result, expected)
result = apply_along_axis(double, 1, m)
assert_(isinstance(result, MyNDArray))
assert_array_equal(result, expected)
def test_preserve_subclass(self):
def double(row):
return row * 2
m = np.matrix([[0, 1], [2, 3]])
result = apply_along_axis(double, 0, m)
assert isinstance(result, np.matrix)
assert_array_equal(
result, np.matrix([[0, 2], [4, 6]])
)
def test_subclass(self):
class MinimalSubclass(np.ndarray):
data = 1
def minimal_function(array):
return array.data
a = np.zeros((6, 3)).view(MinimalSubclass)
assert_array_equal(
apply_along_axis(minimal_function, 0, a), np.array([1, 1, 1])
)
def test_axis_insertion_ma(self):
def f1to2(x):
"""produces an asymmetric non-square matrix from x"""
assert_equal(x.ndim, 1)
res = x[::-1] * x[1:,None]
return np.ma.masked_where(res%5==0, res)
a = np.arange(6*3).reshape((6, 3))
res = apply_along_axis(f1to2, 0, a)
assert_(isinstance(res, np.ma.masked_array))
assert_equal(res.ndim, 3)
assert_array_equal(res[:,:,0].mask, f1to2(a[:,0]).mask)
assert_array_equal(res[:,:,1].mask, f1to2(a[:,1]).mask)
assert_array_equal(res[:,:,2].mask, f1to2(a[:,2]).mask)
def test_axis_insertion(self, cls=np.ndarray):
def f1to2(x):
"""produces an asymmetric non-square matrix from x"""
assert_equal(x.ndim, 1)
return (x[::-1] * x[1:,None]).view(cls)
a2d = np.arange(6*3).reshape((6, 3))
# 2d insertion along first axis
actual = apply_along_axis(f1to2, 0, a2d)
expected = np.stack([
f1to2(a2d[:,i]) for i in range(a2d.shape[1])
], axis=-1).view(cls)
assert_equal(type(actual), type(expected))
assert_equal(actual, expected)
# 2d insertion along last axis
actual = apply_along_axis(f1to2, 1, a2d)
expected = np.stack([
f1to2(a2d[i,:]) for i in range(a2d.shape[0])
], axis=0).view(cls)
assert_equal(type(actual), type(expected))
assert_equal(actual, expected)
# 3d insertion along middle axis
a3d = np.arange(6*5*3).reshape((6, 5, 3))
actual = apply_along_axis(f1to2, 1, a3d)
expected = np.stack([
np.stack([
f1to2(a3d[i,:,j]) for i in range(a3d.shape[0])
], axis=0)
for j in range(a3d.shape[2])
], axis=-1).view(cls)
assert_equal(type(actual), type(expected))
assert_equal(actual, expected)
def test_3d(self):
a = np.arange(27).reshape((3, 3, 3))
assert_array_equal(apply_along_axis(np.sum, 0, a),
[[27, 30, 33], [36, 39, 42], [45, 48, 51]])
def test_simple101(self, level=11):
a = np.ones((10, 101), 'd')
assert_array_equal(
apply_along_axis(len, 0, a), len(a)*np.ones(a.shape[1]))
def test_simple(self):
a = np.ones((20, 10), 'd')
assert_array_equal(
apply_along_axis(len, 0, a), len(a)*np.ones(a.shape[1]))
def test_function_returning_non_scalar(self):
def return_none(x):
pass
a = np.ones(4)
assert_array_equal(apply_along_axis(return_none, 0, a), [None, None, None, None])
def test_function_returning_array(self):
a = np.ones(4)
assert_array_equal(apply_along_axis(np.sqrt, 0, a), [1, 1, 1, 1])
def test_scalar_array(self, cls=np.ndarray):
a = np.ones((6, 3)).view(cls)
res = apply_along_axis(np.sum, 0, a)
assert_(isinstance(res, cls))
assert_array_equal(res, np.array([6, 6, 6]).view(cls))
def test_3d(self):
a = np.arange(27).reshape((3, 3, 3))
assert_array_equal(apply_along_axis(np.sum, 0, a),
[[27, 30, 33], [36, 39, 42], [45, 48, 51]])
def test_simple101(self):
a = np.ones((10, 101), 'd')
assert_array_equal(apply_along_axis(len, 0, a),
len(a) * np.ones(a.shape[1]))
def test_scalar_array(self, cls=np.ndarray):
a = np.ones((6, 3)).view(cls)
res = apply_along_axis(np.sum, 0, a)
assert_(isinstance(res, cls))
assert_array_equal(res, np.array([6, 6, 6]).view(cls))
def test_simple(self):
a = np.ones((20, 10), "d")
assert_array_equal(apply_along_axis(len, 0, a),
len(a) * np.ones(a.shape[1]))
