Source code for lentil.wavefront

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