How can I efficiently "stretch" present values in an array over absent ones

Question

Where 'absent' can mean either nan or np.masked, whichever is easiest to implement this with.

For instance:

>>> from numpy import nan
>>> do_it([1, nan, nan, 2, nan, 3, nan, nan, 4, 3, nan, 2, nan])
array([1, 1, 1, 2, 2, 3, 3, 3, 4, 3, 3, 2, 2])
# each nan is replaced with the first non-nan value before it
>>> do_it([nan, nan, 2, nan])
array([nan, nan, 2, 2])
# don't care too much about the outcome here, but this seems sensible

I can see how you'd do this with a for loop:

def do_it(a):
    res = []
    last_val = nan
    for item in a:
        if not np.isnan(item):
            last_val = item
        res.append(last_val)
    return np.asarray(res)

Is there a faster way to vectorize it?

score 1 · Answer 1 · edited Dec 14 '16 at 19:03

1

cumsumming over an array of flags provides a good way to determine which numbers to write over the NaNs:

def do_it(x):
    x = np.asarray(x)

    is_valid = ~np.isnan(x)
    is_valid[0] = True

    valid_elems = x[is_valid]
    replacement_indices = is_valid.cumsum() - 1
    return valid_elems[replacement_indices]

edited Dec 14 '16 at 19:03

Eric

95,302
53
242
374

answered Dec 14 '16 at 18:11

user2357112

260,549
28
431
505

Hmm, this doesn't work if x is 2d, but I guess that wasn't what I asked for – Eric Dec 14 '16 at 19:04
@Eric: Yeah, I don't know what you'd even want for a 2D input. – user2357112 Dec 14 '16 at 19:06
I'd want it to treat each row independantly as if it were 1d – Eric Dec 14 '16 at 19:17

Benjamin · Answer 2 · 2016-12-14T19:01:32.137

1

Assuming there are no zeros in your data (in order to use numpy.nan_to_num):

b = numpy.maximum.accumulate(numpy.nan_to_num(a))
>>> array([ 1.,  1.,  1.,  2.,  2.,  3.,  3.,  3.,  4.,  4.])
mask = numpy.isnan(a)
a[mask] = b[mask]
>>> array([ 1.,  1.,  1.,  2.,  2.,  3.,  3.,  3.,  4.,  3.])

EDIT: As pointed out by Eric, an even better solution is to replace nans with -inf:

mask = numpy.isnan(a)
a[mask] = -numpy.inf
b = numpy.maximum.accumulate(a)
a[mask] = b[mask]

edited Dec 14 '16 at 19:01

answered Dec 14 '16 at 18:54

Benjamin

11,560
13
70
119

Nice! Replacing `nan` with `-inf` would work here too, right? – Eric Dec 14 '16 at 18:58
@Eric: Indeed, even better solution. – Benjamin Dec 14 '16 at 19:00
Wait, this doesn't work. See my updated test case. You've assumed the values are always increasing, which they're not – Eric Dec 14 '16 at 19:02
@Eric, leaving this for reference. – Benjamin Dec 14 '16 at 19:34
This works great if you throw indices at it instead of values – Eric Dec 14 '16 at 19:39
@Eric; Yes, very nice solution. I couldn't think of how to make it work... How does it compare for speed? – Benjamin Dec 14 '16 at 19:41
Updated my answer with timings. Unlike @user2357112's answer, the speed shouldn't be dependent on the data – Eric Dec 14 '16 at 20:20

Eric · Accepted Answer · 2016-12-14T20:24:42.993

Working from @Benjamin's deleted solution, everything is great if you work with indices

def do_it(data, valid=None, axis=0):
    # normalize the inputs to match the question examples
    data = np.asarray(data)
    if valid is None:
        valid = ~np.isnan(data)

    # flat array of the data values
    data_flat = data.ravel()

    # array of indices such that data_flat[indices] == data
    indices = np.arange(data.size).reshape(data.shape)

    # thanks to benjamin here
    stretched_indices = np.maximum.accumulate(valid*indices, axis=axis)
    return data_flat[stretched_indices]

Comparing solution runtime:

>>> import numpy as np
>>> data = np.random.rand(10000)

>>> %timeit do_it_question(data)
10000 loops, best of 3: 17.3 ms per loop
>>> %timeit do_it_mine(data)
10000 loops, best of 3: 179 µs per loop
>>> %timeit do_it_user(data)
10000 loops, best of 3: 182 µs per loop

# with lots of nans
>>> data[data > 0.25] = np.nan

>>> %timeit do_it_question(data)
10000 loops, best of 3: 18.9 ms per loop
>>> %timeit do_it_mine(data)
10000 loops, best of 3: 177 µs per loop
>>> %timeit do_it_user(data)
10000 loops, best of 3: 231 µs per loop

So both this and @user2357112's solution blow the solution in the question out of the water, but this has the slight edge over @user2357112 when there are high numbers of nans

How can I efficiently "stretch" present values in an array over absent ones

3 Answers3