#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright © 2020 Michael J. Hayford
"""Aberration calculations for (fld, wvl, foc), including focus and image shift
This module refactors some existing ray trace and aberration calculations
in other modules to be expressed for a single field point and wavelength.
The ability to apply focus and image shifts to an already acquired data set
is provided for use interactively and in other performance critical areas.
The following classes are implemented in this module:
- :class:`~.Ray`: trace a single ray
- :class:`~.RayFan`: trace a fan of rays in either the x or y meridian
- :class:`~.RayList`: trace a list of rays from an object point
- :class:`~.RayGrid`: trace a rectilinear grid of rays
All but the `Ray` class are supported by a group of functions to trace the
rays, accumulate the data (trace_*), and refocus (focus_*) the data. A
all-in-one function (eval_*) to trace and apply focus is supplied also.
These are used in the update_data methods of the classes to generate the
ray data.
This module also has functions to calculate chief ray and reference sphere
information as well as functions for calculating the monochromatic PSF of
the model.
.. Created on Sat Feb 22 22:01:56 2020
.. codeauthor: Michael J. Hayford
"""
import numpy as np
from numpy.fft import fftshift, fft2
from scipy.interpolate import interp1d
import rayoptics.optical.model_constants as mc
from rayoptics.raytr import sampler
from rayoptics.raytr import trace
from rayoptics.raytr import traceerror as terr
from rayoptics.raytr import waveabr
# --- Single ray
[docs]class Ray():
"""A ray at the given field and wavelength.
Attributes:
opt_model: :class:`~.OpticalModel` instance
p: relative 2d pupil coordinates
f: index into :class:`~.FieldSpec` or a :class:`~.Field` instance
wl: wavelength (nm) to trace the ray, or central wavelength if None
foc: focus shift to apply to the results
image_pt_2d: base image point. if None, the chief ray is used
image_delta: image offset to apply to image_pt_2d
srf_save:
'single': save the ray data for surface srf_indx
'all': save all of the surface by surface ray data
srf_indx: for single surface retention, the surface index to save
"""
def __init__(self, opt_model, p, f=0, wl=None, foc=None, image_pt_2d=None,
image_delta=None, srf_indx=-1, srf_save='single',
output_filter=None, rayerr_filter=None, color=None):
self.opt_model = opt_model
osp = opt_model.optical_spec
self.pupil = p
self.fld = osp['fov'].fields[f] if isinstance(f, int) else f
self.wvl = osp['wvls'].central_wvl if wl is None else wl
self.foc = osp['focus'].focus_shift if foc is None else foc
self.image_pt_2d = image_pt_2d
self.image_delta = image_delta
self.output_filter = output_filter
self.rayerr_filter = rayerr_filter
self.color = color
self.srf_save = srf_save
self.srf_indx = srf_indx
self.update_data()
[docs] def update_data(self, **kwargs):
"""Trace the ray and calculate transverse aberrations. """
ref_sphere, cr_pkg = trace.setup_pupil_coords(
self.opt_model, self.fld, self.wvl, self.foc,
image_pt=self.image_pt_2d, image_delta=self.image_delta
)
build = kwargs.pop('build', 'rebuild')
if build == 'rebuild':
ray_pkg = trace.trace_safe(
self.opt_model, self.pupil, self.fld, self.wvl,
self.output_filter, self.rayerr_filter,
use_named_tuples=True, **kwargs)
self.ray_seg = ray_pkg[0][self.srf_indx]
if self.srf_save == 'all':
self.ray_pkg = ray_pkg
ray_seg = self.ray_seg
dist = self.foc / ray_seg[mc.d][2]
defocused_pt = ray_seg[mc.p] + dist*ray_seg[mc.d]
reference_image_pt = ref_sphere[0]
self.t_abr = defocused_pt[:2] - reference_image_pt[:2]
return self
# --- Fan of rays
[docs]class RayFan():
"""A fan of rays across the pupil at the given field and wavelength.
Attributes:
opt_model: :class:`~.OpticalModel` instance
f: index into :class:`~.FieldSpec` or a :class:`~.Field` instance
wl: wavelength (nm) to trace the fan, or central wavelength if None
foc: focus shift to apply to the results
image_pt_2d: base image point. if None, the chief ray is used
image_delta: image offset to apply to image_pt_2d
num_rays: number of samples along the fan
xyfan: 'x' or 'y', specifies the axis the fan is sampled on
"""
def __init__(self, opt_model, f=0, wl=None, foc=None, image_pt_2d=None,
image_delta=None, num_rays=21, xyfan='y', output_filter=None,
rayerr_filter=None, color=None, **kwargs):
self.opt_model = opt_model
osp = opt_model.optical_spec
self.fld = osp.field_of_view.fields[f] if isinstance(f, int) else f
self.wvl = osp.spectral_region.central_wvl if wl is None else wl
self.foc = osp.defocus.focus_shift if foc is None else foc
self.image_pt_2d = image_pt_2d
self.image_delta = image_delta
self.num_rays = num_rays
if xyfan == 'x':
self.xyfan = 0
elif xyfan == 'y':
self.xyfan = 1
else:
self.xyfan = int(xyfan)
self.color = color
self.output_filter = output_filter
self.rayerr_filter = rayerr_filter
self.update_data()
def __json_encode__(self):
attrs = dict(vars(self))
del attrs['opt_model']
del attrs['fan_pkg']
return attrs
[docs] def update_data(self, **kwargs):
"""Set the fan attribute to a list of (pupil coords), dx, dy, opd."""
build = kwargs.get('build', 'rebuild')
if build == 'rebuild':
self.fan_pkg = trace_fan(
self.opt_model, self.fld, self.wvl, self.foc, self.xyfan,
image_pt_2d=self.image_pt_2d, image_delta=self.image_delta,
num_rays=self.num_rays,
output_filter=self.output_filter,
rayerr_filter=self.rayerr_filter)
self.fan = focus_fan(self.opt_model, self.fan_pkg,
self.fld, self.wvl, self.foc,
image_pt_2d=self.image_pt_2d,
image_delta=self.image_delta)
return self
[docs]def select_plot_data(fan, xyfan, data_type):
"""Given a fan of data, select the sample points and the resulting data."""
f_x = []
f_y = []
for p, val in fan:
f_x.append(p[xyfan])
f_y.append(val[data_type])
f_x = np.array(f_x)
f_y = np.array(f_y)
return f_x, f_y
[docs]def smooth_plot_data(f_x, f_y, num_points=100):
"""Interpolate fan data points and return a smoothed version."""
interpolator = interp1d(f_x, f_y,
kind='cubic', assume_sorted=True)
x_sample = np.linspace(f_x.min(), f_x.max(), num_points)
y_fit = interpolator(x_sample)
return x_sample, y_fit
[docs]def trace_ray_fan(opt_model, fan_rng, fld, wvl, foc,
output_filter=None, rayerr_filter=None, **kwargs):
"""Trace a fan of rays, according to fan_rng. """
start = np.array(fan_rng[0])
stop = fan_rng[1]
num = fan_rng[2]
step = (stop - start)/(num - 1)
fan = []
for r in range(num):
pupil = np.array(start)
ray_result = trace.trace_safe(opt_model, pupil, fld, wvl,
output_filter, rayerr_filter,
use_named_tuples=True, **kwargs)
if ray_result is not None:
fan.append([pupil[0], pupil[1], ray_result])
start += step
return fan
[docs]def eval_fan(opt_model, fld, wvl, foc, xy,
image_pt_2d=None, image_delta=None, num_rays=21,
output_filter=None, rayerr_filter=None, **kwargs):
"""Trace a fan of rays and evaluate dx, dy, & OPD across the fan."""
fod = opt_model['analysis_results']['parax_data'].fod
ref_sphere, cr_pkg = trace.setup_pupil_coords(opt_model, fld, wvl, foc,
image_pt=image_pt_2d,
image_delta=image_delta)
fld.chief_ray = cr_pkg
fld.ref_sphere = ref_sphere
fan_start = np.array([0., 0.])
fan_stop = np.array([0., 0.])
fan_start[xy] = -1.0
fan_stop[xy] = 1.0
fan_def = [fan_start, fan_stop, num_rays]
fan = trace_ray_fan(opt_model, fan_def, fld, wvl, foc,
output_filter=output_filter,
rayerr_filter=rayerr_filter, **kwargs)
central_wvl = opt_model.optical_spec.spectral_region.central_wvl
convert_to_opd = 1/opt_model.nm_to_sys_units(central_wvl)
def rfc(fi):
pupil_x, pupil_y, ray_pkg = fi
if ray_pkg is not None:
image_pt = ref_sphere[0]
ray = ray_pkg[mc.ray]
dist = foc / ray[-1][mc.d][2]
defocused_pt = ray[-1][mc.p] + dist*ray[-1][mc.d]
t_abr = defocused_pt - image_pt
opdelta = waveabr.wave_abr_full_calc(fod, fld, wvl, foc, ray_pkg,
cr_pkg, ref_sphere)
opd = convert_to_opd*opdelta
return (pupil_x, pupil_y), (t_abr[0], t_abr[1], opd)
else:
return pupil_x, pupil_y, np.NaN
fan_data = [rfc(i) for i in fan]
return fan_data
[docs]def trace_fan(opt_model, fld, wvl, foc, xy,
image_pt_2d=None, image_delta=None, num_rays=21,
output_filter=None, rayerr_filter=None, **kwargs):
"""Trace a fan of rays and precalculate data for rapid refocus later."""
fod = opt_model['analysis_results']['parax_data'].fod
ref_sphere, cr_pkg = trace.setup_pupil_coords(opt_model, fld, wvl, foc,
image_pt=image_pt_2d,
image_delta=image_delta)
fld.chief_ray = cr_pkg
fld.ref_sphere = ref_sphere
""" xy determines whether x (=0) or y (=1) fan """
fan_start = np.array([0., 0.])
fan_stop = np.array([0., 0.])
fan_start[xy] = -1.0
fan_stop[xy] = 1.0
fan_def = [fan_start, fan_stop, num_rays]
fan = trace_ray_fan(opt_model, fan_def, fld, wvl, foc,
output_filter=output_filter,
rayerr_filter=rayerr_filter, **kwargs)
def wpc(fi):
pupil_x, pupil_y, ray_pkg = fi
if ray_pkg is not None and not isinstance(ray_pkg, terr.TraceError):
pre_opd_pkg = waveabr.wave_abr_pre_calc(fod, fld, wvl, foc,
ray_pkg, cr_pkg)
return pre_opd_pkg
else:
return None
upd_fan = [wpc(i) for i in fan]
return fan, upd_fan
[docs]def focus_fan(opt_model, fan_pkg, fld, wvl, foc,
image_pt_2d=None, image_delta=None, **kwargs):
"""Refocus the fan of rays and return the tranverse abr. and OPD."""
fod = opt_model['analysis_results']['parax_data'].fod
fan, upd_fan = fan_pkg
ref_sphere, cr_pkg = trace.setup_pupil_coords(opt_model, fld, wvl, foc,
image_pt=image_pt_2d,
image_delta=image_delta)
central_wvl = opt_model.optical_spec.spectral_region.central_wvl
convert_to_opd = 1/opt_model.nm_to_sys_units(central_wvl)
def rfc(fi, fiu):
pupil_x, pupil_y, ray_pkg = fi
if ray_pkg is not None and not isinstance(ray_pkg, terr.TraceError):
image_pt = ref_sphere[0]
ray = ray_pkg[mc.ray]
dist = foc / ray[-1][mc.d][2]
defocused_pt = ray[-1][mc.p] + dist*ray[-1][mc.d]
t_abr = defocused_pt - image_pt
opdelta = waveabr.wave_abr_calc(fod, fld, wvl, foc,
ray_pkg, cr_pkg, fiu, ref_sphere)
opd = convert_to_opd*opdelta
return (pupil_x, pupil_y), (t_abr[0], t_abr[1], opd)
else:
return pupil_x, pupil_y, np.NaN
fan_data = [rfc(fi, fiu) for fi, fiu in zip(fan, upd_fan)]
return fan_data
# --- List of rays
[docs]class RayList():
"""Container class for a list of rays produced from a list or generator
Attributes:
opt_model: :class:`~.OpticalModel` instance
pupil_gen: (fct, args, kwargs), where:
- fct: a function returning a generator returning a 2d coordinate
- args: passed to fct
- kwargs: passed to fct
pupil_coords: list of 2d coordinates. If None, filled in by calling
pupil_gen.
num_rays: number of samples side of grid. Used only if pupil_coords and
pupil_gen are None.
f: index into :class:`~.FieldSpec` or a :class:`~.Field` instance
wl: wavelength (nm) to trace the fan, or central wavelength if None
foc: focus shift to apply to the results
image_pt_2d: base image point. if None, the chief ray is used
image_delta: image offset to apply to image_pt_2d
apply_vignetting: whether to apply vignetting factors to pupil coords
"""
def __init__(self, opt_model,
pupil_gen=None, pupil_coords=None, num_rays=21,
f=0, wl=None, foc=None, image_pt_2d=None, image_delta=None,
apply_vignetting=True):
self.opt_model = opt_model
osp = opt_model.optical_spec
if pupil_coords is not None and pupil_gen is None:
self.pupil_coords = pupil_coords
self.pupil_gen = None
else:
if pupil_gen is not None:
self.pupil_gen = pupil_gen
else:
grid_start = np.array([-1., -1.])
grid_stop = np.array([1., 1.])
grid_def = [grid_start, grid_stop, num_rays]
self.pupil_gen = (sampler.csd_grid_ray_generator,
(grid_def,), {})
fct, args, kwargs = self.pupil_gen
self.pupil_coords = fct(*args, **kwargs)
self.fld = osp.field_of_view.fields[f] if isinstance(f, int) else f
self.wvl = osp.spectral_region.central_wvl if wl is None else wl
self.foc = osp.defocus.focus_shift if foc is None else foc
self.image_pt_2d = image_pt_2d
self.image_delta = image_delta
self.apply_vignetting = apply_vignetting
self.update_data()
def __json_encode__(self):
attrs = dict(vars(self))
del attrs['opt_model']
del attrs['pupil_gen']
del attrs['pupil_coords']
del attrs['ray_list']
return attrs
[docs] def update_data(self, **kwargs):
build = kwargs.get('build', 'rebuild')
if build == 'rebuild':
if self.pupil_gen:
fct, args, kwargs = self.pupil_gen
self.pupil_coords = fct(*args, **kwargs)
self.ray_list = trace_pupil_coords(
self.opt_model, self.pupil_coords,
self.fld, self.wvl, self.foc,
image_pt_2d=self.image_pt_2d, image_delta=self.image_delta,
apply_vignetting=self.apply_vignetting)
ray_list_data = focus_pupil_coords(
self.opt_model, self.ray_list,
self.fld, self.wvl, self.foc,
image_pt_2d=self.image_pt_2d,
image_delta=self.image_delta)
self.ray_abr = np.rollaxis(ray_list_data, 1)
return self
[docs]def trace_ray_list(opt_model, pupil_coords, fld, wvl, foc,
append_if_none=False,
output_filter=None, rayerr_filter=None,
**kwargs):
"""Trace a list of rays at fld and wvl and return ray_pkgs in a list."""
ray_list = []
for pupil in pupil_coords:
ray_result = trace.trace_safe(opt_model, pupil, fld, wvl,
output_filter, rayerr_filter,
**kwargs)
if ray_result is not None:
ray_list.append([pupil[0], pupil[1], ray_result])
else: # ray outside pupil or failed
if append_if_none:
ray_list.append([pupil[0], pupil[1], None])
return ray_list
[docs]def trace_list_of_rays(opt_model, rays,
output_filter=None, rayerr_filter=None,
**kwargs):
"""Trace a list of rays (pt, dir, wvl) and return ray_pkgs in a list.
Args:
opt_model: :class:`~.OpticalModel` instance
rays: list of (pt0, dir0, wvl)
- pt0: starting point in coords of first interface
- dir0: starting direction cosines in coords of first interface
- wvl: wavelength in nm
output_filter: None, "last", or a callable. See below
**kwargs: kwyword args passed to the trace function
The output_filter keyword argument controls what ray data is returned to
the caller.
- if None, returns the entire traced ray
- if "last", returns the ray data from the last interface
- if a callable, it must take a ray_pkg as an argument and return the
desired data or None
Returns:
A list with an entry for each ray in rays
"""
ray_list = []
for ray in rays:
pt0, dir0, wvl = ray
try:
ray_pkg = trace.trace(opt_model.seq_model, pt0, dir0, wvl, **kwargs)
except terr.TraceError as rayerr:
if rayerr_filter is None:
pass
elif rayerr_filter == 'full':
ray_list.append((ray, rayerr))
elif rayerr_filter == 'summary':
rayerr.ray_pkg = None
ray_list.append((ray, rayerr))
else:
pass
else:
if output_filter is None:
ray_list.append(ray_pkg)
elif output_filter == 'last':
ray, op_delta, wvl = ray_pkg
ray_list.append((ray[-1], op_delta, wvl))
else:
ray_list.append(output_filter(ray_pkg))
return ray_list
[docs]def eval_pupil_coords(opt_model, fld, wvl, foc, image_pt_2d=None,
image_delta=None, num_rays=21, **kwargs):
"""Trace a list of rays and return the transverse abr."""
ref_sphere, cr_pkg = trace.setup_pupil_coords(opt_model, fld, wvl, foc,
image_pt=image_pt_2d,
image_delta=image_delta)
fld.chief_ray = cr_pkg
fld.ref_sphere = ref_sphere
grid_start = np.array([-1., -1.])
grid_stop = np.array([1., 1.])
grid_def = [grid_start, grid_stop, num_rays]
ray_list = trace_ray_list(opt_model, sampler.grid_ray_generator(grid_def),
fld, wvl, foc, check_apertures=True, **kwargs)
def rfc(ri):
pupil_x, pupil_y, ray_pkg = ri
if ray_pkg is not None:
image_pt = ref_sphere[0]
ray = ray_pkg[mc.ray]
dist = foc / ray[-1][mc.d][2]
defocused_pt = ray[-1][mc.p] + dist*ray[-1][mc.d]
t_abr = defocused_pt - image_pt
return t_abr[0], t_abr[1]
else:
return np.NaN
ray_list_data = [rfc(ri) for ri in ray_list]
return np.array(ray_list_data)
[docs]def trace_pupil_coords(opt_model, pupil_coords, fld, wvl, foc,
image_pt_2d=None, image_delta=None, **kwargs):
"""Trace a list of rays and return data needed for rapid refocus."""
ref_sphere, cr_pkg = trace.setup_pupil_coords(opt_model, fld, wvl, foc,
image_pt=image_pt_2d,
image_delta=image_delta)
fld.chief_ray = cr_pkg
fld.ref_sphere = ref_sphere
ray_list = trace_ray_list(opt_model, pupil_coords,
fld, wvl, foc,
check_apertures=True, **kwargs)
return ray_list
[docs]def focus_pupil_coords(opt_model, ray_list, fld, wvl, foc,
image_pt_2d=None, image_delta=None):
"""Given pre-traced rays and a ref. sphere, return the transverse abr."""
ref_sphere, cr_pkg = trace.setup_pupil_coords(opt_model, fld, wvl, foc,
image_pt=image_pt_2d,
image_delta=image_delta)
def rfc(ri):
pupil_x, pupil_y, ray_pkg = ri
if ray_pkg is not None:
image_pt = ref_sphere[0]
ray = ray_pkg[mc.ray]
dist = foc / ray[-1][mc.d][2]
defocused_pt = ray[-1][mc.p] + dist*ray[-1][mc.d]
t_abr = defocused_pt - image_pt
return t_abr[0], t_abr[1]
else:
return np.NaN
ray_list_data = [rfc(ri) for ri in ray_list]
return np.array(ray_list_data)
# --- Square grid of rays
[docs]class RayGrid():
"""Container for a square grid of rays.
Attributes:
opt_model: :class:`~.OpticalModel` instance
f: index into :class:`~.FieldSpec` or a :class:`~.Field` instance
wl: wavelength (nm) to trace the fan, or central wavelength if None
foc: focus shift to apply to the results
image_pt_2d: base image point. if None, the chief ray is used
image_delta: image offset to apply to image_pt_2d
num_rays: number of samples along the side of the grid
"""
def __init__(self, opt_model, f=0, wl=None, foc=None, image_pt_2d=None,
image_delta=None, num_rays=21, value_if_none=np.NaN):
self.opt_model = opt_model
osp = opt_model.optical_spec
self.fld = osp.field_of_view.fields[f] if isinstance(f, int) else f
self.wvl = osp.spectral_region.central_wvl if wl is None else wl
self.foc = osp.defocus.focus_shift if foc is None else foc
self.image_pt_2d = image_pt_2d
self.image_delta = image_delta
self.num_rays = num_rays
self.value_if_none = value_if_none
self.update_data()
def __json_encode__(self):
attrs = dict(vars(self))
del attrs['opt_model']
del attrs['grid_pkg']
return attrs
[docs] def update_data(self, **kwargs):
build = kwargs.get('build', 'rebuild')
if build == 'rebuild':
self.grid_pkg = trace_wavefront(
self.opt_model, self.fld, self.wvl, self.foc,
image_pt_2d=self.image_pt_2d, image_delta=self.image_delta,
num_rays=self.num_rays)
opd = focus_wavefront(self.opt_model, self.grid_pkg,
self.fld, self.wvl, self.foc,
image_pt_2d=self.image_pt_2d,
image_delta=self.image_delta,
value_if_none=self.value_if_none)
self.grid = np.rollaxis(opd, 2)
return self
[docs]def trace_ray_grid(opt_model, grid_rng, fld, wvl, foc, append_if_none=True,
output_filter=None, rayerr_filter=None, **kwargs):
"""Trace a grid of rays at fld and wvl and return ray_pkgs in 2d list."""
start = np.array(grid_rng[0])
stop = grid_rng[1]
num = grid_rng[2]
step = np.array((stop - start)/(num - 1))
grid = []
for i in range(num):
grid_row = []
for j in range(num):
pupil = np.array(start)
ray_result = trace.trace_safe(opt_model, pupil, fld, wvl,
output_filter, rayerr_filter,
apply_vignetting=False, **kwargs)
if ray_result is not None:
grid_row.append([pupil[0], pupil[1], ray_result])
else: # ray outside pupil or failed
if append_if_none:
grid_row.append([pupil[0], pupil[1], None])
start[1] += step[1]
grid.append(grid_row)
start[0] += step[0]
start[1] = grid_rng[0][1]
return grid
[docs]def eval_wavefront(opt_model, fld, wvl, foc, image_pt_2d=None,
image_delta=None, num_rays=21, value_if_none=np.NaN):
"""Trace a grid of rays and evaluate the OPD across the wavefront."""
fod = opt_model['analysis_results']['parax_data'].fod
ref_sphere, cr_pkg = trace.setup_pupil_coords(opt_model, fld, wvl, foc,
image_pt=image_pt_2d,
image_delta=image_delta)
fld.chief_ray = cr_pkg
fld.ref_sphere = ref_sphere
vig_bbox = fld.vignetting_bbox(opt_model['osp']['pupil'])
vig_grid_def = [vig_bbox[0], vig_bbox[1], num_rays]
grid = trace_ray_grid(opt_model, vig_grid_def,
fld, wvl, foc, check_apertures=True)
central_wvl = opt_model.optical_spec.spectral_region.central_wvl
convert_to_opd = 1/opt_model.nm_to_sys_units(central_wvl)
def rfc(gij):
pupil_x, pupil_y, ray_pkg = gij
if ray_pkg is not None:
opdelta = waveabr.wave_abr_full_calc(fod, fld, wvl, foc, ray_pkg,
cr_pkg, ref_sphere)
opd = convert_to_opd*opdelta
return pupil_x, pupil_y, opd
else:
return pupil_x, pupil_y, value_if_none
opd_grid = [[rfc(j) for j in i] for i in grid]
return np.array(opd_grid)
[docs]def trace_wavefront(opt_model, fld, wvl, foc,
image_pt_2d=None, image_delta=None, num_rays=21):
"""Trace a grid of rays and pre-calculate data needed for rapid refocus."""
fod = opt_model['analysis_results']['parax_data'].fod
ref_sphere, cr_pkg = trace.setup_pupil_coords(opt_model, fld, wvl, foc,
image_pt=image_pt_2d,
image_delta=image_delta)
fld.chief_ray = cr_pkg
fld.ref_sphere = ref_sphere
vig_bbox = fld.vignetting_bbox(opt_model['osp']['pupil'])
vig_grid_def = [vig_bbox[0], vig_bbox[1], num_rays]
grid = trace_ray_grid(opt_model, vig_grid_def,
fld, wvl, foc, check_apertures=True)
def wpc(gij):
pupil_x, pupil_y, ray_pkg = gij
if ray_pkg is not None:
pre_opd_pkg = waveabr.wave_abr_pre_calc(fod, fld, wvl, foc,
ray_pkg, cr_pkg)
return pre_opd_pkg
else:
return None
upd_grid = [[wpc(j) for j in i] for i in grid]
return grid, upd_grid
[docs]def focus_wavefront(opt_model, grid_pkg, fld, wvl, foc, image_pt_2d=None,
image_delta=None, value_if_none=np.NaN):
"""Given pre-traced rays and a ref. sphere, return the ray's OPD."""
fod = opt_model['analysis_results']['parax_data'].fod
grid, upd_grid = grid_pkg
ref_sphere, cr_pkg = trace.setup_pupil_coords(opt_model, fld, wvl, foc,
image_pt=image_pt_2d,
image_delta=image_delta)
central_wvl = opt_model.optical_spec.spectral_region.central_wvl
convert_to_opd = 1/opt_model.nm_to_sys_units(central_wvl)
def rfc(gij, uij):
pupil_x, pupil_y, ray_pkg = gij
if ray_pkg is not None:
opdelta = waveabr.wave_abr_calc(fod, fld, wvl, foc,
ray_pkg, cr_pkg, uij, ref_sphere)
opd = convert_to_opd*opdelta
return pupil_x, pupil_y, opd
else:
return pupil_x, pupil_y, value_if_none
refocused_grid = [[rfc(jg, ju) for jg, ju in zip(ig, iu)]
for ig, iu in zip(grid, upd_grid)]
return np.array(refocused_grid)
# --- PSF calculation
[docs]def psf_sampling(n=None, n_pupil=None, n_airy=None):
"""Given 2 of 3 parameters, calculate the third.
Args:
n: The total width of the sampling grid
n_pupil: The sampling across the pupil
n_airy: The sampling across the central peak of the Airy disk
Returns: (n, n_pupil, n_airy)
"""
npa = n, n_pupil, n_airy
i = npa.index(None)
if i == 0:
n = round((n_pupil * n_airy)/2.44)
elif i == 1:
n_pupil = round(2.44*n/n_airy)
elif i == 2:
n_airy = round(2.44*n/n_pupil)
return n, n_pupil, n_airy
[docs]def calc_psf_scaling(pupil_grid, ndim, maxdim):
"""Calculate the input and output grid spacings.
Args:
pupil_grid: A RayGrid instance
ndim: The sampling across the wavefront
maxdim: The total width of the sampling grid
Returns:
delta_x: The linear grid spacing on the entrance pupil
delta_xp: The linear grid spacing on the image plane
"""
opt_model = pupil_grid.opt_model
fod = opt_model['analysis_results']['parax_data'].fod
wl = opt_model.nm_to_sys_units(pupil_grid.wvl)
fill_factor = ndim/maxdim
max_D = 2 * fod.enp_radius / fill_factor
delta_x = max_D / maxdim
C = wl/fod.exp_radius
delta_theta = (fill_factor * C) / 2
ref_sphere_radius = pupil_grid.fld.ref_sphere[2]
delta_xp = delta_theta * ref_sphere_radius
return delta_x, delta_xp
[docs]def calc_psf(wavefront, ndim, maxdim):
"""Calculate the point spread function of wavefront W.
Args:
wavefront: ndim x ndim Numpy array of wavefront errors. No data
condition is indicated by nan
ndim: The sampling across the wavefront
maxdim: The total width of the sampling grid
Returns: AP, the PSF of the input wavefront
"""
maxdim_by_2 = maxdim//2
W = np.zeros([maxdim, maxdim])
nd2 = ndim//2
W[maxdim_by_2-(nd2-1):maxdim_by_2+(nd2+1),
maxdim_by_2-(nd2-1):maxdim_by_2+(nd2+1)] = np.nan_to_num(wavefront)
phase = np.exp(1j*2*np.pi*W)
for i in range(len(phase)):
for j in range(len(phase)):
if phase[i][j] == 1:
phase[i][j] = 0
AP = abs(fftshift(fft2(fftshift(phase))))**2
AP_max = np.nanmax(AP)
AP = AP/AP_max
return AP
[docs]def update_psf_data(pupil_grid, build='rebuild'):
pupil_grid.update_data(build=build)
ndim = pupil_grid.num_rays
maxdim = pupil_grid.maxdim
AP = calc_psf(pupil_grid.grid[2], ndim, maxdim)
return AP