I'm writing this code for a class and am pretty new to python, so I wont be surprised if I'm doing a lot of things wrong. I'm trying to compute the intensity of diffraction pattern of light coming through an aperture, which is given by a Bessel function. I've got the Bessel function itself computed relatively efficiently using numpy arrays, but I've run into a bit of a snag when I try to extend the Bessel function on top of a grid for computing the intensity.
Here's my code:
import matplotlib.pyplot as plt
numDivs = 100
rmax = 1e-6
num = 400
lmda = 500e-9
k = 2*np.pi/lmda
#This function is designed to compute the integral using
# Simpson's Rule. The value N passed in is the number of
# polynomial curve fits to be created. The curve is sliced into
# 2N sections. Designed to work with numpy.
def integrateSimpson(func, lower, upper, N):
delta = (upper-lower)/(2*N)
THETA = np.linspace(lower, upper, 2*N 1)
Y = func(THETA)
Y = Y * (delta / 3)
return Y[0] Y[2*N] 4 * np.sum(Y[1:2*N:2]) 2 * np.sum(Y[2:2*N:2])
#this is the formula for the bessel function
J = lambda m, x: integrateSimpson( lambda theta: np.cos(m*theta-x*np.sin(theta)), 0, np.pi, numDivs)
#this is the formula for the intensity using the above bessel function definition
Intensity = lambda r: (J( 1, k*r)/(k*r))**2
#setting up the grids for computation.
xSpace = np.linspace(-rmax, rmax, num)
ySpace = np.linspace(-rmax, rmax, num)
xx, yy = np.meshgrid(xSpace, ySpace)
grid2 = np.zeros((num,num))
#This is where my problem begins
grid2 = Intensity(np.hypot(xx,yy))
#Set up the image and show it
plt.figure(figsize=(14, 14), dpi=80)
plt.set_cmap('plasma')
plt.imshow(grid2, extent = [-rmax , rmax, -rmax , rmax])
plt.colorbar()
plt.show()
Here is the error code I get:
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
<ipython-input-42-f3614ffc9470> in <module>
21 grid2 = np.zeros((num,num))
22
---> 23 grid2 = Intensity(np.hypot(xx,yy))
24
25 #Set up the image and show it
<ipython-input-42-f3614ffc9470> in <lambda>(r)
11
12 #this is the formula for the intensity using the above bessel function definition
---> 13 Intensity = lambda r: (J( 1, k*r)/(k*r))**2
14
15 #setting up the grid for computation.
<ipython-input-42-f3614ffc9470> in <lambda>(m, x)
8
9 #this is the formula for the bessel function
---> 10 J = lambda m, x: integrateSimpson( lambda theta: np.cos(m*theta-x*np.sin(theta)), 0, np.pi, numDivs)
11
12 #this is the formula for the intensity using the above bessel function definition
<ipython-input-8-1c73b24d5dd3> in integrateSimpson(func, lower, upper, N)
19 delta = (upper-lower)/(2*N)
20 THETA = np.linspace(lower, upper, 2*N 1)
---> 21 Y = func(THETA)
22 Y = Y * (delta / 3)
23
<ipython-input-42-f3614ffc9470> in <lambda>(theta)
8
9 #this is the formula for the bessel function
---> 10 J = lambda m, x: integrateSimpson( lambda theta: np.cos(m*theta-x*np.sin(theta)), 0, np.pi, numDivs)
11
12 #this is the formula for the intensity using the above bessel function definition
ValueError: operands could not be broadcast together with shapes (400,400) (201,)
It looks like the error is that the theta array and the x array are trying to broadcast into each other, but I really don't want them to. Is there a way to set this up which avoid loops but still accomplishes the task? I've fixed the problem for now by backing out of numpy array operations and just using a couple for loops.
The working code which accomplishes what I'm trying to do follows here:
#Type your code here
import matplotlib.pyplot as plt
numDivs = 100
rmax = 1e-6
num = 400
lmda = 500e-9
k = 2*np.pi/lmda
#This function is designed to compute the integral using
# Simpson's Rule. The value N passed in is the number of
# polynomial curve fits to be created. The curve is sliced into
# 2N sections. Designed to work with numpy.
def integrateSimpson(func, lower, upper, N):
delta = (upper-lower)/(2*N)
THETA = np.linspace(lower, upper, 2*N 1)
Y = func(THETA)
Y = Y * (delta / 3)
return Y[0] Y[2*N] 4 * np.sum(Y[1:2*N:2]) 2 * np.sum(Y[2:2*N:2])
#this is the formula for the bessel function
J = lambda m, x: integrateSimpson( lambda theta: np.cos(m*theta-x*np.sin(theta)), 0, np.pi, numDivs)
#this is the formula for the intensity using the above bessel function definition
Intensity = lambda r: (J( 1, k*r)/(k*r))**2
#setting up the grid for computation.
xSpace = np.linspace(-rmax, rmax, num)
ySpace = np.linspace(-rmax, rmax, num)
grid2 = np.zeros((num,num))
#calculate the diffraction pattern. This is an inefficient way to do it for several reasons
# First I'm not utilizing the efficiency of numpy arrays. I ran into a problem where the
# numpy arrays are trying to broadcast into eachother down the line. Not sure how to fix it so
# I reverted to loops. There's also a level of symmetry where I should only have to compute
# one line along r and then rotate it around, instead of computing the entire grid.
# computing the grid takes a few seconds at this density.
for i in range(num):
for j in range(num):
grid2[i,j] = Intensity(np.hypot(xSpace[i],ySpace[j]))
#Set up the image and show it
plt.figure(figsize=(14, 14), dpi=80)
plt.set_cmap('plasma')
plt.imshow(grid2, extent = [-rmax , rmax, -rmax , rmax])
plt.colorbar()
plt.show()
The produced image for the working code:
CodePudding user response：
Broadcasting is actually your friend here, and you can get rid of a lot of temporary arrays by using it. For example, you can set up the grid like this:
x = np.linspace(-rmax, rmax, num)
y = np.linspace(-rmax, rmax, num)
r = np.hypot(x, y[:, None])
image = intensity(r)
For the sake of clarity, and to avoid recomputing unnecessary quantities, I would recommend not using lambdas in this code. Lambdas have a time and a place, but sacrificing legibility for a couple of lines of code is never worth it. In the end, lambdas are just like normal functions in every way that matters except they don't have a meaningful __name__ attribute:
def intensity(r):
r_scaled = k * r
return (bessel(1, r_scaled) / r_scaled)**2
Now to address your actual problem. You have an array r, or in this case r_scaled, that has shape (400, 400). You want to integrate a Bessel function for each element. You are using np.sin and np.cos to generate the elements to integrate. But the problem is that summation is a reduction operation. What dimension are you dropping? The answer is that you need to expand the dimensions of your array. It should become a (400, 400, 201) array, which you then sum back down to (400, 400). Here is an example that uses nested functions to make it a bit easier to read than the mess of lambdas you have.
def bessel(m, x): # M = 1 (scalar), X = r (400, 400)
def func(theta):
return np.cos(m * theta - x[..., None] * np.sin(theta))
return integrate_simpson(func, 0, np.pi, num_divs)
def integrate_simpson(func, lower, upper, n):
n *= 2
delta = (upper - lower) / n
theta = np.linspace(lower, upper, n 1)
y = func(theta)
y *= delta / 3
return y[..., 0] y[..., -1] 4 * y[..., 1:-1:2].sum(axis=-1) 2 * y[..., 2:-1:2].sum(axis=-1)
The major change is that bessel now does math with x[..., None] (i.e. r_scaled[..., None]). ... means "take all existing dimensions, no matter how many there are", and None means np.newaxis: add a unit dimension. The result is a (400, 400, 1) view (no data is copied), which you can then broadcast with your 201-element theta to get a (400, 400, 201) array.
Inside integrate_simpson, you use the broadcasted shape to sum along the last axis. y[..., 0] means "take the first image plane". The ... refers to the leading two indices, 0 is the first element for each pixel, which is exactly what you want. y[..., -1] does the same thing. Numpy, like normal python, supports negative indices. If you want the last element, just use -1. It saves you some calculation, and is more intuitive for a python-savvy reader, which will be you in the future. Hopefully by this point, you can extrapolate the concepts to the remaining two terms in the sum. y[..., 1:-1:2].sum(axis=-1) sums along the last axis: the 201 one that you made using broadcasting. The nice thing about using -1 and ... here is that you don't even need to know how many axes you have. This will work perfectly with no modifications required for next Monday, when you decide you really meant to integrate over 5-dimensional images.
A couple of minor nitpicks:
- I would make
num = 401to see much nicer steps, or setx = np.linspace(-rmax, rmax, num, endpoint=False)(same fory) for a slightly asymmetrical image with the same nice steps. - Python naming convention is CAPS for constants, CamelCase for classes, snake_case for pretty much everything else, including variables, functions, and methods.
CodePudding user response：
Sorry, I just couldn't resist. Let's say you did want to take advantage of the obvious radial symmetry of the problem. That means that you only need to generate the Bessel function for a set of possible radii, and then assign those values to the output image. First find your maximum radius:
r_max_xy = np.hypot(rmax, rmax)
Now generate the integral for just a couple of hundred points across the image. In fact, since you're not doing it for 400 * 400 points, let's splurge and do it for 1000 points:
r_radial = np.linspace(0, r_max_xy, 1000)
image_radial = intensity(r_radial)
Because of how the broadcasting was done in the other answer, intensity(r_radial) will work out of the box. Now you can interpolate the values all over the image. For a first-order approximation, let's just do nearest-neighbor interpolation.
The goal is to compute the index in image_radial to take the image value from. Since r_radial is evenly spaced, you can use simple rounding arithmetic:
r_index = (r / (r_radial.ptp() / (r_radial.size - 1)) 0.5).astype(int)
image = image_radial[r_index]
For a smoother/more accurate interpolation, (a) use more points, and/or (b) add use a better interpolator. Here is an example of the latter:
from scipy.interpolate import Akima1DInterpolator
interp = Akima1DInterpolator(r_radial, image_radial)
image = interp(r)