Python datarray.DataArray类代码示例

def test_swapaxes():
    n_arr = np.random.randn(2,4,3)
    a = DataArray(n_arr, 'xyz')
    b = a.swapaxes('x', 'z')
    c = DataArray(n_arr.transpose(2,1,0), 'zyx')
    yield nt.assert_true, (c==b).all(), 'data not equal in swapaxes test'
    for ax1, ax2 in zip(b.axes, c.axes):
        yield nt.assert_true, ax1==ax2, 'axes not equal in swapaxes test'

def test_squeeze():
    "Test squeeze method"
    d = DataArray(np.random.randn(3, 2, 9), "xyz")
    d2 = d[None, :, None, :, :, None]
    nt.assert_true(d2.shape == (1, 3, 1, 2, 9, 1), "newaxis slicing failed")
    d3 = d.squeeze()
    nt.assert_true(d3.shape == d.shape, "squeezing length-1 dimensions failed")
    nt.assert_true(d3.names == d.names, "Axes got lost in squeeze")

def test_swapaxes():
    n_arr = np.random.randn(2, 4, 3)
    a = DataArray(n_arr, "xyz")
    b = a.swapaxes("x", "z")
    c = DataArray(n_arr.transpose(2, 1, 0), "zyx")
    nt.assert_true((c == b).all(), "data not equal in swapaxes test")
    for ax1, ax2 in zip(b.axes, c.axes):
        nt.assert_true(ax1 == ax2, "axes not equal in swapaxes test")

def test_squeeze():
    "Test squeeze method"
    d = DataArray(np.random.randn(3,2,9), 'xyz')
    d2 = d[None,:,None,:,:,None]
    yield nt.assert_true, d2.shape == (1,3,1,2,9,1), 'newaxis slicing failed'
    d3 = d.squeeze()
    yield nt.assert_true, d3.shape == d.shape, \
          'squeezing length-1 dimensions failed'
    yield nt.assert_true, d3.labels == d.labels, 'Axes got lost in squeeze'

def test_bug44():
    "Bug 44"
    # In instances where axis=None, the operation runs
    # on the flattened array. Here it makes sense to return
    # the op on the underlying np.ndarray.
    A = [[1,2,3],[4,5,6]]
    x = DataArray(A, 'xy').std()
    y = np.std(A)
    nt.assert_equal( x.sum(), y.sum() )

def test_reshape_corners():
    "Test some corner cases for reshape"
    d = DataArray(np.random.randn(3, 4, 5), "xyz")
    d2 = d.reshape(-1)
    nt.assert_true(d2.shape == (60,), "Flattened shape wrong")
    nt.assert_true(type(d2) is np.ndarray, "Flattened type wrong")

    d2 = d.reshape(60)
    nt.assert_true(d2.shape == (60,), "Flattened shape wrong")
    nt.assert_true(type(d2) is np.ndarray, "Flattened type wrong")

def test_reshape_corners():
    "Test some corner cases for reshape"
    d = DataArray(np.random.randn(3,4,5), 'xyz')
    d2 = d.reshape(-1)
    yield nt.assert_true, d2.shape == (60,), 'Flattened shape wrong'
    yield nt.assert_true, type(d2) is np.ndarray, 'Flattened type wrong'

    d2 = d.reshape(60)
    yield nt.assert_true, d2.shape == (60,), 'Flattened shape wrong'
    yield nt.assert_true, type(d2) is np.ndarray, 'Flattened type wrong'

    def test_axes_name_collision(self):
        "Test .axes object for attribute collisions with axis names"
        A = DataArray(np.arange(6).reshape([1, 2, 3]), ("_arr", "_axes", "_namemap"))
        nt.assert_true(A.axes[0] is A.axes("_arr") is A.axes._arr)
        nt.assert_true(A.axes[1] is A.axes("_axes") is A.axes._axes)
        nt.assert_true(A.axes[2] is A.axes("_namemap") is A.axes._namemap)

        # Try to invoke some methods that use these attributes internally
        B = A[np.newaxis, ...]
        nt.assert_equal(B.shape, (1, 1, 2, 3))
        nt.assert_true(np.all(A + A == 2 * A))

def test_flatten_and_ravel():
    "Test the functionality of ravel() and flatten() methods"
    d = DataArray(np.arange(20).reshape(4, 5), "xy")
    df = d.flatten()
    nt.assert_true(type(df) is np.ndarray, "Type error in flatten")
    nt.assert_true(df.shape == (20,), "Wrong shape in flatten")
    df[:4] = 0
    nt.assert_false((d[0, :4] == 0).all(), "Copy not made in flatten")

    dr = d.ravel()
    nt.assert_true(type(dr) is np.ndarray, "Type error in ravel")
    nt.assert_true(dr.shape == (20,), "Wrong shape in ravel")
    dr[:4] = 0
    nt.assert_true((d[0, :4] == 0).all(), "View not made in ravel")

def test_bug34():
    "Bug 34: datetime.date ticks not handled by datarray_to_string"
    from datarray.print_grid import datarray_to_string
    from datetime import date as D
    A = DataArray([[1,2],[3,4]], [('row', ('a', D(2010,1,1))),('col', 'cd')])
    exp_out = """row       col                
--------- -------------------
          c         d        
a                 1         2
2010-01-0         3         4"""
    nt.assert_equal(datarray_to_string(A), exp_out)
    # Output for unsigned integers
    B = A.astype(np.uint32)
    nt.assert_equal(datarray_to_string(B), exp_out)

def test_flatten_and_ravel():
    "Test the functionality of ravel() and flatten() methods"
    d = DataArray(np.arange(20).reshape(4,5), 'xy')
    df = d.flatten()
    yield nt.assert_true, type(df) is np.ndarray, 'Type error in flatten'
    yield nt.assert_true, df.shape == (20,), 'Wrong shape in flatten'
    df[:4] = 0
    yield nt.assert_false, (d[0,:4] == 0).all(), 'Copy not made in flatten'

    dr = d.ravel()
    yield nt.assert_true, type(dr) is np.ndarray, 'Type error in ravel'
    yield nt.assert_true, dr.shape == (20,), 'Wrong shape in ravel'
    dr[:4] = 0
    yield nt.assert_true, (d[0,:4] == 0).all(), 'View not made in ravel'

def test_reshape():
    d = DataArray(np.random.randn(3, 4, 5), "xyz")
    new_shape = (1, 3, 1, 4, 5)
    # Test padding the shape
    d2 = d.reshape(new_shape)
    new_labels = (None, "x", None, "y", "z")
    nt.assert_true(d2.names == new_labels, "Array with inserted dimensions has wrong labels")
    nt.assert_true(d2.shape == new_shape, "New shape wrong")

    # Test trimming the shape
    d3 = d2.reshape(d.shape)
    nt.assert_true(d3.names == d.names, "Array with removed dimensions has wrong labels")
    nt.assert_true(d3.shape == d.shape, "New shape wrong")

    # Test a combo of padding and trimming
    d4 = d2.reshape(3, 4, 1, 5, 1)
    new_labels = ("x", "y", None, "z", None)
    nt.assert_true(d4.names == new_labels, "Array with inserted and removed dimensions has wrong labels")
    nt.assert_true(d4.shape == (3, 4, 1, 5, 1), "New shape wrong")

def test_reshape():
    d = DataArray(np.random.randn(3,4,5), 'xyz')
    new_shape = (1,3,1,4,5)
    # Test padding the shape
    d2 = d.reshape(new_shape)
    new_labels = (None, 'x', None, 'y', 'z')
    yield nt.assert_true, d2.labels == new_labels, \
          'Array with inserted dimensions has wrong labels'
    yield nt.assert_true, d2.shape == new_shape, 'New shape wrong'

    # Test trimming the shape
    d3 = d2.reshape(d.shape)
    yield nt.assert_true, d3.labels == d.labels, \
          'Array with removed dimensions has wrong labels'
    yield nt.assert_true, d3.shape == d.shape, 'New shape wrong'

    # Test a combo of padding and trimming
    d4 = d2.reshape(3,4,1,5,1)
    new_labels = ('x', 'y', None, 'z', None)
    yield nt.assert_true, d4.labels == new_labels, \
          'Array with inserted and removed dimensions has wrong labels'
    yield nt.assert_true, d4.shape == (3,4,1,5,1), 'New shape wrong'

class TestAxesManager:
    def setUp(self):
        self.axes_spec = ('date', ('stocks', ('aapl', 'ibm', 'goog', 'msft')), 'metric')
        self.A = DataArray(np.random.randn(200, 4, 10), axes=self.axes_spec)

    def test_axes_name_collision(self):
        "Test .axes object for attribute collisions with axis names"
        A = DataArray(np.arange(6).reshape([1,2,3]), 
                ('_arr', '_axes', '_namemap'))
        nt.assert_true(A.axes[0] is A.axes('_arr') is A.axes._arr)
        nt.assert_true(A.axes[1] is A.axes('_axes') is A.axes._axes)
        nt.assert_true(A.axes[2] is A.axes('_namemap') is A.axes._namemap)
        
        # Try to invoke some methods that use these attributes internally
        B = A[np.newaxis, ...]
        nt.assert_equal(B.shape, (1,1,2,3))
        nt.assert_true(np.all(A + A == 2*A))

    def test_axes_numeric_access(self):
        for i,spec in enumerate(self.axes_spec):
            try:
                name,labels = spec
            except ValueError:
                name,labels = spec,None
            nt.assert_true(self.A.axes[i] == Axis(name=name, index=i,
                parent_arr=self.A, labels=labels))

    def test_axes_attribute_access(self):
        for spec in self.axes_spec:
            try:
                name,labels = spec
            except ValueError:
                name,labels = spec,None
            nt.assert_true(getattr(self.A.axes, name) is self.A.axes(name))

    def test_equality(self):
        B = DataArray(np.random.randn(200, 4, 10), axes=self.axes_spec)
        nt.assert_true(self.A.axes == B.axes)
        # What if axes differ by labels only?
        D = DataArray(np.random.randn(200, 4, 10), axes=('date', 'stocks', 'metric')) 
        nt.assert_false(self.A.axes == D.axes)

 def setUp(self):
     self.axes_spec = ("date", ("stocks", ("aapl", "ibm", "goog", "msft")), "metric")
     self.A = DataArray(np.random.randn(200, 4, 10), axes=self.axes_spec)