4
\$\begingroup\$

I often perform a lot of reshape(), transpose() or squeeze() operations on numpy arrays for numeric calculations. Often, this obfuscates the code. The initial shape as well as the resulting shape is unknown.

Prior solutions involved:

  • Comments (which may outdate).
  • Tiny functions with somewhat helpful names from_tbf_to_f1tb().

Since experiments change quickly, so do the dimensions. I always enjoy using numpy.einsum(), so I wanted a syntax similar to that, but just for reshapes.

Example:

reshape(A, 't b f -> f 1 b*t')

The test case shows further examples.

How can I simplify my code and make it more readable?

import sys
import unittest
import numpy as np
import nt.testing as tc
from nt.utils.numpy_utils import reshape


def _normalize(op):
    op = op.replace(',', '')
    op = op.replace(' ', '')
    op = ' '.join(c for c in op)
    op = op.replace(' * ', '*')
    op = op.replace('- >', '->')
    return op


def _only_reshape(array, source, target):
    source, target = source.split(), target.replace(' * ', '*').split()
    input_shape = {key: array.shape[index] for index, key in enumerate(source)}

    output_shape = []
    for t in target:
        product = 1
        if not t == '1':
            t = t.split('*')
            for t_ in t:
                product *= input_shape[t_]
        output_shape.append(product)

    return array.reshape(output_shape)


def reshape(array, operation):
    """ This is an experimental version of a generalized reshape.

    See test cases for examples.
    """
    operation = _normalize(operation)

    if '*' in operation.split('->')[0]:
        raise NotImplementedError(
            'Unflatten operation not supported by design. '
            'Actual values for dimensions are not available to this function.'
        )

    # Initial squeeze
    squeeze_operation = operation.split('->')[0].split()
    for axis, op in reversed(list(enumerate(squeeze_operation))):
        if op == '1':
            array = np.squeeze(array, axis=axis)

    # Transpose
    transposition_operation = operation.replace('1', ' ').replace('*', ' ')
    try:
        array = np.einsum(transposition_operation, array)
    except ValueError as e:
        msg = 'op: {}, shape: {}'.format(transposition_operation,
                                         np.shape(array))
        if len(e.args) == 1:
            e.args = (e.args[0]+'\n\n'+msg,)
        else:
            print(msg)
        raise

    # Final reshape
    source = transposition_operation.split('->')[-1]
    target = operation.split('->')[-1]

    return _only_reshape(array, source, target)

Testcase:

T, B, F = 40, 6, 51
A = np.random.uniform(size=(T, B, F))
A2 = np.random.uniform(size=(T, 1, B, F))
A3 = np.random.uniform(size=(T*B*F,))
A4 = np.random.uniform(size=(T, 1, 1, B, 1, F))


class TestReshape(unittest.TestCase):
    def test_noop_comma(self):
        result = reshape(A, 'T,B,F->T,B,F')
        tc.assert_equal(result.shape, (T, B, F))
        tc.assert_equal(result, A)

    def test_noop_space(self):
        result = reshape(A, 'T B F->T B F')
        tc.assert_equal(result.shape, (T, B, F))
        tc.assert_equal(result, A)

    def test_noop_mixed(self):
        result = reshape(A, 'tbf->t, b f')
        tc.assert_equal(result.shape, (T, B, F))
        tc.assert_equal(result, A)

    def test_transpose_comma(self):
        result = reshape(A, 'T,B,F->F,T,B')
        tc.assert_equal(result.shape, (F, T, B))
        tc.assert_equal(result, A.transpose(2, 0, 1))

    def test_transpose_mixed(self):
        result = reshape(A, 't, b, f -> f t b')
        tc.assert_equal(result.shape, (F, T, B))
        tc.assert_equal(result, A.transpose(2, 0, 1))

    def test_broadcast_axis_0(self):
        result = reshape(A, 'T,B,F->1,T,B,F')
        tc.assert_equal(result.shape, (1, T, B, F))
        tc.assert_equal(result, A[None, ...])

    def test_broadcast_axis_2(self):
        result = reshape(A, 'T,B,F->T,B,1,F')
        tc.assert_equal(result.shape, (T, B, 1, F))
        tc.assert_equal(result, A[..., None, :])

    def test_broadcast_axis_3(self):
        result = reshape(A, 'T,B,F->T,B,F,1')
        tc.assert_equal(result.shape, (T, B, F, 1))
        tc.assert_equal(result, A[..., None])

    def test_reshape_comma(self):
        result = reshape(A, 'T,B,F->T,B*F')
        tc.assert_equal(result.shape, (T, B*F))
        tc.assert_equal(result, A.reshape(T, B*F))

    def test_reshape_comma_unflatten(self):
        with tc.assert_raises(NotImplementedError):
            reshape(A3, 't*b*f->t, b, f')

    def test_reshape_comma_unflatten_and_transpose_and_flatten(self):
        with tc.assert_raises(NotImplementedError):
            reshape(A3, 't*b*f->f, t*b')

    def test_reshape_comma_flat(self):
        result = reshape(A, 'T,B,F->T*B*F')
        tc.assert_equal(result.shape, (T*B*F,))
        tc.assert_equal(result, A.ravel())

    def test_reshape_comma_with_singleton_input(self):
        result = reshape(A2, 'T, 1, B, F -> T*B*F')
        tc.assert_equal(result.shape, (T*B*F,))
        tc.assert_equal(result, A2.ravel())

    def test_reshape_comma_with_a_lot_of_singleton_inputs(self):
        result = reshape(A4, 'T, 1, 1, B, 1, F -> T*B*F')
        tc.assert_equal(result.shape, (T*B*F,))
        tc.assert_equal(result, A4.ravel())

    def test_reshape_and_broadcast(self):
        tc.assert_equal(reshape(A, 'T,B,F->T,1,B*F').shape, (T, 1, B*F))
        tc.assert_equal(reshape(A, 'T,B,F->T,1,B*F').ravel(), A.ravel())

    def test_reshape_and_broadcast_many(self):
        result = reshape(A, 'T,B,F->1,T,1,B*F,1')
        tc.assert_equal(result.shape, (1, T, 1, B*F, 1))

    def test_swap_and_reshape(self):
        result = reshape(A, 'T,B,F->T,F*B')
        tc.assert_equal(result.shape, (T, F * B))
        tc.assert_equal(result, A.swapaxes(-1, -2).reshape(T, F * B))

    def test_transpose_and_reshape(self):
        result = reshape(A, 'T,B,F->F,B*T')
        tc.assert_equal(result.shape, (F, B*T))
        tc.assert_equal(result, A.transpose(2, 1, 0).reshape(F, B*T))

    def test_all_comma(self):
        tc.assert_equal(reshape(A, 'T,B,F->F,1,B*T').shape, (F, 1, B*T))

    def test_all_space(self):
        tc.assert_equal(reshape(A, 't b f -> f1b*t').shape, (F, 1, B*T))


suite = unittest.TestLoader().loadTestsFromTestCase(TestReshape)
_ = unittest.TextTestRunner(verbosity=1,stream=sys.stderr).run(suite)

If you prefer reading it as a gist: gist.github.com

\$\endgroup\$
  • \$\begingroup\$ A while back I contributed a patch that fixed the handling of ellipsis in np.einsum. In course of writing that I reverse engineered einsum, focusing on the part the decode the index string. The code is available on my github at github.com/hpaulj/numpy-einsum/blob/master/einsum_py.py I haven't examined your code yet. I don't know if this question fits code-review or not. There aren't a lot of numpy eyes here (compared to StackOverflow). But you might get good general Python comments. An alternative might be a numpy developer's forum. \$\endgroup\$ – hpaulj Dec 2 '16 at 19:50
  • \$\begingroup\$ @hpaulj I am currently working through the code. I did not expect it to be that long. Once I am through, I will reconsider my code and may think about cross-posting it to a NumPy dev forum. \$\endgroup\$ – Lukas Dec 5 '16 at 9:54

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.