0

I have a function that reads in and outputs a 2D array. I want the output to be constant (pi in this case) for every index in the input that equals 0, otherwise I perform some maths on it. E.g:

import numpy as np
import numpy.ma as ma

def my_func(x):

    mask = ma.where(x==0,x)

    # make an array of pi's the same size and shape as the input
    y = np.pi * np.ones(x)

    # psuedo-code bit I can't figure out
    y.not_masked = y**2

    return y 

my_array = [[0,1,2],[1,0,2],[1,2,0]]

result_array = my_func(my_array)

This should give me the following:

result_array = [[3.14, 1, 4],[1, 3.14, 4], [1, 4, 3.14]]

I.e. it has applied y**2 to each element in the 2D list that doesn't equal zero, and replaced all the zeros with pi.

I need this because my function will include division, and I don't know the indexes beforehand. I'm trying to convert a matlab tutorial from a textbook into Python and this function is stumping me!

Thanks

rh1990
  • 880
  • 7
  • 17
  • 32

3 Answers3

4

Just use np.where() directly:

y = np.where(x, x**2, np.pi)

Example:

>>> x = np.asarray([[0,1,2],[1,0,2],[1,2,0]])
>>> y = np.where(x, x**2, np.pi)
>>> print(y)
[[ 3.14159265  1.          4.        ]
 [ 1.          3.14159265  4.        ]
 [ 1.          4.          3.14159265]]
AGN Gazer
  • 8,025
  • 2
  • 27
  • 45
  • 1
    Nice and clean! I have to admit that this is the best solution so far, even though I really like masking. ;) – JE_Muc May 25 '18 at 13:43
  • @Scotty1- Thanks! I think fundamentally both solutions are doing the same thing (there is no other way around it). – AGN Gazer May 25 '18 at 13:45
  • 1
    Yep, but yours is about 35% faster than mine and it is just one line. While the solution using `np.invert` is the slowest and imho most complicated. – JE_Muc May 25 '18 at 13:47
2

Try this:

my_array = np.array([[0,1,2],[1,0,2],[1,2,0]]).astype(float)

def my_func(x):

    mask = x == 0

    x[mask] = np.pi
    x[~mask] = x[~mask]**2  # or some other operation on x...

    return x
JE_Muc
  • 5,403
  • 2
  • 26
  • 41
1

I would suggest rather than using masks you can use a boolean array to achieve what you want. 

def my_func(x):
    #create a boolean matrix, a, that has True where x==0 and
    #False where x!=0 

    a=x==0

    x[a]=np.pi

    #Use np.invert to flip where a is True and False so we can 
    #operate on the non-zero values of the array

    x[~a]=x[~a]**2

    return x #return the transformed array

my_array = np.array([[0.,1.,2.],[1.,0.,2.],[1.,2.,0.]])

result_array = my_func(my_array)

this gives the output: 

array([[ 3.14159265,  1.        ,  4.        ],
       [ 1.        ,  3.14159265,  4.        ],
       [ 1.        ,  4.        ,  3.14159265]])

Notice that I passed to the function an numpy array specifically, originally you passed a list and that will give problems when you attempt to do mathematical operations. Also notice I defined the array with 1. rather than just 1, in order to make sure it was an array of floats rather than integers, because if it is an array of integers when you set values equal to pi it will truncate to 3.

Perhaps it would be good to add a piece to the function to check the dtype of the input argument and see if it is a numpy array rather than a list or other object, and also to make sure it contains floats, and if not you can adjust accordingly.

EDIT: Change to using ~a rather than invert(a) as per Scotty1's suggestion.

JJR4
  • 436
  • 3
  • 7
  • 1
    `np.invert` is about 30% slover than `~`, thus I'd recommend using `~` as in my answer whereever it is possible. – JE_Muc May 25 '18 at 13:41