CODEWITHSUNDEEP
pythonnumpyscipy
James Downs
James Downs

Reputation: 43

Distances to positions on a square integer lattice

I want to calculate the distance to every other position in an integer lattice from it's centre and the number of positions at each distance. I'm currently using the following code to calculate this: 

x = numpy.arange(-10, 11, 1)  
[X, Y] = numpy.meshgrid(x, x)  

R = numpy.sqrt(X**2+Y**2)  
R2 = numpy.ndarray.flatten(R)  
R3 = numpy.unique(R2)  
r = R3[1:] # excludes the 0  
Nr = numpy.zeros(numpy.size(r))  

for i in range(numpy.size(r)):  
    Nr[i] = numpy.count_nonzero(R2 == r[i]

This tells me that the possible distances are 1, sqrt2, 2, sqrt5 etc.
It also tells me their is 4x1, 4xsqrt2, 4x2, 8xsqrt5 etc.

As this is a common problem in physics I was wondering if there is a function from a library such as numpy or scipy which could return these values more easily.

Upvotes: 1

Views: 312

Answers (3)

Divakar
Divakar

Reputation: 221524

The lattice is centered at (0,0). So it's symmetric on the four quadrants. So, we can use this restriction to our advantage as we could compute those required unique distances and counts for one quadrant and multiply those counts by 4 to simulate for all four quadrants.

So, let's say we use the first quadrant (upper right quad). We would skip the elements on the (y = 0) line, because otherwise with the multiplication by 4 for simulating on all four quadrants would result in duplicating results. Additionally, this way we won't have to exclude the first element, as done in the original post.

Thus, an implementation would be -

N = 11 # Lattice size
xa, ya = np.ogrid[0:N,1:N] # x's:0:N, y's:1:N 
unq_dists, count = np.unique(np.sqrt(xa**2 + ya**2), return_counts=1)
count = count*4

For further performance boost, we could use np.unique on the squared summations and then use np.sqrt on the unique ones. The idea is to perform the slow square-root computation on smaller unique set, like so -

unq_dists, count = np.unique(xa**2 + ya**2, return_counts=1)
unq_dists = np.sqrt(unq_dists)

Upvotes: 1

user6655984
user6655984

Reputation:

Your code can be made more efficient in a few ways. You are first forming two large matrices and then square the terms. It would be better to square first. Also, forming meshgrid only to perform outer addition is unnecessary: there is numpy.add.outer. Lastly, the loop you have is made unnecessary by the return_counts=True option in numpy.unique. (Which, incidentally, flattens the array itself, so you don't have to.) So the code is shortened to three lines.

x = numpy.arange(-10, 11, 1)
R = numpy.sqrt(numpy.add.outer(x**2, x**2))
r, Nr = numpy.unique(R, return_counts=True)

(And if you want to exclude 0 distance, return r[1:] and Nr[1:])

Upvotes: 0

DYZ
DYZ

Reputation: 57033

After you flattened the array, convert it to a Pandas Series and count unique values:

distances = pandas.Series(numpy.ndarray.flatten(R))
distances.value_counts()
# 9.219544     16
# 8.062258     16
# 5.000000     12
#10.000000     12
# 7.071068     12
# 2.236068      8
# ....

Upvotes: 0

Related Questions