The (Very) Short Introduction to EasyTS#
Author: Hannu Parviainen Edited: 9 August 2024
This notebook is a less verbose version of the (not so) short introduction into transmission spectroscopy with EasyTS, showing how a typical transmission spectroscopy analysis would look like without all the comments and more useful to be used as a template for a new analysis.
[1]:
%run ../setup_multiprocessing.py
[2]:
%matplotlib inline
[3]:
from multiprocessing import Pool
from xarray import load_dataset
from matplotlib.pyplot import subplots, setp
from numpy import geomspace, linspace, newaxis, diff, zeros
from easyts import EasyTS, TSData
Data preparation#
[4]:
def read_data(fname, name=""):
with load_dataset(fname) as ds:
return TSData(time=ds.time.values, wavelength=ds.wavelength.values, fluxes=ds.flux.values, errors=ds.error.values, name=name)
[5]:
d1 = read_data('data/nirHiss_order_1.h5', "WASP-39b JWST NIRISS Order 1")
d1.remove_outliers()
d2 = read_data('data/nirHiss_order_2.h5', "WASP-39b JSWT NIRISS Order 2")
d2.remove_outliers()
db = d2.bin_wavelength(r = 100) + d1.bin_wavelength(r = 100)
db.errors[:,:] = diff(db.fluxes[:, :140]).std(1)[:,newaxis]
fig = db.plot(figsize=(6, 7))
[6]:
ts = EasyTS('01a_lowres_power2', ldmodel='power-2', data=db, nk=50, nldc=10, nthreads=1)
ts.set_prior('tc', 'NP', 2459694.286, 0.003)
ts.set_prior('p', 'NP', 4.05487, 1e-5)
ts.set_radius_ratio_prior('UP', 0.14, 0.15)
ts.set_ldtk_prior(teff=(5327, 139), logg=(4.38, 0.09), metal=(-0.01, 0.1), uncertainty_multiplier=10)
[7]:
ts.ps
[7]:
[ 0 |G| rho U(a = 0.1, b = 25.0) [ 0.00 .. inf],
1 |G| tc N(μ = 2459694.286, σ = 0.003) [ -inf .. inf],
2 |G| p N(μ = 4.05487, σ = 1e-05) [ 0.00 .. inf],
3 |G| b U(a = 0.0, b = 1.0) [ 0.00 .. inf],
4 |G| secw N(μ = 0.0, σ = 1e-05) [ -1.00 .. 1.00],
5 |G| sesw N(μ = 0.0, σ = 1e-05) [ -1.00 .. 1.00],
6 |G| ldc1_00.57277 N(μ = 0.755, σ = 0.018) [ -inf .. inf],
7 |G| ldc2_00.57277 N(μ = 0.835, σ = 0.033) [ -inf .. inf],
8 |G| ldc1_00.82264 N(μ = 0.61, σ = 0.016) [ -inf .. inf],
9 |G| ldc2_00.82264 N(μ = 0.711, σ = 0.027) [ -inf .. inf],
10 |G| ldc1_01.07250 N(μ = 0.546, σ = 0.014) [ -inf .. inf],
11 |G| ldc2_01.07250 N(μ = 0.632, σ = 0.022) [ -inf .. inf],
12 |G| ldc1_01.32236 N(μ = 0.566, σ = 0.016) [ -inf .. inf],
13 |G| ldc2_01.32236 N(μ = 0.533, σ = 0.02) [ -inf .. inf],
14 |G| ldc1_01.57223 N(μ = 0.631, σ = 0.019) [ -inf .. inf],
15 |G| ldc2_01.57223 N(μ = 0.358, σ = 0.013) [ -inf .. inf],
16 |G| ldc1_01.82209 N(μ = 0.61, σ = 0.016) [ -inf .. inf],
17 |G| ldc2_01.82209 N(μ = 0.321, σ = 0.01) [ -inf .. inf],
18 |G| ldc1_02.07195 N(μ = 0.55, σ = 0.013) [ -inf .. inf],
19 |G| ldc2_02.07195 N(μ = 0.349, σ = 0.009) [ -inf .. inf],
20 |G| ldc1_02.32182 N(μ = 0.49, σ = 0.013) [ -inf .. inf],
21 |G| ldc2_02.32182 N(μ = 0.386, σ = 0.012) [ -inf .. inf],
22 |G| ldc1_02.57168 N(μ = 0.435, σ = 0.012) [ -inf .. inf],
23 |G| ldc2_02.57168 N(μ = 0.416, σ = 0.014) [ -inf .. inf],
24 |G| ldc1_02.82154 N(μ = 0.414, σ = 0.011) [ -inf .. inf],
25 |G| ldc2_02.82154 N(μ = 0.402, σ = 0.013) [ -inf .. inf],
26 |G| k_00.57277 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
27 |G| k_00.61867 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
28 |G| k_00.66456 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
29 |G| k_00.71045 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
30 |G| k_00.75635 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
31 |G| k_00.80224 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
32 |G| k_00.84813 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
33 |G| k_00.89403 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
34 |G| k_00.93992 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
35 |G| k_00.98581 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
36 |G| k_01.03171 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
37 |G| k_01.07760 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
38 |G| k_01.12349 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
39 |G| k_01.16939 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
40 |G| k_01.21528 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
41 |G| k_01.26117 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
42 |G| k_01.30707 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
43 |G| k_01.35296 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
44 |G| k_01.39885 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
45 |G| k_01.44475 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
46 |G| k_01.49064 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
47 |G| k_01.53653 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
48 |G| k_01.58243 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
49 |G| k_01.62832 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
50 |G| k_01.67421 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
51 |G| k_01.72011 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
52 |G| k_01.76600 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
53 |G| k_01.81189 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
54 |G| k_01.85779 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
55 |G| k_01.90368 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
56 |G| k_01.94957 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
57 |G| k_01.99547 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
58 |G| k_02.04136 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
59 |G| k_02.08725 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
60 |G| k_02.13315 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
61 |G| k_02.17904 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
62 |G| k_02.22493 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
63 |G| k_02.27083 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
64 |G| k_02.31672 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
65 |G| k_02.36261 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
66 |G| k_02.40850 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
67 |G| k_02.45440 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
68 |G| k_02.50029 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
69 |G| k_02.54618 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
70 |G| k_02.59208 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
71 |G| k_02.63797 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
72 |G| k_02.68386 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
73 |G| k_02.72976 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
74 |G| k_02.77565 U(a = 0.14, b = 0.15) [ 0.00 .. inf],
75 |G| k_02.82154 U(a = 0.14, b = 0.15) [ 0.00 .. inf]]
Customize the radius ratio knot locations#
[8]:
ts.set_radius_ratio_knots(geomspace(*ts.wavelength[[0, -1]], 50))
ts.add_radius_ratio_knots(linspace(0.768-0.015, 0.768+0.015, 5))
ts.set_radius_ratio_prior('UP', 0.14, 0.15)
ts.plot_setup();
Fit the white light curve#
[9]:
ts.fit_white()
fig = ts.plot_white()
Normalise the baseline#
[10]:
ts.normalize_baseline(1)
fig = ts.plot_baseline()
Set up multiprocessing#
[11]:
def lnpostf(pv):
return ts.lnposterior(pv)
[12]:
pool = Pool(8)
Fit the transmission spectrum#
[16]:
ts.fit(niter=1500, npop=250, pool=pool, lnpost=lnpostf)
[17]:
fig = ts.plot_fit(result='fit', figsize=(13,8),
res_args=dict(pmin=5, pmax=95),
trs_args=dict(xscale='log', ylim=(2.0,2.25), xticks=[0.6, 0.88, 1.16, 1.44, 1.72, 2.0, 2.30, 2.8]))
MCMC sampling#
[18]:
ts.sample(100, thin=10, repeats=3, pool=pool, lnpost=lnpostf)
[49]:
import astropy.io.fits as pf
import codecs
import json
import pickle
from easyts.ldtkld import LDTkLD
from astropy.table import Table
def save(self, overwrite: bool = False) -> None:
"""Saves the EasyTS analysis to a FITS file.
Parameters
----------
overwrite : bool, optional
Flag indicating whether to overwrite an existing file with the same name.
"""
pri = pf.PrimaryHDU()
pri.header['name'] = self.name
pri.header['t0'] = self.transit_center
pri.header['t14'] = self.transit_duration
pri.header['ndgroups'] = self.data.ngroups
pr = pf.ImageHDU(name='priors')
priors = [pickle.dumps(p) for p in self.ps]
pr.header['priors'] = json.dumps(codecs.encode(pickle.dumps(priors), "base64").decode())
if isinstance(self._tsa.ldmodel, LDTkLD):
ldm = self._tsa.ldmodel
pri.header['ldmodel'] = 'ldtk'
pri.header['ldtkld'] = json.dumps(codecs.encode(pickle.dumps((ldm.sc.filters, ldm.sc.teff, ldm.sc.logg,
ldm.sc.metal, ldm.dataset)), "base64").decode())
else:
pri.header['ldmodel'] = self._tsa.ldmodel
k_knots = pf.ImageHDU(self._tsa.k_knots, name='k_knots')
ld_knots = pf.ImageHDU(self._tsa.ld_knots, name='ld_knots')
hdul = pf.HDUList([pri, k_knots, ld_knots, pr])
for i, d in enumerate(self.data.data):
flux = pf.ImageHDU(d.fluxes, name=f'flux_{i}')
flux.header['name'] = d.name
ferr = pf.ImageHDU(d.errors, name=f'ferr_{i}')
wave = pf.ImageHDU(d.wavelength, name=f'wavelength_{i}')
time = pf.ImageHDU(d.time, name=f'time_{i}')
hdul.extend([time, wave, flux, ferr])
if self._tsa.de is not None:
de = pf.BinTableHDU(Table(self._tsa._de_population, names=self.ps.names), name='DE')
de.header['npop'] = self._tsa.de.n_pop
de.header['ndim'] = self._tsa.de.n_par
de.header['imin'] = self._tsa.de.minimum_index
hdul.append(de)
if self._tsa.sampler is not None:
mc = pf.BinTableHDU(Table(self._tsa.sampler.flatchain, names=self.ps.names), name='MCMC')
mc.header['npop'] = self._tsa.sampler.nwalkers
mc.header['ndim'] = self._tsa.sampler.ndim
hdul.append(mc)
hdul.writeto(f"{self.name}.fits", overwrite=True)
save(ts)
[55]:
from easyts.tsdata import TSDataSet
from pytransit.param import ParameterSet
def load_model(fname, name: str | None = None):
"""Loads an EasyTS analysis from a FITS file.
Parameters
----------
fname : str
The name of the savefile.
name : str, optional
The name of the new EasyTS model. If not provided, the original analysis name will be used.
Returns
-------
a : EasyTS
An EasyTS analysis.
Raises
------
IOError
If there is an error while opening or reading the file.
ValueError
If the file format is invalid or does not match the expected format.
"""
with pf.open(fname) as hdul:
d = []
for i in range(hdul[0].header['NDGROUPS']):
print(hdul[f'FLUX_{i}'].header)
d.append(TSData(hdul[f'TIME_{i}'].data.astype('d'), hdul[f'WAVELENGTH_{i}'].data.astype('d'),
hdul[f'FLUX_{i}'].data.astype('d'), hdul[f'FERR_{i}'].data.astype('d'),
name=hdul[f'FLUX_{i}'].header['NAME']))
data = TSDataSet(d)
if hdul[0].header['LDMODEL'] == 'ldtk':
filters, teff, logg, metal, dataset = pickle.loads(codecs.decode(json.loads(hdul[0].header['LDTKLD']).encode(), "base64"))
ldm = LDTkLD(filters, teff, logg, metal, dataset=dataset)
else:
ldm = hdul[0].header['LDMODEL']
a = EasyTS(name or hdul[0].header['NAME'], ldmodel=ldm, data=data)
a.set_radius_ratio_knots(hdul['K_KNOTS'].data.astype('d'))
a.set_limb_darkening_knots(hdul['LD_KNOTS'].data.astype('d'))
a.transit_center = hdul[0].header['T0']
a.transit_duration = hdul[0].header['T14']
if a.transit_duration is not None:
a.ootmask = abs(a.time - a.transit_center) > 0.502 * a.transit_duration
priors = pickle.loads(codecs.decode(json.loads(hdul['PRIORS'].header['PRIORS']).encode(), "base64"))
a._tsa.ps = ParameterSet([pickle.loads(p) for p in priors])
a._tsa.ps.freeze()
if 'DE' in hdul:
a._tsa._de_population = Table(hdul['DE'].data).to_pandas().values
a._tsa._de_imin = hdul['DE'].header['IMIN']
if 'MCMC' in hdul:
npop = hdul['MCMC'].header['NPOP']
ndim = hdul['MCMC'].header['NDIM']
a._tsa._mc_chains = Table(hdul['MCMC'].data).to_pandas().values.reshape([npop, -1, ndim])
return a
ts = load_model("01a_lowres_power2.fits")
XTENSION= 'IMAGE ' / Image extension BITPIX = -64 / array data type NAXIS = 2 / number of array dimensions NAXIS1 = 518 NAXIS2 = 91 PCOUNT = 0 / number of parameters GCOUNT = 1 / number of groups EXTNAME = 'FLUX_0 ' / extension name NAME = 'WASP-39b JSWT NIRISS Order 2' END
XTENSION= 'IMAGE ' / Image extension BITPIX = -64 / array data type NAXIS = 2 / number of array dimensions NAXIS1 = 518 NAXIS2 = 118 PCOUNT = 0 / number of parameters GCOUNT = 1 / number of groups EXTNAME = 'FLUX_1 ' / extension name NAME = 'WASP-39b JWST NIRISS Order 1' END
[57]:
ts.plot_fit()
[57]:
[22]:
ts.sample(1000, thin=100, repeats=3, pool=pool, lnpost=lnpostf)
[21]:
fig = ts.plot_fit(result='mcmc', figsize=(13,7),
res_args=dict(pmin=5, pmax=95),
trs_args=dict(xscale='log', ylim=(2.0,2.25), xticks=[0.6, 0.88, 1.16, 1.44, 1.72, 2.0, 2.30, 2.8]))
©2024 Hannu Parviainen