def test_fit_selfconsist(): """ Tests that the masses we get from orbitize! are what we expeect. Note that this does not test for correctness. """ # generate planet b orbital parameters b_params = [1, 0, 0, 0, 0, 0.5] tau_ref_epoch = 0 mass_b = 0.001 # Msun m0 = 1 # Msun plx = 1 # mas # generate planet c orbital parameters # at 0.3 au, and starts on the opposite side of the star relative to b c_params = [0.3, 0, 0, np.pi, 0, 0.5] mass_c = 0.002 # Msun mtot_c = m0 + mass_b + mass_c mtot_b = m0 + mass_b period_b = np.sqrt(b_params[0]**3 / mtot_b) period_c = np.sqrt(c_params[0]**3 / mtot_c) epochs = np.linspace(0, period_b * 365.25, 20) + tau_ref_epoch # the full period of b, MJD # comptue Keplerian orbit of b ra_model_b, dec_model_b, vz_model = kepler.calc_orbit( epochs, b_params[0], b_params[1], b_params[2], b_params[3], b_params[4], b_params[5], plx, mtot_b, mass_for_Kamp=m0, tau_ref_epoch=tau_ref_epoch) # comptue Keplerian orbit of c ra_model_c, dec_model_c, vz_model_c = kepler.calc_orbit( epochs, c_params[0], c_params[1], c_params[2], c_params[3], c_params[4], c_params[5], plx, mtot_c, tau_ref_epoch=tau_ref_epoch) # perturb b due to c ra_model_b_orig = np.copy(ra_model_b) dec_model_b_orig = np.copy(dec_model_b) # the sign is positive b/c of 2 negative signs cancelling. ra_model_b += mass_c / m0 * ra_model_c dec_model_b += mass_c / m0 * dec_model_c # # perturb c due to b # ra_model_c += mass_b/m0 * ra_model_b_orig # dec_model_c += mass_b/m0 * dec_model_b_orig # generate some fake measurements to fit to. Make it with b first t = table.Table([ epochs, np.ones(epochs.shape, dtype=int), ra_model_b, 0.00001 * np.ones(epochs.shape, dtype=int), dec_model_b, 0.00001 * np.ones(epochs.shape, dtype=int) ], names=[ "epoch", "object", "raoff", "raoff_err", "decoff", "decoff_err" ]) # add c for eps, ra, dec in zip(epochs, ra_model_c, dec_model_c): t.add_row([eps, 2, ra, 0.000001, dec, 0.000001]) filename = os.path.join(orbitize.DATADIR, "multiplanet_fake_2planettest.csv") t.write(filename, overwrite=True) # create the orbitize system and generate model predictions using the standard 1 body model for b, and the 2 body model for b and c astrom_dat = read_input.read_file(filename) sys = system.System(2, astrom_dat, m0, plx, tau_ref_epoch=tau_ref_epoch, fit_secondary_mass=True) # fix most of the orbital parameters to make the dimensionality a bit smaller sys.sys_priors[sys.param_idx['ecc1']] = b_params[1] sys.sys_priors[sys.param_idx['inc1']] = b_params[2] sys.sys_priors[sys.param_idx['aop1']] = b_params[3] sys.sys_priors[sys.param_idx['pan1']] = b_params[4] sys.sys_priors[sys.param_idx['ecc2']] = c_params[1] sys.sys_priors[sys.param_idx['inc2']] = c_params[2] sys.sys_priors[sys.param_idx['aop2']] = c_params[3] sys.sys_priors[sys.param_idx['pan2']] = c_params[4] sys.sys_priors[sys.param_idx['m1']] = priors.LogUniformPrior( mass_b * 0.01, mass_b * 100) sys.sys_priors[sys.param_idx['m2']] = priors.LogUniformPrior( mass_c * 0.01, mass_c * 100) n_walkers = 30 samp = sampler.MCMC(sys, num_temps=1, num_walkers=n_walkers, num_threads=1) # should have 8 parameters assert samp.num_params == 6 # start walkers near the true location for the orbital parameters np.random.seed(123) # planet b samp.curr_pos[:, 0] = np.random.normal(b_params[0], 0.01, n_walkers) # sma samp.curr_pos[:, 1] = np.random.normal(b_params[-1], 0.01, n_walkers) # tau # planet c samp.curr_pos[:, 2] = np.random.normal(c_params[0], 0.01, n_walkers) # sma samp.curr_pos[:, 3] = np.random.normal(c_params[-1], 0.01, n_walkers) # tau # we will make a fairly broad mass starting position samp.curr_pos[:, 4] = np.random.uniform(mass_b * 0.25, mass_b * 4, n_walkers) samp.curr_pos[:, 5] = np.random.uniform(mass_c * 0.25, mass_c * 4, n_walkers) samp.curr_pos[0, 4] = mass_b samp.curr_pos[0, 5] = mass_c samp.run_sampler(n_walkers * 50, burn_steps=600) res = samp.results print(np.median(res.post[:, sys.param_idx['m1']]), np.median(res.post[:, sys.param_idx['m2']])) assert np.median(res.post[:, sys.param_idx['sma1']]) == pytest.approx( b_params[0], abs=0.01) assert np.median(res.post[:, sys.param_idx['sma2']]) == pytest.approx( c_params[0], abs=0.01) assert np.median(res.post[:, sys.param_idx['m2']]) == pytest.approx( mass_c, abs=0.5 * mass_c)
def __init__(self, num_secondary_bodies, data_table, stellar_mass, plx, mass_err=0, plx_err=0, restrict_angle_ranges=None, tau_ref_epoch=58849, fit_secondary_mass=False, results=None): self.num_secondary_bodies = num_secondary_bodies self.sys_priors = [] self.labels = [] self.results = [] self.fit_secondary_mass = fit_secondary_mass self.tau_ref_epoch = tau_ref_epoch # # Group the data in some useful ways # self.data_table = data_table # Creates a copy of the input in case data_table needs to be modified self.input_table = self.data_table.copy() # List of arrays of indices corresponding to each body self.body_indices = [] # List of arrays of indices corresponding to epochs in RA/Dec for each body self.radec = [] # List of arrays of indices corresponding to epochs in SEP/PA for each body self.seppa = [] # List of index arrays corresponding to each rv for each body self.rv = [] radec_indices = np.where(self.data_table['quant_type'] == 'radec') seppa_indices = np.where(self.data_table['quant_type'] == 'seppa') rv_indices = np.where(self.data_table['quant_type'] == 'rv') # save indicies for all of the ra/dec, sep/pa measurements for convenience self.all_radec = radec_indices self.all_seppa = seppa_indices for body_num in np.arange(self.num_secondary_bodies + 1): self.body_indices.append( np.where(self.data_table['object'] == body_num)) self.radec.append( np.intersect1d(self.body_indices[body_num], radec_indices)) self.seppa.append( np.intersect1d(self.body_indices[body_num], seppa_indices)) self.rv.append( np.intersect1d(self.body_indices[body_num], rv_indices)) # we should track the influence of the planet(s) on each other/the star if we are not fitting massless planets and # we are not fitting relative astrometry of just a single body self.track_planet_perturbs = self.fit_secondary_mass and \ ((len(self.radec[1]) + len(self.seppa[1]) + len(self.rv[1]) < len(data_table)) or \ (self.num_secondary_bodies > 1)) if restrict_angle_ranges: angle_upperlim = np.pi else: angle_upperlim = 2. * np.pi # # Set priors for each orbital element # for body in np.arange(num_secondary_bodies): # Add semimajor axis prior self.sys_priors.append(priors.LogUniformPrior(0.001, 1e7)) self.labels.append('sma{}'.format(body + 1)) # Add eccentricity prior self.sys_priors.append(priors.UniformPrior(0., 1.)) self.labels.append('ecc{}'.format(body + 1)) # Add inclination angle prior self.sys_priors.append(priors.SinPrior()) self.labels.append('inc{}'.format(body + 1)) # Add argument of periastron prior self.sys_priors.append(priors.UniformPrior(0., 2. * np.pi)) self.labels.append('aop{}'.format(body + 1)) # Add position angle of nodes prior self.sys_priors.append(priors.UniformPrior(0., angle_upperlim)) self.labels.append('pan{}'.format(body + 1)) # Add epoch of periastron prior. self.sys_priors.append(priors.UniformPrior(0., 1.)) self.labels.append('tau{}'.format(body + 1)) # # Set priors on total mass and parallax # self.labels.append('plx') if plx_err > 0: self.sys_priors.append(priors.GaussianPrior(plx, plx_err)) else: self.sys_priors.append(plx) # checking for rv data to include appropriate rv priors: if len(self.rv[0]) > 0 and self.fit_secondary_mass: self.sys_priors.append(priors.UniformPrior( -5, 5)) # gamma prior in km/s self.labels.append('gamma') self.sys_priors.append(priors.LogUniformPrior( 1e-4, 0.05)) # jitter prior in km/s self.labels.append('sigma') if self.fit_secondary_mass: for body in np.arange(num_secondary_bodies) + 1: self.sys_priors.append(priors.LogUniformPrior( 1e-6, 2)) # in Solar masses for now self.labels.append('m{}'.format(body)) self.labels.append('m0') else: self.labels.append('mtot') # still need to append m0/mtot, even though labels are appended above if mass_err > 0: self.sys_priors.append(priors.GaussianPrior( stellar_mass, mass_err)) else: self.sys_priors.append(stellar_mass) # add labels dictionary for parameter indexing self.param_idx = dict(zip(self.labels, np.arange(len(self.labels))))
def __init__(self, num_secondary_bodies, data_table, stellar_mass, plx, mass_err=0, plx_err=0, restrict_angle_ranges=None, tau_ref_epoch=58849, fit_secondary_mass=False, results=None): self.num_secondary_bodies = num_secondary_bodies self.sys_priors = [] self.labels = [] self.results = [] self.fit_secondary_mass = fit_secondary_mass self.tau_ref_epoch = tau_ref_epoch self.restrict_angle_ranges = restrict_angle_ranges # # Group the data in some useful ways # self.data_table = data_table # Creates a copy of the input in case data_table needs to be modified self.input_table = self.data_table.copy() # Rob: check if instrument column is other than default. If other than default, then separate data table into n number of instruments. # gather list of instrument: list_instr = self.data_table['instruments'][name of instrument] # List of arrays of indices corresponding to each body # instruments = np.unique(self.data_table['instruments']) gives a list of unique names self.body_indices = [] # List of arrays of indices corresponding to epochs in RA/Dec for each body self.radec = [] # List of arrays of indices corresponding to epochs in SEP/PA for each body self.seppa = [] # List of index arrays corresponding to each rv for each body self.rv = [] # instr1_tbl = np.where(self.data_table['instruments'] == list_instr[0]) # loop through the indices per input_table: # rv_indices = np.where(instr1_tbl['quant_type'] == 'rv') # ... return the parameter labels for each table # ... self.fit_astrometry = True radec_indices = np.where(self.data_table['quant_type'] == 'radec') seppa_indices = np.where(self.data_table['quant_type'] == 'seppa') if len(radec_indices[0]) == 0 and len(seppa_indices[0]) == 0: self.fit_astrometry = False rv_indices = np.where(self.data_table['quant_type'] == 'rv') # defining all indices to loop through the unique rv instruments to get different offsets and jitters instrument_list = np.unique(self.data_table['instrument']) inst_indices_all = [] for inst in instrument_list: inst_indices = np.where(self.data_table['instrument'] == inst) inst_indices_all.append(inst_indices) # defining indices for unique instruments in the data table self.rv_instruments = np.unique( self.data_table['instrument'][rv_indices]) self.rv_inst_indices = [] for inst in self.rv_instruments: inst_indices = np.where(self.data_table['instrument'] == inst) self.rv_inst_indices.append(inst_indices) # astrometry instruments same for radec and seppa: self.astr_instruments = np.unique( self.data_table['instrument'][np.where( self.data_table['quant_type'] != 'rv')]) # save indicies for all of the ra/dec, sep/pa measurements for convenience self.all_radec = radec_indices self.all_seppa = seppa_indices for body_num in np.arange(self.num_secondary_bodies + 1): self.body_indices.append( np.where(self.data_table['object'] == body_num)) self.radec.append( np.intersect1d(self.body_indices[body_num], radec_indices)) self.seppa.append( np.intersect1d(self.body_indices[body_num], seppa_indices)) self.rv.append( np.intersect1d(self.body_indices[body_num], rv_indices)) # we should track the influence of the planet(s) on each other/the star if we are not fitting massless planets and # we are not fitting relative astrometry of just a single body self.track_planet_perturbs = self.fit_secondary_mass and \ ((len(self.radec[1]) + len(self.seppa[1]) + len(self.rv[1]) < len(data_table)) or \ (self.num_secondary_bodies > 1)) if restrict_angle_ranges: angle_upperlim = np.pi else: angle_upperlim = 2. * np.pi # # Set priors for each orbital element # for body in np.arange(num_secondary_bodies): # Add semimajor axis prior self.sys_priors.append(priors.LogUniformPrior(0.001, 1e7)) self.labels.append('sma{}'.format(body + 1)) # Add eccentricity prior self.sys_priors.append(priors.UniformPrior(0., 1.)) self.labels.append('ecc{}'.format(body + 1)) # Add inclination angle prior self.sys_priors.append(priors.SinPrior()) self.labels.append('inc{}'.format(body + 1)) # Add argument of periastron prior self.sys_priors.append(priors.UniformPrior(0., 2. * np.pi)) self.labels.append('aop{}'.format(body + 1)) # Add position angle of nodes prior self.sys_priors.append(priors.UniformPrior(0., angle_upperlim)) self.labels.append('pan{}'.format(body + 1)) # Add epoch of periastron prior. self.sys_priors.append(priors.UniformPrior(0., 1.)) self.labels.append('tau{}'.format(body + 1)) # # Set priors on total mass and parallax # self.labels.append('plx') if plx_err > 0: self.sys_priors.append(priors.GaussianPrior(plx, plx_err)) else: self.sys_priors.append(plx) # checking for rv data to include appropriate rv priors: if len(self.rv[0]) > 0 and self.fit_secondary_mass: # Rob and Lea: # for instrument in rv_instruments: # add gamma and sigma for each and label each unique gamma and sigma per instrument name (gamma+instrument1, ...) for instrument in self.rv_instruments: self.sys_priors.append(priors.UniformPrior( -5, 5)) # gamma prior in km/s self.labels.append('gamma_{}'.format(instrument)) self.sys_priors.append(priors.LogUniformPrior( 1e-4, 0.05)) # jitter prior in km/s self.labels.append('sigma_{}'.format(instrument)) if self.fit_secondary_mass: for body in np.arange(num_secondary_bodies) + 1: self.sys_priors.append(priors.LogUniformPrior( 1e-6, 2)) # in Solar masses for now self.labels.append('m{}'.format(body)) self.labels.append('m0') else: self.labels.append('mtot') # still need to append m0/mtot, even though labels are appended above if mass_err > 0: self.sys_priors.append(priors.GaussianPrior( stellar_mass, mass_err)) else: self.sys_priors.append(stellar_mass) # add labels dictionary for parameter indexing self.param_idx = dict(zip(self.labels, np.arange(len(self.labels))))
def __init__(self, num_secondary_bodies, data_table, stellar_mass, plx, mass_err=0, plx_err=0, restrict_angle_ranges=None, tau_ref_epoch=58849, fit_secondary_mass=False, results=None): self.num_secondary_bodies = num_secondary_bodies self.sys_priors = [] self.labels = [] self.results = [] self.fit_secondary_mass = fit_secondary_mass self.tau_ref_epoch = tau_ref_epoch # # Group the data in some useful ways # self.data_table = data_table # Creates a copy of the input in case data_table needs to be modified self.input_table = self.data_table.copy() # List of arrays of indices corresponding to each body self.body_indices = [] # List of arrays of indices corresponding to epochs in RA/Dec for each body self.radec = [] # List of arrays of indices corresponding to epochs in SEP/PA for each body self.seppa = [] radec_indices = np.where(self.data_table['quant_type'] == 'radec') seppa_indices = np.where(self.data_table['quant_type'] == 'seppa') for body_num in np.arange(self.num_secondary_bodies + 1): self.body_indices.append( np.where(self.data_table['object'] == body_num)) self.radec.append( np.intersect1d(self.body_indices[body_num], radec_indices)) self.seppa.append( np.intersect1d(self.body_indices[body_num], seppa_indices)) if restrict_angle_ranges: angle_upperlim = np.pi else: angle_upperlim = 2. * np.pi # # Set priors for each orbital element # for body in np.arange(num_secondary_bodies): # Add semimajor axis prior self.sys_priors.append(priors.LogUniformPrior(0.001, 1e7)) self.labels.append('sma{}'.format(body + 1)) # Add eccentricity prior self.sys_priors.append(priors.UniformPrior(0., 1.)) self.labels.append('ecc{}'.format(body + 1)) # Add inclination angle prior self.sys_priors.append(priors.SinPrior()) self.labels.append('inc{}'.format(body + 1)) # Add argument of periastron prior self.sys_priors.append(priors.UniformPrior(0., 2. * np.pi)) self.labels.append('aop{}'.format(body + 1)) # Add position angle of nodes prior self.sys_priors.append(priors.UniformPrior(0., angle_upperlim)) self.labels.append('pan{}'.format(body + 1)) # Add epoch of periastron prior. self.sys_priors.append(priors.UniformPrior(0., 1.)) self.labels.append('tau{}'.format(body + 1)) # # Set priors on total mass and parallax # self.labels.append('plx') if plx_err > 0: self.sys_priors.append(priors.GaussianPrior(plx, plx_err)) else: self.sys_priors.append(plx) if self.fit_secondary_mass: for body in np.arange(num_secondary_bodies): self.sys_priors.append(priors.LogUniformPrior( 1e-6, 1)) # in Solar masses for now self.labels.append('m{}'.format(body)) self.labels.append('m0') else: self.labels.append('mtot') # still need to append m0/mtot, even though labels are appended above if mass_err > 0: self.sys_priors.append(priors.GaussianPrior( stellar_mass, mass_err)) else: self.sys_priors.append(stellar_mass) # add labels dictionary for parameter indexing self.param_idx = dict(zip(self.labels, np.arange(len(self.labels))))