#wv_index_list = [self.input_psfs_wvs.index(wv) for wv in wvs]
__author__ = 'jruffio'
import multiprocessing as mp
import ctypes
import numpy as np
import pyklip.spectra_management as spec
import os
import itertools
from pyklip.fmlib.nofm import NoFM
import pyklip.fm as fm
from scipy import interpolate
from copy import copy
import astropy.io.fits as pyfits
from time import time
import scipy.ndimage as ndimage
debug = False
[docs]
class MatchedFilter(NoFM):
"""
Matched filter with forward modelling.
"""
def __init__(self, inputs_shape,
numbasis,
input_psfs,
input_psfs_wvs,
spectrallib,
save_per_sector = None,
datatype="float",
fakes_sepPa_list = None,
disable_FM = None,
true_fakes_pos = None,
ref_center = None,
flipx = None,
rm_edge = None,
edge_threshold = None,
planet_radius = None,
background_width = None,
save_bbfm = None,
save_fm = None,
save_fmout_path = None):
'''
Defining the forward model matched filter parameters
Args:
inputs_shape: shape of the inputs numpy array. Typically (N, y, x)
numbasis: 1d numpy array consisting of the number of basis vectors to use
input_psfs: the psf of the image. A numpy array with shape (wv, y, x)
input_psfs_wvs: the wavelegnths that correspond to the input psfs
spectrallib: if not None, a list of spectra in raw DN units. The spectra should:
- have the total flux of the star, ie correspond to a contrast of 1.
- represent the total flux of the PSF and not the simply peak value.
- be multiplied by the atmospheric and instrumental transmission.
- have the same size as the number of images in the dataset.
save_per_sector: If not None, should be a filename where the fmout array will be saved after each sector.
(Caution: huge file!! easily tens of Gb.)
datatype: datatype to be used for the numpy arrays: "double" or "float" (default).
fakes_sepPa_list: [(sep_pix1,pa1),(sep_pix2,pa2),...].
List of separations and pas for the simulated planets in the data.
If not None, it will only calculated the matched filter at the position of the fakes and skip the rest.
disable_FM: Disable the calculation of the forward model in the code.
The unchanged original PSF will be used instead. (Default False)
true_fakes_pos: If True and fakes_only is True, calculate the forward model at the exact position of the
fakes and not at the center of the pixels. (Default False)
ref_center: reference center to which all the images are aligned. It should be the same center as the one
used in fm.parallelized.
flipx: Determines whether a relfection about the x axis is necessary to rotate image North-up East left.
Should match the same attribute in the instrument class.
rm_edge: When True (default), remove image edges to avoid edge effect. When there is more than 25% of NaNs
in the projection of the FM model on the data, the result of the projection is set to NaNs right away.
edge_threshold: if rm_edge is true, defines the fraction of pixels that needs to be finite in the aperture
around the planet for the projection to be calculated.
planet_radius: Radius of the aperture to be used for the matched filter (pick something of the order of the
2xFWHM)
background_width: Half the width of the arc in which the local standard deviation will be calculated.
save_bbfm: If true, saves the broadband forward models
savebfm: If true, saves all the forward models. This is only tractable for extremely small dataset.
save_fmout_path: If not None, path to a directory where to save fmout at the end.
'''
# allocate super class
super(MatchedFilter, self).__init__(inputs_shape, np.array(numbasis))
self.supports_rdi = True # flag needed while not all fm classes support RDI
self.save_bbfm = save_bbfm
self.save_fm = save_fm
if (self.save_bbfm or self.save_fm) and (np.size(numbasis) != 1 or len(spectrallib) != 1):
raise ValueError("save_bbfm or save_fm only work when np.size(numbasis)==1 and len(self.spectrallib)==1 ")
if (self.save_bbfm and self.save_fm):
raise ValueError("save_bbfm or save_fm cannot be used at the same time ")
self.save_fmout_path = save_fmout_path
if rm_edge is not None:
self.rm_edge = rm_edge
else:
self.rm_edge = True
if edge_threshold is None:
self.edge_threshold = 0.75
else:
self.edge_threshold = edge_threshold
if true_fakes_pos is None:
self.true_fakes_pos = False
else:
self.true_fakes_pos = true_fakes_pos
if datatype=="double":
self.data_type = ctypes.c_double
elif datatype=="float":
self.data_type = ctypes.c_float
if save_per_sector is not None:
self.fmout_dir = save_per_sector
self.save_raw_fmout = True
else:
self.save_raw_fmout = False
self.N_numbasis = np.size(numbasis)
self.ny = self.inputs_shape[1]
self.nx = self.inputs_shape[2]
self.N_frames = self.inputs_shape[0]
self.fakes_sepPa_list = fakes_sepPa_list
if disable_FM is None:
self.disable_FM = False
else:
self.disable_FM = disable_FM
self.inputs_shape = self.inputs_shape
self.input_psfs_wvs = list(np.array(input_psfs_wvs,dtype=self.data_type))
# Make sure the total flux of each PSF is unity for all wavelengths
# So the peak value won't be unity.
self.input_psfs = input_psfs/np.nansum(input_psfs,axis=(1,2))[:,None,None]
numwv_psf,ny_psf,nx_psf = self.input_psfs.shape
self.spectrallib = spectrallib
self.N_spectra = len(self.spectrallib)
# create bounds for PSF stamp size
self.row_m = int(np.floor(ny_psf/2.0)) # row_minus
self.row_p = int(np.ceil(ny_psf/2.0)) # row_plus
self.col_m = int(np.floor(nx_psf/2.0)) # col_minus
self.col_p = int(np.ceil(nx_psf/2.0)) # col_plus
self.psf_centx_notscaled = {}
self.psf_centy_notscaled = {}
self.curr_pa_fk = {}
self.curr_sep_fk = {}
x_psf_grid, y_psf_grid = np.meshgrid(np.arange(nx_psf * 1.)-nx_psf//2,np.arange(ny_psf* 1.)-ny_psf//2)
psfs_func_list = []
self.input_psfs[np.where(np.isnan(self.input_psfs))] = 0
import warnings
with warnings.catch_warnings():
warnings.simplefilter("ignore")
for wv_index in range(numwv_psf):
model_psf = self.input_psfs[wv_index, :, :]
psf_func = interpolate.LSQBivariateSpline(x_psf_grid.ravel(),y_psf_grid.ravel(),model_psf.ravel(),x_psf_grid[0,0:nx_psf-1]+0.5,y_psf_grid[0:ny_psf-1,0]+0.5)
psfs_func_list.append(psf_func)
self.psfs_func_list = psfs_func_list
ny_PSF,nx_PSF = input_psfs.shape[1:]
if ny_PSF < 8 or nx_PSF < 8:
raise Exception("PSF cube is too small. It needs a stamp width bigger than 8 pixels.")
stamp_PSF_x_grid, stamp_PSF_y_grid = np.meshgrid(np.arange(0,nx_PSF,1)-nx_PSF//2,np.arange(0,ny_PSF,1)-ny_PSF//2)
self.stamp_PSF_mask = np.ones((ny_PSF,nx_PSF))
r_PSF_stamp = np.sqrt((stamp_PSF_x_grid)**2 +(stamp_PSF_y_grid)**2)
if planet_radius is not None:
self.planet_radius = planet_radius
else:
self.planet_radius = int(np.round(np.min([ny_PSF,nx_PSF])*7./20.))
self.stamp_PSF_mask[np.where(r_PSF_stamp < self.planet_radius)] = np.nan
# self.stamp_PSF_mask[np.where(r_PSF_stamp < 4.)] = np.nan
if background_width is not None:
self.background_width = background_width
else:
self.background_width = np.min([ny_PSF,nx_PSF])//2
self.bbfm_mask = np.ones((self.planet_radius*2,self.planet_radius*2))
stamp_bbfm_x_grid, stamp_bbfm_y_grid = np.meshgrid(np.arange(0,self.planet_radius*2,1)-self.planet_radius,
np.arange(0,self.planet_radius*2,1)-self.planet_radius)
r_bbfm_stamp = abs((stamp_bbfm_x_grid) +(stamp_bbfm_y_grid)*1j)
self.bbfm_mask[np.where(r_bbfm_stamp < self.planet_radius)] = np.nan
if ref_center is not None:
self.ref_center = ref_center
# create some parameters for a blank canvas to draw psfs on
nx = inputs_shape[2]
ny = inputs_shape[1]
self.x_grid, self.y_grid = np.meshgrid(np.arange(nx * 1.)-ref_center[0], np.arange(ny * 1.)-ref_center[1])
self.r_grid = abs(self.x_grid +self.y_grid*1j)
self.pa_grid = np.arctan2( -self.x_grid,self.y_grid) % (2.0 * np.pi)
if flipx is not None:
sign = -1.
if flipx:
sign = 1.
self.th0_grid = np.arctan2(sign*self.x_grid,self.y_grid)
# def alloc_interm(self, max_sector_size, numsciframes):
# """Allocates shared memory array for intermediate step
#
# Intermediate step is allocated for a sector by sector basis
#
# Args:
# max_sector_size: number of pixels in this sector. Max because this can be variable. Stupid rotating sectors
#
# Returns:
# interm: mp.array to store intermediate products from one sector in
# interm_shape:shape of interm array (used to convert to numpy arrays)
#
# """
# if self.save_bbfm is not None:
# interm_size = self.planet_radius*self.planet_radius*self.ny*self.nx
# interm = mp.Array(ctypes.c_double, interm_size)
# interm_shape = [self.ny,self.nx,self.planet_radius,self.planet_radius]
#
# return interm, interm_shape
# else:
# return None,None
[docs]
def alloc_fmout(self, output_img_shape):
"""
Allocates shared memory for the output of the shared memory
Args:
output_img_shape: Not used
Returns:
fmout: mp.array to store auxilliary data in
fmout_shape: shape of auxilliary array = (3,self.N_spectra,self.N_numbasis,self.N_frames,self.ny,self.nx)
The 3 is for saving the different term of the matched filter:
0: dot product
1: square of the norm of the model
2: Local estimated variance of the data
3: Number of pixels used in the matched filter
"""
if self.save_bbfm:
fmout_size = 4*self.N_spectra*self.N_numbasis*self.N_frames*self.ny*self.nx + \
2*self.planet_radius*2*self.planet_radius*self.ny*self.nx
fmout = mp.Array(self.data_type, fmout_size)
fmout_shape = (4*self.N_spectra*self.N_numbasis*self.N_frames*self.ny*self.nx + \
2*self.planet_radius*2*self.planet_radius*self.ny*self.nx,1)
elif self.save_fm:
fmout_size = 4*self.N_spectra*self.N_numbasis*self.N_frames*self.ny*self.nx + \
self.N_frames*2*self.planet_radius*2*self.planet_radius*self.ny*self.nx
fmout = mp.Array(self.data_type, fmout_size)
fmout_shape = (4*self.N_spectra*self.N_numbasis*self.N_frames*self.ny*self.nx + \
self.N_frames*2*self.planet_radius*2*self.planet_radius*self.ny*self.nx,1)
else:
fmout_size = 4*self.N_spectra*self.N_numbasis*self.N_frames*self.ny*self.nx
fmout = mp.Array(self.data_type, fmout_size)
fmout_shape = (4,self.N_spectra,self.N_numbasis,self.N_frames,self.ny,self.nx)
return fmout, fmout_shape
[docs]
def skip_section(self, radstart, radend, phistart, phiend,flipx=True):
"""
Returns a boolean indicating if the section defined by (radstart, radend, phistart, phiend) should be skipped.
When True is returned the current section in the loop in klip_parallelized() is skipped.
Args:
radstart: minimum radial distance of sector [pixels]
radend: maximum radial distance of sector [pixels]
phistart: minimum azimuthal coordinate of sector [radians]
phiend: maximum azimuthal coordinate of sector [radians]
flipx: if True, flip x coordinate in final image
Returns:
Boolean: False so by default it never skips.
"""
margin_sep = np.sqrt(2)/2.
margin_phi = np.sqrt(2)/(2*radstart)
if self.fakes_sepPa_list is not None:
skipSectionBool = True
for sep_it,pa_it in self.fakes_sepPa_list:
if flipx:
paend= ((-phistart + np.pi/2.)% (2.0 * np.pi))
pastart = ((-phiend + np.pi/2.)% (2.0 * np.pi))
else:
pastart = ((phistart - np.pi/2.)% (2.0 * np.pi))
paend= ((phiend - np.pi/2.)% (2.0 * np.pi))
# Normal case when there are no 2pi wrap
if pastart < paend:
if (radstart-margin_sep<=sep_it<=radend+margin_sep) and ((pa_it%360)/180.*np.pi >= pastart-margin_phi) & ((pa_it%360)/180.*np.pi < paend+margin_phi):
skipSectionBool = False
break
# 2 pi wrap case
else:
if (radstart-margin_sep<=sep_it<=radend+margin_sep) and (((pa_it%360)/180.*np.pi >= pastart-margin_phi) | ((pa_it%360)/180.*np.pi < paend+margin_phi)):
skipSectionBool = False
break
else:
skipSectionBool = False
return skipSectionBool
[docs]
def fm_from_eigen(self, klmodes=None, evals=None, evecs=None, input_img_shape=None, input_img_num=None,
ref_psfs_indicies=None, section_ind=None,section_ind_nopadding=None, aligned_imgs=None, pas=None,
wvs=None, radstart=None, radend=None, phistart=None, phiend=None, padding=None,IOWA = None, ref_center=None,
parang=None, ref_wv=None, numbasis=None, fmout=None, perturbmag=None,klipped=None, flipx=True,
rdi_psfs=None, **kwargs):
"""
Calculate and project the FM at every pixel of the sector. Store the result in fmout.
Args:
klmodes: unpertrubed KL modes
evals: eigenvalues of the covariance matrix that generated the KL modes in ascending order
(lambda_0 is the 0 index) (shape of [nummaxKL])
evecs: corresponding eigenvectors (shape of [p, nummaxKL])
input_img_shape: 2-D shape of inpt images ([ysize, xsize])
input_img_num: index of sciece frame
ref_psfs_indicies: array of indicies for each reference PSF
section_ind: array indicies into the 2-D x-y image that correspond to this section.
Note needs be called as section_ind[0]
pas: array of N parallactic angles corresponding to N reference images [degrees]
wvs: array of N wavelengths of those referebce images
radstart: radius of start of segment
radend: radius of end of segment
phistart: azimuthal start of segment [radians]
phiend: azimuthal end of segment [radians]
padding: amount of padding on each side of sector
IOWA: tuple (IWA,OWA) where IWA = Inner working angle and OWA = Outer working angle both in pixels.
It defines the separation interva in which klip will be run.
ref_center: center of image
numbasis: array of KL basis cutoffs
parang: parallactic angle of input image [DEGREES]
ref_wv: wavelength of science image
fmout: numpy output array for FM output. Shape is (N, y, x, b)
klipped: array of shape (p,b) that is the PSF subtracted data for each of the b KLIP basis
cutoffs. If numbasis was an int, then sub_img_row_selected is just an array of length p
rdi_psfs: array of M RDI reference images in this sector.
kwargs: any other variables that we don't use but are part of the input
"""
if hasattr(self,"ref_center"):
if (self.ref_center[0] != ref_center[0]) or (self.ref_center[1] != ref_center[1]):
raise ValueError("ref_center needs to be the same when defining the matchedFilter class and calling parallelized")
if np.size(numbasis) != 1:
raise ValueError("Numbasis should only have a single element. e.g. numbasis = [30]. numbasis = [10,20,30] is not accepted.")
if self.save_bbfm:
fmout1 = fmout[:4*self.N_spectra*self.N_numbasis*self.N_frames*self.ny*self.nx]
fmout1.shape = (4,self.N_spectra,self.N_numbasis,self.N_frames,self.ny,self.nx)
fmout2 = fmout[4*self.N_spectra*self.N_numbasis*self.N_frames*self.ny*self.nx::]
fmout2.shape = (self.ny,self.nx,2*self.planet_radius,2*self.planet_radius)
elif self.save_fm:
fmout1 = fmout[:4*self.N_spectra*self.N_numbasis*self.N_frames*self.ny*self.nx]
fmout1.shape = (4,self.N_spectra,self.N_numbasis,self.N_frames,self.ny,self.nx)
fmout2 = fmout[4*self.N_spectra*self.N_numbasis*self.N_frames*self.ny*self.nx::]
fmout2.shape = (self.N_frames,self.ny,self.nx,2*self.planet_radius,2*self.planet_radius)
else:
fmout1 = fmout
ref_wv = ref_wv.astype(self.data_type)
sci = aligned_imgs[input_img_num, section_ind[0]]
refs = aligned_imgs[ref_psfs_indicies, :]
refs = refs[:, section_ind[0]]
ref_rdi_indices = np.zeros(np.size(ref_psfs_indicies)) # only used to track RDI frames. 1 if RDI.
# handle the RDI case
if rdi_psfs is not None:
# there are RDI frames
refs = np.append(refs, rdi_psfs, axis=0)
pas = np.append(pas, np.nan * np.ones(rdi_psfs.shape[0]))
wvs = np.append(wvs, np.nan * np.ones(rdi_psfs.shape[0]))
ref_rdi_indices = np.append(ref_rdi_indices, np.ones(rdi_psfs.shape[0]))
# Calculate the PA,sep 2D map
if hasattr(self,"x_grid") and hasattr(self,"y_grid"):
x_grid, y_grid = self.x_grid,self.y_grid
else:
x_grid, y_grid = np.meshgrid(np.arange(self.nx * 1.)-ref_center[0], np.arange(self.ny * 1.)-ref_center[1])
x_grid=x_grid.astype(self.data_type)
y_grid=y_grid.astype(self.data_type)
# Define the masks for where the planet is and the background.
if hasattr(self,"r_grid"):
r_grid = self.r_grid
else:
r_grid = np.sqrt((x_grid)**2 + (y_grid)**2)
if hasattr(self,"pa_grid"):
pa_grid = self.pa_grid
else:
pa_grid = np.arctan2( -x_grid,y_grid) % (2.0 * np.pi)
if flipx:
paend= ((-phistart + np.pi/2.)% (2.0 * np.pi))
pastart = ((-phiend + np.pi/2.)% (2.0 * np.pi))
else:
pastart = ((phistart - np.pi/2.)% (2.0 * np.pi))
paend= ((phiend - np.pi/2.)% (2.0 * np.pi))
# Normal case when there are no 2pi wrap
if pastart < paend:
where_section = np.where((r_grid >= radstart) & (r_grid < radend) & (pa_grid >= pastart) & (pa_grid < paend))
# 2 pi wrap case
else:
where_section = np.where((r_grid >= radstart) & (r_grid < radend) & ((pa_grid >= pastart) | (pa_grid < paend)))
# Get a list of the PAs and sep of the PA,sep map falling in the current section
r_list = r_grid[where_section]
pa_list = pa_grid[where_section]
x_list = x_grid[where_section]
y_list = y_grid[where_section]
row_id_list = where_section[0]
col_id_list = where_section[1]
# Only select pixel with fakes if needed
if self.fakes_sepPa_list is not None:
r_list_tmp = []
pa_list_tmp = []
row_id_list_tmp = []
col_id_list_tmp = []
for sep_it,pa_it in self.fakes_sepPa_list:
x_it = sep_it*np.cos(np.radians(90+pa_it))
y_it = sep_it*np.sin(np.radians(90+pa_it))
dist_list = np.sqrt((x_list-x_it)**2+(y_list-y_it)**2)
min_id = np.nanargmin(dist_list)
min_dist = dist_list[min_id]
if min_dist < np.sqrt(2)/2.:
if self.true_fakes_pos:
r_list_tmp.append(sep_it)
pa_list_tmp.append(np.radians(pa_it))
else:
r_list_tmp.append(r_list[min_id])
pa_list_tmp.append(pa_list[min_id])
row_id_list_tmp.append(row_id_list[min_id])
col_id_list_tmp.append(col_id_list[min_id])
r_list = r_list_tmp
pa_list = pa_list_tmp
row_id_list = row_id_list_tmp
col_id_list = col_id_list_tmp
greenboard = np.zeros((self.ny,self.nx))
x_bbfm, y_bbfm = np.meshgrid(np.arange(self.ny, dtype=float), np.arange(self.nx, dtype=float))
# flip x if needed to get East left of North
if flipx is True:
x_bbfm = ref_center[0] - (x_bbfm - ref_center[0])
# do rotation. CW rotation formula to get a CCW of the image
angle_rad = np.radians(parang)
cosa = np.cos(angle_rad)
sina = np.sin(angle_rad)
xp = (x_bbfm-ref_center[0])*cosa + (y_bbfm-ref_center[1])*sina + ref_center[0]
yp = -(x_bbfm-ref_center[0])*sina + (y_bbfm-ref_center[1])*cosa + ref_center[1]
# create bounds for PSF stamp size
self.bbfm_m = int(np.floor(self.planet_radius))
self.bbfm_p = int(np.ceil(self.planet_radius))
# Loop over the input template spectra and the number of KL modes in numbasis
for spec_id,N_KL_id in itertools.product(range(self.N_spectra),range(self.N_numbasis)):
# Calculate the projection of the FM and the klipped section for every pixel in the section.
# 1/ Inject a fake at one pa and sep in the science image
# 2/ Inject the corresponding planets at the same PA and sep in the reference images remembering that the
# references rotate.
# 3/ Calculate the perturbation of the KL modes
# 4/ Calculate the FM
# 5/ Calculate dot product (matched filter)
for sep_fk,pa_fk,row_id,col_id in zip(r_list,np.rad2deg(pa_list),row_id_list,col_id_list):
# print(sep_fk,pa_fk,row_id,col_id)
# 1/ Inject a fake at one pa and sep in the science image
try:
model_sci,mask = self.generate_model_sci(input_img_shape, section_ind, parang, ref_wv,
radstart, radend, phistart, phiend, padding, ref_center,
parang, ref_wv,sep_fk,pa_fk, flipx)
except:
# print("skip",sep_fk,pa_fk)
continue
# Normalize the science image according to the spectrum. the model is normalize to unit contrast,
model_sci = model_sci*self.spectrallib[spec_id][input_img_num]
where_fk = np.where(mask==2)[0]
where_background = np.where(mask>=1)[0] # Caution: it includes where the fake is...
where_background_strict = np.where(mask==1)[0]
# JB hack trying to include covariances
# if 0:
# section_ind_y,section_ind_x = np.unravel_index(section_ind,(self.ny,self.nx))
# section_ind_y,section_ind_x = section_ind_y[0][where_fk],section_ind_x[0][where_fk]
# distances = np.sqrt((section_ind_x[:,None]-section_ind_x[None,:])**2+(section_ind_y[:,None]-section_ind_y[None,:])**2)
# from scipy.interpolate import interp1d
# corr_profile = interp1d([0,1,np.sqrt(2),2],[1,0.2,-0.1,0],bounds_error=False, fill_value=0)
# spatial_covar =corr_profile(distances)
# inv_spatial_covar = np.linalg.inv(spatial_covar)
#
# # import matplotlib.pyplot as plt
# # plt.subplot(1,3,1)
# # plt.imshow(distances,interpolation="nearest")
# # plt.subplot(1,3,2)
# # plt.imshow(spatial_covar,interpolation="nearest")
# # plt.subplot(1,3,3)
# # plt.imshow(np.linalg.inv(spatial_covar),interpolation="nearest")
# # plt.show()
# # exit()
if np.size(klipped[where_fk,N_KL_id]) == 0:
quitnow = True
elif self.rm_edge and float(np.sum(np.isfinite(klipped[where_fk,N_KL_id])))/float(np.size(klipped[where_fk,N_KL_id]))<=self.edge_threshold:
quitnow = True
else:
quitnow = False
if quitnow:
fmout1[0,spec_id,N_KL_id,input_img_num,row_id,col_id] = np.nan
fmout1[1,spec_id,N_KL_id,input_img_num,row_id,col_id] = np.nan
fmout1[2,spec_id,N_KL_id,input_img_num,row_id,col_id] = np.nan
fmout1[3,spec_id,N_KL_id,input_img_num,row_id,col_id] = np.nan
continue
# 2/ Inject the corresponding planets at the same PA and sep in the reference images remembering that the
# references rotate.
if not self.disable_FM:
models_ref = self.generate_models(input_img_shape, section_ind, pas, wvs, radstart, radend,
phistart, phiend, padding, ref_center, parang, ref_wv,
sep_fk, pa_fk, flipx, ref_rdi_indices)
# Normalize the models with the spectrum. the model is normalize to unit contrast,
input_spectrum = self.spectrallib[spec_id][ref_psfs_indicies]
if rdi_psfs is not None:
input_spectrum = np.append(input_spectrum, np.zeros(rdi_psfs.shape[0]), axis=0)
models_ref = models_ref * input_spectrum[:, None]
# 3/ Calculate the perturbation of the KL modes
# using original Kl modes and reference models, compute the perturbed KL modes.
# Spectrum is already in the model, that's why we use perturb_specIncluded(). (Much faster)
delta_KL = fm.perturb_specIncluded(evals, evecs, klmodes, refs, models_ref)
# 4/ Calculate the FM: calculate postklip_psf using delta_KL
# postklip_psf has unit broadband contrast
model_sci_fk = model_sci[where_fk]
delta_KL_kl = delta_KL[:,where_fk]
klmodes_fk = klmodes[:,where_fk]
postklip_psf = calculate_fm_opti(delta_KL, klmodes,sci, model_sci_fk,delta_KL_kl,klmodes_fk)
# postklip_psf, oversubtraction, selfsubtraction = fm.calculate_fm(delta_KL, klmodes, numbasis,
# sci, model_sci, inputflux=None)
else:
#if one doesn't want the FM
if np.size(numbasis) == 1:
postklip_psf = model_sci[None,:]
else:
# Mh, is this actually working? shouldn't I use tile?
postklip_psf = model_sci
# 5/ Calculate dot product (matched filter)
# fmout_shape = (3,self.N_spectra,self.N_numbasis,self.N_frames,self.ny,self.nx)
# The 3 is for saving the different term of the matched filter:
# 0: dot product
# 1: square of the norm of the model
# 2: Local estimated variance of the data
sky = np.nanmean(klipped[where_background_strict,N_KL_id])
# postklip_psf[N_KL_id,where_fk] = postklip_psf[N_KL_id,where_fk]-np.mean(postklip_psf[N_KL_id,where_background])
# Subtract local sky background to the klipped image
klipped_sub = klipped[where_fk,N_KL_id]-sky
# klipped_sub_finite = np.where(np.isfinite(klipped_sub))
# klipped_sub_nan = np.where(np.isnan(klipped_sub))
postklip_psf[N_KL_id,np.where(np.isnan(klipped_sub))[0]] = np.nan
postklip_psf_fk = postklip_psf[N_KL_id,:]
dot_prod = np.nansum(klipped_sub*postklip_psf_fk)
model_norm = np.nansum(postklip_psf_fk*postklip_psf_fk)
klipped_rm_pl = copy(klipped[:,N_KL_id]) -sky
klipped_rm_pl[where_fk] -= (dot_prod/model_norm)*postklip_psf_fk
klipped_rm_pl_bkg = klipped_rm_pl[where_background]
if self.rm_edge and (float(np.sum(np.isfinite(klipped_rm_pl_bkg)))/float(np.size(klipped_rm_pl_bkg))<=self.edge_threshold):
variance = np.nan
npix = np.nan
else:
variance = np.nanvar(klipped_rm_pl_bkg)
npix = np.sum(np.isfinite(klipped_rm_pl_bkg))
# JB hack trying to include covariances
# else:
# dot_prod = np.nansum(klipped_sub*postklip_psf_fk)
# model_norm = np.nansum(postklip_psf_fk*postklip_psf_fk)
# klipped_rm_pl = copy(klipped[:,N_KL_id]) -sky
# klipped_rm_pl[where_fk] -= (dot_prod/model_norm)*postklip_psf_fk
# klipped_rm_pl_bkg = klipped_rm_pl[where_background]
# if self.rm_edge and (float(np.sum(np.isfinite(klipped_rm_pl_bkg)))/float(np.size(klipped_rm_pl_bkg))<=self.edge_threshold):
# variance = np.nan
# npix = np.nan
# else:
# variance = np.nanvar(klipped_rm_pl_bkg)
# npix = np.sum(np.isfinite(klipped_rm_pl_bkg))
#
# inv_spatial_covar = inv_spatial_covar/variance
# # print(klipped_sub.shape,inv_spatial_covar.shape)
# # print(np.dot(klipped_sub,inv_spatial_covar).shape)
# dot_prod = np.nansum(np.dot(klipped_sub,inv_spatial_covar)*postklip_psf_fk)
# model_norm = np.nansum(np.dot(postklip_psf_fk,inv_spatial_covar)*postklip_psf_fk)
# variance = 1.0
# # exit()
fmout1[0,spec_id,N_KL_id,input_img_num,row_id,col_id] = dot_prod
fmout1[1,spec_id,N_KL_id,input_img_num,row_id,col_id] = model_norm
fmout1[2,spec_id,N_KL_id,input_img_num,row_id,col_id] = variance
fmout1[3,spec_id,N_KL_id,input_img_num,row_id,col_id] = npix
if self.save_bbfm:
greenboard.shape = [input_img_shape[0] * input_img_shape[1]]
greenboard[section_ind[0][where_fk]] = postklip_psf[N_KL_id,:]
greenboard.shape = [input_img_shape[0],input_img_shape[1]]
rot_stamp = ndimage.map_coordinates(greenboard,
[yp[(row_id-self.bbfm_m):(row_id+self.bbfm_p),
(col_id-self.bbfm_m):(col_id+self.bbfm_p)].ravel(),
xp[(row_id-self.bbfm_m):(row_id+self.bbfm_p),
(col_id-self.bbfm_m):(col_id+self.bbfm_p)].ravel()],
cval=np.nan)
rot_stamp.shape = [2*self.planet_radius,2*self.planet_radius]
fmout2[row_id,col_id,:,:] = fmout2[row_id,col_id,:,:]+rot_stamp
elif self.save_fm:
greenboard.shape = [input_img_shape[0] * input_img_shape[1]]
greenboard[section_ind[0][where_fk]] = postklip_psf[N_KL_id,:]
greenboard.shape = [input_img_shape[0],input_img_shape[1]]
greenboard_pad = np.pad(greenboard,((self.bbfm_m,self.bbfm_p),(self.bbfm_m,self.bbfm_p)),mode="constant",constant_values=np.nan)
mystamp = copy(greenboard_pad[(row_id):(row_id+self.bbfm_m+self.bbfm_p),
(col_id):(col_id+self.bbfm_m+self.bbfm_p)])
# print((row_id-self.bbfm_m),(row_id+self.bbfm_p),
# (col_id-self.bbfm_m),(col_id+self.bbfm_p))
# import matplotlib.pyplot as plt
# plt.subplot(1,2,1)
# plt.imshow(greenboard, interpolation="nearest")
# plt.subplot(1,2,2)
# plt.imshow(mystamp, interpolation="nearest")
# plt.show()
fmout2[input_img_num,row_id,col_id,:,:] = mystamp
if 0:
print(dot_prod,model_norm,variance,npix)
# Plot sector, klipped and FM model for debug only
if 0: # and row_id>=10 and col_id > 5 # and np.nansum(klipped[where_fk,N_KL_id]) != 0:
# plt.figure(1)
# blackboard = np.zeros((self.ny,self.nx))
# blackboard.shape = [input_img_shape[0] * input_img_shape[1]]
# blackboard[section_ind] = klipped
# blackboard.shape = [input_img_shape[0],input_img_shape[1]]
# plt.imshow(blackboard,interpolation="nearest")
# plt.colorbar()
# plt.figure(2)
# for k in range(numbasis[0]):
# blackboard = np.zeros((self.ny,self.nx))
# blackboard.shape = [input_img_shape[0] * input_img_shape[1]]
# blackboard[section_ind] = klmodes[k,:]
# blackboard.shape = [input_img_shape[0],input_img_shape[1]]
# plt.subplot(1,numbasis[0],k+1)
# plt.imshow(blackboard[::-1,:],interpolation="nearest")
# plt.title("KL{0}".format(k))
# plt.colorbar()
# plt.show()
#if 0:
# print(klipped_sub)
# print(np.isfinite(klipped_sub))
# print(np.size(klipped_sub))
# print(float(np.sum(np.isfinite(klipped_sub)))/float(np.size(klipped_sub)))
# print(float(np.sum(np.isfinite(klipped[where_background,N_KL_id])))/float(np.size(klipped[where_background,N_KL_id])))
print(sep_fk,pa_fk,row_id,col_id)
print(dot_prod,model_norm,variance,npix)
# print(np.nanmean(klipped-sky),sky,dot_prod,model_norm,np.nanmean((dot_prod/model_norm)*postklip_psf[N_KL_id,:]))
# print(klipped.shape,postklip_psf[N_KL_id,:].shape)
# print(float(np.sum(np.isfinite(klipped_rm_pl[where_background]))),float(np.size(klipped_rm_pl[where_background])))
blackboard1 = np.zeros((self.ny,self.nx))
blackboard2 = np.zeros((self.ny,self.nx))
blackboard3 = np.zeros((self.ny,self.nx))
#print(section_ind)
plt.figure(1)
plt.subplot(1,3,1)
blackboard1.shape = [input_img_shape[0] * input_img_shape[1]]
blackboard1[section_ind] = mask
blackboard1[section_ind] = blackboard1[section_ind] + 1
blackboard1.shape = [input_img_shape[0],input_img_shape[1]]
plt.imshow(blackboard1,interpolation="nearest")
plt.colorbar()
plt.subplot(1,3,2)
blackboard2.shape = [input_img_shape[0] * input_img_shape[1]]
# blackboard2[section_ind[0][where_fk]] = klipped[where_fk,N_KL_id]
blackboard2[section_ind[0]] = klipped#klipped_rm_pl
blackboard2.shape = [input_img_shape[0],input_img_shape[1]]
plt.imshow(blackboard2,interpolation="nearest")
plt.colorbar()
plt.subplot(1,3,3)
blackboard3.shape = [input_img_shape[0] * input_img_shape[1]]
blackboard3[section_ind[0][where_fk]] = postklip_psf[N_KL_id,:]
blackboard3.shape = [input_img_shape[0],input_img_shape[1]]
plt.imshow(blackboard3,interpolation="nearest")
plt.colorbar()
#print(klipped[where_fk,N_KL_id])
#print(postklip_psf[N_KL_id,where_fk])
# print(np.sum(klipped[where_fk,N_KL_id]*postklip_psf[N_KL_id,where_fk]))
# print(np.sum(postklip_psf[N_KL_id,where_fk]*postklip_psf[N_KL_id,where_fk]))
# print(np.sum(klipped[where_fk,N_KL_id]*klipped[where_fk,N_KL_id]))
plt.show()
# exit()
[docs]
def fm_end_sector(self, interm_data=None, fmout=None, sector_index=None,
section_indicies=None):
"""
Save the fmout object at the end of each sector if save_per_sector was defined when initializing the class.
"""
#fmout_shape = (3,self.N_spectra,self.N_numbasis,self.N_frames,self.ny,self.nx)
if self.save_raw_fmout:
if self.save_bbfm or self.save_fm:
fmout1 = fmout[:4*self.N_spectra*self.N_numbasis*self.N_frames*self.ny*self.nx]
fmout1.shape = (4,self.N_spectra,self.N_numbasis,self.N_frames,self.ny,self.nx)
else:
fmout1 = fmout
hdu = pyfits.PrimaryHDU(fmout)
hdulist = pyfits.HDUList([hdu])
hdulist.writeto(self.fmout_dir,clobber=True)
return
[docs]
def save_fmout(self, dataset, fmout, outputdir, fileprefix, numbasis, klipparams=None, calibrate_flux=False,
spectrum=None, pixel_weights=1):
"""
Saves the fmout data to disk following the instrument's savedata function
Args:
dataset: Instruments.Data instance. Will use its dataset.savedata() function to save data
fmout: the fmout data passed from fm.klip_parallelized which is passed as the output of cleanup_fmout
outputdir: output directory
fileprefix: the fileprefix to prepend the file name
numbasis: KL mode cutoffs used
klipparams: string with KLIP-FM parameters
calibrate_flux: if True, flux calibrate the data (if applicable)
spectrum: if not None, the spectrum to weight the data by. Length same as dataset.wvs
pixel_weights: weights for each pixel for weighted mean. [Not used in matched filter]
"""
if self.save_bbfm:
fmout1 = fmout[:4*self.N_spectra*self.N_numbasis*self.N_frames*self.ny*self.nx]
fmout1.shape = (4,self.N_spectra,self.N_numbasis,self.N_frames,self.ny,self.nx)
fmout2 = fmout[4*self.N_spectra*self.N_numbasis*self.N_frames*self.ny*self.nx::]
fmout2.shape = (self.ny,self.nx,2*self.planet_radius,2*self.planet_radius)
elif self.save_fm:
fmout1 = fmout[:4*self.N_spectra*self.N_numbasis*self.N_frames*self.ny*self.nx]
fmout1.shape = (4,self.N_spectra,self.N_numbasis,self.N_frames,self.ny,self.nx)
fmout2 = fmout[4*self.N_spectra*self.N_numbasis*self.N_frames*self.ny*self.nx::]
fmout2.shape = (self.N_frames,self.ny,self.nx,2*self.planet_radius,2*self.planet_radius)
else:
fmout1 = fmout
if self.save_fmout_path is not None:
hdu = pyfits.PrimaryHDU(fmout1)
hdulist = pyfits.HDUList([hdu])
suffix = "fmout"
try:
hdulist.writeto(os.path.join(self.save_fmout_path,fileprefix+'-'+suffix+'.fits'), overwrite=True)
except TypeError:
hdulist.writeto(os.path.join(self.save_fmout_path,fileprefix+'-'+suffix+'.fits'), clobber=True)
#fmout_shape = (3,self.N_spectra,self.N_numbasis,self.N_frames,self.ny,self.nx)
fmout1[np.where(fmout1==0)] = np.nan
# hdu = pyfits.PrimaryHDU(fmout1)
# hdulist = pyfits.HDUList([hdu])
# hdulist.writeto(outputdir+os.path.sep+'fmout1_mvt6_KL30.fits',clobber=True)
# The mf.MatchedFilter class calculate the projection of the FM on the data for each pixel and images.
# The final combination to form the cross cross correlation, matched filter and contrast maps is done right
# here.
FMCC_map = np.nansum(fmout1[0,:,:,:,:,:],axis=2) \
/ np.sqrt(np.nansum(fmout1[1,:,:,:,:,:],axis=2))
FMCC_map[np.where(FMCC_map==0)]=np.nan
self.FMCC_map = FMCC_map
FMMF_map = np.nansum(fmout1[0,:,:,:,:,:]/fmout1[2,:,:,:,:,:],axis=2) \
/ np.sqrt(np.nansum(fmout1[1,:,:,:,:,:]/fmout1[2,:,:,:,:,:],axis=2))
FMMF_map[np.where(FMMF_map==0)]=np.nan
self.FMMF_map = FMMF_map
contrast_map = np.nansum(fmout1[0,:,:,:,:,:]/fmout1[2,:,:,:,:,:],axis=2) \
/ np.nansum(fmout1[1,:,:,:,:,:]/fmout1[2,:,:,:,:,:],axis=2)
contrast_map[np.where(contrast_map==0)]=np.nan
self.contrast_map = contrast_map
self.FMNpix_map = np.nansum(fmout1[3,:,:,:,:,:],axis=2)
self.metricMap = [self.FMMF_map,self.FMCC_map,self.contrast_map]
for k in range(np.size(self.numbasis)):
# Save the outputs (matched filter, shape map and klipped image) as fits files
suffix = "FMMF-KL{0}".format(self.numbasis[k])
dataset.savedata(outputdir+os.path.sep+fileprefix+'-'+suffix+'.fits',
self.FMMF_map[0,k,:,:],
filetype=suffix)
suffix = "FMCont-KL{0}".format(self.numbasis[k])
dataset.savedata(outputdir+os.path.sep+fileprefix+'-'+suffix+'.fits',
self.contrast_map[0,k,:,:],
filetype=suffix)
suffix = "FMCC-KL{0}".format(self.numbasis[k])
dataset.savedata(outputdir+os.path.sep+fileprefix+'-'+suffix+'.fits',
self.FMCC_map[0,k,:,:],
filetype=suffix)
suffix = "FMN_pix-KL{0}".format(self.numbasis[k])
dataset.savedata(outputdir+os.path.sep+fileprefix+'-'+suffix+'.fits',
self.FMNpix_map[0,k,:,:],
filetype=suffix)
if self.save_bbfm:
suffix = "BBFM-KL{0}".format(self.numbasis[k])
dataset.savedata(outputdir+os.path.sep+fileprefix+'-'+suffix+'.fits',
fmout2,
filetype=suffix)
elif self.save_fm:
suffix = "AllFM-KL{0}".format(self.numbasis[k])
dataset.savedata(outputdir+os.path.sep+fileprefix+'-'+suffix+'.fits',
fmout2,
filetype=suffix)
return
[docs]
def generate_model_sci(self, input_img_shape, section_ind, pa, wv, radstart, radend, phistart, phiend, padding, ref_center, parang, ref_wv,sep_fk,pa_fk, flipx):
"""
Generate model PSFs at the correct location of this segment of the science image denotated by its wv and
parallactic angle.
Args:
input_img_shape: 2-D shape of inpt images ([ysize, xsize])
section_ind: array indicies into the 2-D x-y image that correspond to this section.
Note needs be called as section_ind[0]
pa: parallactic angle of the science image [degrees]
wv: wavelength of the science image
radstart: radius of start of segment (not used)
radend: radius of end of segment (not used)
phistart: azimuthal start of segment [radians] (not used)
phiend: azimuthal end of segment [radians] (not used)
padding: amount of padding on each side of sector
ref_center: center of image
parang: parallactic angle of input image [DEGREES] (not used)
ref_wv: wavelength of science image
sep_fk: separation of the planet to be injected.
pa_fk: position angle of the planet to be injected.
flipx: if True, flip x coordinate in final image
Return: (models, mask)
models: vector of size p where p is the number of pixels in the segment
mask: vector of size p where p is the number of pixels in the segment
if pixel == 1: arc shape where to calculate the standard deviation
if pixel == 2: 7 pixels disk around the position of the planet.
"""
# create some parameters for a blank canvas to draw psfs on
nx = input_img_shape[1]
ny = input_img_shape[0]
if hasattr(self,"x_grid") and hasattr(self,"y_grid"):
x_grid, y_grid = self.x_grid,self.y_grid
else:
x_grid, y_grid = np.meshgrid(np.arange(nx * 1.)-ref_center[0], np.arange(ny * 1.)-ref_center[1])
numwv, ny_psf, nx_psf = self.input_psfs.shape
# a blank img array of write model PSFs into
whiteboard = np.zeros((ny,nx))
# grab PSF given wavelength
wv_index = spec.find_nearest(self.input_psfs_wvs,wv)[1]
sign = -1.
if flipx:
sign = 1.
# The trigonometric calculation are save in a dictionary to avoid calculating them many times.
recalculate_trig = False
if pa not in self.psf_centx_notscaled:
recalculate_trig = True
else:
if pa_fk != self.curr_pa_fk[pa] or sep_fk != self.curr_sep_fk[pa]:
recalculate_trig = True
if recalculate_trig: # we could actually store the values for the different pas too...
# flipx requires the opposite rotation
self.psf_centx_notscaled[pa] = sep_fk * np.cos(np.radians(90. - sign*pa_fk - pa))
self.psf_centy_notscaled[pa] = sep_fk * np.sin(np.radians(90. - sign*pa_fk - pa))
self.curr_pa_fk[pa] = pa_fk
self.curr_sep_fk[pa] = sep_fk
psf_centx = (ref_wv/wv) * self.psf_centx_notscaled[pa]
psf_centy = (ref_wv/wv) * self.psf_centy_notscaled[pa]
# create a coordinate system for the image that is with respect to the model PSF
# round to nearest pixel and add offset for center
l = int(round(psf_centx + ref_center[0]))
k = int(round(psf_centy + ref_center[1]))
# recenter coordinate system about the location of the planet
# x_vec_stamp_centered = x_grid[0, (l-col_m):(l+col_p)]-psf_centx
# y_vec_stamp_centered = y_grid[(k-row_m):(k+row_p), 0]-psf_centy
x_vec_stamp_centered = x_grid[0, np.max([(l-self.col_m),0]):np.min([(l+self.col_p),nx])]-psf_centx
y_vec_stamp_centered = y_grid[np.max([(k-self.row_m),0]):np.min([(k+self.row_p),ny]), 0]-psf_centy
# rescale to account for the align and scaling of the refernce PSFs
# e.g. for longer wvs, the PSF has shrunk, so we need to shrink the coordinate system
x_vec_stamp_centered /= (ref_wv/wv)
y_vec_stamp_centered /= (ref_wv/wv)
# use intepolation spline to generate a model PSF and write to temp img
# whiteboard[(k-row_m):(k+row_p), (l-col_m):(l+col_p)] = \
# self.psfs_func_list[wv_index[0]](x_vec_stamp_centered,y_vec_stamp_centered).transpose()
whiteboard[np.max([(k-self.row_m),0]):np.min([(k+self.row_p),ny]), np.max([(l-self.col_m),0]):np.min([(l+self.col_p),nx])] = \
self.psfs_func_list[wv_index](x_vec_stamp_centered,y_vec_stamp_centered).transpose()
# write model img to output (segment is collapsed in x/y so need to reshape)
whiteboard.shape = [input_img_shape[0] * input_img_shape[1]]
segment_with_model = copy(whiteboard[section_ind])
whiteboard.shape = [input_img_shape[0],input_img_shape[1]]
# Define the masks for where the planet is and the background.
if hasattr(self,"th0_grid") and hasattr(self,"r_grid"):
r_grid = self.r_grid
th_grid = (self.th0_grid-sign*np.radians(pa))% (2.0 * np.pi)
else:
r_grid = abs(x_grid +y_grid*1j)
th_grid = (np.arctan2(sign*x_grid,y_grid)-sign*np.radians(pa))% (2.0 * np.pi)
w = self.background_width
thstart = (np.radians(pa_fk)- float(w)/sep_fk) % (2.0 * np.pi) # -(2*np.pi-np.radians(pa))
thend = (np.radians(pa_fk) + float(w)/sep_fk) % (2.0 * np.pi) # -(2*np.pi-np.radians(pa))
# thstart = (np.radians(pa_fk)- 2*float(w)/sep_fk) % (2.0 * np.pi) # -(2*np.pi-np.radians(pa))
#thend = (np.radians(pa_fk) + 2*float(w)/sep_fk) % (2.0 * np.pi) # -(2*np.pi-np.radians(pa))
if thstart < thend:
where_mask = np.where((r_grid>=(sep_fk-w)) & (r_grid<(sep_fk+w)) & (th_grid >= thstart) & (th_grid < thend))
else:
where_mask = np.where((r_grid>=(sep_fk-w)) & (r_grid<(sep_fk+w)) & ((th_grid >= thstart) | (th_grid < thend)))
whiteboard[where_mask] = 1
#TODO check the modification I did to these lines
whiteboard = np.pad(whiteboard,((self.row_m,self.row_p),(self.col_m,self.col_p)),mode="constant",constant_values=0)
# whiteboard[(k-row_m):(k+row_p), (l-col_m):(l+col_p)][np.where(np.isnan(self.stamp_PSF_mask))]=2
whiteboard[(k):(k+self.row_m+self.row_p), (l):(l+self.col_m+self.col_p)][np.where(np.isnan(self.stamp_PSF_mask))]=2
whiteboard = np.ascontiguousarray(whiteboard[self.row_m:self.row_m+input_img_shape[0],self.col_m:self.col_m+input_img_shape[1]])
whiteboard.shape = [input_img_shape[0] * input_img_shape[1]]
mask = whiteboard[section_ind]
# create a canvas to place the new PSF in the sector on
if 0:#np.size(np.where(mask==2)[0])==0: 296
print(pa,pa_fk)
print(thstart,thend)
whiteboard[section_ind] = whiteboard[section_ind] + 0.5
whiteboard.shape = (input_img_shape[0], input_img_shape[1])
blackboard = np.zeros((ny,nx))
blackboard.shape = [input_img_shape[0] * input_img_shape[1]]
blackboard[section_ind] = segment_with_model
blackboard.shape = [input_img_shape[0],input_img_shape[1]]
plt.figure(1)
plt.subplot(1,3,1)
im = plt.imshow(whiteboard)
plt.colorbar(im)
plt.subplot(1,3,2)
im = plt.imshow(blackboard)
plt.colorbar(im)
plt.subplot(1,3,3)
im = plt.imshow(np.degrees(th_grid))
plt.colorbar(im)
plt.show()
return segment_with_model,mask
[docs]
def generate_models(self, input_img_shape, section_ind, pas, wvs, radstart, radend, phistart, phiend, padding, ref_center, parang, ref_wv,sep_fk,pa_fk,flipx, rdi_indices):
"""
Generate model PSFs at the correct location of this segment for each image denotated by its wv and parallactic
angle.
Args:
input_img_shape: 2-D shape of inpt images ([ysize, xsize])
section_ind: array indicies into the 2-D x-y image that correspond to this section.
Note needs be called as section_ind[0]
pas: array of N parallactic angles corresponding to N images [degrees]
wvs: array of N wavelengths of those images
radstart: radius of start of segment (not used)
radend: radius of end of segment (not used)
phistart: azimuthal start of segment [radians] (not used)
phiend: azimuthal end of segment [radians] (not used)
padding: amount of padding on each side of sector
ref_center: center of image
parang: parallactic angle of input image [DEGREES] (not used)
ref_wv: wavelength of science image
sep_fk: separation of the planet to be injected.
pa_fk: position angle of the planet to be injected.
flipx: if True, flip x coordinate in final image
Return:
models: array of size (N, p) where p is the number of pixels in the segment
"""
# create some parameters for a blank canvas to draw psfs on
nx = input_img_shape[1]
ny = input_img_shape[0]
try:
x_grid, y_grid = self.x_grid,self.y_grid
except:
x_grid, y_grid = np.meshgrid(np.arange(nx * 1.)-ref_center[0], np.arange(ny * 1.)-ref_center[1])
sign = -1.
if flipx:
sign = 1.
# a blank img array of write model PSFs into
whiteboard = np.zeros((ny,nx))
models = []
#print(self.input_psfs.shape)
for pa, wv, is_rdi in zip(pas, wvs, rdi_indices):
# if RDI, generate a blank canvas since we assume there is no astrophysical signal in the RDI frames
if is_rdi:
models.append(np.zeros(np.size(section_ind)))
continue
# grab PSF given wavelength
wv_index = [spec.find_nearest(self.input_psfs_wvs,wv)[1]]
# find center of psf
# to reduce calculation of sin and cos, see if it has already been calculated before
recalculate_trig = False
if pa not in self.psf_centx_notscaled:
recalculate_trig = True
else:
#print(self.psf_centx_notscaled[pa],pa)
if pa_fk != self.curr_pa_fk[pa] or sep_fk != self.curr_sep_fk[pa]:
recalculate_trig = True
if recalculate_trig: # we could actually store the values for the different pas too...
self.psf_centx_notscaled[pa] = sep_fk * np.cos(np.radians(90. - sign*pa_fk - pa))
self.psf_centy_notscaled[pa] = sep_fk * np.sin(np.radians(90. - sign*pa_fk - pa))
self.curr_pa_fk[pa] = pa_fk
self.curr_sep_fk[pa] = sep_fk
psf_centx = (ref_wv/wv) * self.psf_centx_notscaled[pa]
psf_centy = (ref_wv/wv) * self.psf_centy_notscaled[pa]
# create a coordinate system for the image that is with respect to the model PSF
# round to nearest pixel and add offset for center
l = int(round(psf_centx + ref_center[0]))
k = int(round(psf_centy + ref_center[1]))
# recenter coordinate system about the location of the planet
x_vec_stamp_centered = x_grid[0, (l-self.col_m):(l+self.col_p)]-psf_centx
y_vec_stamp_centered = y_grid[(k-self.row_m):(k+self.row_p), 0]-psf_centy
# rescale to account for the align and scaling of the refernce PSFs
# e.g. for longer wvs, the PSF has shrunk, so we need to shrink the coordinate system
x_vec_stamp_centered /= (ref_wv/wv)
y_vec_stamp_centered /= (ref_wv/wv)
# use intepolation spline to generate a model PSF and write to temp img
whiteboard[(k-self.row_m):(k+self.row_p), (l-self.col_m):(l+self.col_p)] = \
self.psfs_func_list[wv_index[0]](x_vec_stamp_centered,y_vec_stamp_centered).transpose()
# write model img to output (segment is collapsed in x/y so need to reshape)
whiteboard.shape = [input_img_shape[0] * input_img_shape[1]]
segment_with_model = copy(whiteboard[section_ind])
whiteboard.shape = [input_img_shape[0],input_img_shape[1]]
models.append(segment_with_model)
# create a canvas to place the new PSF in the sector on
if 0:
blackboard = np.zeros((ny,nx))
blackboard.shape = [input_img_shape[0] * input_img_shape[1]]
blackboard[section_ind] = segment_with_model
blackboard.shape = [input_img_shape[0],input_img_shape[1]]
plt.figure(1)
plt.subplot(1,2,1)
im = plt.imshow(whiteboard)
plt.colorbar(im)
plt.subplot(1,2,2)
im = plt.imshow(blackboard)
plt.colorbar(im)
plt.show()
whiteboard[(k-self.row_m):(k+self.row_p), (l-self.col_m):(l+self.col_p)] = 0.0
return np.array(models)
[docs]
def calculate_fm_opti(delta_KL, original_KL, sci, model_sci_fk,delta_KL_fk, original_KL_fk):
r"""
Optimized version for calculate_fm() (if numbas) for a single numbasis.
Calculate what the PSF looks up post-KLIP using knowledge of the input PSF, assumed spectrum of the science target,
and the partially calculated KL modes (\Delta Z_k^\lambda in Laurent's paper). If inputflux is None,
the spectral dependence has already been folded into delta_KL_nospec (treat it as delta_KL).
Note: if inputflux is None and delta_KL_nospec has three dimensions (ie delta_KL_nospec was calculated using
pertrurb_nospec() or perturb_nospec_modelsBased()) then only klipped_oversub and klipped_selfsub are returned.
Besides they will have an extra first spectral dimension.
Args:
delta_KL_nospec: perturbed KL modes but without the spectral info. delta_KL = spectrum x delta_Kl_nospec.
Shape is (numKL, wv, pix). If inputflux is None, delta_KL_nospec = delta_KL
orignal_KL: unpertrubed KL modes (array of size [numbasis, numpix])
sci: array of size p representing the science data
model_sci: array of size p corresponding to the PSF of the science frame
input_spectrum: array of size wv with the assumed spectrum of the model
If delta_KL_nospec does NOT include a spectral dimension or if inputflux is not None:
Returns:
fm_psf: array of shape (b,p) showing the forward modelled PSF
Skipped if inputflux = None, and delta_KL_nospec has 3 dimensions.
klipped_oversub: array of shape (b, p) showing the effect of oversubtraction as a function of KL modes
klipped_selfsub: array of shape (b, p) showing the effect of selfsubtraction as a function of KL modes
Note: psf_FM = model_sci - klipped_oversub - klipped_selfsub to get the FM psf as a function of K Lmodes
(shape of b,p)
If inputflux = None and if delta_KL_nospec include a spectral dimension:
Returns:
klipped_oversub: Sum(<S|KL>KL) with klipped_oversub.shape = (1,Npix)
klipped_selfsub: Sum(<N|DKL>KL) + Sum(<N|KL>DKL) with klipped_selfsub.shape = (1,N_lambda or N_ref,N_pix)
"""
# remove means and nans from science image
sci_mean_sub = (sci - np.nanmean(sci))[None,:]
sci_mean_sub[np.where(np.isnan(sci_mean_sub))] =0
# science PSF models, ready for FM
# /!\ JB: If subtracting the mean. It should be done here. not in klip_math since we don't use model_sci there.
model_sci_mean_sub = model_sci_fk[None,:] # should be subtracting off the mean?
model_sci_mean_sub[np.where(np.isnan(model_sci_mean_sub))] =0
# Forward model the PSF
# 3 terms: 1 for oversubtracton (planet attenauted by speckle KL modes),
# and 2 terms for self subtraction (planet signal leaks in KL modes which get projected onto speckles)
#
# Klipped = N-Sum(<N|KL>KL) + S-Sum(<S|KL>KL) - Sum(<N|DKL>KL) - Sum(<N|KL>DKL)
# With N = noise/speckles (science image)
# S = signal/planet model
# KL = KL modes
# DKL = perturbation of the KL modes/Delta_KL
#
# sci_mean_sub_rows.shape = (1,N_pix)
# model_sci_mean_sub_rows.shape = (1,N_pix)
# original_KL.shape = (max_basis,N_pix)
# delta_KL.shape = (max_basis,N_pix)
oversubtraction_inner_products = np.dot(model_sci_mean_sub, original_KL_fk.T)
selfsubtraction_1_inner_products = np.dot(sci_mean_sub, delta_KL.T)
selfsubtraction_2_inner_products = np.dot(sci_mean_sub, original_KL.T)
# oversubtraction_inner_products = (1,numbasis)
klipped_oversub = np.dot(oversubtraction_inner_products, original_KL_fk)
# selfsubtraction_1_inner_products = (1,numbasis)
# selfsubtraction_2_inner_products = (1,numbasis)
klipped_selfsub = np.dot(selfsubtraction_1_inner_products, original_KL_fk) + \
np.dot(selfsubtraction_2_inner_products, delta_KL_fk)
return model_sci_fk[None,:] - klipped_oversub - klipped_selfsub