Figuring out the structure of a shape tuple

Question

I have some code that uses a multidimensional look-up table (LUT) to do interpolation. The typical application is a color space conversion, where 3D inputs (RGB) are converted to 4D (CMYK), but the code is rather general. The look-up table is a numpy array, which will generally have a shape like (4, 17, 17, 17). Here, 4 is the number of output dimensions, o, 17 is the number of grid points along each axis, g, which must be the same for all, and 3 (the number of 17s in the shape tuple) is the number of input dimensions, i. There is also lut_axis, 0 in this case, that indicates the position of o in the shape tuple. There is an additional requirement for g to be either a power of two (1, 2, 4, 8, 16...) or a power of two plus one (2, 3, 5, 9, 17...).

I am currently checking that the lut and lut_axis I receive as inputs are compatible with the above requirements. But in most cases there is only one value of lut_axis that would yield a correct result, so I think I could improve the code by only using lut_axis when the shape of the LUT doesn't give away what it should be, and else ignoring it.

My first attempt was this logical spaghetti code, that gets things done, but which I will be unable to figure out a month from today :

OPTION 1

from operator import and_

def process_lut_shape(lut_shape, default_axis=0) :
    i = len(lut_shape) - 1
    values = sorted(set(lut_shape))
    counts = [lut_shape.count(j) for j in values]
    count_of_values = len(values)
    figured_lut_out = False
    if count_of_values > 2 or i < 1 :
        figured_lut_out = -1
    elif count_of_values == 2 :
        is_pow_of_two = [j - 2**(j.bit_length() - 1) in (0, 1) for j in values]
        other_is_one = [counts[(j + 1) % 2] for j in xrange(2)]
        could_be_g = map(and_, is_pow_of_two, other_is_one)
        candidates = could_be_g.count(True)
        if candidates == 0 :
            figured_lut_out = -1
        elif candidates == 1 :
            idx = could_be_g.index(True)
            g = values[idx]
            o = values[(idx + 1) % 2]
            lut_axis = lut_shape.index(o)
            figured_out_lut = True
    if figured_lut_out == False :
        o = lut_shape[default_axis]
        g = (o if count_of_values == 1 else
             values[(values.index(o) + 1) % 2])
        lut_axis = default_axis
    elif figured_out_lut == -1 :
        msg = 'Array of shape {0} cannot be a lut.'
        msg = msg.format(str(lut_shape))
        raise ValueError(msg)
    return lut_axis, i, o, g

lut_axis, i, o, g = process_lut_shape(lut_shape, default_axis)

I think that, since the LUTs that are reasonable to expect will never have more than few dimensions, typically 4 or 5, it may be much clearer to brute-force my way through this:

OPTION 2

def check_lut_shape(lut_shape, lut_axis) :
    i = len(lut_shape) - 1
    if i < 1 :
        return False, None
    o = lut_shape[lut_axis]
    g = lut_shape[(lut_axis + 1) % (i + 1)]
    if g - 2**(g.bit_length() - 1) not in (0, 1) :
        return False, None
    shape = (g,) * lut_axis + (o,) + (g,) * (i - lut_axis)
    if any(lut_shape[j] != shape[j] for j in xrange(i + 1)) :
        return False, None
    return True, (i, o, g)

def find_lut_axis(lut_shape, default_lut_axis=0) :
    func = lambda j : check_lut_shape(lut_shape, j)[0]
    could_be_lut_axis = map(func, xrange(len(lut_shape)))
    candidates = could_be_lut_axis.count(True)
    if candidates == 0 :
        msg = 'Array of shape {0} cannot be a lut.'
        msg = msg.format(str(lut_shape))
        raise ValueError(msg)
    elif candidates == 1 :
        return could_be_lut_axis.index(True)
    # I don't think the following can ever happen...
    elif could_be_lut_axis[default_lut_axis] != True :
        msg = 'Cannot determine lut_axis for array of shape {0}.'
        msg = msg.format(str(lut_shape))
        raise ValueError(msg)
    return default_lut_axis

lut_axis = find_lut_axis(lut_shape, default_lut_axis)
i, o, g = check_lut_shape(lut_shape, lut_axis)[1]

There's one call too many to check_lut_shape to fetch the values of i, o and g, but since the code is already not an example of efficiency, the focus should be on readability, and I find it is easier to understand like this.

I would feel more comfortable if I could put together a more elegant version of option 1, but since I don't see how to go about that, I am more and more leaning to something along the lines of option 2. Any comments on which way to go, or improvements over the code above are more than welcome.

Answer 1

Firstly, you could use some tests. Here are a few I cooked up:

assert_equals( process_lut_shape( (4, 17, 17, 17) ) , (0, 3, 4, 17) )
assert_equals( process_lut_shape( (14, 16, 16) ) , (0, 2, 14, 16) )
assert_raises( ValueError, process_lut_shape, (14, 15, 15) )
assert_equals( process_lut_shape( (1, 16, 16) ) , (0, 2, 1, 16) )
assert_equals( process_lut_shape( (16, 16, 3) ) , (2, 2, 3, 16) )
assert_raises( ValueError, process_lut_shape, (14,) )
assert_raises( ValueError, process_lut_shape, () )

If you try these tests, you'll find your function doesn't actually work. It uses figured_lut_out and figured_out_lut interchangeably.

def process_lut_shape(lut_shape, default_axis=0) :
    i = len(lut_shape) - 1
    values = sorted(set(lut_shape))

Why sort it?

    counts = [lut_shape.count(j) for j in values]
    count_of_values = len(values)
    figured_lut_out = False

Flag variables are delayed gotos. They make code hard to read. Instead of setting a flag to do something later, just do it.

    if count_of_values > 2 or i < 1 :

Using i here is confusing, I'd use len(lut_shape)

        figured_lut_out = -1

Using -1 to create a trinary value is just confusing.

    elif count_of_values == 2 :
        is_pow_of_two = [j - 2**(j.bit_length() - 1) in (0, 1) for j in values]

Somewhat confusingly named because its not really just checking for a power of 2

        other_is_one = [counts[(j + 1) % 2] for j in xrange(2)]
        could_be_g = map(and_, is_pow_of_two, other_is_one)

I think this is too clever, simply writing it out without using loops would probably be clearer

        candidates = could_be_g.count(True)
        if candidates == 0 :
            figured_lut_out = -1
        elif candidates == 1 :
            idx = could_be_g.index(True)
            g = values[idx]
            o = values[(idx + 1) % 2]
            lut_axis = lut_shape.index(o)
            figured_out_lut = True

You already have all the data here, just return it.

    if figured_lut_out == False :
        o = lut_shape[default_axis]
        g = (o if count_of_values == 1 else
             values[(values.index(o) + 1) % 2])

That feels very awkward.

        lut_axis = default_axis
    elif figured_out_lut == -1 :
        msg = 'Array of shape {0} cannot be a lut.'
        msg = msg.format(str(lut.shape))

There is no reason to use str when doing format. raise ValueError(msg)

def check_lut_shape(lut_shape, lut_axis) :
    i = len(lut_shape) - 1
    if i < 1 :

Why check i? Isn't len(lut_shape) more clear? return False, None

You don't need to return two values, just return None.

    o = lut_shape[lut_axis]
    g = lut_shape[(lut_axis + 1) % (i + 1)]

Again use len not i

    if g - 2**(g.bit_length() - 1) not in (0, 1) :
        return False, None
    shape = (g,) * lut_axis + (o,) + (g,) * (i - lut_axis)
    if any(lut_shape[j] != shape[j] for j in xrange(i + 1)) :

just use if lut_shape != shape. However, reconstructing the shape to compare it feels a little awkward.

        return False, None
    return True, (i, o, g)

def find_lut_axis(lut_shape, default_lut_axis=0) :
    func = lambda j : check_lut_shape(lut_shape, j)[0]
    could_be_lut_axis = map(func, xrange(len(lut_shape)))

Use could_be_lut_axis = [check_lut_shape(lut_shape, j)[0] for j in xrange(len(lut_shape))]

    candidates = could_be_lut_axis.count(True)
    if candidates == 0 :
        msg = 'Array of shape {0} cannot be a lut.'
        msg = msg.format(str(lut_shape))
        raise ValueError(msg)
    elif candidates == 1 :
        return could_be_lut_axis.index(True)
    # I don't think the following can ever happen...
    elif could_be_lut_axis[default_lut_axis] != True :
        msg = 'Cannot determine lut_axis for array of shape {0}.'
        msg = msg.format(str(lut_shape))
        raise ValueError(msg)
    return default_lut_axis

Here's my approach:

def process_lut_shape(lut_shape, default_axis=0):

    if len(lut_shape) < 2:
        # anything less then two long cannot be a shape
        raise NoneLutShapeError(lut_shape)
    elif len(lut_shape) == 2:
        # given two, either could be the lut_axis
        # try the default if it can work
        # otherwise try the other one
        if is_almost_power_of_two(lut_shape[default_axis]):
            lut_axis = default_axis
        else:
            lut_axis = 1 - default_axis
    else: # len(lut_shape) >= 3
        counts = Counter(lut_shape)
        if len(counts) == 1:
            # since all axis are the same, we can't distinguish between them
            lut_axis = default_axis
        elif len(counts) > 2:
            # can't have more then 2 distinct values in the shape
            raise NoneLutShapeError(lut_shape)
        else:
            # the least common count, with three elements should be the one
            lut_value, count = counts.most_common()[-1]
            if count != 1:
                raise NoneLutShapeError(lut_shape)
            lut_axis = lut_shape.index(lut_value)

    g = lut_shape[(lut_axis + 1) % len(lut_shape)]
    if not is_almost_power_of_two(g):
        raise NoneLutShapeError(g)
    return lut_axis, len(lut_shape) - 1, lut_shape[lut_axis], g

However, I'd really recommend you not write this function at all. This function is guessing. As the zen of python says:

In the face of ambiguity, refuse the temptation to guess.

Given (16, 16, 16) you have no way of knowing which axis was intended to be the lut_axis. So I really really recommend not trying to guess. You'll just do the wrong thing oddly in certain circumstances. Even trying to figure it out when its not ambigious is problematic because it'll still become ambigious at times and your code will work differently. As the Zen of Python also says:

Explicit is better than implicit.

Make your users explicitly pass the lut_axis. You can't do the correct thing for all inputs, so don't create false expectations by doing the correct thing for some inputs.