CODEWITHSUNDEEP
pythonnumpy
giorgi nguyen
giorgi nguyen

Reputation: 166

Loop over elements in a numpy array of arbitrary dimension

I have a code like the following:

def infball_proj(mu, beta):
    newmu = np.zeros(mu.shape)
    if len(mu.shape) == 2:
        for i in range(mu.shape[0]):
            for j in range(mu.shape[1]):
                if np.abs(mu[i,j]) > beta:
                    newmu[i,j] = np.sign(mu[i,j]) * beta
                else:
                    newmu[i,j] = mu[i,j]
        return newmu
    elif len(mu.shape) == 1:
        for i in range(mu.shape[0]):
            if np.abs(mu[i]) > beta:
                newmu[i] = np.sign(mu[i]) * beta
            else:
                newmu[i] = mu[i]
        return newmu

Is there a smarter way to do this so I don't have to write the 2 different cases? It would be nice if I could have a version that scales to an arbitrary dimension (i.e. numbers of axes).

Upvotes: 0

Views: 61

Answers (2)

Ken T
Ken T

Reputation: 2553

mu[np.abs(mu)>beta] = np.sign(mu[np.abs(mu)>beta]) * beta

np.abs(mu)>beta will create a boolean array which can then be used for boolean indexing.

The LHS mu[np.abs(mu)>beta] will return a view of the elements being selected by the boolean indexing and can be assigned to the value your want, that is, the RHS.

REMEMBER: Try to avoid for-loop of NumPy array as it is very inefficient.

Upvotes: 1

Marat
Marat

Reputation: 15738

Something like this should do the job:

newmu = np.where(np.abs(mu) > beta, np.sign(mu) * beta, mu)

Or, if I get the logic right,

newmu = np.minimum(np.abs(mu), beta) * np.sign(mu)

Upvotes: 1

Related Questions