upsAnd_rejection_sampler_no_parallel.py

import numpy as np
import orbitize
from orbitize import hipparcos, gaia, read_input, system
import pandas

import compute_sep  # directly from Sarah
import astropy
from astropy.table import Table
from astropy.io import ascii
from astropy.time import Time
from datetime import datetime

radvel_datapath = './data/9826_rv_main_binned.txt'
radvel_resultspath = './fits/results/9826_radvel_chains.csv'
save_orbitize_datapath = './data/orbitize_basis_data.csv'

mass = (1.29, 0.04)                 # from Rosenthal et al. 2021
parallax = (74.1940, 0.2083)        # from Gaia EDR3 https://simbad.cds.unistra.fr/simbad/sim-ref?bibcode=2020yCat.1350....0G 
n_planets = 3

iad_path = './data/H007513.d'
gost_path = './data/gost_22.4.3_806543_2025-08-29-15-24-50_ups_And.csv'
hip_num = 7513

n_samples = 50_000
posterior_savepath = './RJ_results_no_parallel.csv'

def convert_data_RadToOrb(radvel_path, save_path=False):
    ''' 
    Function to convert radvel data file to orbitize data file and save it 
    in the specified directory.

    args:
        radvel_path (str): full path to radvel data .txt file
        save_name (bool or str): path to optionally save the new orbitize-friendly data file (.csv)

    return:
        data (astropy.table.Table): data with new column names
    '''

    data = ascii.read(radvel_path)
    # rename columns to orbitize labels
    data.rename_columns(names=('time', 'mnvel', 'errvel'), new_names=('epoch', 'rv', 'rv_err'))
    
    # convert from JD to MJD
    for row in data:
        if (row['epoch'] > 2400000.5):
            mjd_epoch = Time(row['epoch'], format='jd').mjd         # use Astropy to convert from jd to mjd
            row['epoch'] = mjd_epoch
            
    # add column for object number
    data.add_column(np.zeros(len(data['epoch']), dtype=int), name='object', index=1)
    
    # save the data
    data.write(save_path, format='csv', overwrite=True)

    return data

def convert_results_RadToOrb(radvel_path:str, radvel_data, mass:tuple, parallax:tuple, n_planets:int):
    ''' 
    Function to convert radvel results file to orbitize results.post shape (orbits, params).

    args:
        radvel_path (str): full path to the radvel results file (.csv)
        radvel_data (astropy.table.Table): the radvel data table object
        mass (tuple, float): tuple of (stellar mass, mass error)
        parallax (tuple, float: tuple of (stellar plx, plx error)
        n_planets (int): the number of planetary companions in the system
    returns:
        orbitize_post (np.ndarray): the posterior data array size NxM for N orbits and M parameters
    '''
    df = pandas.read_csv(radvel_path)

    epochs = Time(radvel_data['epoch'].value, format='jd')
    basis_string = 'per tc secosw sesinw k'
    m0, m0_err = mass
    plx, plx_err = parallax

    # place to store each planet posterior array
    planet_param_arr = []
    planet_mass_arr = []
    star_mass_arr = []
    plx_arr = []

    for planet in np.arange(n_planets)+1:
        # pl_str = str(planet)
        seps, df_orb = compute_sep.compute_sep(df, epochs, basis_string, m0, m0_err, plx, plx_err, n_planets=n_planets, pl_num=planet)
        
        # use column names to restack the posterior
        planet_param_arr.append(df_orb[['sma', 'ecc', 'inc_rad', 'omega_pl_rad', 'lan_rad', 'tau_58849']].to_numpy())
        planet_mass_arr.append(df_orb['mp'].to_numpy().reshape((len(df_orb)),1))
        star_mass_arr.append(df_orb['m_st'].to_numpy().reshape((len(df_orb)),1))
        plx_arr.append(df_orb['plx'].to_numpy().reshape((len(df_orb)),1))

        # add in these terms to match the default orbitize gaia param_idx
        gamma_term = np.ones(shape=(len(df_orb),1)) * np.nan
        sigma_term = np.ones(shape=(len(df_orb),1)) * np.nan

    # add together in one np.ndarray
    param_stack = np.hstack((planet_param_arr))
    int_post = np.hstack((param_stack, plx_arr[0], gamma_term, sigma_term))
    mass_stack = np.hstack((planet_mass_arr))

    orbitize_post = np.hstack((int_post, mass_stack, star_mass_arr[0]))

    return orbitize_post

def create_system_object(IAD_path:str, gost_path:str, rv_datapath:str, hip_num:str, 
                         num_secondary_bodies:int, mass:tuple, parallax:tuple):
    ''' 
    Function to create the orbitize.system.System object from gaia+hipparcos data files.

    args:
        IAD_path (str): 
        gost_path (str):
        rv_datapath (str):
        hip_num (str):
        num_secondary_bodies (int):
        mass (tuple, flaot):
        parallax (tuple, float):
        
    returns:
        system_object (orbitize.system.System)
    '''
    
    hipparcos_lnprob = hipparcos.HipparcosLogProb(IAD_path, int(hip_num), num_secondary_bodies=num_secondary_bodies)
    hgca_lnprob = gaia.HGCALogProb(hip_id=int(hip_num), hiplogprob=hipparcos_lnprob, gost_filepath=gost_path)

    data_table = read_input.read_file(rv_datapath)

    m0, m0_err = mass
    plx, plx_err = parallax

    system_object = system.System(
        num_secondary_bodies,
        data_table,
        m0, 
        plx,
        mass_err=m0_err,
        plx_err=plx_err,
        tau_ref_epoch=58849,
        fit_secondary_mass=True,
        gaia=hgca_lnprob)

    return system_object

def draw_samples(posterior, sample_size):
    '''
    Function to draw random samples from the orbitize posteriors.

    args:
        posterior (NxM np.ndarray): for N orbits and M parameters
        sample_size (int): number of samples to draw from the posteriors
    
    returns:
        sample_posterior (SxM np.ndarray): for S sample size and M parameters
    '''
    # create random indices
    generator = np.random.default_rng(seed=None)
    sample_posterior = generator.choice(posterior, size=sample_size, replace=True, axis=0)

    return sample_posterior

def compute_chi2_prob_and_accept(orbitize_system, posterior):
    '''
    Function to 
    1) compute proper motion anomally models from RV posteriors (at the gaia+hipp data epochs)
    2) assess goodness of fit (chi2_prob) with the HGCA data
    3) accept or reject orbits

    args:
        orbitize_system (orbitize.system.System object): to call .compute_all_orbits from
        posterior (np.array): len M array (M parameters) from which to draw model parameters ** assumes one orbit input at a time **

    returns:
        accepted (list): N x 2 shape of N accepted orbits with returned posterior/ orbit params and the accepted chi2_prob value 
        rejected (list): N x 2 shape of N rejected orbits with returned posterior/ orbit params and the rejected chi2_prob value 
    '''

    if type(orbitize_system) != orbitize.system.System:
        raise TypeError('orbitize_system must be an orbitize.system.System object')
    
    # determine which epochs to compute models for
    gaiahip_epochs = astropy.time.Time(np.append(orbitize_system.gaia.hipparcos_epoch, orbitize_system.gaia.gaia_epoch),
                format="decimalyear").mjd

    accepted = []

    for post in posterior:
        ra_model, dec_model, rv_model = orbitize_system.compute_all_orbits(
            post, epochs=gaiahip_epochs
        )
        ra_star, dec_star = ra_model[:,0,:], dec_model[:,0,:]

        # compute normalization term to subtract later
        norm_term = orbitize.lnlike.chi2_norm_term(orbitize_system.gaia.dpm_err, orbitize_system.gaia.dpm_corr)

        # compute the lnlikes and correct for normalization
            # this is doing: data - model --> where data is the gaia-hipp 
        lnlike = orbitize_system.gaia.compute_lnlike(ra_star, dec_star,                 # orbitize_system.gaia == orbitize.gaia.HGCALogProb here (defined when creating the system object)
                                        samples=0, param_idx=orbitize_system.param_idx) # samples param not actually used in the function
        chi2 = lnlike - norm_term
        chi2_prob = np.exp(chi2)

        # compute a random threshold and compare
        threshold = np.random.random(size=1)
        if chi2_prob > threshold:
            # add the chi2 value to the end of the posterior
            accepted_post = np.append(post, chi2_prob)
            # add the updated posterior array to the accepted orbits list
            accepted.append(accepted_post)

    return np.array(accepted)     # has shape N accepted x M params

## functions to set up sampler/ call other functions ##

def prepare_sampler(radvel_datapath:str, radvel_resultspath:str, IAD_path:str, 
                    gost_path:str, mass:tuple, parallax:tuple, hip_num:int,
                    n_planets:int, save_datapath:str):
    ''' 
    Note: make sure you run the IAD test before this step to check for quality of data.
    Function to prepare the RV data and results for the rejection sampler, by calling
    functions defined earlier in this file.

    args:
        radvel_datapath (str):
        radvel_resultspath (str):
        IAD_path (str):
        gost_path (str):
        mass (tuple, float):
        parallax (tuple, float):
        hip_num (str): Hipparcos identifier number
        n_planets (int): number of secondary bodies in the system
        save_datapath (str): 

    return:
        posterior (np.ndarray):
        system_object (orbitize.system.System):
    '''
    # convert the data file to an orbitize readable array
    radvel_data = convert_data_RadToOrb(radvel_datapath, save_path=save_datapath)
    # create an orbitize-shaped posterior
    orbitize_post = convert_results_RadToOrb(radvel_resultspath, radvel_data, 
                                             mass=mass, parallax=parallax, n_planets=n_planets)
    
    # create a system object with the gaia data
    system_object = create_system_object(IAD_path=IAD_path, gost_path=gost_path, rv_datapath=save_datapath, 
                                         hip_num=hip_num, num_secondary_bodies=n_planets, mass=mass, parallax=parallax)
    
    return orbitize_post, system_object

def run_rejection_sampler(n_samples:int, posterior, system_object):
    ''' 
    Function to run all steps of the rejection sampler.

    args:
        n_samples (int): the number of orbits in the sample
        posterior (np.ndarray):
        system_object (orbitize.system.System):
        
    return:
        chi2_out (list shape (n_batches, n_samples, N_accepted, 2))
    '''

    # gather a random sample from the posterior the length of samples
    sample_posterior = draw_samples(posterior=posterior, sample_size=n_samples)

    # compute chi2, reject or accept
    accepted = compute_chi2_prob_and_accept(orbitize_system=system_object, posterior=sample_posterior)

    return accepted

if __name__ == "__main__":
    # prepare the sampler inputs
    full_posterior, system_object = prepare_sampler(radvel_datapath=radvel_datapath, radvel_resultspath=radvel_resultspath,
                                                    IAD_path=iad_path, gost_path=gost_path, mass=mass,
                                                    parallax=parallax, hip_num=hip_num, n_planets=n_planets, save_datapath=save_orbitize_datapath)

    print()
    print(f'Beginning Rejection Sampling:   total sample size: {n_samples}.\n')
    start = datetime.now()

    # run the rejection sampler
    accepted_orbits = run_rejection_sampler(n_samples=n_samples, posterior=full_posterior, system_object=system_object)
    end = datetime.now()

    # compute how long it took
    duration_seconds = (end-start).total_seconds()
    if duration_seconds > (60*60):
        duration = duration_seconds/60/60 # hr
        print(f'Completed Rejection Sampling:   duration {duration:.2f} [hrs].\n')
    elif duration_seconds > 60:
        duration = duration_seconds/60 # min
        print(f'Completed Rejection Sampling:   duration {duration:.2f} [mins].\n')
    else:
        print(f'Completed Rejection Sampling:   duration {duration_seconds:.2f} [sec].\n')
    
    # save the posteriors for the accepted orbits in a csv file

    # save to an astropy table
    system_object.param_idx['chi2'] = 0 # add the accepted chi2 key
    names = system_object.param_idx.keys()
    accepted_table = Table(accepted_orbits, names=names)
    
    print(f'Results file contains {len(accepted_orbits)} accepted orbits:   saving to "{posterior_savepath}"\n')

    # save to path
    accepted_table.write(posterior_savepath, format='csv', overwrite=True)
    print(f'Results saved. Sampler run is complete!\n')