Perform Astrometry and Photometry on Outer Planet Companion

This notebook will show how to forward model the PSF of the planet through KLIP (since KLIP distorts the planet signal), and then we will use MCMC to measure the posterior location/flux of the planet and Gaussian processes to model the correlated noise in the data. We call this technique Bayesian KLIP Astrometry (BKA) and you can read more about it in the pyKLIP docs: http://pyklip.readthedocs.io/en/latest/bka.html.

import os
import astropy.io.fits as fits
import numpy as np
import scipy
import scipy.ndimage as ndi
import matplotlib.pylab as plt
import pandas as pd
%matplotlib inline
%config InlineBackend.figure_format = 'retina'

import pyklip.klip
import pyklip.fakes as fakes
import pyklip.fm as fm
import pyklip.instruments.Instrument as Instrument
import pyklip.fmlib.fmpsf as fmpsf
import pyklip.fitpsf as fitpsf

First, we must read in the data

Just as is done in the basic KLIP reduction notebook, we’ll read in our two roll angles. Unfortunately, RDI is not yet supported with BKA, so we won’t use the reference star data, but we will use the unocculted star for flux calibration and for our PSF model.

filtername = "f300m"
datadir = "datadir/""
roll1 = "roll1.fits"
roll2 = "roll2.fits"
ref_cube = "ref_cube.fits"
unocculted_psf = "unocculted.fits"

# read in roll 1
with fits.open(f"{datadir}{roll1}") as hdulist:
    roll1_cube = hdulist[0].data

# read in roll 2
with fits.open(f"{datadir}{roll2}") as hdulist:
    roll2_cube = hdulist[0].data

# read in ref star
with fits.open(f"{datadir}{ref_cube}") as hdulist:
    ref_cube = hdulist[0].data

# read in unocculted PSF
with fits.open(f"{datadir}{unocculted}") as hdulist:
    psf_cube = hdulist[0].data

fig = plt.figure(figsize=(10,4))

ax1 = fig.add_subplot(1,4,1)
ax1.imshow(roll1_cube[0], interpolation="nearest", cmap="inferno")
ax1.invert_yaxis()
ax1.set_title("Roll 1")

ax2 = fig.add_subplot(1,4,2)
ax2.imshow(roll2_cube[0], interpolation="nearest", cmap="inferno")
ax2.invert_yaxis()
ax2.set_title("Roll 2")

ax3 = fig.add_subplot(1,4,3)
ax3.imshow(ref_cube[0], interpolation="nearest", cmap="inferno")
ax3.invert_yaxis()
ax3.set_title("Reference Star")

ax4 = fig.add_subplot(1,4,4)
ax4.imshow(psf_cube[0], interpolation="nearest", cmap="inferno")
ax4.invert_yaxis()
ax4.set_title("Unocculted PSF")

(6, 101, 101)
(6, 101, 101)
(54, 101, 101)
(6, 101, 101)

Text(0.5, 1.0, 'Unocculted PSF')

Create the Generic Dataset

We’ll combine the two roll angles to create a dataset.

# combine the two rows
full_seq = np.concatenate([roll1_cube, roll2_cube], axis=0)

# two rolls are offset 10 degrees, this is the right sign (trust me)
pas = np.append([0 for _ in range(roll1_cube.shape[0])], [10 for _ in range(roll2_cube.shape[0])])

# for each image, the (x,y) center where the star is is just the center of the image
centers = np.array([np.array(frame.shape)/2. for frame in full_seq])

# give it some names, just in case we want to refer to them
filenames = np.append(["roll1_{0}".format(i) for i in range(roll1_cube.shape[0])],
                      ["roll2_{0}".format(i) for i in range(roll1_cube.shape[0])])

# create the GenericData object. This will standardize the data for pyKLIP
dataset = Instrument.GenericData(full_seq, centers, IWA=4, parangs=pas, filenames=filenames)
dataset.flipx = False # get the right handedness of the data

Center on the unocculted PSF

So that we have a nicely centered model of the star we can use for PSF modelling

# collapse reference psf in time
psf_frame = np.nanmean(psf_cube, axis=0)

# find the centroid
bestfit = fakes.gaussfit2d(psf_frame, 71, 30, searchrad=3, guessfwhm=2, guesspeak=1, refinefit=True)

psf_xcen, psf_ycen = bestfit[2:4]
print(psf_xcen, psf_ycen)

# recenter PSF to that location
x, y = np.meshgrid(np.arange(-20,20.1,1), np.arange(-20,20.1,1))
x += psf_xcen
y += psf_ycen

psf_stamp = scipy.ndimage.map_coordinates(psf_frame, [y,x])

plt.figure()
plt.imshow(psf_stamp, interpolation="nearest", cmap="inferno")
plt.gca().invert_yaxis()
plt.title("Reference PSF")

70.45855749153525 29.8360856033007

Text(0.5, 1.0, 'Reference PSF')

KLIP-FM for forward modelling of outer source

Since pyKLIP oversubtraction can distort our planet position and flux, the following steps will go through forward modeling the planet PSF.

Correcting for coronagraphic throughput

To improve the accuracy of our forward model, we can also account for the effect of JWST’s coronograph on light transmission. When initializing the fmpsf.FMPlanetPSF class, we can provide the optional argument ‘field_dependent_correction’ which accepts a user provided function to correct for coronagraphic throughput. The signature of the user provided function is (input_stamp: numpy.ndarray, input_dx: numpy.ndarray, input_dy: numpy.ndarray) -> numpy.ndarray. All the numpy arrays in the signature should be 2D. As an example of how this would be done, we’ll use the transmission profile of the MASK210 coronagraph (obtained from the Occulting Masks section of this NIRCAM webpage: https://jwst-docs.stsci.edu/near-infrared-camera/nircam-instrumentation/nircam-coronagraphic-occulting-masks-and-lyot-stops).

# Read in the transmission profile csv
mask210 = pd.read_csv("MASK210R.csv", names = ["rad_dist", "trans"])

# Create the throughput correction function
def transmission_corrected(input_stamp, input_dx, input_dy):

    """
    Args:
        input_stamp (array): 2D array of the region surrounding the fake planet injection site
        input_dx (array): 2D array specifying the x distance of each stamp pixel from the center
        input_dy (array): 2D array specifying the y distance of each stamp pixel from the center

    Returns:
        output_stamp (array): 2D array of the throughput corrected planet injection site.
        """

    # Calculate the distance of each pixel in the input stamp from the center
    distance_from_center = np.sqrt((input_dx) ** 2 + (input_dy) ** 2)

    # Interpolate to find the transmission value for each pixel in the input stamp (we need to turn the columns into arrays so np.interp can accept them)
    distance = np.array(mask210["rad_dist"])
    transmission = np.array(mask210["trans"])

    # Calculate the distance in pixels
    distance_pix = mask210['rad_dist']/0.063

    trans_at_dist = np.interp(distance_from_center, distance_pix, transmission)

    # Reshape the interpolated array to have the same dimensions as the input stamp
    transmission_stamp = trans_at_dist.reshape(input_stamp.shape)

    # Make the throughput correction
    output_stamp = transmission_stamp * input_stamp

    return output_stamp

Now, we’ll set up guesses for the location and flux of the planet. Then we can run the code that generates the forward model of the planet PSF using the unocculted PSF of the star. Note that RDI is not yet working with KLIP-FM, but should be soon (and should make fitting the planet even better)

We’ll get a warning about the spline fitting of the PSF being rank deficient. This should be fine.

# setup FM guesses
numbasis = np.array([1, 3, 10]) # KL basis cutoffs you want to try
guess_dx = 9 # in pxiels (positive is to the left)
guess_dy = 4 # in pixels (positive is up)
guesssep = np.sqrt(guess_dx**2 + guess_dy**2) # estimate of separation in pixels
guesspa = np.degrees(np.arctan2(guess_dx, guess_dy)) # estimate of position angle, in degrees
guessflux = 1e-4 # estimated contrast
guessspec = np.array([1]) # braodband, so don't need to guess spectrum

# initialize the FM Planet PSF class
fm_class = fmpsf.FMPlanetPSF(dataset.input.shape, numbasis, guesssep, guesspa, guessflux, np.array([psf_stamp]),
                             np.unique(dataset.wvs), spectrallib_units="contrast", spectrallib=[guessspec], field_dependent_correction = None)

/Users/jeaadams/opt/anaconda3/lib/python3.7/site-packages/scipy/interpolate/fitpack2.py:1133: UserWarning:
The coefficients of the spline returned have been computed as the
minimal norm least-squares solution of a (numerically) rank deficient
system (deficiency=255). If deficiency is large, the results may be
inaccurate. Deficiency may strongly depend on the value of eps.
  warnings.warn(message)

Perform PSF Subtraction

# PSF subtraction parameters
# You should change these to be suited to your data!
outputdir = "./" # where to write the output files
prefix = "pyklipfm-b-ADI-k50m1" # fileprefix for the output files
annulus_bounds = [[guesssep-20, guesssep+20]] # one annulus centered on the planet
subsections = 1 # we are not breaking up the annulus
padding = 0 # we are not padding our zones
movement = 1 # basically, we want to use the other roll angle for ADI.

# run KLIP-FM
import pyklip.fm as fm
fm.klip_dataset(dataset, fm_class, outputdir=outputdir, fileprefix=prefix, numbasis=numbasis,
                annuli=annulus_bounds, subsections=subsections, padding=padding, movement=movement, maxnumbasis=50)

Begin align and scale images for each wavelength
Align and scale finished
Starting KLIP for sector 1/1 with an area of 2475.2879453098526 pix^2
Time spent on last sector: 0s
Time spent since beginning: 0s
First sector: Can't predict remaining time

Closing threadpool
Writing KLIPed Images to directory /Users/jeaadams/JWST-ERS-Pipeline/notebooks

Load in the forward model and the data and create the fitting object.

KLIP-FM produces both the KLIPed data and the forward model, so that the data and model are reduced with the exact same parameters for consistency.

The FMAstrometry object fits for both the astrometry and the photometry of the planet. We specify the fitting box size and extract out stamps of the planet in both the data and forward model fot fitting.

output_prefix = os.path.join(outputdir, prefix)
with fits.open(output_prefix + "-fmpsf-KLmodes-all.fits") as fm_hdu:
    # get FM frame, use KL=7
    fm_frame = fm_hdu[0].data[1]
    fm_centx = fm_hdu[0].header['PSFCENTX']
    fm_centy = fm_hdu[0].header['PSFCENTY']

with fits.open(output_prefix + "-klipped-KLmodes-all.fits") as data_hdu:
    # get data_stamp frame, use KL=7
    data_frame = data_hdu[0].data[1]
    data_centx = data_hdu[0].header["PSFCENTX"]
    data_centy = data_hdu[0].header["PSFCENTY"]

fitboxsize = 17
fma = fitpsf.FMAstrometry(guesssep, guesspa, fitboxsize)

# generate FM stamp
# padding should be greater than 0 so we don't run into interpolation problems
fma.generate_fm_stamp(fm_frame, [fm_centx, fm_centy], padding=5)

# generate data_stamp stamp
# not that dr=4 means we are using a 4 pixel wide annulus to sample the noise for each pixel
# exclusion_radius excludes all pixels less than that distance from the estimated location of the planet
fma.generate_data_stamp(data_frame, [data_centx, data_centy], dr=4, exclusion_radius=10)

/Users/jeaadams/opt/anaconda3/lib/python3.7/site-packages/numpy/lib/nanfunctions.py:1667: RuntimeWarning: Degrees of freedom <= 0 for slice.
  keepdims=keepdims)

# Let's plot our model!
plt.imshow(fm_frame, cmap = 'inferno')

<matplotlib.image.AxesImage at 0x7fbf81004f50>

# And also the data
plt.imshow(data_frame, cmap = 'inferno')

<matplotlib.image.AxesImage at 0x7fbf5162bc90>

Set up the Gaussian process kernel, initialize priors, run MCMC

We will specify that we are using a Matern \(\nu=3/2\) kernel as the analytical form to model the correlated noise in the data. This kernel has wider tails than a standard squared exponential kernel. Basically, with this kernel, we specify how much correlation there is between two pixels as a function of the distance between these two pixels. We will keep one free parameter, the correlation length, which we will fit with MCMC. Longer correlation lengths means that pixels separated by a given distance will be more correlated.

We also will initialize the bounds for the priors. We will give 1 pixel wiggle room for the planet’s true location, and a factor or 10 wiggle room for the planet’s flux (and also for the correlation length of the Gaussian process).

To run MCMC, we will use the affine invariant sampler that comes with the emcee package. This could take several minutes.

# set kernel
corr_len_guess = 3. # in pixels, our guess for the correlation length
corr_len_label = r"$l$" # label for this variable.
fma.set_kernel("matern32", [corr_len_guess], [corr_len_label])

# set prior boundson parameters
x_range = 1.0 # pixels
y_range = 1.0 # pixels
flux_range = 1. # flux can vary by an order of magnitude
corr_len_range = 1. # between 0.3 and 30
fma.set_bounds(x_range, y_range, flux_range, [corr_len_range])

# run MCMC fit
fma.fit_astrometry(nwalkers=100, nburn=200, nsteps=800, numthreads=2)

Running burn in
Burn in finished. Now sampling posterior
MCMC sampler has finished

Plot the fits

First we will plot the MCMC walkers to make sure they are converged

fig = plt.figure(figsize=(10,8))

# grab the chains from the sampler
chain = fma.sampler.chain

# plot RA offset
ax1 = fig.add_subplot(511)
ax1.plot(chain[:,:,0].T, '-', color='k', alpha=0.3)
ax1.set_xlabel("Steps")
ax1.set_ylabel(r"$\Delta$ RA")

# plot Dec offset
ax2 = fig.add_subplot(512)
ax2.plot(chain[:,:,1].T, '-', color='k', alpha=0.3)
ax2.set_xlabel("Steps")
ax2.set_ylabel(r"$\Delta$ Dec")

# plot flux scaling
ax3 = fig.add_subplot(513)
ax3.plot(chain[:,:,2].T, '-', color='k', alpha=0.3)
ax3.set_xlabel("Steps")
ax3.set_ylabel(r"$\alpha$")

# plot hyperparameters.. we only have one for this example: the correlation length
ax4 = fig.add_subplot(514)
ax4.plot(chain[:,:,3].T, '-', color='k', alpha=0.3)
ax4.set_xlabel("Steps")
ax4.set_ylabel(r"$l$")

Text(0, 0.5, '$l$')

Plot the comparison between the data and the best fit model

fig = plt.figure()
fig = fma.best_fit_and_residuals(fig=fig)

Plot the posterior plots for the parameters

fig = fma.make_corner_plot()

Finally, we can read off the planet’s location and flux

print("\nPlanet Raw RA offset is {0} +/- {1}, Raw Dec offset is {2} +/- {3}".format(fma.raw_RA_offset.bestfit, fma.raw_RA_offset.error,
                                                                                    fma.raw_Dec_offset.bestfit, fma.raw_Dec_offset.error))

print("Planet Flux is {0} +/- {1}".format(fma.raw_flux.bestfit*guessflux, fma.raw_flux.error*guessflux))

Planet Raw RA offset is 9.296172270077676 +/- 0.05087799532961057, Raw Dec offset is 3.8443487808025054 +/- 0.05675071335087978
Planet Flux is 5.469367576971451e-05 +/- 3.4749672931897125e-06