Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merging gs-inject-source into develop #203

Merged
merged 1 commit into from
Jul 30, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
255 changes: 252 additions & 3 deletions spaceKLIP/imagetools.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
from jwst.pipeline import Detector1Pipeline, Image2Pipeline, Coron3Pipeline
import jwst.datamodels
from pyklip import parallelized
from scipy.ndimage import gaussian_filter, median_filter, fourier_shift
from scipy.ndimage import gaussian_filter, median_filter, fourier_shift, rotate
from scipy.ndimage import shift as spline_shift
from scipy.optimize import leastsq, minimize
from skimage.registration import phase_cross_correlation
Expand All @@ -33,8 +33,15 @@
from webbpsf_ext import robust
from webbpsf_ext.coords import dist_image
from webbpsf_ext.webbpsf_ext_core import _transmission_map

from tqdm.auto import trange
import copy
import pyklip.fakes as fakes
from spaceKLIP.psf import get_offsetpsf
from spaceKLIP.pyklippipeline import get_pyklip_filepaths
from pyklip.instruments.JWST import JWSTData
from webbpsf.constants import JWST_CIRCUMSCRIBED_DIAMETER
from astropy.io import fits
from spaceKLIP.starphot import get_stellar_magnitudes, read_spec_file

import logging
log = logging.getLogger(__name__)
Expand Down Expand Up @@ -1524,7 +1531,249 @@ def hpf(self,
self.database.update_obs(key, j, fitsfile, maskfile)

pass


def inject_companions(self,
companions,
starfile,
spectral_type,
output_dir,
highpass=False,
subdir='test',
date='auto',
kwargs={}):
'''
Function to inject synthetic PSFs into a set of frames loaded from a dataset, and save the new frames with the
injected companion.

Parameters (some may need to be adjusted!)
----------
companions : list of list of three float, optional
List of companions to be injected.
For each companion, there should be a three element list containing
[RA offset (arcsec), Dec offset (arcsec), contrast].
raw_dataset : pyKLIP dataset
A pyKLIP dataset which companions will be injected into and KLIP
will be performed on.
injection_psf : 2D-array
The PSF of the companion to be injected.
injection_seps : 1D-array
List of separations to inject companions at (pixels).
injection_pas : 1D-array
List of position angles to inject companions at (degrees).
injection_spacing : int, None
Spacing between companions injected in a single image. If companions
are too close then it can pollute the recovered flux. Set to 'None'
to inject only one companion at a time (pixels).
injection_fluxes : 1D-array
Same size as injection_seps, units should correspond to the image
units. This is the *peak* flux of the injection.
true_companions : list of list of three float, optional
List of real companions to be masked before computing the raw contrast.
For each companion, there should be a three element list containing
[RA offset (pixels), Dec offset (pixels), mask radius (pixels)].
The default is None.

Returns
-------
None
'''

# Check input.
if not isinstance(companions[0], list):
if len(companions) == 3:
companions = [companions]
for i in range(len(companions)):
if len(companions[i]) != 3:
raise UserWarning('There should be three elements for each companion in the companions list')

Ncompanions = len(companions)
for _, key in enumerate(self.database.obs.keys()):
ww_type=list(self.database.obs[key]['TYPE'])
list_of_injected = []
all_injected = False

log.info('--> Concatenation ' + key)
#######################################################################################################################
filepaths, psflib_filepaths = get_pyklip_filepaths(self.database, key)
raw_dataset = JWSTData(filepaths, psflib_filepaths)

for ww in range(len(ww_type)):
# Read input files and store values that we just want to save in the output_dir
fitsfile = self.database.obs[key]['FITSFILE'][ww]
data, erro, pxdq, head_pri, head_sci, is2d, imshifts, maskoffs = ut.read_obs(fitsfile)
maskfile = self.database.obs[key]['MASKFILE'][ww]
mask = ut.read_msk(maskfile)
crpix1 = self.database.obs[key]['CRPIX1'][ww]
crpix2 = self.database.obs[key]['CRPIX2'][ww]
head, tail = os.path.split(fitsfile)

# Write FITS file and PSF mask.
head_sci['CRPIX1'] = crpix1
head_sci['CRPIX2'] = crpix2

# Inject only into SCI type data
if ww_type[ww] == 'SCI':
# Convert the host star brightness from vegamag to MJy. Use an
# unocculted model PSF whose integrated flux is normalized to
# one in order to obtain the theoretical peak count of the
# star.
filt = self.database.obs[key]['FILTER'][ww]

# Get stellar magnitudes and filter zero points.
mstar, fzero, fzero_si = get_stellar_magnitudes(starfile, spectral_type,
self.database.obs[key]['INSTRUME'][ww], return_si=True,
output_dir=output_dir,
**kwargs) # vegamag, Jy, erg/cm^2/s/A
# Compute the pixel area in steradian.
pxsc_arcsec = self.database.obs[key]['PIXSCALE'][ww] # arcsec
pxsc_rad = pxsc_arcsec / 3600. / 180. * np.pi # rad
pxar = pxsc_rad ** 2 # sr

# Set output directory.
output_dir = os.path.join(self.database.output_dir, subdir)
if not os.path.exists(output_dir):
os.makedirs(output_dir)

# Make a copy of the dataset
dataset = copy.deepcopy(raw_dataset)

apername = self.database.obs[key]['APERNAME'][ww]
if 'planetfile' not in kwargs.keys() or kwargs['planetfile'] is None:
if starfile is not None and starfile.endswith('.txt'):
sed = read_spec_file(starfile)
else:
sed = None
else:
sed = read_spec_file(kwargs['planetfile'])
ww_sci = np.where(self.database.obs[key]['TYPE'] == 'SCI')[0]
if date is not None:
if date == 'auto':
date = fits.getheader(self.database.obs[key]['FITSFILE'][ww_sci[0]], 0)['DATE-BEG']

offsetpsf_func = JWST_PSF(apername,
filt,
date=date,
fov_pix=65,
oversample=2,
sp=sed,
use_coeff=False)

# Offset PSF that is not affected by the coronagraphic
# mask, but only the Lyot stop.
psf_no_coronmsk = offsetpsf_func.psf_off
log.info('--> Injecting companions, writing FITS files and updating spaceKLIP database: ')

# Loop over companions
for i in trange(Ncompanions):
# Initial guesses for the fit parameters.
guess_dx = companions[i][0] / pxsc_arcsec # pix
guess_dy = companions[i][1] / pxsc_arcsec # pix
guess_sep = np.sqrt(guess_dx ** 2 + guess_dy ** 2) # pix
guess_pa = np.rad2deg(np.arctan2(guess_dx, guess_dy)) # deg
guess_flux = companions[i][2] # contrast

roll_ref = self.database.obs[key]['ROLL_REF'][ww] # deg

# Get shift between star and coronagraphic mask
# position. If positive, the coronagraphic mask center
# is to the left/bottom of the star position.
_, _, _, _, _, _, _, maskoffs = ut.read_obs(self.database.obs[key]['FITSFILE'][ww])

# NIRCam.
if maskoffs is not None:
mask_xoff = -maskoffs[:, 0] # pix
mask_yoff = -maskoffs[:, 1] # pix

# Need to rotate by the roll angle (CCW) and flip
# the x-axis so that positive RA is to the left.
mask_raoff = -(mask_xoff * np.cos(np.deg2rad(roll_ref)) - mask_yoff * np.sin(
np.deg2rad(roll_ref))) # pix
mask_deoff = mask_xoff * np.sin(np.deg2rad(roll_ref)) + mask_yoff * np.cos(
np.deg2rad(roll_ref)) # pix

# Compute the true offset between the companion and
# the coronagraphic mask center.
sim_dx = guess_dx - mask_raoff # pix
sim_dy = guess_dy - mask_deoff # pix
sim_sep = np.sqrt(sim_dx ** 2 + sim_dy ** 2) * pxsc_arcsec # arcsec
sim_pa = np.rad2deg(np.arctan2(sim_dx, sim_dy)) # deg

# Take median of observation. Typically, each
# dither position is a separate observation.
sim_sep = np.median(sim_sep)
sim_pa = np.median(sim_pa)

# Otherwise.
else:
sim_sep = np.sqrt(guess_dx ** 2 + guess_dy ** 2) * pxsc_arcsec # arcsec
sim_pa = np.rad2deg(np.arctan2(guess_dx, guess_dy)) # deg

# Generate offset PSF for this roll angle. Do not add
# the V3Yidl angle as it has already been added to the
# roll angle by spaceKLIP. This is only for estimating
# the coronagraphic mask throughput!
offsetpsf_coronmsk = offsetpsf_func.gen_psf([sim_sep, sim_pa],
mode='rth',
PA_V3=roll_ref,
do_shift=False,
quick=True,
addV3Yidl=False)

# Coronagraphic mask throughput is not incorporated
# into the flux calibration of the JWST pipeline so
# that the companion flux from the detector pixels will
# be underestimated. Therefore, we need to scale the
# model offset PSF to account for the coronagraphic
# mask throughput (it becomes fainter). Compute scale
# factor by comparing a model PSF with and without
# coronagraphic mask.
scale_factor = np.sum(offsetpsf_coronmsk) / np.sum(psf_no_coronmsk)

# Normalize model offset PSF to a total integrated flux
# of 1 at infinity. Generates a new webbpsf model with
# PSF normalization set to 'exit_pupil'.
offsetpsf = offsetpsf_func.gen_psf([sim_sep, sim_pa],
mode='rth',
PA_V3=roll_ref,
do_shift=False,
quick=False,
addV3Yidl=False,
normalize='exit_pupil')

# Normalize model offset PSF by the flux of the star.
mcomp = mstar[filt] -2.5*np.log10(guess_flux)
offsetpsf *= fzero[filt] / 10 ** (mcomp / 2.5) / 1e6 / pxar # MJy/sr

## Question: am I injecting companions too bright?! Shouldn't I rescale them to the distance
## of the target? Why doesen't work if I try.
# parallax_mas=12.1549
# parallax_arcseconds = parallax_mas / 1000
# distance_parsecs = 1 / parallax_arcseconds
# offsetpsf *= (1 / distance_parsecs) ** 2


# Apply scale factor to incorporate the coronagraphic
# mask througput.
offsetpsf *= scale_factor

# Injected PSF needs to be a 3D array that matches dataset
inj_psf_3d = np.array([offsetpsf for k in range(dataset.input.shape[0])])

# Inject the PSF
fakes.inject_planet(frames=dataset.input, centers=dataset.centers, inputflux=inj_psf_3d,
astr_hdrs=dataset.wcs, radius=guess_sep, pa=guess_pa, stampsize=65)

data = dataset.input


# Write FITS file and PSF mask.
fitsfile = ut.write_obs(fitsfile, output_dir, data, erro, pxdq, head_pri, head_sci, is2d,
imshifts, maskoffs)
maskfile = ut.write_msk(maskfile, mask, fitsfile)

# Update spaceKLIP database.
self.database.update_obs(key,ww, fitsfile, maskfile, crpix1=crpix1, crpix2=crpix2)

def update_nircam_centers(self):
"""
Determine offset between SIAF reference pixel position and true mask
Expand Down