CODEWITHSUNDEEP
pythonarraysimageloopsnumpy
astrohuman
astrohuman

Reputation: 45

Avoiding loops over 2D array using numpy

I am stuck trying to figure out how to avoid doing 2 loops while doing some analysis on an image. Its a square image (same dimension in the x and y). The current way its set up its that for each pixel, in the array "stars", the first coordinate is the x position, the second coordinate is the y position, and the third coordinate is the brightness of the pixel. I know there must be a far far superior way to set this up but every thing I try just works when i=j (giving me essentially only the diagonal of this matrix). How do I do this better?!

def firstmoment(stars,bins):
    xsum=0.0
    ysum=0.0
    itot=0.0
    for i in range(0, bins):
        for j in range(0, bins):
            xsum=xsum+stars[i][j][0]*stars[i][j][2]
            ysum=ysum+stars[i][j][1]*stars[i][j][2]
            itot=itot+stars[i][j][2]
    x0=xsum/itot
    y0=ysum/itot
    return x0, y0, itot

Upvotes: 0

Views: 356

Answers (1)

John Zwinck
John Zwinck

Reputation: 249394

Let's just focus on the x part, ignoring y because it works the same way.

def firstmoment(stars, bins):
    xsum = 0.0
    itot = 0.0
    for i in range(bins):
        for j in range(bins):
            xsum += stars[i][j][0] * stars[i][j][2]
            itot += stars[i][j][2]
    x0 = xsum / itot
    return x0, itot

stars is a 3D array indexed as [x, y, luminance]. At the end, itot is the sum of the luminance values, which is trivially stars[:bins,:bins,2].sum(). Using a similar transformation for xsum, we get:

def firstmoment(stars, bins):
    xsum = (stars[:bins,:bins,0] * stars[:bins,:bins,2]).sum()
    itot = stars[:bins,:bins,2].sum()
    x0 = xsum / itot
    return x0, itot

Upvotes: 1

Related Questions