Skip to content

Commit

Permalink
Adding function to perform MCMC fitting with a psf convolved by a gau…
Browse files Browse the repository at this point in the history 
…ssian
  • Loading branch information
@strampelligiovanni
strampelligiovanni committed Jul 26, 2024
1 parent 27e8f6e commit 6afaf4f
Showing 1 changed file with 252 additions and 3 deletions.
255 changes: 252 additions & 3 deletions spaceKLIP/imagetools.py
View file
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

0 comments on commit 6afaf4f

Please sign in to comment.