Skip to content

Commit

Permalink
Refactor tests to use fixtures
Browse files Browse the repository at this point in the history
  • Loading branch information
@andykee
andykee committed Mar 18, 2024
1 parent 98f73ea commit 20ac09d
Show file tree
Hide file tree
Showing 4 changed files with 30 additions and 232 deletions.
20 changes: 0 additions & 20 deletions lentil/propagate.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -263,26 +263,6 @@ def _mask_shift(x, threshold=0):
return rc_extent - rc_full, cc_extent - cc_full


def _overlap(field_shape, field_shift, output_shape):
# Return True if there's any overlap between a shifted field and the
# output shape
output_shape = np.asarray(output_shape)
field_shape = np.asarray(field_shape)
field_shift = np.asarray(field_shift)

# Output coordinates of the upper left corner of the shifted data array
field_shifted_ul = (output_shape / 2) - (field_shape / 2) + field_shift

if field_shifted_ul[0] > output_shape[0]:
return False
if field_shifted_ul[0] + field_shape[0] < 0:
return False
if field_shifted_ul[1] > output_shape[1]:
return False
if field_shifted_ul[1] + field_shape[1] < 0:
return False
return True

def _propagate_ptype(ptype, method='fraunhofer'):
if method == 'fraunhofer':
if ptype not in (lentil.pupil, lentil.image):
Expand Down
2 changes: 1 addition & 1 deletion tests/fixtures/pupil.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ def pupil():

def _pupil(focal_length, diameter, shape, radius, coeffs=None):

amplitude = lentil.circle(shape, radius)
amplitude = lentil.normalize_power(lentil.circle(shape, radius))
pixelscale = diameter/(2*radius)

if coeffs is not None:
Expand Down
190 changes: 29 additions & 161 deletions tests/test_propagate.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,167 +5,30 @@
import lentil


def airy(diameter, focal_length, wavelength, pixelscale, length, oversample=1):
# https://en.wikipedia.org/wiki/Airy_disk#Mathematical_Formulation
f_number = focal_length / diameter
length *= oversample

c = (length - 1) / 2
x = np.arange(length, dtype=float)
x -= c

q = x * (pixelscale / oversample)
X = (np.pi * q) / (wavelength * f_number)

# if length is odd, the center value will be zero which will throw a
# divide by zero error. To avoid this, we'll set any zeros to machine
# epsilon (eps)
X[X == 0] = np.finfo(X.dtype).eps

return (2 * scipy.special.jn(1, X) / X) ** 2


def airy2(diameter, focal_length, wavelength, pixelscale, shape, oversample=1):
# https://en.wikipedia.org/wiki/Airy_disk#Mathematical_Formulation

f_number = focal_length / diameter

shape = np.asarray(shape)
shape *= oversample

c = (shape - 1) / 2
y, x = np.indices(shape, dtype=float)

x -= c[1]
y -= c[0]

x *= (pixelscale[0] / oversample)
y *= (pixelscale[1] / oversample)

q = np.sqrt(x ** 2 + y ** 2)
X = (np.pi * q) / (wavelength * f_number)

# if length is odd, the center value will be zero which will throw a
# divide by zero error. To avoid this, we'll set any zeros to machine
# epsilon (eps)
X[X == 0] = np.finfo(X.dtype).eps

return (2 * scipy.special.jn(1, X) / X) ** 2


def test_airy():
diameter = 1
focal_length = 10
f_number = focal_length / diameter
wave = 650e-9
pixelscale = 5e-6 / 1000 # divide by 1000 here to fake oversampling
length = 2 ** 14 # an appropriately large number
oversample = 1

airy1 = airy(diameter, focal_length, wave, pixelscale, length, oversample)

c = (length - 1) / 2
x = np.arange(length, dtype=float)
x -= c
x *= pixelscale

airy_interp = scipy.interpolate.interp1d(x, airy1)

# TODO: move this in to the documentation and just point to that as a reference here
# derivation of airy FWHM expression
# ----------------------------------
# the half-maximum of the central airy disk occurs at x_hwhm ~= 1.61633 [1]
# since x = (pi*q)/(wave*f_number), q = (x * wave * f_number)/pi where
# q is the radial distance from the optics axis in the observation
# (or focal) plane
#
# substituting in x_hwhm,
# q_hwhm = (1.61633/pi) * wave * f_number
#
# because we are more interested in FWHM, we multiply by 2
# q_fwhm = 2*(1.61633/pi) * wave * f_number = 1.028987 * wave * f_number
#
# Refs:
# [1] https://en.wikipedia.org/wiki/Airy_disk#Mathematical_Formulation

# verify the FWHM occurs at 1.029 * wave * f_number
j1_hwhm = 1.61633
fwhm = 2 * (j1_hwhm / np.pi) * wave * f_number # actual linear distance
assert (np.abs(airy_interp(fwhm / 2) - 0.5) < 1e-5)

# verify the first five zeros (as listed on Wikipedia) occur in the
# right places
j1_zeros = np.asarray([3.8317, 7.0156, 10.1735, 13.3237, 16.4706])
for j1_zero in j1_zeros:
assert (airy_interp(j1_zero / np.pi * wave * f_number) < 1e-7)


def test_airy2():
# evaluate the 2D Airy code against the 1D Airy code. since we've done
# some pretty exhaustive testing of the 1D Airy code, this simple
# comparison should be sufficient

diameter = 1
focal_length = 10
wave = 650e-9
pixelscale = 5e-6 / 10 # divide by 10 here to fake oversampling
npix = 511
oversample = 1
slice_index = npix // 2

a1 = airy(diameter, focal_length, wave, pixelscale, npix, oversample)
a2 = airy2(diameter, focal_length, wave, (pixelscale, pixelscale), (npix, npix), oversample)

res = a2[:, slice_index] - a1
assert (np.all(np.abs(res) < 1e-9))


class TiltPupil(lentil.Pupil):
def __init__(self, npix, diameter=1, coeffs=None):

amplitude = lentil.normalize_power(lentil.circle((npix, npix), npix//2))
mask = lentil.circle((npix, npix), npix//2, antialias=False)

if coeffs is None:
coeffs = 5e-6 * np.random.uniform(low=-0.5, high=0.5, size=3)
coeffs[0] = 0

opd = lentil.zernike_compose(mask=mask, coeffs=coeffs)

super().__init__(focal_length=10,
pixelscale=1/npix,
amplitude=amplitude,
opd=opd,
mask=mask,
diameter=diameter)

self.coeffs = coeffs


def test_amplitude_normalize_power():
p = TiltPupil(npix=256)
def test_amplitude_normalize_power(pupil):
p = pupil(focal_length=10, diameter=1, shape=(256,256), radius=120, coeffs=None)
w = lentil.Wavefront(wavelength=650e-9)
w *= p
w = lentil.propagate_dft(w, shape=(64,64), pixelscale=5e-6, oversample=1)
psf = w.intensity
assert (np.sum(psf) <= 1) and (np.sum(psf) >= 0.95)


def test_amplitude_normalize_power_oversample():
p = TiltPupil(npix=256)
def test_amplitude_normalize_power_oversample(pupil):
p = pupil(focal_length=10, diameter=1, shape=(256,256), radius=120, coeffs=None)
w = lentil.Wavefront(wavelength=650e-9)
w *= p
w = lentil.propagate_dft(w, shape=(64,64), pixelscale=5e-6, oversample=2)
psf = w.intensity
assert (np.sum(psf) <= 1) and (np.sum(psf) >= 0.95)


def test_propagate_airy():
def test_propagate_airy(airy2, pupil):
# we set oversample = 1 and apply an "oversampling" factor of 25 to pixelscale
# so that we can more easily enforce an odd npix. having an odd npix enables
# comparison between numerical propagations via dft2() and the airy2() function
# since their DC pixels will be coincident.
p = TiltPupil(npix=512, diameter=1, coeffs=[0])
p = pupil(focal_length=10, diameter=1, shape=(512,512), radius=250, coeffs=None)

psf_airy = airy2(diameter=p.diameter, focal_length=p.focal_length, wavelength=650e-9,
pixelscale=[5e-6, 5e-6], shape=(511, 511), oversample=1)
Expand All @@ -179,9 +42,10 @@ def test_propagate_airy():
assert np.all(np.isclose(psf, psf_airy, atol=1e-3))


def test_propagate_fit_tilt():

p = TiltPupil(npix=256)
def test_propagate_fit_tilt(pupil):
coeffs = 5e-6 * np.random.uniform(low=-0.5, high=0.5, size=3)
coeffs[0] = 0
p = pupil(focal_length=10, diameter=1, shape=(256,256), radius=120, coeffs=coeffs)

w = lentil.Wavefront(650e-9)
w = p.multiply(w)
Expand All @@ -202,14 +66,17 @@ def test_propagate_fit_tilt():
assert np.all(delta <= 1e-10)


def test_propagate_tilt_analytic():
def test_propagate_tilt_analytic(pupil):
oversample = 10
pixelscale = 5e-6
npix = np.array([64, 64])
pupil = TiltPupil(npix=256, diameter=1)

coeffs = 5e-6 * np.random.uniform(low=-0.5, high=0.5, size=3)
coeffs[0] = 0
p = pupil(focal_length=10, diameter=1, shape=(256,256), radius=120, coeffs=coeffs)

w = lentil.Wavefront(650e-9)
w = pupil.multiply(w)
w = p.multiply(w)
w = lentil.propagate_dft(w, shape=npix, pixelscale=pixelscale, oversample=oversample)
psf = w.intensity

Expand All @@ -220,26 +87,29 @@ def test_propagate_tilt_analytic():
centroid = np.asarray(lentil.util.centroid(psf))
shift = centroid - center

pupil_tilt = pupil.coeffs[1:3]
pupil_tilt = coeffs[1:3]

# The factor of 4 gets you from RMS to PV
analytic_shift = ((pupil_tilt/pupil.diameter)*pupil.focal_length/pixelscale) * oversample * 4
analytic_shift = ((pupil_tilt/p.diameter)*p.focal_length/pixelscale) * oversample * 4

# Analytic shift is in terms of (x, y) but we really want (r, c) to match the centroid
analytic_shift = analytic_shift[1], -analytic_shift[0]

assert np.all((np.abs(shift - analytic_shift)/oversample) < 0.2)


def test_propagate_tilt_fit_tilt_analytic():
def test_propagate_tilt_fit_tilt_analytic(pupil):
oversample = 10
pixelscale = 5e-6
npix = np.array([64, 64])
pupil = TiltPupil(npix=256, diameter=1)

coeffs = 5e-6 * np.random.uniform(low=-0.5, high=0.5, size=3)
coeffs[0] = 0
p = pupil(focal_length=10, diameter=1, shape=(256,256), radius=120, coeffs=coeffs)

pupil.fit_tilt()
p.fit_tilt()
w = lentil.Wavefront(650e-9)
w = w * pupil
w = w * p
w = lentil.propagate_dft(w, shape=npix, pixelscale=pixelscale, oversample=oversample)
psf = w.intensity

Expand All @@ -250,24 +120,22 @@ def test_propagate_tilt_fit_tilt_analytic():
centroid = np.asarray(lentil.util.centroid(psf))
shift = centroid - center

pupil_tilt = pupil.coeffs[1:3]
pupil_tilt = coeffs[1:3]

# The factor of 4 gets you from RMS to PV
analytic_shift = ((pupil_tilt/pupil.diameter)*pupil.focal_length/pixelscale) * oversample * 4
analytic_shift = ((pupil_tilt/p.diameter)*p.focal_length/pixelscale) * oversample * 4

# Analytic shift is in terms of (x, y) but we really want (r, c) to match the centroid
analytic_shift = analytic_shift[1], -analytic_shift[0]

assert np.all((np.abs(shift - analytic_shift)/oversample) < 0.2)


def test_propagate_resample():
def test_propagate_resample(pupil):

amp = lentil.circle((256, 256), 120)
coeffs = np.random.uniform(low=-200e-9, high=200e-9, size=6)
opd = lentil.zernike_compose(amp, coeffs)

p = lentil.Pupil(focal_length=10, pixelscale=1 / 240, amplitude=amp, opd=opd)
p = pupil(focal_length=10, diameter=1, shape=(256,256), radius=120, coeffs=coeffs)
w = lentil.Wavefront(650e-9)
w *= p
wi = lentil.propagate_dft(w, shape=(64,64), pixelscale=5e-6, oversample=10)
Expand Down
50 changes: 0 additions & 50 deletions tests/test_propagate_segmented.py
View file

This file was deleted.

0 comments on commit 20ac09d

Please sign in to comment.