Skip to content

poor performance of multidimensional masked median #4683

Closed

Description

@juliantaylor

due to the use of apply_over_axis masked median operations gets incredibly slow when computing them along a small dimension of a large array:

In [3]: d = np.random.uniform(size=(200, 200, 50))
In [4]: %timeit np.median(d,axis=2)
10 loops, best of 3: 40.1 ms per loop
In [5]: %timeit np.nanmedian(d,axis=2)
1 loops, best of 3: 1.32 s per loop
In [6]: %timeit np.ma.median(d,axis=2)
1 loops, best of 3: 4.89 s per loop

a way out is using bottlenecks nanmedian which is on par with our non-masked performance (and better for the slow axis) but I think we should think of ways to improve this in numpy.

I can think of following improvements for median:

  • implement apply_over_axis in C, should be relatively easy but will only give us ~50% boost at best. it should still be worthwhile as there are likely more users of that function out there.
  • extend partition to broadcast the index, then one could compute the right index from the mask and do everything in one partition call. Probably worthwhile but it will make the partition interface more complicated and would probably imply rewritting our pretty ugly sorting code.
  • probably the simplest would be add a check on the size of the axis and use sort instead of partition as currently done in ma.median but then instead of selecting via apply_along_axis, one selects the elements to average (or not) via fancy indexing (always select the middle two and if the non-masked count is odd duplicate one of the elements so mean() does nothing)
  • add a MaskedPartition C-Api function
  • see if we can integrate bottleneck directly somehow

any more options I'm might be overlooking?

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions