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

Explicitly set reference epoch, save chi2 to object, and fixes plotting for stationary tracks #47

Merged
merged 7 commits into from
Jan 4, 2024
Merged
6 changes: 4 additions & 2 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -14,10 +14,12 @@ bgstar_emcee.py
.vscode*
*beta_pic*
*SAO*
*sao*
*yses2*
*YSES2*
.pyproject.toml.sw
.pytest_cache/
tests/test_changes.py
tests/HD_131399A_*
tests/dynesty.save
tests/yses2*
tests/YSES_2*
tests/gaia_query_*
127 changes: 106 additions & 21 deletions backtracks/backtracks.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,13 @@
from typing import Optional, Tuple

import astropy.units as u
from scipy.optimize import minimize
import dynesty
import novas.compat as novas
import numpy as np
import pandas as pd

from astropy.coordinates import SkyCoord
from astropy.coordinates import SkyCoord, Distance
from astropy.table import Table
from astropy.io import fits
from astropy.time import Time
Expand All @@ -34,7 +35,7 @@
from schwimmbad import MPIPool

from backtracks.utils import pol2car, transform_gengamm, transform_normal, transform_uniform, utc2tt, HostStarPriors
from backtracks.plotting import diagnostic, neighborhood, plx_prior, posterior, trackplot
from backtracks.plotting import diagnostic, neighborhood, plx_prior, posterior, trackplot, stationtrackplot


# Set the Gaia data release to DR3
Expand Down Expand Up @@ -111,7 +112,14 @@ def __init__(self, target_name: str, candidate_file: str, nearby_window: float =
self.decs = astrometry[2]
self.decserr = astrometry[4]
self.rho = astrometry[5]


# choose time of reference observation of relative candidate position (defaults to the first observation)
if 'ref_epoch_idx' in kwargs:
self.ref_epoch_idx = kwargs['ref_epoch_idx']
self.ref_epoch = self.epochs[self.ref_epoch_idx]
else:
self.ref_epoch_idx = 0
self.ref_epoch = self.epochs[self.ref_epoch_idx]

if 'query_file' in kwargs and kwargs['query_file'] is not None:
self.gaia_id = int(Path(kwargs['query_file']).stem.split('_')[-1])
Expand Down Expand Up @@ -162,6 +170,7 @@ def __init__(self, target_name: str, candidate_file: str, nearby_window: float =
if kwargs['rv_host_method'].lower() == 'uniform':
self.rv_host_method='uniform'
self.rv_lower,self.rv_upper=kwargs['rv_host_params'] # e.g., rv_host_params=(-10,10)
self.radvelo = np.mean([self.rv_lower,self.rv_upper])
elif kwargs['rv_host_method'].lower() == 'normal':
self.rv_host_method='normal'
self.radvelo,self.sig_rv=kwargs['rv_host_params'] # e.g., rv_host_params=(10,0.5)
Expand All @@ -180,16 +189,16 @@ def __init__(self, target_name: str, candidate_file: str, nearby_window: float =
[ra_dec_corr*sig_ra*sig_dec,sig_dec**2,sig_dec*sig_pmra*dec_pmra_corr,sig_dec*sig_pmdec*dec_pmdec_corr,sig_dec*sig_parallax*dec_parallax_corr],
[ra_pmra_corr*sig_ra*sig_pmra,dec_pmra_corr*sig_pmra*sig_dec,sig_pmra**2,sig_pmra*sig_pmdec*pmra_pmdec_corr,sig_pmra*sig_parallax*parallax_pmra_corr],[sig_pmdec*sig_ra*ra_pmdec_corr,sig_pmdec*sig_dec*dec_pmdec_corr,sig_pmdec*sig_pmra*pmra_pmdec_corr,sig_pmdec**2,sig_pmdec*sig_parallax*parallax_pmdec_corr],[sig_parallax*sig_ra*ra_parallax_corr,sig_parallax*sig_dec*dec_parallax_corr,sig_parallax*sig_pmra*parallax_pmra_corr,sig_parallax*sig_pmdec*parallax_pmdec_corr,sig_parallax**2]])
self.HostStarPriors = HostStarPriors(self.host_mean,self.host_cov)
# initial estimate for background star scenario (best guesses)
# (manually looked at approximate offset from star with cosine compensation for
# dependence of RA on declination (RA is a singularity at the declination poles))
# this offset is calculated wrong (should use spherical offsets to) but it kinda works because the uniform prior is bigger than this. not all of these parameters are still in use
self.ra0 = self.rao-(self.ras[0])/1000/3600/np.cos(self.deco/180.*np.pi)
self.dec0 = self.deco-(self.decs[0])/1000/3600
self.pmra0 = self.pmrao # mas/yr
self.pmdec0 = self.pmdeco # mas/yr
self.par0 = self.paro/10 # mas
self.radvel0 = 0 # km/s

# set inital guess for position at gaia epoch
self.set_initial_position()

# create parameter set for stationary case
# params = ra, dec, pmra, pmdec, par, ra_host, dec_host, pmra_host, pmdec_host, par_host, rv_host
self.stationary_params = [self.ra0, self.dec0, 0, 0, 0, self.rao, self.deco, self.pmrao, self.pmdeco, self.paro, self.radvelo]
# Compute useful chi2 value
self.stationary_loglike = self.loglike(self.stationary_params)
self.stationary_chi2_red = -2.*self.stationary_loglike/((2*(len(self.epochs)-1))-self.ndim)

jd_start, jd_end, number = ephem_open()
print('[BACKTRACK INFO]: Opened ephemeris file')
Expand All @@ -215,6 +224,48 @@ def __init__(self, target_name: str, candidate_file: str, nearby_window: float =

# this converts the Epoch from the Gaia ref_epoch (2016 for DR3) to 2000 following ICRS

def set_initial_position(self):
# initial estimate for background star scenario (best guesses)
print(f'[BACKTRACK INFO]: Estimating candidate position if stationary in RA,Dec @ {self.gaia_epoch} from observation #'+str(self.ref_epoch_idx))
# we'll do a rough estimate using astropy, then minimize the distance between
# the novas projection and the specified observation using scipy's BFGS with
# RA,DEC @ Gaia epoch as free parameters.

# educated guess with SkyCoord transformations
init_host_coord = SkyCoord(ra=self.rao*u.deg, dec=self.deco*u.deg, distance=Distance(parallax=self.paro*u.mas), pm_ra_cosdec=self.pmrao*u.mas/u.yr, pm_dec=self.pmdeco*u.mas/u.yr, obstime=Time(self.gaia_epoch,format='jyear',scale='tcb'))
# propagate host position from Gaia epoch (2016.0 for DR3) to reference observation epoch
init_host_coord_at_obs = init_host_coord.apply_space_motion(Time(self.ref_epoch, format='jd'))
# apply measurement of relative astrometry to get candidate in absolute coordinates at time of observation
init_cand_coord_at_obs = init_host_coord_at_obs.spherical_offsets_by(self.ras[self.ref_epoch_idx] * u.mas, self.decs[self.ref_epoch_idx] * u.mas)
init_cand_coord = SkyCoord(ra=init_cand_coord_at_obs.ra, dec=init_cand_coord_at_obs.dec, distance=Distance(parallax=self.paro*u.mas), pm_ra_cosdec=self.pmrao*u.mas/u.yr, pm_dec=self.pmdeco*u.mas/u.yr, obstime=Time(self.ref_epoch, format='jd'))
# propagate candidate position from epoch of observation to Gaia epoch
tdiff = Time(self.ref_epoch, format='jd').decimalyear-self.gaia_epoch
init_cand_coord_at_gaia = init_cand_coord.apply_space_motion(Time(self.gaia_epoch+tdiff,format='jyear',scale='tcb'))
ra0 = init_cand_coord_at_gaia.ra.value
dec0 = init_cand_coord_at_gaia.dec.value

def dummy_loglike(radec):
ra, dec = radec
dummy_param = ra, dec, 0, 0, 0, self.rao, self.deco, self.pmrao, self.pmdeco, self.paro, self.radvelo
xs, ys = self.radecdists([utc2tt(self.ref_epoch)], dummy_param)

distance = np.linalg.norm(np.array([self.ras[self.ref_epoch_idx],self.decs[self.ref_epoch_idx]])-np.array([xs[0],ys[0]]))

return distance

# minimize distance between these points (bounds are 5 mas, tolerance might need adjusting?)
bound = 5e-5
init_result = minimize(dummy_loglike, np.array([ra0,dec0]), tol=1e-15, method='Nelder-Mead', bounds=((ra0-bound, ra0+bound), (dec0-bound, dec0+bound)))

# not sure which of these is used anymore
self.ra0 = init_result.x[0]
self.dec0 = init_result.x[1]
self.pmra0 = self.pmrao # mas/yr
self.pmdec0 = self.pmdeco # mas/yr
self.par0 = self.paro/10 # mas
self.radvel0 = 0 # km/s


def query_astrometry(self, nearby_window: float = 0.5):
# resolve target in simbad
target_result_table = Simbad.query_object(self.target_name)
Expand Down Expand Up @@ -529,41 +580,47 @@ def fit(self, dlogz=0.5, npool=4, dynamic=False, nlive=200, mpi_pool=False, resu
self.run_median = np.median(samples, axis=0)
self.run_quant = np.quantile(samples, [0.1587, 0.8413], axis=0)

# Compute useful chi2 value
self.median_loglike = self.loglike(self.run_median)
self.median_chi2_red = -2.*self.median_loglike/((2*len(self.epochs))-self.ndim)

return self.results

def save_results(self, fileprefix='./'):
save_dict = {'med': self.run_median, 'quant': self.run_quant, 'results': self.results}
target_label = self.target_name.replace(' ','_')
file_name = f'{fileprefix}{target_label}_dynestyrun_results.pkl'
print('[BACKTRACK INFO]: Saving results to {}'.format(file_name))
pickle.dump(save_dict, open(file_name, "wb"))

def load_results(self, fileprefix: str = './'):
target_label = self.target_name.replace(' ', '_')
file_name = f'{fileprefix}{target_label}_dynestyrun_results.pkl'
print('[BACKTRACK INFO]: Loading results from {}'.format(file_name))
save_dict = pickle.load(open(file_name, "rb"))
self.run_median = save_dict['med']
self.run_quant = save_dict['quant']
self.results = save_dict['results']

# recompute useful chi2 value
self.median_loglike = self.loglike(self.run_median)
self.median_chi2_red = -2.*self.median_loglike/((2*len(self.epochs))-self.ndim)

def generate_plots(
self,
days_backward: float = 3.*365.,
days_forward: float = 3.*365.,
step_size: float = 10.,
ref_epoch: Optional[Tuple[int, int, int]] = None,
# ref_epoch: Optional[Tuple[int, int, int]] = None,
plot_radec: bool = False,
plot_stationary: bool = False,
fileprefix: str = './',
filepost: str = '.pdf',
) -> Tuple[Figure, Figure, Figure, Figure, Figure]:
"""
"""
if ref_epoch is None:
mean_epoch = np.floor(np.mean(self.epochs))
ref_epoch = Time(mean_epoch, format='jd')

else:
ref_epoch = Time(f'{ref_epoch[0]}-{ref_epoch[1]}-{ref_epoch[2]}T12:00:00', format='isot')
print('[BACKTRACK INFO]: Generating Plots')
ref_epoch = Time(self.ref_epoch, format='jd')

fig_track = trackplot(
self,
Expand All @@ -581,5 +638,33 @@ def generate_plots(
fig_diag = diagnostic(self, fileprefix=fileprefix, filepost=filepost)
fig_prior = plx_prior(self, fileprefix=fileprefix, filepost=filepost)
fig_hood = neighborhood(self, fileprefix=fileprefix, filepost=filepost)

print('[BACKTRACK INFO]: Plots saved to {}'.format(fileprefix))
return fig_track, fig_post, fig_diag, fig_prior, fig_hood

def generate_stationary_plot(
self,
days_backward: float = 3.*365.,
days_forward: float = 3.*365.,
step_size: float = 10.,
# ref_epoch: Optional[Tuple[int, int, int]] = None,
plot_radec: bool = False,
fileprefix: str = './',
filepost: str = '.pdf',
) -> Figure:
"""
"""
print('[BACKTRACK INFO]: Generating Stationary plot')
ref_epoch = Time(self.ref_epoch, format='jd')

fig_track = stationtrackplot(
self,
ref_epoch=ref_epoch.jd,
days_backward=days_backward,
days_forward=days_forward,
step_size=step_size,
plot_radec=plot_radec,
fileprefix=fileprefix,
filepost=filepost
)
print('[BACKTRACK INFO]: Stationary plot saved to {}'.format(fileprefix))
return fig_track
Loading
Loading