Loop over elements in a numpy array of arbitrary dimension

Question

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).

Marat · Accepted Answer · 2020-08-24 16:15:42Z

1

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)

answered Aug 24, 2020 at 16:15

Marat

15.9k3 gold badges44 silver badges53 bronze badges

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

1 Comment

Eric Hodonsky Over a year ago

Some commentary on your answer would help with clarity. Please translate why. Also, you should elaborate on what "get the logic right" means as a difference from the related above.

Ken T · Accepted Answer · 2020-08-27 06:28:26Z

1

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.

edited Aug 27, 2020 at 6:28

answered Aug 24, 2020 at 19:52

Ken T

2,5631 gold badge27 silver badges35 bronze badges

4 Comments

Eric Hodonsky Over a year ago

Some commentary on your answer would help with clarity

Ken T Over a year ago

@EricHodonsky Does it really help if I add “Something like this should the job” in my answer? I doubt.

Eric Hodonsky Over a year ago

Unfortunately flip doesn't get you very far in the internet community. You have a nice score at time of this comment of 1,339. That means you've been here for a while, you've helped others, and been helped. When prompted for clarity, it is in the answer, not that you have answered at all. Please describe WHY you have written the code this way, and HOW it answers their question. Remember, SO Q&A isn't just for the OP and whom ever is selected as the correct answer. Now to answer your flip, in triage I am only presented with one answer, I don't see them all.

Ken T Over a year ago

@EricHodonsky You are correct int his. I agree. More clarification has been given

Collectives™ on Stack Overflow

Loop over elements in a numpy array of arbitrary dimension

2 Answers 2

1 Comment

4 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

4 Comments

Your Answer

Sign up or log in

Post as a guest

Related