def coadd_long_exposures(self):
"""Coadd the interim long exposures."""
# Iterate over long exposure images
for index, file in enumerate(self.long_exp_files):
# Read data
with fits.open(os.path.join(self.tmp_dir, file)) as hdu_list:
tmp_image = hdu_list[0].data
if self.var_ext is not None and self.var_ext in hdu_list:
tmp_image_var = hdu_list[self.var_ext].data
else:
tmp_image_var = None
# Co-add reconstructions and var images
self.image += alignment.pad_array(tmp_image, self.pad_vectors[index], mode=self.mode,
reference_image_pad_vector=self.reference_pad_vector)
if tmp_image_var is not None:
self.var += alignment.pad_array(tmp_image_var, self.pad_vectors[index], mode=self.mode,
reference_image_pad_vector=self.reference_pad_vector)
# Update out_file
self.out_file.data = self.image
if self.var_ext is not None and self.var is not None:
self.out_file.update_extension(ext_name=self.var_ext, data=self.var)
return self.image, self.var
def test_pad_array(self):
pad_vectors, ref_pad_vector = alignment.get_pad_vectors(
self.shifts,
cube_mode=False,
return_reference_image_pad_vector=True)
padded = alignment.pad_array(np.ones(self.image_shape),
pad_vectors[1],
mode='same',
reference_image_pad_vector=ref_pad_vector)
imshow(padded)
pad_vectors, ref_pad_vector = alignment.get_pad_vectors(
self.shifts,
cube_mode=True,
return_reference_image_pad_vector=True)
for pad_vector in pad_vectors:
padded = alignment.pad_array(
np.ones(self.cube_shape),
pad_vector=pad_vector,
mode='same',
reference_image_pad_vector=ref_pad_vector)
def coadd_fft(self):
"""Co-add the Fourier transforms of the image and PSF frames.
Returns:
fourier_image (np.ndarray, dtype=np.comlex128):
Fourier-transformed object reconstruction.
"""
# Padding and Fourier transforming the images
logger.info("Padding the images and PSFs...")
for file_index in trange(len(self.in_files), desc="Processing files"):
# Open PSF and image files
psf_cube = fits.getdata(self.psf_files[file_index])
image_cube = fits.getdata(
os.path.join(self.in_dir, self.in_files[file_index]))
n_frames = image_cube.shape[0]
for frame_index in trange(n_frames,
desc="Fourier transforming frames"):
# Padding and transforming the image
img = pad_array(
array=image_cube[frame_index],
pad_vector=self.pad_vectors[file_index],
mode=self.mode,
reference_image_pad_vector=self.reference_image_pad_vector)
f_img = fftshift(fft2(img))
# Padding and Fourier transforming PSF
psf = psf_cube[frame_index]
psf = np.pad(
psf,
self.psf_pad_vector,
mode='constant',
)
f_psf = fftshift(fft2(psf))
# Co-adding for the average
self.enumerator += np.multiply(f_img, np.conjugate(f_psf))
self.denominator += np.abs(np.square(f_psf))
# Compute the object:
# Note that by this division implicitly does averaging. By this implicit summing up of enumerator and
# denominator, this computation is cheaper in terms of memory usage
self.fourier_image = np.divide(self.enumerator, self.denominator)
return self.fourier_image
def coadd_frames(cube, var_cube=None, box=None):
"""Compute the simple shift-and-add (SSA) reconstruction of a data cube.
This function uses the SSA algorithm to coadd frames of a cube. If provided, this function coadds the variances
within a var cube considering the exact same shifts.
Args:
cube (np.ndarray, ndim=3):
Data cube which is integrated along the zero-th axis.
var_cube (np.ndarray, ndim=3, optional):
Data cube of variances which is integrated along the zero-th axis with the same shifts as the cube.
box (Box object, optional):
Constraining the search for the intensity peak to the specified box. Searching the full frames if not
provided.
Returns:
coadded (np.ndarray, ndim=2):
SSA-integrated frames of the input cube.
var_coadded (np.ndarray, ndim=2):
SSA-integrated variances of the input cube or the variance map itself if provided as a 2D cube.
"""
if not isinstance(cube, np.ndarray):
raise SpecklepyTypeError('coadd_frames()',
argname='cube',
argtype=type(cube),
expected='np.ndarray')
if cube.ndim is not 3:
raise SpecklepyValueError('coadd_frames()',
argname='cube.ndim',
argvalue=cube.ndim,
expected='3')
if var_cube is not None:
if not isinstance(var_cube, np.ndarray):
raise SpecklepyTypeError('coadd_frames()',
argname='var_cube',
argtype=type(var_cube),
expected='np.ndarray')
if var_cube.ndim == cube.ndim and var_cube.shape != cube.shape:
raise SpecklepyValueError('coadd_frames()',
argname='var_cube.shape',
argvalue=str(var_cube.shape),
expected=str(cube.shape))
elif var_cube.ndim == cube.ndim - 1:
if var_cube.shape[0] != cube.shape[1] or var_cube.shape[
1] != cube.shape[2]:
raise SpecklepyValueError('coadd_frames()',
argname='var_cube.shape',
argvalue=str(var_cube.shape),
expected=str(cube.shape))
# Compute shifts
peak_indizes = np.zeros((cube.shape[0], 2), dtype=int)
for index, frame in enumerate(cube):
if box is not None:
frame = box(frame)
peak_indizes[index] = np.array(np.unravel_index(
np.argmax(frame, axis=None), frame.shape),
dtype=int)
# Compute shifts from indizes
peak_indizes = peak_indizes.transpose()
xmean, ymean = np.mean(np.array(peak_indizes), axis=1)
xmean = int(xmean)
ymean = int(ymean)
shifts = np.array([xmean - peak_indizes[0], ymean - peak_indizes[1]])
shifts = shifts.transpose()
# Shift frames and add to coadded
coadded = np.zeros(cube[0].shape)
pad_vectors, ref_pad_vector = alignment.get_pad_vectors(
shifts, cube_mode=False, return_reference_image_pad_vector=True)
for index, frame in enumerate(cube):
coadded += alignment.pad_array(
frame,
pad_vectors[index],
mode='same',
reference_image_pad_vector=ref_pad_vector)
# Coadd variance cube (if not an image itself)
if var_cube is not None:
if var_cube.ndim == 3:
var_coadded = np.zeros(coadded.shape)
for index, frame in enumerate(var_cube):
var_coadded += alignment.pad_array(
frame,
pad_vectors[index],
mode='same',
reference_image_pad_vector=ref_pad_vector)
elif var_cube.ndim == 2:
var_coadded = var_cube
else:
raise RuntimeError(
f"var_cube has unexpected shape: {var_cube.shape}")
else:
var_coadded = None
return coadded, var_coadded
def ssa(files,
mode='same',
reference_file=None,
outfile=None,
in_dir=None,
tmp_dir=None,
lazy_mode=True,
box_indexes=None,
debug=False,
**kwargs):
"""Compute the SSA reconstruction of a list of files.
The simple shift-and-add (SSA) algorithm makes use of the structure of typical speckle patterns, i.e.
short-exposure point-spread functions (PSFs). These show multiple peaks resembling the diffraction-limited PSF of
coherent fractions within the telescope aperture. Under good conditions or on small telescopes, there is typically
one largest coherent atmospheric cell and therefore, speckle PSFs typically show one major intensity peak. The
algorithm makes use of this fact and identifies the emission peak in a given observation frame, assuming that this
always belongs to the same star, and aligns all frames on the coordinate of the emission peak.
See Bates & Cady (1980) for references.
Args:
files (list or array_like):
List of complete paths to the fits files that shall be considered for the SSA reconstruction.
mode (str):
Name of the reconstruction mode: In 'same' mode, the reconstruction covers the same field of view of the
reference file. In 'full' mode, every patch of the sky that is covered by at least one frame will be
contained in the final reconstruction.
reference_file (str, int, optional):
Path to a reference file or index of the file in files, relative to which the shifts are computed. See
specklepy.core.aligment.get_shifts for details. Default is 0.
outfile (specklepy.io.recfile, optional):
Object to write the result to, if provided.
in_dir (str, optional):
Path to the files. `None` is substituted by an empty string.
tmp_dir (str, optional):
Path of a directory in which the temporary results are stored in.
lazy_mode (bool, optional):
Set to False, to enforce the alignment of a single file with respect to the reference file. Default is True.
box_indexes (list, optional):
Constraining the search for the intensity peak to the specified box. Searching the full frames if not
provided.
debug (bool, optional):
Show debugging information. Default is False.
Returns:
reconstruction (np.ndarray):
The image reconstruction. The size depends on the mode argument.
"""
logger.info("Starting SSA reconstruction...")
# Check parameters
if not isinstance(files, (list, np.ndarray)):
if isinstance(files, str):
files = [files]
else:
raise SpecklepyTypeError('ssa()',
argname='files',
argtype=type(files),
expected='list')
if isinstance(mode, str):
if mode not in ['same', 'full', 'valid']:
raise SpecklepyValueError('ssa()',
argname='mode',
argvalue=mode,
expected="'same', 'full' or 'valid'")
else:
raise SpecklepyTypeError('ssa()',
argname='mode',
argtype=type(mode),
expected='str')
if reference_file is None:
reference_file = files[0]
elif isinstance(reference_file, int):
reference_file = files[reference_file]
elif not isinstance(reference_file, str):
raise SpecklepyTypeError('ssa()',
argname='reference_file',
argtype=type(reference_file),
expected='str or int')
if outfile is None:
pass
elif isinstance(outfile, str):
outfile = ReconstructionFile(files=files,
filename=outfile,
cards={"RECONSTRUCTION": "SSA"})
elif isinstance(outfile, ReconstructionFile):
pass
else:
raise SpecklepyTypeError('ssa()',
argname='outfile',
argtype=type(outfile),
expected='str')
if in_dir is None:
in_dir = ''
reference_file = os.path.join(in_dir, reference_file)
if tmp_dir is not None:
if isinstance(tmp_dir, str) and not os.path.isdir(tmp_dir):
os.makedirs(tmp_dir)
if not isinstance(lazy_mode, bool):
raise SpecklepyTypeError('ssa()',
argname='lazy_mode',
argtype=type(lazy_mode),
expected='bool')
if box_indexes is not None:
box = Box(box_indexes)
else:
box = None
if 'variance_extension_name' in kwargs.keys():
var_ext = kwargs['variance_extension_name']
else:
var_ext = 'VAR'
if debug:
logger.setLevel('DEBUG')
logger.handlers[0].setLevel('DEBUG')
logger.info("Set logging level to DEBUG")
# Align reconstructions if multiple files are provided
if lazy_mode and len(files) == 1:
# Do not align just a single file
with fits.open(os.path.join(in_dir, files[0])) as hdu_list:
cube = hdu_list[0].data
if var_ext in hdu_list:
var_cube = hdu_list[var_ext].data
else:
var_cube = None
reconstruction, reconstruction_var = coadd_frames(
cube, var_cube=var_cube, box=box)
else:
# Compute temporary reconstructions of the individual cubes
tmp_files = []
for index, file in enumerate(files):
with fits.open(os.path.join(in_dir, file)) as hdu_list:
cube = hdu_list[0].data
if var_ext in hdu_list:
var_cube = hdu_list[var_ext].data
logger.debug(
f"Found variance extension {var_ext} in file {file}")
else:
logger.debug(
f"Did not find variance extension {var_ext} in file {file}"
)
var_cube = None
tmp, tmp_var = coadd_frames(cube, var_cube=var_cube, box=box)
if debug:
imshow(box(tmp), norm='log')
tmp_file = os.path.basename(file).replace(".fits", "_ssa.fits")
tmp_file = os.path.join(tmp_dir, tmp_file)
logger.info(
"Saving interim SSA reconstruction of cube to {}".format(
tmp_file))
tmp_file_object = Outfile(tmp_file, data=tmp, verbose=True)
# Store variance of temporary reconstruction
if tmp_var is not None:
tmp_file_object.new_extension(var_ext, data=tmp_var)
del tmp_var
tmp_files.append(tmp_file)
# Align tmp reconstructions and add up
file_shifts, image_shape = alignment.get_shifts(
tmp_files,
reference_file=reference_file,
return_image_shape=True,
lazy_mode=True)
pad_vectors, ref_pad_vector = alignment.get_pad_vectors(
file_shifts,
cube_mode=(len(image_shape) == 3),
return_reference_image_pad_vector=True)
# Iterate over file-wise reconstructions
reconstruction = None
reconstruction_var = None
for index, file in enumerate(tmp_files):
# Read data
with fits.open(file) as hdu_list:
tmp_image = hdu_list[0].data
if var_ext in hdu_list:
tmp_image_var = hdu_list[var_ext].data
else:
tmp_image_var = None
# Initialize or co-add reconstructions and var images
if reconstruction is None:
reconstruction = alignment.pad_array(
tmp_image,
pad_vectors[index],
mode=mode,
reference_image_pad_vector=ref_pad_vector)
if tmp_image_var is not None:
reconstruction_var = alignment.pad_array(
tmp_image_var,
pad_vectors[index],
mode=mode,
reference_image_pad_vector=ref_pad_vector)
else:
reconstruction += alignment.pad_array(
tmp_image,
pad_vectors[index],
mode=mode,
reference_image_pad_vector=ref_pad_vector)
if tmp_image_var is not None:
reconstruction_var += alignment.pad_array(
tmp_image_var,
pad_vectors[index],
mode=mode,
reference_image_pad_vector=ref_pad_vector)
logger.info("Reconstruction finished...")
# Save the result to an Outfile
if outfile is not None:
outfile.data = reconstruction
if reconstruction_var is not None:
outfile.new_extension(name=var_ext, data=reconstruction_var)
# Return reconstruction (and the variance map if computed)
if reconstruction_var is not None:
return reconstruction, reconstruction_var
return reconstruction
def __init__(self, in_files, psf_files, shifts, mode='same', in_dir=None):
""" Initialize a FourierObject instance.
Args:
in_files (list):
List of paths of the input files.
psf_files (list):
List of paths of the PSF files.
shifts (list):
List of integer shifts between the files.
mode (str, optional):
Define the size of the output image as 'same' to the reference image or expanding to include the 'full'
covered field. Default is 'same'.
in_dir (str, optional):
Path to the input files.
"""
# Assert that there are the same number of inFiles and psfFiles, which should be the case after running the
# holography function.
if not len(in_files) == len(psf_files):
raise ValueError(
f"The number of input files ({len(in_files)}) and PSF files ({len(psf_files)}) do not "
f"match!")
self.in_files = in_files
self.psf_files = psf_files
self.shifts = shifts
# Check whether mode is supported
if mode not in ['same', 'full', 'valid']:
raise SpecklepyValueError(
'FourierObject',
argname='mode',
argvalue=mode,
expected="either 'same', 'full', or 'valid'")
self.mode = mode
if in_dir is None:
self.in_dir = ''
else:
self.in_dir = in_dir
# Extract padding vectors for images and reference image
logger.info("Initializing padding vectors")
# files_contain_data_cubes = fits.getdata(in_files[0]).ndim == 3
self.pad_vectors, self.reference_image_pad_vector = get_pad_vectors(
shifts=shifts,
cube_mode=False,
return_reference_image_pad_vector=True)
file_index = 0
image_pad_vector = self.pad_vectors[file_index]
# Get example image frame, used as final image size
image_file = in_files[file_index]
logger.info(f"\tUsing example image frame from {image_file}")
img = fits.getdata(os.path.join(
self.in_dir, image_file))[0] # Remove time axis padding
img = pad_array(
array=img,
pad_vector=image_pad_vector,
mode=mode,
reference_image_pad_vector=self.reference_image_pad_vector)
logger.info(f"\tShift: {shifts[file_index]}")
logger.info(f"\tShape: {img.shape}")
# Get example PSF frame
psf_file = psf_files[file_index]
logger.info(f"\tUsing example PSF frame from {psf_file}")
psf = fits.getdata(psf_file)[0]
logger.info(f"\tShape: {psf.shape}")
# Estimate the padding vector for the f_psf frames to have the same xy-extent as f_img
dx = img.shape[0] - psf.shape[0]
dy = img.shape[1] - psf.shape[1]
psf_pad_vector = ((dx // 2, int(np.ceil(dx / 2))),
(dy // 2, int(np.ceil(dy / 2))))
logger.info(f"\tPad_width for PSFs: {psf_pad_vector}")
# Apply padding to PSF frame
psf = np.pad(
psf,
psf_pad_vector,
mode='constant',
)
if not img.shape == psf.shape:
raise ValueError(
f"The Fourier transformed images and PSFs have different shape, {img.shape} and "
f"{psf.shape}. Something went wrong with the padding!")
self.psf_pad_vector = psf_pad_vector
# Initialize the enumerator, denominator and Fourier object attributes
self.enumerator = np.zeros(img.shape, dtype='complex128')
self.denominator = np.zeros(img.shape, dtype='complex128')
self.fourier_image = np.zeros(img.shape, dtype='complex128')