Numpy arctan2 of multidimensional array

Question

I am trying to shape up some code that was written to take single float values, so it works fine using 1D (and eventually 2D) numpy.arrays as input.

Striped down to a minimal example the function looks like this (no this example is not doing anything useful, but if do_math and do_some_more_math are removed, it will produce exactly the described behavior):

def do_complicated_math(r, g, b):
    rgb = numpy.array([r, g, b])

    # Math! No change in array shape. To run example just comment out.
    rgb = do_math(rgb)

    m_2 = numpy.array([[rgb[0], 0, 0], [0, rgb[1], 0], [0, 0, rgb[2]]])

    # Get additional matrices needed for transformation.
    # These are actually predefined 3x3 float arrays
    m_1 = numpy.ones((3, 3))
    m_3 = numpy.ones((3, 3))

    # Transform the rgb array
    rgb_transformed = m_1.dot(m_2).dot(m_3).dot(rgb)

    # More math! No change in array shape. To run example just comment out.
    rgb_transformed = do_some_more_math(rgb_transformed)

    # Almost done just one more thing...
    return numpy.arctan2(rgb_transformed, rgb_transformed)

# Works fine
do_complicated_math(1, 1, 1)

# Fails
x = numpy.ones(6)
do_complicated_math(x, x, x)

This function works fine as long, as as r, g and b are individual numbers, however, if they are given as numpy.array (e.g., in order to transform multiple rgb values at once) the numpy.arctan2 throws the following exception:

Traceback (most recent call last):
  (...) line 32, in do_complicated_math
    numpy.arctan2(rgb_transformed, rgb_transformed)
AttributeError: 'numpy.ndarray' object has no attribute 'arctan2'

I haven't found any definitive answer as to what this is trying to tell me. arctan2 seems to work fine is used with multidimensional arrays like this:

numpy.arctan2(numpy.ones((3,4,5)), numpy.ones((3,4,5)))

So I assume the problem has to be somewhere in how m_2 is created, or how the multiplications of m_1, m_2, m_3 and rgb get propagated, but I can't seem to figure out just where it breaks.

Answer 1

The problem is that rgb_transformed is no longer a standard numpy array when you pass it to arctan2, it has become an object array:

print rgb_transformed
"""[[array([ 9.,  9.,  9.,  9.,  9.,  9.])
  array([ 9.,  9.,  9.,  9.,  9.,  9.])
  array([ 9.,  9.,  9.,  9.,  9.,  9.])
  array([ 9.,  9.,  9.,  9.,  9.,  9.])
  array([ 9.,  9.,  9.,  9.,  9.,  9.])
  array([ 9.,  9.,  9.,  9.,  9.,  9.])]
 [array([ 9.,  9.,  9.,  9.,  9.,  9.])
  array([ 9.,  9.,  9.,  9.,  9.,  9.])
  array([ 9.,  9.,  9.,  9.,  9.,  9.])
  array([ 9.,  9.,  9.,  9.,  9.,  9.])
  array([ 9.,  9.,  9.,  9.,  9.,  9.])
  array([ 9.,  9.,  9.,  9.,  9.,  9.])]
 [array([ 9.,  9.,  9.,  9.,  9.,  9.])
  array([ 9.,  9.,  9.,  9.,  9.,  9.])
  array([ 9.,  9.,  9.,  9.,  9.,  9.])
  array([ 9.,  9.,  9.,  9.,  9.,  9.])
  array([ 9.,  9.,  9.,  9.,  9.,  9.])
  array([ 9.,  9.,  9.,  9.,  9.,  9.])]]"""
print rgb_transformed.shape
#(3, 6)
print rgb_transformed.dtype
#object

So the problem it simpler than I thought:

This line:

m_2 = numpy.array([[rgb[0], 0, 0], [0, rgb[1], 0], [0, 0, rgb[2]]])
print m_2
#array([[array([ 1.,  1.,  1.,  1.,  1.,  1.]), 0, 0],
#       [0, array([ 1.,  1.,  1.,  1.,  1.,  1.]), 0],
#       [0, 0, array([ 1.,  1.,  1.,  1.,  1.,  1.])]], dtype=object)

Here object arrays are created, propagating through the rest of the code.

EDIT

To get around this issue you likely need to broadcast your arrays slightly differently. Basically change the outer dimension to reflect the changing rgb values. Disclaimer: I don't have a good way to verify the result of this in the context of your question, so treat the output with due care.

import numpy as np

def do_complicated_math(r, g, b):
    rgb = np.array([r, g, b])

    # create a transposed version of the m_2 array
    m_2 = np.zeros((r.size,3,3))
    for ii,ar in enumerate(rgb):
        m_2[:,ii][:,ii][:] = ar
    m_1 = np.ones((3, 3))
    m_3 = np.ones((3, 3))

    rgb_transformed = m_1.dot(m_2).dot(m_3).dot(rgb)

    print rgb_transformed
    return np.arctan2(rgb_transformed, rgb_transformed)

x = np.ones(6)
do_complicated_math(x, x, x)                                                                                                                        

r = np.array([0.2,0.3,0.1])
g = np.array([1.0,1.0,0.2])
b = np.array([0.3,0.3,0.3])
do_complicated_math(r, g, b)

This will work for arrays as input only, but adding handling for single values as input should be trivial.