import copy
import numpy as np
import lentil.field
from lentil.field import Field
import lentil.fourier
import lentil.helper
[docs]class Wavefront:
"""A class representing a monochromatic wavefront.
Parameters
----------
wavelength : float
Wavelength in meters
pixelscale : float, optional
Physical sampling of wavefront
shape : (2,) array_like, optional
Wavefront shape. If ``shape`` is None (default), the wavefront is
assumed to be infinite (broadcastable to any shape).
data : list_like, optional
Wavefront data. Default is [1+0j] (a plane wave).
focal_length : float or np.inf
Wavefront focal length. A plane wave (default) has an infinite focal
length (``np.inf``).
"""
__slots__ = ('wavelength', '_pixelscale', 'focal_length',
'data', 'shape', 'planetype')
def __init__(self, wavelength, pixelscale=None, shape=None, planetype=None,
data=None, focal_length=None):
self.wavelength = wavelength
self._pixelscale = () if pixelscale is None else lentil.sanitize_shape(pixelscale)
self.shape = () if shape is None else shape
if data is None:
self.data = [Field(data=1, pixelscale=pixelscale, offset=[0, 0], tilt=[])]
else:
self.data = [*data]
self.focal_length = focal_length if focal_length else np.inf
self.planetype = planetype
def __mul__(self, plane):
return plane.multiply(self, inplace=False)
def __imul__(self, plane):
return plane.multiply(self, inplace=True)
def __rmul__(self, other):
return self.__mul__(other)
@property
def pixelscale(self):
"""Physical sampling of the wavefront"""
return self._pixelscale
@pixelscale.setter
def pixelscale(self, value):
self._pixelscale = lentil.sanitize_shape(value)
@property
def field(self):
"""Wavefront complex field"""
out = np.zeros(self.shape, dtype=complex)
for field in self.data:
out = lentil.field.insert(field, out)
return out
@property
def intensity(self):
"""Wavefront intensity"""
out = np.zeros(self.shape, dtype=float)
for field in lentil.field.reduce(*self.data):
out = lentil.field.insert(field, out, intensity=True)
return out
[docs] def copy(self):
return copy.deepcopy(self)
[docs] def insert(self, out, weight=1):
"""Directly insert wavefront intensity data into an output array.
This method can avoid repeatedly allocating large arrays of zeros
when accumulating :attr:`intensity`.
Parameters
----------
out : ndarray
Array to insert wavefront data into
weight : float
Scale factor applied to wavefront data
Returns
-------
out : ndarray
Array with wavefront data inserted into it at the appropriate location
"""
for field in lentil.field.reduce(*self.data):
out = lentil.field.insert(field, out, intensity=True, weight=weight)
return out
[docs] def propagate_image(self, pixelscale, npix, npix_prop=None, oversample=2,
inplace=True):
"""Propagate the Wavefront from a Pupil to an Image plane using
Fraunhofer diffraction.
Parameters
----------
pixelscale : float or (2,) float
Physical sampling of output (image) plane. If a single value is supplied,
the output is assumed to be uniformly sampled in both x and y.
npix : int or (2,) tuple of ints
Shape of output plane.
npix_prop : int or (2,) tuple of ints, optional
Shape of propagation output plane. If None (default),
``npix_prop = npix``. If ``npix_prop != npix``, the propagation
result is placed in the appropriate location in the output plane.
npix_prop cannot be larger than npix.
oversample : int, optional
Number of times to oversample the output plane. Default is 2.
inplace : bool, optional
If True (default) the wavefront is propagated in-place, otherwise
a copy is created and propagated.
Returns
-------
wavefront : :class:`~lentil.Wavefront`
A Wavefront propagated to the specified image plane
"""
if self.planetype != 'pupil':
raise ValueError("Wavefront must have planetype 'pupil'")
npix = np.asarray(lentil.sanitize_shape(npix))
npix_prop = npix if npix_prop is None else np.asarray(lentil.sanitize_shape(npix_prop))
prop_shape = npix_prop * oversample
dx = self.pixelscale
du = np.asarray(lentil.sanitize_shape(pixelscale))
z = self.focal_length
data = self.data
if inplace:
out = self
out.data = []
out.pixelscale = du / oversample
out.shape = npix * oversample
out.focal_length = np.inf
out.planetype = 'image'
else:
out = Wavefront(wavelength=self.wavelength, data=[],
pixelscale=du/oversample, shape=npix*oversample,
planetype='image')
for field in data:
# compute the field shift from any embedded tilts. note the return value
# is specified in terms of (r, c)
shift = field.shift(z=z, wavelength=self.wavelength,
pixelscale=du, oversample=oversample,
indexing='ij')
fix_shift = np.fix(shift)
dft_shift = shift - fix_shift
if _overlap(prop_shape, fix_shift, out.shape):
alpha = lentil.helper.dft_alpha(dx=dx, du=du,
wave=self.wavelength, z=z,
oversample=oversample)
data = lentil.fourier.dft2(f=field.data, alpha=alpha,
npix=prop_shape, shift=dft_shift,
offset=field.offset, unitary=True)
out.data.append(Field(data=data, pixelscale=du/oversample,
offset=fix_shift))
return out
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