Python- Numpy array in function with conditions

Question

I have a numpy array (say xs), for which I am writing a function to create another array (say ys), which has the same values as of xs until the first half of xs and two times of xs in the remaining half. for example, if xs=[0,1,2,3,4,5], the required output is [0,1,2,6,8,10]

I wrote the following function:

import numpy as np
xs=np.arange(0,6,1)

def step(xs):
    ys1=np.array([]);ys2=np.array([])
    if xs.all() <=2:
        ys1=xs
    else:
        ys2=xs*2
    return np.concatenate((ys1,ys2))

print(xs,step(xs))

Which produces output: `array([0., 1., 2., 3., 4., 5.]), ie the second condition is not executed. Does anybody know how to fix it? Thanks in advance.

if xs.all() <=2: is True so ys1 is assigned xs and ys2 remains an empty array. On concatenation, you just get ys1 back which is the same as xs — Sheldore
– Sheldore, Commented Aug 5, 2018 at 13:51

jpp · Accepted Answer · 2018-08-05 12:56:40Z

2

You can use vectorised operations instead of Python-level iteration. With the below method, we first copy the array and then multiply the second half of the array by 2.

import numpy as np

xs = np.arange(0,6,1)

def step(xs):
    arr = xs.copy()
    arr[int(len(arr)/2):] *= 2
    return arr

print(xs, step(xs))

[0 1 2 3 4 5] [ 0  1  2  6  8 10]

answered Aug 5, 2018 at 12:56

jpp

166k37 gold badges301 silver badges362 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

hyl · Accepted Answer · 2018-08-06 07:44:17Z

2

import numpy as np

xs=np.arange(0,6,1)

def f(a):
    it = np.nditer([a, None])
    for x, y in it:
        y[...] = x if x <= 2 else x * 2
    return it.operands[1]

print(f(xs))

[ 0  1  2  6  8 10]

Sorry, I did not find your bug, but I felt it can be implemented differently.

edited Aug 6, 2018 at 7:44

answered Aug 5, 2018 at 12:44

hyl

3691 silver badge7 bronze badges

Collectives™ on Stack Overflow

Python- Numpy array in function with conditions

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related