def __init__(self, nhex=15, narr=200, extrasupportindex=None):
"""
Input: nzern: number of Noll Zernikes to use in the fit
Input: narr: the live pupil array size you want to use
Sets up list of poly's and grids & support grids
Makes coordinate grid for rho and phi and circular support mask
Calculates 'overlap integrals' (covariance matrix) of the Zernike polynomials on your grid and array size
Calculates the inverse of this matrix, so it's 'ready to fit' your incoming array
"""
self.narr = narr
self.nhex = nhex # tbd - allowed numbers from Pascal's Triangle sum(n) starting from n=1, viz. n(n+1)/2
self.grid = (N.indices((self.narr, self.narr), dtype=N.float) -
self.narr // 2) / (float(self.narr) * 0.5)
self.grid_rho = (self.grid**2.0).sum(0)**0.5
self.grid_phi = N.arctan2(self.grid[0], self.grid[1])
self.grid_mask = self.grid_rho <= 1
self.grid_outside = self.grid_rho > 1
if extrasupportindex is not None:
self.grid_mask[extrasupportindex] = 0
self.grid_outside[extrasupportindex] = 1
# Compute list of explicit Zernike polynomials and keep them around for fitting
self.hex_list = z.hexike_basis(nterms=self.nhex, npix=self.narr)
self.hex_list = z.hexike_basis(nterms=self.nhex, npix=self.narr)
# Force hexikes to be unit standard deviation over hex mask
for h, hfunc in enumerate(self.hex_list):
if h > 0:
self.hex_list[h] = (
hfunc / hfunc[self.grid_mask].std()) * self.grid_mask
else:
self.hex_list[0] = hfunc * self.grid_mask
#### Write out a cube of the hexikes in here
if 0:
self.stack = N.zeros((self.nhex, self.narr, self.narr))
for ii in range(len(self.hex_list)):
self.stack[ii, :, :] = self.hex_list[ii]
# Calculate covariance between all Zernike polynomials
self.cov_mat = N.array([[N.sum(hexi * hexj) for hexi in self.hex_list]
for hexj in self.hex_list])
self.grid_mask = z.hex_aperture(npix=self.narr) == 1
self.grid_outside = self.grid_mask == False
# Invert covariance matrix using SVD
self.cov_mat_in = N.linalg.pinv(self.cov_mat)
def _test_cross_hexikes(testj=4, nterms=10, npix=500):
"""Verify the functions are orthogonal, by taking the
integrals of a given Hexike times N other ones.
This is a helper function for test_cross_hexike.
Parameters :
--------------
testj : int
Index of the Zernike polynomial to test against the others
nterms : int
Test that polynomial against those from 1 to this N
npix : int
Size of array to use for this test
"""
hexike_basis = zernike.hexike_basis(nterms=nterms, npix=npix)
test_hexike = hexike_basis[testj - 1]
assert np.sum(
np.isfinite(test_hexike)) > 0, "Hexike calculation failure; all NaNs."
for idx, hexike_array in enumerate(hexike_basis):
j = idx + 1
if j == testj or j == 1:
continue # discard piston term and self
prod = hexike_array * test_hexike
wg = np.where(np.isfinite(prod))
cross_sum = np.abs(prod[wg].sum())
# Threshold was originally 1e-9, but we ended up getting 1.19e-9 on some machines (not always)
# this seems acceptable, so relaxing criteria slightly
assert cross_sum < 2e-9, (
"orthogonality failure, Sum[Hexike(j={}) * Hexike(j={})] = {} (> 2e-9)"
.format(j, testj, cross_sum))
def fit_hexikes_to_surface(self, surface, choosemodes=False):
"""
Input: surface: input surface to be fit (2D array)
Output: zcoeffs: 1d vector of coefficients of the fit (self.nzern in length)
Output: rec_wf: the 'recovered wavefront' - i.e. the fitted zernikes, in same array size as surface
Output: res_wf: surface - rec_wf, i.e. the residual error in the fit
"""
# Calculate the inner product of each Zernike mode with the test surface
wf_hex_inprod = N.array(
[N.sum(surface * hexi) for hexi in self.hex_list])
# Given the inner product vector of the test wavefront with Zernike basis,
# calculate the Zernike polynomial coefficients
hcoeffs = N.dot(self.cov_mat_in, wf_hex_inprod)
# Reconstruct (e.g. wavefront) surface from Zernike components
hexikes = z.hexike_basis(nterms=len(hcoeffs), npix=self.narr)
if choosemodes is not False:
hcoeffs[choosemodes == 0] = 0
rec_wf = sum(val * hexikes[i] for (i, val) in enumerate(hcoeffs))
if 0:
print "First 10 recovered Hernike coeffts:", hcoeffs[:10]
print "Standard deviation of fit is %.3e" % (
surface * self.grid_mask - rec_wf)[self.grid_mask].std()
return hcoeffs, rec_wf, (surface - rec_wf) * self.grid_mask
def _test_cross_hexikes(testj=4, nterms=10, npix=500):
"""Verify the functions are orthogonal, by taking the
integrals of a given Hexike times N other ones.
Parameters :
--------------
testj : int
Index of the Zernike polynomial to test against the others
nterms : int
Test that polynomial against those from 1 to this N
npix : int
Size of array to use for this test
"""
hexike_basis = zernike.hexike_basis(nterms=nterms, npix=npix)
test_hexike = hexike_basis[testj - 1]
for idx, hexike_array in enumerate(hexike_basis):
j = idx + 1
if j == testj or j == 1:
continue # discard piston term and self
prod = hexike_array * test_hexike
wg = np.where(np.isfinite(prod))
cross_sum = np.abs(prod[wg].sum())
assert cross_sum < 1e-9, (
"orthogonality failure, Sum[Hexike(j={}) * Hexike(j={})] = {} (> 1e-9)".format(
j, testj, cross_sum)
)
def _test_cross_hexikes(testj=4, nterms=10, npix=500):
"""Verify the functions are orthogonal, by taking the
integrals of a given Hexike times N other ones.
This is a helper function for test_cross_hexike.
Parameters :
--------------
testj : int
Index of the Zernike polynomial to test against the others
nterms : int
Test that polynomial against those from 1 to this N
npix : int
Size of array to use for this test
"""
hexike_basis = zernike.hexike_basis(nterms=nterms, npix=npix)
test_hexike = hexike_basis[testj - 1]
assert np.sum(np.isfinite(test_hexike)) > 0, "Hexike calculation failure; all NaNs."
for idx, hexike_array in enumerate(hexike_basis):
j = idx + 1
if j == testj or j == 1:
continue # discard piston term and self
prod = hexike_array * test_hexike
wg = np.where(np.isfinite(prod))
cross_sum = np.abs(prod[wg].sum())
# Threshold was originally 1e-9, but we ended up getting 1.19e-9 on some machines (not always)
# this seems acceptable, so relaxing criteria slightly
assert cross_sum < 2e-9, (
"orthogonality failure, Sum[Hexike(j={}) * Hexike(j={})] = {} (> 2e-9)".format(
j, testj, cross_sum)
)
def projectionfilter(data,
nterms=None,
bases=None,
npix=None,
basis_type='Zernike',
outside='nan',
basis_kwds={}):
""" Filtering reconstructed band structure using orthogonal polynomial approximation.
**Parameters**\n
data: 2D array
Band dispersion in 2D to filter.
nterms: int | None
Number of terms.
bases: 3D array | None
Bases for decomposition.
npix: int | None
Size (number of pixels) in one direction of each basis term.
basis_type: str | 'Zernike'
Type of basis to use for filtering.
outside: numeric/str | 'nan'
Values to fill for regions outside the Brillouin zone boundary.
basis_kwds: dictionary | {}
Keywords for basis generator (see `poppy.zernike.hexike_basis()` if hexagonal Zernike polynomials are used).
"""
nterms = int(nterms)
# Generate basis functions
if bases is None:
if basis_type == 'Zernike':
bases = ppz.hexike_basis(nterms=nterms, npix=npix, **basis_kwds)
# Decompose into the given basis
coeffs = decomposition_hex2d(data,
bases=bases,
baxis=0,
nterms=nterms,
basis_type=basis_type,
ret='coeffs')
# Reconstruct the smoothed version of the energy band
recon = reconstruction_hex2d(coeffs,
bases=bases,
baxis=0,
npix=npix,
basis_type=basis_type,
ret='band')
if outside == 'nan':
recon = u.to_masked(recon, val=0)
return recon, coeffs
elif outside == 0:
return recon, coeffs
def decomposition_hex2d(band,
bases=None,
baxis=0,
nterms=100,
basis_type='Zernike',
ret='coeffs'):
""" Decompose energy band in 3D momentum space using the orthogonal polynomials in a hexagon.
**Parameters**\n
band: 2D array
2D electronic band structure.
bases: 3D array | None
Matrix composed of bases to decompose into.
baxis: int | 0
Axis of the basis index.
nterms: int | 100
Number of basis terms.
basis_type: str | 'Zernike'
Type of basis to use.
ret: str | 'coeffs'
Options for the return values.
"""
nbr, nbc = band.shape
if nbr != nbc:
raise ValueError('Input band surface should be square!')
if bases is None:
if basis_type == 'Zernike':
bases = ppz.hexike_basis(nterms=nterms,
npix=nbr,
vertical=True,
outside=0)
elif basis_type == 'Fourier':
raise NotImplementedError
else:
raise NotImplementedError
else:
if baxis != 0:
bases = np.moveaxis(bases, baxis, 0)
nbas, nbasr, nbasc = bases.shape
band_flat = band.reshape((band.size, ))
coeffs = np.linalg.pinv(bases.reshape(
(nbas, nbasr * nbasc))).T.dot(band_flat)
if ret == 'coeffs':
return coeffs
elif ret == 'all':
return coeffs, bases
def reconstruction_hex2d(coeffs,
bases=None,
baxis=0,
npix=256,
basis_type='Zernike',
ret='band'):
""" Reconstruction of energy band in 3D momentum space using orthogonal polynomials
and the term-wise coefficients.
**Parameters**\n
coeffs: 1D array
Polynomial coefficients to use in reconstruction.
bases: 3D array | None
Matrix composed of bases to decompose into.
baxis: int | 0
Axis of the basis index.
npix: int | 256
Number of pixels along one side in the square image.
basis_type: str | 'Zernike'
Type of basis to use.
ret: str | 'band'
Options for the return values.
"""
coeffs = coeffs.ravel()
nterms = coeffs.size
if bases is None:
if basis_type == 'Zernike':
bases = ppz.hexike_basis(nterms=nterms,
npix=npix,
vertical=True,
outside=0)
elif basis_type == 'Fourier':
raise NotImplementedError
else:
raise NotImplementedError
else:
if baxis != 0:
bases = np.moveaxis(bases, baxis, 0)
nbas, nbasr, nbasc = bases.shape
band_recon = bases.reshape((nbas, nbasr * nbasc)).T.dot(coeffs).reshape(
(nbasr, nbasc))
if ret == 'band':
return band_recon
elif ret == 'all':
return band_recon, bases
def basisgen(self,
nterms,
npix,
vertical=True,
outside=0,
basis_type='Zernike'):
""" Generate polynomial bases for energy band synthesis.
"""
if basis_type == 'Zernike':
self.bases = ppz.hexike_basis(nterms=nterms,
npix=npix,
vertical=vertical,
outside=outside)
def hexmask(hexdiag=128,
imside=256,
image=None,
padded=False,
margins=[],
pad_top=None,
pad_bottom=None,
pad_left=None,
pad_right=None,
vertical=True,
outside='nan',
ret='mask',
**kwargs):
""" Generate a hexagonal mask. To use the function, either the argument ``imside`` or ``image`` should be
given. The image padding on four sides could be specified with either ``margins`` altogether or separately
with the individual arguments ``pad_xxx``. For the latter, at least two independent padding values are needed.
**Parameters**\n
hexdiag: int | 128
Number of pixels along the hexagon's diagonal.
imside: int | 256
Number of pixels along the side of the (square) reference image.
image: 2D array | None
2D reference image to construct the mask for. If the reference (image) is given, each side
of the generated mask is at least that of the smallest dimension of the reference.
padded: bool | False
Option to pad the image (need to set to True to enable the margins).
margins: list/tuple | []
Margins of the image [top, bottom, left, right]. Overrides the `pad_xxx` arguments.
pad_top, pad_bottom, pad_left, pad_right : int, int, int, int | None, None, None, None
Number of padded pixels on each of the four sides of the image.
vertical: bool | True
Option to align the diagonal of the hexagon with the vertical image axis.
outside: numeric/str | 'nan'
Pixel value outside the masked region.
ret: str | 'mask'
Return option ('mask', 'masked_image', 'all').
"""
if image is not None:
imshape = image.shape
minside = min(imshape)
mask = ppz.hexike_basis(nterms=1, npix=minside, vertical=vertical)[0,
...]
else:
imshape = kwargs.pop('imshape', (imside, imside))
mask = ppz.hexike_basis(nterms=1, npix=hexdiag, vertical=vertical)[0,
...]
# Use a padded version of the original mask
if padded == True:
# Padding image margins on all sides
if len(margins) == 4:
top, bottom, left, right = margins
else:
# Total padding pixel numbers along horizontal and vertical directions
padsides = np.abs(np.asarray(imshape) - hexdiag)
top, bottom = u.nonneg_sum_decomposition(a=pad_top,
b=pad_bottom,
absum=padsides[0])
left, right = u.nonneg_sum_decomposition(a=pad_left,
b=pad_right,
absum=padsides[1])
mask = np.pad(mask, ((top, bottom), (left, right)),
mode='constant',
constant_values=np.nan)
if outside == 0:
mask = np.nan_to_num(mask)
if ret == 'mask':
return mask
elif ret == 'masked_image':
return mask * image
elif ret == 'all':
margins = [top, bottom, left, right]
return mask, margins
def analytical_model(zernike_pol, coef, cali=False):
"""
:param zernike_pol:
:param coef:
:param cali: bool; True if we already have calibration coefficients to use. False if we still need to create them.
:return:
"""
#-# Parameters
dataDir = os.path.join(CONFIG_INI.get('local', 'local_data_path'),
'active')
telescope = CONFIG_INI.get('telescope', 'name')
nb_seg = CONFIG_INI.getint(telescope, 'nb_subapertures')
tel_size_m = CONFIG_INI.getfloat(telescope, 'diameter') * u.m
real_size_seg = CONFIG_INI.getfloat(
telescope, 'flat_to_flat'
) # in m, size in meters of an individual segment flatl to flat
size_seg = CONFIG_INI.getint(
'numerical',
'size_seg') # pixel size of an individual segment tip to tip
wvln = CONFIG_INI.getint(telescope, 'lambda') * u.nm
inner_wa = CONFIG_INI.getint(telescope, 'IWA')
outer_wa = CONFIG_INI.getint(telescope, 'OWA')
tel_size_px = CONFIG_INI.getint(
'numerical', 'tel_size_px') # pupil diameter of telescope in pixels
im_size_pastis = CONFIG_INI.getint(
'numerical', 'im_size_px_pastis') # image array size in px
sampling = CONFIG_INI.getfloat('numerical', 'sampling') # sampling
size_px_tel = tel_size_m / tel_size_px # size of one pixel in pupil plane in m
px_sq_to_rad = (size_px_tel * np.pi / tel_size_m) * u.rad
zern_max = CONFIG_INI.getint('zernikes', 'max_zern')
sz = CONFIG_INI.getint('numerical', 'im_size_lamD_hcipy')
# Create Zernike mode object for easier handling
zern_mode = util.ZernikeMode(zernike_pol)
#-# Mean subtraction for piston
if zernike_pol == 1:
coef -= np.mean(coef)
#-# Generic segment shapes
if telescope == 'JWST':
# Load pupil from file
pupil = fits.getdata(
os.path.join(dataDir, 'segmentation', 'pupil.fits'))
# Put pupil in randomly picked, slightly larger image array
pup_im = np.copy(pupil) # remove if lines below this are active
#pup_im = np.zeros([tel_size_px, tel_size_px])
#lim = int((pup_im.shape[1] - pupil.shape[1])/2.)
#pup_im[lim:-lim, lim:-lim] = pupil
# test_seg = pupil[394:,197:315] # this is just so that I can display an individual segment when the pupil is 512
# test_seg = pupil[:203,392:631] # ... when the pupil is 1024
# one_seg = np.zeros_like(test_seg)
# one_seg[:110, :] = test_seg[8:, :] # this is the centered version of the individual segment for 512 px pupil
# Creat a mini-segment (one individual segment from the segmented aperture)
mini_seg_real = poppy.NgonAperture(
name='mini', radius=real_size_seg
) # creating real mini segment shape with poppy
#test = mini_seg_real.sample(wavelength=wvln, grid_size=flat_diam, return_scale=True) # fix its sampling with wavelength
mini_hdu = mini_seg_real.to_fits(wavelength=wvln,
npix=size_seg) # make it a fits file
mini_seg = mini_hdu[
0].data # extract the image data from the fits file
elif telescope == 'ATLAST':
# Create mini-segment
pupil_grid = hcipy.make_pupil_grid(dims=tel_size_px,
diameter=real_size_seg)
focal_grid = hcipy.make_focal_grid(
pupil_grid, sampling, sz, wavelength=wvln.to(
u.m).value) # fov = lambda/D radius of total image
prop = hcipy.FraunhoferPropagator(pupil_grid, focal_grid)
mini_seg_real = hcipy.hexagonal_aperture(circum_diameter=real_size_seg,
angle=np.pi / 2)
mini_seg_hc = hcipy.evaluate_supersampled(
mini_seg_real, pupil_grid, 4
) # the supersampling number doesn't really matter in context with the other numbers
mini_seg = mini_seg_hc.shaped # make it a 2D array
# Redefine size_seg if using HCIPy
size_seg = mini_seg.shape[0]
# Make stand-in pupil for DH array
pupil = fits.getdata(
os.path.join(dataDir, 'segmentation', 'pupil.fits'))
pup_im = np.copy(pupil)
#-# Generate a dark hole mask
#TODO: simplify DH generation and usage
dh_area = util.create_dark_hole(
pup_im, inner_wa, outer_wa, sampling
) # this might become a problem if pupil size is not same like pastis image size. fine for now though.
if telescope == 'ATLAST':
dh_sz = util.zoom_cen(dh_area, sz * sampling)
#-# Import information form segmentation script
Projection_Matrix = fits.getdata(
os.path.join(dataDir, 'segmentation', 'Projection_Matrix.fits'))
vec_list = fits.getdata(
os.path.join(dataDir, 'segmentation', 'vec_list.fits')) # in pixels
NR_pairs_list = fits.getdata(
os.path.join(dataDir, 'segmentation', 'NR_pairs_list_int.fits'))
# Figure out how many NRPs we're dealing with
NR_pairs_nb = NR_pairs_list.shape[0]
#-# Chose whether calibration factors to do the calibraiton with
if cali:
filename = 'calibration_' + zern_mode.name + '_' + zern_mode.convention + str(
zern_mode.index)
ck = fits.getdata(
os.path.join(dataDir, 'calibration', filename + '.fits'))
else:
ck = np.ones(nb_seg)
coef = coef * ck
#-# Generic coefficients
# the coefficients in front of the non redundant pairs, the A_q in eq. 13 in Leboulleux et al. 2018
generic_coef = np.zeros(
NR_pairs_nb
) * u.nm * u.nm # setting it up with the correct units this will have
for q in range(NR_pairs_nb):
for i in range(nb_seg):
for j in range(i + 1, nb_seg):
if Projection_Matrix[i, j, 0] == q + 1:
generic_coef[q] += coef[i] * coef[j]
#-# Constant sum and cosine sum - calculating eq. 13 from Leboulleux et al. 2018
if telescope == 'JWST':
i_line = np.linspace(-im_size_pastis / 2., im_size_pastis / 2.,
im_size_pastis)
tab_i, tab_j = np.meshgrid(i_line, i_line)
cos_u_mat = np.zeros(
(int(im_size_pastis), int(im_size_pastis), NR_pairs_nb))
elif telescope == 'ATLAST':
i_line = np.linspace(-(2 * sz * sampling) / 2.,
(2 * sz * sampling) / 2., (2 * sz * sampling))
tab_i, tab_j = np.meshgrid(i_line, i_line)
cos_u_mat = np.zeros((int((2 * sz * sampling)), int(
(2 * sz * sampling)), NR_pairs_nb))
# Calculating the cosine terms from eq. 13.
# The -1 with each NR_pairs_list is because the segment names are saved starting from 1, but Python starts
# its indexing at zero, so we have to make it start at zero here too.
for q in range(NR_pairs_nb):
# cos(b_q <dot> u): b_q with 1 <= q <= NR_pairs_nb is the basis of NRPS, meaning the distance vectors between
# two segments of one NRP. We can read these out from vec_list.
# u is the position (vector) in the detector plane. Here, those are the grids tab_i and tab_j.
# We need to calculate the dot product between all b_q and u, so in each iteration (for q), we simply add the
# x and y component.
cos_u_mat[:, :, q] = np.cos(
px_sq_to_rad *
(vec_list[NR_pairs_list[q, 0] - 1, NR_pairs_list[q, 1] - 1, 0] *
tab_i) + px_sq_to_rad *
(vec_list[NR_pairs_list[q, 0] - 1, NR_pairs_list[q, 1] - 1, 1] *
tab_j)) * u.dimensionless_unscaled
sum1 = np.sum(
coef**2
) # sum of all a_{k,l} in eq. 13 - this works only for single Zernikes (l fixed), because np.sum would sum over l too, which would be wrong.
if telescope == 'JWST':
sum2 = np.zeros(
(int(im_size_pastis), int(im_size_pastis))
) * u.nm * u.nm # setting it up with the correct units this will have
elif telescope == 'ATLAST':
sum2 = np.zeros(
(int(2 * sz * sampling), int(2 * sz * sampling))) * u.nm * u.nm
for q in range(NR_pairs_nb):
sum2 = sum2 + generic_coef[q] * cos_u_mat[:, :, q]
#-# Local Zernike
if telescope == 'JWST':
# Generate a basis of Zernikes with the mini segment being the support
isolated_zerns = zern.hexike_basis(nterms=zern_max,
npix=size_seg,
rho=None,
theta=None,
vertical=False,
outside=0.0)
# Calculate the Zernike that is currently being used and put it on one single subaperture, the result is Zer
# Apply the currently used Zernike to the mini-segment.
if zernike_pol == 1:
Zer = np.copy(mini_seg)
elif zernike_pol in range(2, zern_max - 2):
Zer = np.copy(mini_seg)
Zer = Zer * isolated_zerns[zernike_pol - 1]
# Fourier Transform of the Zernike - the global envelope
mf = mft.MatrixFourierTransform()
ft_zern = mf.perform(Zer, im_size_pastis / sampling, im_size_pastis)
elif telescope == 'ATLAST':
isolated_zerns = hcipy.make_zernike_basis(num_modes=zern_max,
D=real_size_seg,
grid=pupil_grid,
radial_cutoff=False)
Zer = hcipy.Wavefront(mini_seg_hc * isolated_zerns[zernike_pol - 1],
wavelength=wvln.to(u.m).value)
# Fourier transform the Zernike
ft_zern = prop(Zer)
#-# Final image
if telescope == 'JWST':
# Generating the final image that will get passed on to the outer scope, I(u) in eq. 13
intensity = np.abs(ft_zern)**2 * (sum1.value + 2. * sum2.value)
elif telescope == 'ATLAST':
intensity = ft_zern.intensity.shaped * (sum1.value + 2. * sum2.value)
# PASTIS is only valid inside the dark hole, so we cut out only that part
if telescope == 'JWST':
tot_dh_im_size = sampling * (outer_wa + 3)
intensity_zoom = util.zoom_cen(
intensity, tot_dh_im_size
) # zoom box is (owa + 3*lambda/D) wide, in terms of lambda/D
dh_area_zoom = util.zoom_cen(dh_area, tot_dh_im_size)
dh_psf = dh_area_zoom * intensity_zoom
elif telescope == 'ATLAST':
dh_psf = dh_sz * intensity
"""
# Create plots.
plt.subplot(1, 3, 1)
plt.imshow(pupil, origin='lower')
plt.title('JWST pupil and diameter definition')
plt.plot([46.5, 464.5], [101.5, 409.5], 'r-') # show how the diagonal of the pupil is defined
plt.subplot(1, 3, 2)
plt.imshow(mini_seg, origin='lower')
plt.title('JWST individual mini-segment')
plt.subplot(1, 3, 3)
plt.imshow(dh_psf, origin='lower')
plt.title('JWST dark hole')
plt.show()
"""
# dh_psf is the image of the dark hole only, the pixels outside of it are zero
# intensity is the entire final image
return dh_psf, intensity
