在堆叠的条形图中定义每个条形的底部

Question

I have a two-dimensional array absolute_heights of shape (2, 6) . I'd like to define a new two-dimensional array bottoms of shape (2, 6) that holds 0 at each position i unless

1) The sign of absolute_heights[0, i] - absolute_heights[1, i] matches that of absolute_heights[0, i] , in which case bottoms[0, i] should be set to absolute_heights[1, i] .

2) #1 is false, in which case bottoms[1, i] should be set to absolute_heights[0, i] .

Below is a for loop that achieves this:

def _get_bottoms(absolute_heights):
    """Define the bottom of each bar in a stacked bar graph.

    Parameters
    ----------
    absolute_heights : np.array
      The absolute height of each bar.  Stacking of the bars is along
      the first axis of this array.

    Returns
    -------
    bottoms : np.array
      The absolute height of the bar in each stack that is closest to
      zero.

    """
    bottoms = np.zeros((2, 6))
    for i, diff in enumerate(absolute_heights[0, :] - absolute_heights[1, :]):
        if np.sign(diff) == np.sign(absolute_heights[0, i]):
            bottoms[0, i] = absolute_heights[1, i]
        else:
            bottoms[1, i] = absolute_heights[0, i]
    return bottoms

Is there a more efficient way of doing this in numpy ?

Answer 1

You could use boolean indexing to avoid the for loop:

def _get_bottoms(absolute_heights):
    bottoms = np.zeros((2,6))
    diff = absolute_heights[0, :] - absolute_heights[1, :]
    i = np.sign(diff) == np.sign(absolute_heights[0, :])
    bottoms[0, i] = absolute_heights[1, i]
    bottoms[1, ~i] = absolute_heights[0, ~i]
    return bottoms

In this function i is a boolean array indicating whether the signs match (essentially your if statement). Inverting the boolean values with ~i gives the array for the else statement.

Answer 2

Another solution using np.where

b = np.where(np.sign(ah[0,:]) == np.sign(ah[0,:] - ah[1,:]), ah[1,:], 0.)
b2 = np.where(np.sign(ah[0,:]) != np.sign(ah[0,:] - ah[1,:]), ah[0,:], 0.)
np.vstack((b2,b))

Unlikely to be substantially faster than the one mentioned above, possibly - slightly more readable.

np.where takes an array of bool conditions and then uses the first argument ( ah[1,:] above) if the condition is True else second argument( ah[0,:] above).