Broadcast across pandas MultiIndex level even if index values happen to agree

Viewed 63

This one has me stumped. I have two pd.Series s and t as follows:

Common  Level s
Foo     a          1
        b          2
Name: s, dtype: int64
Common  Level t
Foo     A          10
        B          20
Name: t, dtype: int64

pandas lets me add these and broadcasts across the common level 'Common'

Input:

s + t

Output:

Common  Level s  Level t
Foo     a        A          11
                 B          21
        b        A          12
                 B          22
dtype: int64

Consider now another pd.Series u where the index labels happen to agree with those of s

Common  Level u
Foo     a          100
        b          200
Name: u, dtype: int64

In other words, we have (s.index.values == u.index.values).all() returns True. Because of this, pandas no longer broadcasts

Input:

s + u

Output:

Common  Level s
Foo     a          101
        b          202
dtype: int64

even though s.index.names and u.index.names disagree.

Lastly, if the order is changed but not the labels, such as for v:

Common  Level v
Foo     b          1000
        a          2000
Name: v, dtype: int64

so that s.index.values and v.index.values don't agree outright, then broadcasting happens.

Input:

s + v

Output:

Common  Level s  Level v
Foo     a        b          1001
                 a          2001
        b        b          1002
                 a          2002
dtype: int64

My question: How can I add s and u such that pandas still broadcasts? (For my particular application, I am actually interested in elementwise-and s & u, not the sum s + u.)


Code

s = pd.Series([1, 2],
              index=pd.MultiIndex.from_tuples(
                  [('Foo', 'a'), ('Foo', 'b')],
                  names=['Common', 'Level s']), name='s')

t = pd.Series([10, 20],
              index=pd.MultiIndex.from_tuples(
                  [('Foo', 'A'), ('Foo', 'B')],
                  names=['Common', 'Level t']), name='t')

u = pd.Series([100, 200],
              index=pd.MultiIndex.from_tuples(
                  [('Foo', 'a'), ('Foo', 'b')],
                  names=['Common', 'Level u']), name='u')

v = pd.Series([1000, 2000],
              index=pd.MultiIndex.from_tuples(
                  [('Foo', 'b'), ('Foo', 'a')],
                  names=['Common', 'Level v']), name='v')
1 Answers

The method that is called to normalise indexes of Series and DataFrames for broadcasting is align. This method is called internally to get the new indexes, we can see this with:

left, right = s.align(t, join='outer', level=None, copy=False)

# left
Common  Level s  Level t
Foo     a        A          1
                 B          1
        b        A          2
                 B          2
Name: s, dtype: int64
# right
Common  Level s  Level t
Foo     a        A          10
                 B          20
        b        A          10
                 B          20
Name: t, dtype: int64

Notice that this call will produce the "non-broadcasted" values when the indexes are equal since outer join produces a single level:

left, right = s.align(u, join='outer', level=None, copy=False)

# left
Common  Level s
Foo     a          1
        b          2
Name: s, dtype: int64

# right
Common  Level u
Foo     a          100
        b          200
Name: u, dtype: int64

If we want to force the levels to generate we can use the branch from _align_series for when indexes are non-equal:

join_index, lidx, ridx = self.index.join(
    other.index, how=join, level=level, return_indexers=True
) 
left = self._reindex_indexer(join_index, lidx, copy)
right = other._reindex_indexer(join_index, ridx, copy)

We can use Index.join and _reindex_indexer to create the aligned Series:

join_index, lidx, ridx = s.index.join(
    u.index, how='outer', level=None, return_indexers=True
)
left = s._reindex_indexer(join_index, lidx, copy=False)
right = u._reindex_indexer(join_index, ridx, copy=False)

*Note: we are using a private method as there is no equivalent reindexer in the public API.

Now that we have aligned Series we can just do:

left + right

To get the results:

Common  Level s  Level u
Foo     a        a          101
                 b          201
        b        a          102
                 b          202
dtype: int64

If we wanted to avoid using the private method, we could also iloc to select values using index locations then set_axis to overwrite with the join index with appropriate labels:

join_index, lidx, ridx = s.index.join(
    u.index, how='outer', level=None, return_indexers=True
)
left = s.iloc[lidx].set_axis(join_index)
right = u.iloc[ridx].set_axis(join_index)
left + right
Common  Level s  Level u
Foo     a        a          101
                 b          201
        b        a          102
                 b          202
dtype: int64
Related