def test_orbit_e99(): """ Test a highly eccentric orbit (ecc=0.99). Again validate against James Graham's orbit code """ # sma, ecc, inc, argp, lan, tau, plx, mtot orbital_params = np.array([10, 0.99, 3, 0.5, 1.5, 0.3, 50, 1.5]) epochs = np.array([1000, 1101.4]) raoffs, deoffs, vzs = kepler.calc_orbit(epochs, orbital_params[0], orbital_params[1], orbital_params[2], orbital_params[3], orbital_params[4], orbital_params[5], orbital_params[6], orbital_params[7], tau_ref_epoch=0) true_raoff = [-589.45575, -571.48432] true_deoff = [-447.32217, -437.68456] true_vz = [.39208876, .42041953] for meas, truth in zip(raoffs, true_raoff): assert truth == pytest.approx(meas, abs=threshold) for meas, truth in zip(deoffs, true_deoff): assert truth == pytest.approx(meas, abs=threshold) for meas, truth in zip(vzs, true_vz): assert truth == pytest.approx(meas, abs=1e-8)
def test_orbit_with_mass(): """ Test a orbit where we specify the mass of the body too. This will change the radial velocity, which normally assumes the body is a test particle We will test two equal mass bodies, which will reduce the RV signal by 2, compared to the RV signal of a massless particle in a system with the same total mass. """ # sma, ecc, inc, argp, lan, tau, plx, mtot orbital_params = np.array([10, 0.99, 3, 0.5, 1.5, 0.3, 50, 1.5]) epochs = np.array([1000, 1101.4]) raoffs, deoffs, vzs = kepler.calc_orbit(epochs, orbital_params[0], orbital_params[1], orbital_params[2], orbital_params[3], orbital_params[4], orbital_params[5], orbital_params[6], orbital_params[7], mass_for_Kamp=orbital_params[7] / 2) true_raoff = [-589.45575, -571.48432] true_deoff = [-447.32217, -437.68456] true_vz = [.39208876 / 2, .42041953 / 2] for meas, truth in zip(raoffs, true_raoff): assert truth == pytest.approx(meas, abs=threshold) for meas, truth in zip(deoffs, true_deoff): assert truth == pytest.approx(meas, abs=threshold) for meas, truth in zip(vzs, true_vz): assert truth == pytest.approx(meas, abs=1e-8)
def test_orbit_scalar(): """ Test orbitize.kepler.calc_orbit() with scalar values """ sma = 10 ecc = 0.3 inc = 3 argp = 0.5 lan = 1.5 tau = 0.3 plx = 50 mtot = 1.5 epochs = 1000 raoffs, deoffs, vzs = kepler.calc_orbit(epochs, sma, ecc, inc, argp, lan, tau, plx, mtot, tau_ref_epoch=0) true_raoff = 152.86786 true_deoff = -462.91038 true_vz = .86448656 assert true_raoff == pytest.approx(raoffs, abs=threshold) assert true_deoff == pytest.approx(deoffs, abs=threshold) assert true_vz == pytest.approx(vzs, abs=1e-8)
def test_orbit_e03(): """ Test orbitize.kepler.calc_orbit() by comparing this code to the output of James Graham's code which has been used in many published papers. Note that orbitize currently uses Rob De Rosa's eccentricity solver. Pretty standard orbit with ecc = 0.3 """ # sma, ecc, inc, argp, lan, tau, plx, mtot orbital_params = np.array([10, 0.3, 3, 0.5, 1.5, 0.3, 50, 1.5]) epochs = np.array([1000, 1101.4]) raoffs, deoffs, vzs = kepler.calc_orbit( epochs, orbital_params[0], orbital_params[1], orbital_params[2], orbital_params[3], orbital_params[4], orbital_params[5], orbital_params[6], orbital_params[7]) true_raoff = [152.86786, 180.39408] #mas true_deoff = [-462.91038, -442.0127] true_vz = [.86448656, .97591289] for meas, truth in zip(raoffs, true_raoff): assert truth == pytest.approx(meas, abs=threshold) for meas, truth in zip(deoffs, true_deoff): assert truth == pytest.approx(meas, abs=threshold) for meas, truth in zip(vzs, true_vz): assert truth == pytest.approx(meas, abs=1e-8)
def test_orbit_e03_array(): """ Test orbitize.kepler.calc_orbit() with a standard orbit with ecc = 0.3 and an array of keplerian input """ # sma, ecc, inc, argp, lan, tau, plx, mtot sma = np.array([10, 10, 10]) ecc = np.array([0.3, 0.3, 0.3]) inc = np.array([3, 3, 3]) argp = np.array([0.5, 0.5, 0.5]) lan = np.array([1.5, 1.5, 1.5]) tau = np.array([0.3, 0.3, 0.3]) plx = np.array([50, 50, 50]) mtot = np.array([1.5, 1.5, 1.5]) epochs = np.array([1000, 1101.4]) raoffs, deoffs, vzs = kepler.calc_orbit(epochs, sma, ecc, inc, argp, lan, tau, plx, mtot) true_raoff = np.array([[152.86786, 152.86786, 152.86786], [180.39408, 180.39408, 180.39408]]) true_deoff = np.array([[-462.91038, -462.91038, -462.91038], [-442.0127, -442.0127, -442.0127]]) true_vz = np.array([[.86448656, .86448656, .86448656], [.97591289, .97591289, .97591289]]) for ii in range(0, 3): for meas, truth in zip(raoffs[:, ii], true_raoff[:, ii]): assert truth == pytest.approx(meas, abs=threshold) for meas, truth in zip(deoffs[:, ii], true_deoff[:, ii]): assert truth == pytest.approx(meas, abs=threshold) for meas, truth in zip(vzs[:, ii], true_vz[:, ii]): assert truth == pytest.approx(meas, abs=1e-8)
def compute_model(self, params_arr): """ Compute model predictions for an array of fitting parameters. Args: params_arr (np.array of float): RxM array of fitting parameters, where R is the number of parameters being fit, and M is the number of orbits we need model predictions for. Must be in the same order documented in ``System()`` above. If M=1, this can be a 1d array. Returns: np.array of float: Nobsx2xM array model predictions. If M=1, this is a 2d array, otherwise it is a 3d array. """ if len(params_arr.shape) == 1: model = np.zeros((len(self.data_table), 2)) else: model = np.zeros((len(self.data_table), 2, params_arr.shape[1])) for body_num in np.arange(self.num_secondary_bodies)+1: epochs = self.data_table['epoch'][self.body_indices[body_num]] sma = params_arr[body_num-1] ecc = params_arr[body_num] inc = params_arr[body_num+1] argp = params_arr[body_num+2] lan = params_arr[body_num+3] tau = params_arr[body_num+4] plx = params_arr[6*self.num_secondary_bodies] if self.fit_secondary_mass: # mass of secondary bodies are in order from -1-num_bodies until -2 in order. mass = params_arr[-1-self.num_secondary_bodies+(body_num-1)] else: mass = None mtot = params_arr[-1] raoff, decoff, vz = kepler.calc_orbit( epochs, sma, ecc, inc, argp, lan, tau, plx, mtot, mass=mass, tau_ref_epoch=self.tau_ref_epoch ) if len(raoff[self.radec[body_num]]) > 0: # (prevent empty array dimension errors) model[self.radec[body_num], 0] = raoff[self.radec[body_num]] model[self.radec[body_num], 1] = decoff[self.radec[body_num]] if len(raoff[self.seppa[body_num]]) > 0: sep, pa = radec2seppa( raoff[self.seppa[body_num]], decoff[self.seppa[body_num]] ) model[self.seppa[body_num], 0] = sep model[self.seppa[body_num], 1] = pa return model
def gen_data(): ''' Simulates radial velocity and astrometric data for a test system. Returns: (rvs,rv_epochs): Tuple of generated radial velocity measurements (rvs) and their corresponding measurement epochs (rv_epochs) (ras,decs,astro_epochs): Tuple containing simulated astrometric measurements (ras, decs) and the corresponding measurement epochs (astro_epochs) ''' #set parameters for the synthetic data sma=1 inc=np.pi/2 ecc=0.2 aop=np.pi/4 pan=np.pi/4 tau=0.4 plx=50 mass_for_kamp=0.1 mtot=1.1 #epochs and errors for rv rv_epochs=np.linspace(51544,52426,200) #epochs and errors for astrometry astro_epochs=np.linspace(51500,52500,10) astro_err=0 #generate rv trend rvset=kepler.calc_orbit(rv_epochs,sma,ecc,inc,aop,pan,tau,plx,mtot,mass_for_Kamp=mass_for_kamp) rvs=rvset[2] #generate predictions for astrometric epochs astro_set=kepler.calc_orbit(astro_epochs,sma,ecc,inc,aop,pan,tau,plx,mtot,mass_for_Kamp=mass_for_kamp) ras,decs=astro_set[0],astro_set[1] #return model generations return (rvs,rv_epochs),(ras,decs,astro_epochs)
def test_2planet_nomass(): """ Compare multiplanet to rebound for planets with mass. """ # generate a planet orbit mjup = u.Mjup.to(u.Msun) mass_b = 0 * mjup mass_c = 0 * mjup params = np.array([ 10, 0.1, np.radians(45), np.radians(45), np.radians(45), 0.5, 3, 0.1, np.radians(45), np.radians(190), np.radians(45), 0.2, 1, mass_b, mass_c, 1.5 - mass_b - mass_c ]) tau_ref_epoch = 0 epochs = np.linspace(0, 365.25 * 4, 100) + tau_ref_epoch # nearly the full period, MJD # doesn't matter that this is right, just needs to be the same size. below doesn't include effect of c # just want to generate some measurements of plaent b to test compute model b_ra_model, b_dec_model, b_vz_model = kepler.calc_orbit( epochs, params[0], params[1], params[2], params[3], params[4], params[5], params[-2], params[-1], tau_ref_epoch=tau_ref_epoch) # generate some fake measurements of planet b, just to feed into system.py to test bookkeeping t = table.Table([ epochs, np.ones(epochs.shape, dtype=int), b_ra_model, np.zeros(b_ra_model.shape), b_dec_model, np.zeros(b_dec_model.shape) ], names=[ "epoch", "object", "raoff", "raoff_err", "decoff", "decoff_err" ]) filename = os.path.join(orbitize.DATADIR, "rebound_2planet.csv") t.write(filename) # create the orbitize system and generate model predictions using the ground truth astrom_dat = read_input.read_file(filename) sys = system.System(2, astrom_dat, params[-1], params[-2], tau_ref_epoch=tau_ref_epoch, fit_secondary_mass=True) # generate measurement radec_orbitize, _ = sys.compute_model(params) b_ra_orb = radec_orbitize[:, 0] b_dec_orb = radec_orbitize[:, 1] # now project the orbit with rebound b_manom = basis.tau_to_manom(epochs[0], params[0], params[-1] + params[-3], params[5], tau_ref_epoch) c_manom = basis.tau_to_manom(epochs[0], params[0 + 6], params[-1] + params[-2], params[5 + 6], tau_ref_epoch) sim = rebound.Simulation() sim.units = ('yr', 'AU', 'Msun') # add star sim.add(m=params[-1]) # add two planets sim.add(m=mass_c, a=params[0 + 6], e=params[1 + 6], M=c_manom, omega=params[3 + 6], Omega=params[4 + 6] + np.pi / 2, inc=params[2 + 6]) sim.add(m=mass_b, a=params[0], e=params[1], M=b_manom, omega=params[3], Omega=params[4] + np.pi / 2, inc=params[2]) ps = sim.particles sim.move_to_com() # Use Wisdom Holman integrator (fast), with the timestep being < 5% of inner planet's orbital period sim.integrator = "ias15" sim.dt = ps[1].P / 1000. # integrate and measure star/planet separation b_ra_reb = [] b_dec_reb = [] for t in epochs: sim.integrate(t / 365.25) b_ra_reb.append(-(ps[2].x - ps[0].x)) # ra is negative x b_dec_reb.append(ps[2].y - ps[0].y) b_ra_reb = np.array(b_ra_reb) b_dec_reb = np.array(b_dec_reb) diff_ra = b_ra_reb - b_ra_orb / params[6 * 2] diff_dec = b_dec_reb - b_dec_orb / params[6 * 2] # should be as good as the one planet case assert np.all(np.abs(diff_ra) / (params[0] * params[6 * 2]) < 1e-9) assert np.all(np.abs(diff_dec) / (params[0] * params[6 * 2]) < 1e-9)
def compute_all_orbits(self, params_arr, epochs=None, comp_rebound=False): """ Calls orbitize.kepler.calc_orbit and optionally accounts for multi-body interactions. Also computes total quantities like RV (without jitter/gamma) Args: params_arr (np.array of float): RxM array of fitting parameters, where R is the number of parameters being fit, and M is the number of orbits we need model predictions for. Must be in the same order documented in ``System()`` above. If M=1, this can be a 1d array. epochs (np.array of float): epochs (in mjd) at which to compute orbit predictions. comp_rebound (bool, optional): A secondary optional input for use of N-body solver Rebound; by default, this will be set to false and a Kepler solver will be used instead. Returns: tuple: raoff (np.array of float): N_epochs x N_bodies x N_orbits array of RA offsets from barycenter at each epoch. decoff (np.array of float): N_epochs x N_bodies x N_orbits array of Dec offsets from barycenter at each epoch. vz (np.array of float): N_epochs x N_bodies x N_orbits array of radial velocities at each epoch. """ if epochs is None: epochs = self.data_table['epoch'] n_epochs = len(epochs) if len(params_arr.shape) == 1: n_orbits = 1 else: n_orbits = params_arr.shape[1] ra_kepler = np.zeros((n_epochs, self.num_secondary_bodies + 1, n_orbits)) # N_epochs x N_bodies x N_orbits dec_kepler = np.zeros( (n_epochs, self.num_secondary_bodies + 1, n_orbits)) ra_perturb = np.zeros( (n_epochs, self.num_secondary_bodies + 1, n_orbits)) dec_perturb = np.zeros( (n_epochs, self.num_secondary_bodies + 1, n_orbits)) vz = np.zeros((n_epochs, self.num_secondary_bodies + 1, n_orbits)) # mass/mtot used to compute each Keplerian orbit will be needed later to compute perturbations if self.track_planet_perturbs: masses = np.zeros((self.num_secondary_bodies + 1, n_orbits)) mtots = np.zeros((self.num_secondary_bodies + 1, n_orbits)) if comp_rebound or self.use_rebound: sma = params_arr[self.sma_indx] ecc = params_arr[self.ecc_indx] inc = params_arr[self.inc_indx] argp = params_arr[self.aop_indx] lan = params_arr[self.pan_indx] tau = params_arr[self.tau_indx] plx = params_arr[self.basis.standard_basis_idx['plx']] if self.fit_secondary_mass: m_pl = params_arr[self.mpl_idx] m0 = params_arr[self.basis.param_idx['m0']] mtot = m0 + sum(m_pl) else: m_pl = np.zeros(self.num_secondary_bodies) # if not fitting for secondary mass, then total mass must be stellar mass mtot = params_arr[self.basis.param_idx['mtot']] raoff, deoff, vz = nbody.calc_orbit( epochs, sma, ecc, inc, argp, lan, tau, plx, mtot, tau_ref_epoch=self.tau_ref_epoch, m_pl=m_pl, output_star=True) else: for body_num in np.arange(self.num_secondary_bodies) + 1: sma = params_arr[self.basis.standard_basis_idx['sma{}'.format( body_num)]] ecc = params_arr[self.basis.standard_basis_idx['ecc{}'.format( body_num)]] inc = params_arr[self.basis.standard_basis_idx['inc{}'.format( body_num)]] argp = params_arr[self.basis.standard_basis_idx['aop{}'.format( body_num)]] lan = params_arr[self.basis.standard_basis_idx['pan{}'.format( body_num)]] tau = params_arr[self.basis.standard_basis_idx['tau{}'.format( body_num)]] plx = params_arr[self.basis.standard_basis_idx['plx']] if self.fit_secondary_mass: # mass of secondary bodies are in order from -1-num_bodies until -2 in order. mass = params_arr[self.basis.standard_basis_idx[ 'm{}'.format(body_num)]] m0 = params_arr[self.basis.standard_basis_idx['m0']] # For what mtot to use to calculate central potential, we should use the mass enclosed in a sphere with r <= distance of planet. # We need to select all planets with sma < this planet. all_smas = params_arr[self.sma_indx] within_orbit = np.where(all_smas <= sma) outside_orbit = np.where(all_smas > sma) all_pl_masses = params_arr[self.secondary_mass_indx] inside_masses = all_pl_masses[within_orbit] mtot = np.sum(inside_masses) + m0 else: m_pl = np.zeros(self.num_secondary_bodies) # if not fitting for secondary mass, then total mass must be stellar mass mass = None m0 = None mtot = params_arr[self.basis.standard_basis_idx['mtot']] if self.track_planet_perturbs: masses[body_num] = mass mtots[body_num] = mtot # solve Kepler's equation raoff, decoff, vz_i = kepler.calc_orbit( epochs, sma, ecc, inc, argp, lan, tau, plx, mtot, mass_for_Kamp=m0, tau_ref_epoch=self.tau_ref_epoch) # raoff, decoff, vz are scalers if the length of epochs is 1 if len(epochs) == 1: raoff = np.array([raoff]) decoff = np.array([decoff]) vz_i = np.array([vz_i]) # add Keplerian ra/deoff for this body to storage arrays ra_kepler[:, body_num, :] = np.reshape(raoff, (n_epochs, n_orbits)) dec_kepler[:, body_num, :] = np.reshape(decoff, (n_epochs, n_orbits)) vz[:, body_num, :] = np.reshape(vz_i, (n_epochs, n_orbits)) # vz_i is the ith companion radial velocity if self.fit_secondary_mass: vz0 = np.reshape( vz_i * -(mass / m0), (n_epochs, n_orbits) ) # calculating stellar velocity due to ith companion vz[:, 0, :] += vz0 # adding stellar velocity and gamma # if we are fitting for the mass of the planets, then they will perturb the star # add the perturbation on the star due to this planet on the relative astrometry of the planet that was measured # We are superimposing the Keplerian orbits, so we can add it linearly, scaled by the mass. # Because we are in Jacobi coordinates, for companions, we only should model the effect of planets interior to it. # (Jacobi coordinates mean that separation for a given companion is measured relative to the barycenter of all interior companions) if self.track_planet_perturbs: for body_num in np.arange(self.num_secondary_bodies + 1): if body_num > 0: # for companions, only perturb companion orbits at larger SMAs than this one. sma = params_arr[self.basis.standard_basis_idx[ 'sma{}'.format(body_num)]] all_smas = params_arr[self.sma_indx] outside_orbit = np.where(all_smas > sma)[0] which_perturb_bodies = outside_orbit + 1 # the planet will also perturb the star which_perturb_bodies = np.append([0], which_perturb_bodies) else: # for the star, what we are measuring is its position relative to the system barycenter # so we want to account for all of the bodies. which_perturb_bodies = np.arange( self.num_secondary_bodies + 1) for other_body_num in which_perturb_bodies: # skip itself since the the 2-body problem is measuring the planet-star separation already if (body_num == other_body_num) | (body_num == 0): continue ## NOTE: we are only handling astrometry right now (TODO: integrate RV into this) # this computes the perturbation on the other body due to the current body # star is perturbed in opposite direction if other_body_num == 0: ra_perturb[:, other_body_num, :] -= ( masses[body_num] / mtots[body_num]) * ra_kepler[:, body_num, :] dec_perturb[:, other_body_num, :] -= ( masses[body_num] / mtots[body_num]) * dec_kepler[:, body_num, :] else: ra_perturb[:, other_body_num, :] += ( masses[body_num] / mtots[body_num]) * ra_kepler[:, body_num, :] dec_perturb[:, other_body_num, :] += ( masses[body_num] / mtots[body_num]) * dec_kepler[:, body_num, :] raoff = ra_kepler + ra_perturb deoff = dec_kepler + dec_perturb if self.fitting_basis == 'XYZ': # Find and filter out unbound orbits bad_orbits = np.where(np.logical_or(ecc >= 1., ecc < 0.))[0] if (bad_orbits.size != 0): raoff[:, :, bad_orbits] = np.inf deoff[:, :, bad_orbits] = np.inf vz[:, :, bad_orbits] = np.inf return raoff, deoff, vz else: return raoff, deoff, vz else: return raoff, deoff, vz
def compute_sep(df, epochs, basis, m0, m0_err, plx, plx_err, n_planets=1, pl_num=1): """ Computes a sky-projected angular separation posterior given a RadVel-computed DataFrame. Args: df (pd.DataFrame): Radvel-computed posterior (in any orbital basis) epochs (np.array of astropy.time.Time): epochs at which to compute separations basis (str): basis string of input posterior (see radvel.basis.BASIS_NAMES` for the full list of possibilities). m0 (float): median of primary mass distribution (assumed Gaussian). m0_err (float): 1sigma error of primary mass distribution (assumed Gaussian). plx (float): median of parallax distribution (assumed Gaussian). plx_err: 1sigma error of parallax distribution (assumed Gaussian). n_planets (int): total number of planets in RadVel posterior pl_num (int): planet number used in RadVel fits (e.g. a RadVel label of 'per1' implies `pl_num` == 1) Example: >> df = pandas.read_csv('sample_radvel_chains.csv.bz2', index_col=0) >> epochs = astropy.time.Time([2022, 2024], format='decimalyear') >> seps, df_orb = compute_sep( df, epochs, 'per tc secosw sesinw k', 0.82, 0.02, 312.22, 0.47 ) Returns: tuple of: np.array of size (len(epochs) x len(df)): sky-projected angular separations [mas] at each input epoch pd.DataFrame: corresponding orbital posterior in orbitize basis """ myBasis = Basis(basis, n_planets) df = myBasis.to_synth(df) chain_len = len(df) tau_ref_epoch = 58849 # convert RadVel posteriors -> orbitize posteriors m_st = np.random.normal(m0, m0_err, size=chain_len) semiamp = df['k{}'.format(pl_num)].values per_day = df['per{}'.format(pl_num)].values period_yr = per_day / 365.25 ecc = df['e{}'.format(pl_num)].values msini = (Msini(semiamp, per_day, m_st, ecc, Msini_units='Earth') * (u.M_earth / u.M_sun).to('')) cosi = (2. * np.random.random(size=chain_len)) - 1. inc = np.arccos(cosi) m_pl = msini / np.sin(inc) mtot = m_st + m_pl sma = (period_yr**2 * mtot)**(1 / 3) omega_st_rad = df['w{}'.format(pl_num)].values omega_pl_rad = omega_st_rad + np.pi parallax = np.random.normal(plx, plx_err, size=chain_len) lan = np.random.random_sample(size=chain_len) * 2. * np.pi tp_mjd = df['tp{}'.format(pl_num)].values - 2400000.5 tau = tp_to_tau(tp_mjd, tau_ref_epoch, period_yr) # compute projected separation in mas raoff, deoff, _ = calc_orbit(epochs.mjd, sma, ecc, inc, omega_pl_rad, lan, tau, parallax, mtot, tau_ref_epoch=tau_ref_epoch) seps = np.sqrt(raoff**2 + deoff**2) df_orb = pd.DataFrame(np.transpose( [sma, ecc, inc, omega_pl_rad, lan, tau, parallax, m_st, m_pl]), columns=[ 'sma', 'ecc', 'inc_rad', 'omega_pl_rad', 'lan_rad', 'tau_58849', 'plx', 'm_st', 'mp' ]) return seps, df_orb
def plot_orbits(self, object_to_plot=1, start_mjd=51544., num_orbits_to_plot=100, num_epochs_to_plot=100, square_plot=True, show_colorbar=True, cmap=cmap, sep_pa_color='lightgrey', sep_pa_end_year=2025.0, cbar_param='epochs', mod180=False): """ Plots one orbital period for a select number of fitted orbits for a given object, with line segments colored according to time Args: object_to_plot (int): which object to plot (default: 1) start_mjd (float): MJD in which to start plotting orbits (default: 51544, the year 2000) num_orbits_to_plot (int): number of orbits to plot (default: 100) num_epochs_to_plot (int): number of points to plot per orbit (default: 100) square_plot (Boolean): Aspect ratio is always equal, but if square_plot is True (default), then the axes will be square, otherwise, white space padding is used show_colorbar (Boolean): Displays colorbar to the right of the plot [True] cmap (matplotlib.cm.ColorMap): color map to use for making orbit tracks (default: modified Purples_r) sep_pa_color (string): any valid matplotlib color string, used to set the color of the orbit tracks in the Sep/PA panels (default: 'lightgrey'). sep_pa_end_year (float): decimal year specifying when to stop plotting orbit tracks in the Sep/PA panels (default: 2025.0). cbar_param (string): options are the following: epochs, sma1, ecc1, inc1, aop1, pan1, tau1. Number can be switched out. Default is epochs. mod180 (Bool): if True, PA will be plotted in range [180, 540]. Useful for plotting short arcs with PAs that cross 360 deg during observations (default: False) Return: ``matplotlib.pyplot.Figure``: the orbit plot if input is valid, ``None`` otherwise (written): Henry Ngo, Sarah Blunt, 2018 Additions by Malena Rice, 2019 """ if Time(start_mjd, format='mjd').decimalyear >= sep_pa_end_year: raise Exception('start_mjd keyword date must be less than sep_pa_end_year keyword date.') with warnings.catch_warnings(): warnings.simplefilter('ignore', ErfaWarning) dict_of_indices = { 'sma': 0, 'ecc': 1, 'inc': 2, 'aop': 3, 'pan': 4, 'tau': 5, 'plx':6 } if cbar_param == 'epochs': pass elif cbar_param[0:3] in dict_of_indices: try: object_id = np.int(cbar_param[3:]) except ValueError: object_id = 1 index = dict_of_indices[cbar_param[0:3]] + 6*(object_id-1) else: raise Exception('Invalid input; acceptable inputs include epochs, sma1, ecc1, inc1, aop1, pan1, tau1, sma2, ecc2, ...') # Split the 2-D post array into series of 1-D arrays for each orbital parameter num_objects, remainder = np.divmod(self.post.shape[1],6) if object_to_plot > num_objects: return None sma = self.post[:,dict_of_indices['sma']] ecc = self.post[:,dict_of_indices['ecc']] inc = self.post[:,dict_of_indices['inc']] aop = self.post[:,dict_of_indices['aop']] pan = self.post[:,dict_of_indices['pan']] tau = self.post[:,dict_of_indices['tau']] plx = self.post[:,dict_of_indices['plx']] # Then, get the other parameters if 'mtot' in self.labels: mtot = self.post[:,-1] elif 'm0' in self.labels: m0 = self.post[:,-1] mplanet = self.post[:,-2] mtot = m0 + mplanet # Select random indices for plotted orbit if num_orbits_to_plot > len(sma): num_orbits_to_plot = len(sma) choose = np.random.randint(0, high=len(sma), size=num_orbits_to_plot) raoff = np.zeros((num_orbits_to_plot, num_epochs_to_plot)) deoff = np.zeros((num_orbits_to_plot, num_epochs_to_plot)) epochs = np.zeros((num_orbits_to_plot, num_epochs_to_plot)) # Loop through each orbit to plot and calcualte ra/dec offsets for all points in orbit # Need this loops since epochs[] vary for each orbit, unless we want to just plot the same time period for all orbits for i in np.arange(num_orbits_to_plot): orb_ind = choose[i] # Compute period (from Kepler's third law) period = np.sqrt(4*np.pi**2.0*(sma*u.AU)**3/(consts.G*(mtot*u.Msun))) period = period.to(u.day).value # Create an epochs array to plot num_epochs_to_plot points over one orbital period epochs[i,:] = np.linspace(start_mjd, float(start_mjd+period[orb_ind]), num_epochs_to_plot) # Calculate ra/dec offsets for all epochs of this orbit raoff0, deoff0, _ = kepler.calc_orbit( epochs[i,:], sma[orb_ind], ecc[orb_ind], inc[orb_ind], aop[orb_ind], pan[orb_ind], tau[orb_ind], plx[orb_ind], mtot[orb_ind], tau_ref_epoch=self.tau_ref_epoch ) raoff[i,:] = raoff0 deoff[i,:] = deoff0 # Create a linearly increasing colormap for our range of epochs if cbar_param != 'epochs': cbar_param_arr = self.post[:,index] norm = mpl.colors.Normalize(vmin=np.min(cbar_param_arr), vmax=np.max(cbar_param_arr)) norm_yr = mpl.colors.Normalize(vmin=np.min(cbar_param_arr), vmax=np.max(cbar_param_arr)) elif cbar_param == 'epochs': norm = mpl.colors.Normalize(vmin=np.min(epochs), vmax=np.max(epochs[-1,:])) norm_yr = mpl.colors.Normalize( vmin=np.min(Time(epochs,format='mjd').decimalyear), vmax=np.max(Time(epochs,format='mjd').decimalyear) ) # Create figure for orbit plots fig = plt.figure(figsize=(14,6)) ax = plt.subplot2grid((2, 14), (0, 0), rowspan=2, colspan=6) # Plot each orbit (each segment between two points coloured using colormap) for i in np.arange(num_orbits_to_plot): points = np.array([raoff[i,:], deoff[i,:]]).T.reshape(-1,1,2) segments = np.concatenate([points[:-1], points[1:]], axis=1) lc = LineCollection( segments, cmap=cmap, norm=norm, linewidth=1.0 ) if cbar_param != 'epochs': lc.set_array(np.ones(len(epochs[0]))*cbar_param_arr[i]) elif cbar_param == 'epochs': lc.set_array(epochs[i,:]) ax.add_collection(lc) # modify the axes if square_plot: adjustable_param='datalim' else: adjustable_param='box' ax.set_aspect('equal', adjustable=adjustable_param) ax.set_xlabel('$\\Delta$RA [mas]') ax.set_ylabel('$\\Delta$Dec [mas]') ax.locator_params(axis='x', nbins=6) ax.locator_params(axis='y', nbins=6) ax.invert_xaxis() # To go to a left-handed coordinate system # add colorbar if show_colorbar: cbar_ax = fig.add_axes([0.47, 0.15, 0.015, 0.7]) # xpos, ypos, width, height, in fraction of figure size cbar = mpl.colorbar.ColorbarBase(cbar_ax, cmap=cmap, norm=norm_yr, orientation='vertical', label=cbar_param) # plot sep/PA zoom-in panels ax1 = plt.subplot2grid((2, 14), (0, 9), colspan=6) ax2 = plt.subplot2grid((2, 14), (1, 9), colspan=6) ax2.set_ylabel('PA [$^{{\\circ}}$]') ax1.set_ylabel('$\\rho$ [mas]') ax2.set_xlabel('Epoch') epochs_seppa = np.zeros((num_orbits_to_plot, num_epochs_to_plot)) for i in np.arange(num_orbits_to_plot): orb_ind = choose[i] epochs_seppa[i,:] = np.linspace( start_mjd, Time(sep_pa_end_year, format='decimalyear').mjd, num_epochs_to_plot ) # Calculate ra/dec offsets for all epochs of this orbit raoff0, deoff0, _ = kepler.calc_orbit( epochs_seppa[i,:], sma[orb_ind], ecc[orb_ind], inc[orb_ind], aop[orb_ind], pan[orb_ind], tau[orb_ind], plx[orb_ind], mtot[orb_ind], tau_ref_epoch=self.tau_ref_epoch ) raoff[i,:] = raoff0 deoff[i,:] = deoff0 yr_epochs = Time(epochs_seppa[i,:],format='mjd').decimalyear plot_epochs = np.where(yr_epochs <= sep_pa_end_year)[0] yr_epochs = yr_epochs[plot_epochs] seps, pas = orbitize.system.radec2seppa(raoff[i,:], deoff[i,:], mod180=mod180) plt.sca(ax1) plt.plot(yr_epochs, seps, color=sep_pa_color) plt.sca(ax2) plt.plot(yr_epochs, pas, color=sep_pa_color) ax1.locator_params(axis='x', nbins=6) ax1.locator_params(axis='y', nbins=6) ax2.locator_params(axis='x', nbins=6) ax2.locator_params(axis='y', nbins=6) return fig
def compute_model(self, params_arr): """ Compute model predictions for an array of fitting parameters. Args: params_arr (np.array of float): RxM array of fitting parameters, where R is the number of parameters being fit, and M is the number of orbits we need model predictions for. Must be in the same order documented in ``System()`` above. If M=1, this can be a 1d array. Returns: np.array of float: Nobsx2xM array model predictions. If M=1, this is a 2d array, otherwise it is a 3d array. """ if len(params_arr.shape) == 1: model = np.zeros((len(self.data_table), 2)) # Rob and Lea: adding gamma zeros gamma = np.zeros((len(self.data_table), 2)) jitter = np.zeros((len(self.data_table), 2)) else: model = np.zeros((len(self.data_table), 2, params_arr.shape[1])) gamma = np.zeros((len(self.data_table), 2, params_arr.shape[1])) jitter = np.zeros((len(self.data_table), 2, params_arr.shape[1])) if self.track_planet_perturbs: radec_perturb = np.zeros(model.shape) if len(self.rv[0]) > 0 and self.fit_secondary_mass: # looping through instruments to get the gammas for rv_idx in range(len(self.rv_instruments)): gamma[self.rv_inst_indices[rv_idx], 0] = params_arr[6 * self.num_secondary_bodies + 1 + 2 * rv_idx] # km/s jitter[self.rv_inst_indices[rv_idx], 0] = params_arr[6 * self.num_secondary_bodies + 2 + 2 * rv_idx] gamma[self.rv_inst_indices[rv_idx], 1] = np.nan jitter[self.rv_inst_indices[rv_idx], 1] = np.nan # need to put planetary rv later total_rv0 = 0 # If we're not fitting rv, then we don't regard the total rv and will not use this for body_num in np.arange(self.num_secondary_bodies) + 1: # we're going to compute at all epochs for convenience of indexing right now # self.radec, and self.seppa index into the entire data table, not just the values for a particular body epochs = self.data_table['epoch'] #[self.body_indices[body_num]] startindex = 6 * (body_num - 1) sma = params_arr[startindex] ecc = params_arr[startindex + 1] inc = params_arr[startindex + 2] argp = params_arr[startindex + 3] lan = params_arr[startindex + 4] tau = params_arr[startindex + 5] plx = params_arr[6 * self.num_secondary_bodies] if self.fit_secondary_mass: # mass of secondary bodies are in order from -1-num_bodies until -2 in order. mass = params_arr[-1 - self.num_secondary_bodies + (body_num - 1)] m0 = params_arr[-1] # For what mtot to use to calculate central potential, we should use the mass enclosed in a sphere with r <= distance of planet. # We need to select all planets with sma < this planet. all_smas = params_arr[0:6 * self.num_secondary_bodies:6] within_orbit = np.where(all_smas <= sma) outside_orbit = np.where(all_smas > sma) all_pl_masses = params_arr[-1 - self.num_secondary_bodies:-1] inside_masses = all_pl_masses[within_orbit] mtot = np.sum(inside_masses) + m0 else: # if not fitting for secondary mass, then total mass must be stellar mass mass = None m0 = None mtot = params_arr[-1] # i = 1,2,3... (companion index) raoff, decoff, vz_i = kepler.calc_orbit( epochs, sma, ecc, inc, argp, lan, tau, plx, mtot, mass_for_Kamp=m0, tau_ref_epoch=self.tau_ref_epoch, tau_warning=False) # raoff, decoff, vz are scalers if the length of epochs is 1. # Jason is too lazy to figure out how to make it return a one element array without breaking everything else # so hard code it here to convert them into 1-element numpy arrays. if len(epochs) == 1: raoff = np.array([raoff]) decoff = np.array([decoff]) vz = np.array([vz_i]) # vz_i is the ith companion radial velocity if len(self.rv[0]) > 0 and self.fit_secondary_mass: vz0 = vz_i * -( mass / m0 ) # calculating stellar velocity due to ith companion total_rv0 = total_rv0 + vz0 # Adding stellar velocity and gamma # for the model points that correspond to this planet's orbit, add the model prediction # RA/Dec if len(self.radec[body_num] ) > 0: # (prevent empty array dimension errors) model[self.radec[body_num], 0] = raoff[self.radec[body_num]] model[self.radec[body_num], 1] = decoff[self.radec[body_num]] # Sep/PA if len(self.seppa[body_num]) > 0: sep, pa = radec2seppa(raoff, decoff) model[self.seppa[body_num], 0] = sep[self.seppa[body_num]] model[self.seppa[body_num], 1] = pa[self.seppa[body_num]] # RV if len(self.rv[body_num]) > 0: model[self.rv[body_num], 0] = vz_i[self.rv[body_num]] model[self.rv[body_num], 1] = np.nan # for the other epochs, if we are fitting for the mass of the planets, then they will perturb the star # add the perturbation on the star due to this planet on the relative astrometry of the planet that was measured # We are superimposing the Keplerian orbits, so we can add it linearly, scaled by the mass. # Because we are in Jacobi coordinates, for companions, we only should model the effect of planets interior to it. # (Jacobi coordinates mean that separation for a given companion is measured relative to the barycenter of all interior companions) if self.track_planet_perturbs: if body_num > 0: # for companions, only perturb companion orbits at larger SMAs than this one. # note the +1, since the 0th planet is body_num 1. which_perturb_bodies = outside_orbit[0] + 1 else: # for the star, what we are measuring is it's position relative to the system barycenter # so we want to account for all of the bodies. which_perturb_bodies = range(self.num_secondary_bodies + 1) for other_body_num in which_perturb_bodies: # skip itself since the the 2-body problem is measuring the planet-star separation already if (body_num == other_body_num) | (body_num == 0): continue ## NOTE: we are only handling ra/dec and sep/pa right now ## TODO: integrate RV into this if len(self.radec[other_body_num]) > 0: radec_perturb[self.radec[other_body_num], 0] += -(mass / mtot) * raoff[ self.radec[other_body_num]] radec_perturb[self.radec[other_body_num], 1] += -(mass / mtot) * decoff[ self.radec[other_body_num]] if len(self.seppa[other_body_num]) > 0: radec_perturb[self.seppa[other_body_num], 0] += -(mass / mtot) * raoff[ self.seppa[other_body_num]] radec_perturb[self.seppa[other_body_num], 1] += -(mass / mtot) * decoff[ self.seppa[other_body_num]] if len(self.rv[0]) > 0 and self.fit_secondary_mass: if len(total_rv0[self.rv[0]]) > 0: model[self.rv[0], 0] = total_rv0[self.rv[0]] model[self.rv[0], 1] = np.nan # nans only for rv indices # add the effects of other planets on the measured astrometry if self.track_planet_perturbs: for body_num in range(self.num_secondary_bodies + 1): if len(self.radec[body_num]) > 0: model[self.radec[body_num]] -= radec_perturb[ self.radec[body_num]] if len(self.seppa[body_num]) > 0: # for seppa, add the perturbations in radec space and convert back ra_unperturb, dec_unperturb = seppa2radec( model[self.seppa[body_num], 0], model[self.seppa[body_num], 1]) ra_perturb = ra_unperturb - radec_perturb[ self.seppa[body_num], 0] dec_perturb = dec_unperturb - radec_perturb[ self.seppa[body_num], 1] sep_perturb, pa_perturb = radec2seppa( ra_perturb, dec_perturb) model[self.seppa[body_num], 0] = sep_perturb model[self.seppa[body_num], 1] = pa_perturb if self.fit_secondary_mass: return model + gamma, jitter else: return model, jitter
def plot_orbits(results, object_to_plot=1, start_mjd=51544., num_orbits_to_plot=100, num_epochs_to_plot=100, square_plot=True, show_colorbar=True, cmap=cmap, sep_pa_color='lightgrey', sep_pa_end_year=2025.0, cbar_param='Epoch [year]', mod180=False, rv_time_series=False, plot_astrometry=True, plot_astrometry_insts=False, fig=None): """ Plots one orbital period for a select number of fitted orbits for a given object, with line segments colored according to time Args: object_to_plot (int): which object to plot (default: 1) start_mjd (float): MJD in which to start plotting orbits (default: 51544, the year 2000) num_orbits_to_plot (int): number of orbits to plot (default: 100) num_epochs_to_plot (int): number of points to plot per orbit (default: 100) square_plot (Boolean): Aspect ratio is always equal, but if square_plot is True (default), then the axes will be square, otherwise, white space padding is used show_colorbar (Boolean): Displays colorbar to the right of the plot [True] cmap (matplotlib.cm.ColorMap): color map to use for making orbit tracks (default: modified Purples_r) sep_pa_color (string): any valid matplotlib color string, used to set the color of the orbit tracks in the Sep/PA panels (default: 'lightgrey'). sep_pa_end_year (float): decimal year specifying when to stop plotting orbit tracks in the Sep/PA panels (default: 2025.0). cbar_param (string): options are the following: 'Epoch [year]', 'sma1', 'ecc1', 'inc1', 'aop1', 'pan1', 'tau1', 'plx. Number can be switched out. Default is Epoch [year]. mod180 (Bool): if True, PA will be plotted in range [180, 540]. Useful for plotting short arcs with PAs that cross 360 deg during observations (default: False) rv_time_series (Boolean): if fitting for secondary mass using MCMC for rv fitting and want to display time series, set to True. plot_astrometry (Boolean): set to True by default. Plots the astrometric data. plot_astrometry_insts (Boolean): set to False by default. Plots the astrometric data by instruments. fig (matplotlib.pyplot.Figure): optionally include a predefined Figure object to plot the orbit on. Most users will not need this keyword. Return: ``matplotlib.pyplot.Figure``: the orbit plot if input is valid, ``None`` otherwise (written): Henry Ngo, Sarah Blunt, 2018 Additions by Malena Rice, 2019 """ if Time(start_mjd, format='mjd').decimalyear >= sep_pa_end_year: raise ValueError('start_mjd keyword date must be less than sep_pa_end_year keyword date.') if object_to_plot > results.num_secondary_bodies: raise ValueError("Only {0} secondary bodies being fit. Requested to plot body {1} which is out of range".format(results.num_secondary_bodies, object_to_plot)) if object_to_plot == 0: raise ValueError("Plotting the primary's orbit is currently unsupported. Stay tuned.") with warnings.catch_warnings(): warnings.simplefilter('ignore', ErfaWarning) data = results.data[results.data['object'] == object_to_plot] possible_cbar_params = [ 'sma', 'ecc', 'inc', 'aop' 'pan', 'tau', 'plx' ] if cbar_param == 'Epoch [year]': pass elif cbar_param[0:3] in possible_cbar_params: index = results.param_idx[cbar_param] else: raise Exception( "Invalid input; acceptable inputs include 'Epoch [year]', 'plx', 'sma1', 'ecc1', 'inc1', 'aop1', 'pan1', 'tau1', 'sma2', 'ecc2', ...)" ) # Select random indices for plotted orbit num_orbits = len(results.post[:, 0]) if num_orbits_to_plot > num_orbits: num_orbits_to_plot = num_orbits choose = np.random.randint(0, high=num_orbits, size=num_orbits_to_plot) # Get posteriors from random indices standard_post = [] if results.sampler_name == 'MCMC': # Convert the randomly chosen posteriors to standard keplerian set for i in np.arange(num_orbits_to_plot): orb_ind = choose[i] param_set = np.copy(results.post[orb_ind]) standard_post.append(results.basis.to_standard_basis(param_set)) else: # For OFTI, posteriors are already converted for i in np.arange(num_orbits_to_plot): orb_ind = choose[i] standard_post.append(results.post[orb_ind]) standard_post = np.array(standard_post) sma = standard_post[:, results.standard_param_idx['sma{}'.format(object_to_plot)]] ecc = standard_post[:, results.standard_param_idx['ecc{}'.format(object_to_plot)]] inc = standard_post[:, results.standard_param_idx['inc{}'.format(object_to_plot)]] aop = standard_post[:, results.standard_param_idx['aop{}'.format(object_to_plot)]] pan = standard_post[:, results.standard_param_idx['pan{}'.format(object_to_plot)]] tau = standard_post[:, results.standard_param_idx['tau{}'.format(object_to_plot)]] plx = standard_post[:, results.standard_param_idx['plx']] # Then, get the other parameters if 'mtot' in results.labels: mtot = standard_post[:, results.standard_param_idx['mtot']] elif 'm0' in results.labels: m0 = standard_post[:, results.standard_param_idx['m0']] m1 = standard_post[:, results.standard_param_idx['m{}'.format(object_to_plot)]] mtot = m0 + m1 raoff = np.zeros((num_orbits_to_plot, num_epochs_to_plot)) deoff = np.zeros((num_orbits_to_plot, num_epochs_to_plot)) vz_star = np.zeros((num_orbits_to_plot, num_epochs_to_plot)) epochs = np.zeros((num_orbits_to_plot, num_epochs_to_plot)) # Loop through each orbit to plot and calcualte ra/dec offsets for all points in orbit # Need this loops since epochs[] vary for each orbit, unless we want to just plot the same time period for all orbits for i in np.arange(num_orbits_to_plot): # Compute period (from Kepler's third law) period = np.sqrt(4*np.pi**2.0*(sma*u.AU)**3/(consts.G*(mtot*u.Msun))) period = period.to(u.day).value # Create an epochs array to plot num_epochs_to_plot points over one orbital period epochs[i, :] = np.linspace(start_mjd, float( start_mjd+period[i]), num_epochs_to_plot) # Calculate ra/dec offsets for all epochs of this orbit raoff0, deoff0, _ = kepler.calc_orbit( epochs[i, :], sma[i], ecc[i], inc[i], aop[i], pan[i], tau[i], plx[i], mtot[i], tau_ref_epoch=results.tau_ref_epoch ) raoff[i, :] = raoff0 deoff[i, :] = deoff0 # Create a linearly increasing colormap for our range of epochs if cbar_param != 'Epoch [year]': cbar_param_arr = results.post[:, index] norm = mpl.colors.Normalize(vmin=np.min(cbar_param_arr), vmax=np.max(cbar_param_arr)) norm_yr = mpl.colors.Normalize(vmin=np.min( cbar_param_arr), vmax=np.max(cbar_param_arr)) elif cbar_param == 'Epoch [year]': min_cbar_date = np.min(epochs) max_cbar_date = np.max(epochs[-1, :]) # if we're plotting orbital periods greater than 1,000 yrs, limit the colorbar dynamic range if max_cbar_date - min_cbar_date > 1000 * 365.25: max_cbar_date = min_cbar_date + 1000 * 365.25 norm = mpl.colors.Normalize(vmin=min_cbar_date, vmax=max_cbar_date) norm_yr = mpl.colors.Normalize( vmin=Time(min_cbar_date, format='mjd').decimalyear, vmax=Time(max_cbar_date, format='mjd').decimalyear ) # Before starting to plot rv data, make sure rv data exists: rv_indices = np.where(data['quant_type'] == 'rv') if rv_time_series and len(rv_indices) == 0: warnings.warn("Unable to plot radial velocity data.") rv_time_series = False # Create figure for orbit plots if fig is None: fig = plt.figure(figsize=(14, 6)) if rv_time_series: fig = plt.figure(figsize=(14, 9)) ax = plt.subplot2grid((3, 14), (0, 0), rowspan=2, colspan=6) else: fig = plt.figure(figsize=(14, 6)) ax = plt.subplot2grid((2, 14), (0, 0), rowspan=2, colspan=6) else: plt.set_current_figure(fig) if rv_time_series: ax = plt.subplot2grid((3, 14), (0, 0), rowspan=2, colspan=6) else: ax = plt.subplot2grid((2, 14), (0, 0), rowspan=2, colspan=6) astr_inds=np.where((~np.isnan(data['quant1'])) & (~np.isnan(data['quant2']))) astr_epochs=data['epoch'][astr_inds] radec_inds = np.where(data['quant_type'] == 'radec') seppa_inds = np.where(data['quant_type'] == 'seppa') sep_data, sep_err=data['quant1'][seppa_inds],data['quant1_err'][seppa_inds] pa_data, pa_err=data['quant2'][seppa_inds],data['quant2_err'][seppa_inds] if len(radec_inds[0] > 0): sep_from_ra_data, pa_from_dec_data = orbitize.system.radec2seppa( data['quant1'][radec_inds], data['quant2'][radec_inds] ) num_radec_pts = len(radec_inds[0]) sep_err_from_ra_data = np.empty(num_radec_pts) pa_err_from_dec_data = np.empty(num_radec_pts) for j in np.arange(num_radec_pts): sep_err_from_ra_data[j], pa_err_from_dec_data[j], _ = orbitize.system.transform_errors( np.array(data['quant1'][radec_inds][j]), np.array(data['quant2'][radec_inds][j]), np.array(data['quant1_err'][radec_inds][j]), np.array(data['quant2_err'][radec_inds][j]), np.array(data['quant12_corr'][radec_inds][j]), orbitize.system.radec2seppa ) sep_data = np.append(sep_data, sep_from_ra_data) sep_err = np.append(sep_err, sep_err_from_ra_data) pa_data = np.append(pa_data, pa_from_dec_data) pa_err = np.append(pa_err, pa_err_from_dec_data) # For plotting different astrometry instruments if plot_astrometry_insts: astr_colors = ('#FF7F11', '#11FFE3', '#14FF11', '#7A11FF', '#FF1919') astr_symbols = ('*', 'o', 'p', 's') ax_colors = itertools.cycle(astr_colors) ax_symbols = itertools.cycle(astr_symbols) astr_data = data[astr_inds] astr_insts = np.unique(data[astr_inds]['instrument']) # Indices corresponding to each instrument in datafile astr_inst_inds = {} for i in range(len(astr_insts)): astr_inst_inds[astr_insts[i]]=np.where(astr_data['instrument']==astr_insts[i].encode())[0] # Plot each orbit (each segment between two points coloured using colormap) for i in np.arange(num_orbits_to_plot): points = np.array([raoff[i, :], deoff[i, :]]).T.reshape(-1, 1, 2) segments = np.concatenate([points[:-1], points[1:]], axis=1) lc = LineCollection( segments, cmap=cmap, norm=norm, linewidth=1.0 ) if cbar_param != 'Epoch [year]': lc.set_array(np.ones(len(epochs[0]))*cbar_param_arr[i]) elif cbar_param == 'Epoch [year]': lc.set_array(epochs[i, :]) ax.add_collection(lc) if plot_astrometry: ra_data,dec_data=orbitize.system.seppa2radec(sep_data,pa_data) # Plot astrometry along with instruments if plot_astrometry_insts: for i in range(len(astr_insts)): ra = ra_data[astr_inst_inds[astr_insts[i]]] dec = dec_data[astr_inst_inds[astr_insts[i]]] ax.scatter(ra, dec, marker=next(ax_symbols), c=next(ax_colors), zorder=10, s=60, label=astr_insts[i]) else: ax.scatter(ra_data, dec_data, marker='*', c='#FF7F11', zorder=10, s=60) # modify the axes if square_plot: adjustable_param = 'datalim' else: adjustable_param = 'box' ax.set_aspect('equal', adjustable=adjustable_param) ax.set_xlabel('$\\Delta$RA [mas]') ax.set_ylabel('$\\Delta$Dec [mas]') ax.locator_params(axis='x', nbins=6) ax.locator_params(axis='y', nbins=6) ax.invert_xaxis() # To go to a left-handed coordinate system # plot sep/PA and/or rv zoom-in panels if rv_time_series: ax1 = plt.subplot2grid((3, 14), (0, 8), colspan=6) ax2 = plt.subplot2grid((3, 14), (1, 8), colspan=6) ax3 = plt.subplot2grid((3, 14), (2, 0), colspan=14, rowspan=1) ax2.set_ylabel('PA [$^{{\\circ}}$]') ax1.set_ylabel('$\\rho$ [mas]') ax3.set_ylabel('RV [km/s]') ax3.set_xlabel('Epoch') ax2.set_xlabel('Epoch') plt.subplots_adjust(hspace=0.3) else: ax1 = plt.subplot2grid((2, 14), (0, 9), colspan=6) ax2 = plt.subplot2grid((2, 14), (1, 9), colspan=6) ax2.set_ylabel('PA [$^{{\\circ}}$]') ax1.set_ylabel('$\\rho$ [mas]') ax2.set_xlabel('Epoch') if plot_astrometry_insts: ax1_colors = itertools.cycle(astr_colors) ax1_symbols = itertools.cycle(astr_symbols) ax2_colors = itertools.cycle(astr_colors) ax2_symbols = itertools.cycle(astr_symbols) epochs_seppa = np.zeros((num_orbits_to_plot, num_epochs_to_plot)) for i in np.arange(num_orbits_to_plot): epochs_seppa[i, :] = np.linspace( start_mjd, Time(sep_pa_end_year, format='decimalyear').mjd, num_epochs_to_plot ) # Calculate ra/dec offsets for all epochs of this orbit if rv_time_series: raoff0, deoff0, _ = kepler.calc_orbit( epochs_seppa[i, :], sma[i], ecc[i], inc[i], aop[i], pan[i], tau[i], plx[i], mtot[i], tau_ref_epoch=results.tau_ref_epoch, mass_for_Kamp=m0[i] ) raoff[i, :] = raoff0 deoff[i, :] = deoff0 else: raoff0, deoff0, _ = kepler.calc_orbit( epochs_seppa[i, :], sma[i], ecc[i], inc[i], aop[i], pan[i], tau[i], plx[i], mtot[i], tau_ref_epoch=results.tau_ref_epoch ) raoff[i, :] = raoff0 deoff[i, :] = deoff0 yr_epochs = Time(epochs_seppa[i, :], format='mjd').decimalyear seps, pas = orbitize.system.radec2seppa(raoff[i, :], deoff[i, :], mod180=mod180) plt.sca(ax1) plt.plot(yr_epochs, seps, color=sep_pa_color) plt.sca(ax2) plt.plot(yr_epochs, pas, color=sep_pa_color) # Plot sep/pa instruments if plot_astrometry_insts: for i in range(len(astr_insts)): sep = sep_data[astr_inst_inds[astr_insts[i]]] pa = pa_data[astr_inst_inds[astr_insts[i]]] epochs = astr_epochs[astr_inst_inds[astr_insts[i]]] plt.sca(ax1) plt.scatter(Time(epochs,format='mjd').decimalyear,sep,s=10,marker=next(ax1_symbols),c=next(ax1_colors),zorder=10,label=astr_insts[i]) plt.sca(ax2) plt.scatter(Time(epochs,format='mjd').decimalyear,pa,s=10,marker=next(ax2_symbols),c=next(ax2_colors),zorder=10) plt.sca(ax1) plt.legend(title='Instruments', bbox_to_anchor=(1.3, 1), loc='upper right') else: plt.sca(ax1) plt.scatter(Time(astr_epochs,format='mjd').decimalyear,sep_data,s=10,marker='*',c='purple',zorder=10) plt.sca(ax2) plt.scatter(Time(astr_epochs,format='mjd').decimalyear,pa_data,s=10,marker='*',c='purple',zorder=10) if rv_time_series: rv_data = results.data[results.data['object'] == 0] rv_data = rv_data[rv_data['quant_type'] == 'rv'] # switch current axis to rv panel plt.sca(ax3) # get list of rv instruments insts = np.unique(rv_data['instrument']) if len(insts) == 0: insts = ['defrv'] # get gamma/sigma labels and corresponding positions in the posterior gams=['gamma_'+inst for inst in insts] if isinstance(results.labels,list): labels=np.array(results.labels) else: labels=results.labels # get the indices corresponding to each gamma within results.labels gam_idx=[np.where(labels==inst_gamma)[0] for inst_gamma in gams] # indices corresponding to each instrument in the datafile inds={} for i in range(len(insts)): inds[insts[i]]=np.where(rv_data['instrument']==insts[i].encode())[0] # choose the orbit with the best log probability best_like=np.where(results.lnlike==np.amax(results.lnlike))[0][0] med_ga=[results.post[best_like,i] for i in gam_idx] # Get the posteriors for this index and convert to standard basis best_post = results.basis.to_standard_basis(results.post[best_like].copy()) # Get the masses for the best posteriors: best_m0 = best_post[results.standard_param_idx['m0']] best_m1 = best_post[results.standard_param_idx['m{}'.format(object_to_plot)]] best_mtot = best_m0 + best_m1 # colour/shape scheme scheme for rv data points clrs=('#0496FF','#372554','#FF1053','#3A7CA5','#143109') symbols=('o','^','v','s') ax3_colors = itertools.cycle(clrs) ax3_symbols = itertools.cycle(symbols) # get rvs and plot them for i,name in enumerate(inds.keys()): inst_data=rv_data[inds[name]] rvs=inst_data['quant1'] epochs=inst_data['epoch'] epochs=Time(epochs, format='mjd').decimalyear rvs-=med_ga[i] rvs -= best_post[results.param_idx[gams[i]]] plt.scatter(epochs,rvs,s=5,marker=next(ax3_symbols),c=next(ax3_colors),label=name,zorder=5) if len(inds.keys()) == 1 and 'defrv' in inds.keys(): pass else: plt.legend() # calculate the predicted rv trend using the best orbit _, _, vz = kepler.calc_orbit( epochs_seppa[0, :], best_post[results.standard_param_idx['sma{}'.format(object_to_plot)]], best_post[results.standard_param_idx['ecc{}'.format(object_to_plot)]], best_post[results.standard_param_idx['inc{}'.format(object_to_plot)]], best_post[results.standard_param_idx['aop{}'.format(object_to_plot)]], best_post[results.standard_param_idx['pan{}'.format(object_to_plot)]], best_post[results.standard_param_idx['tau{}'.format(object_to_plot)]], best_post[results.standard_param_idx['plx']], best_mtot, tau_ref_epoch=results.tau_ref_epoch, mass_for_Kamp=best_m0 ) vz=vz*-(best_m1)/np.median(best_m0) # plot rv trend plt.plot(Time(epochs_seppa[0, :],format='mjd').decimalyear, vz, color=sep_pa_color) # add colorbar if show_colorbar: if rv_time_series: # Create an axes for colorbar. The position of the axes is calculated based on the position of ax. # You can change x1.0.05 to adjust the distance between the main image and the colorbar. # You can change 0.02 to adjust the width of the colorbar. cbar_ax = fig.add_axes( [ax.get_position().x1+0.005, ax.get_position().y0, 0.02, ax.get_position().height]) cbar = mpl.colorbar.ColorbarBase( cbar_ax, cmap=cmap, norm=norm_yr, orientation='vertical', label=cbar_param) else: # xpos, ypos, width, height, in fraction of figure size cbar_ax = fig.add_axes([0.47, 0.15, 0.015, 0.7]) cbar = mpl.colorbar.ColorbarBase( cbar_ax, cmap=cmap, norm=norm_yr, orientation='vertical', label=cbar_param) ax1.locator_params(axis='x', nbins=6) ax1.locator_params(axis='y', nbins=6) ax2.locator_params(axis='x', nbins=6) ax2.locator_params(axis='y', nbins=6) return fig
""" A script to test some edge cases for the calc_orbit() function. We threw a bunch of weird inputs in and now everything is returning nans. Try to figure out which input or inputs is causing it to return nan. If you have time, update the code to do some error checking to alert the user the input is invalid. """ import numpy as np import orbitize.kepler as kepler ra, dec, rv = kepler.calc_orbit(np.array([1000, 1e6, -12]), 10.0, 0.1, -1, 1000, 0, 0.5, 100, -99) print(ra, dec, rv) assert np.all(np.isfinite(ra)) assert np.all(np.isfinite(dec)) assert np.all(np.isfinite(rv))
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 test_1planet(): """ Sanity check that things agree for 1 planet case """ # generate a planet orbit sma = 1 ecc = 0.1 inc = np.radians(45) aop = np.radians(45) pan = np.radians(45) tau = 0.5 plx = 1 mtot = 1 tau_ref_epoch = 0 mjup = u.Mjup.to(u.Msun) mass_b = 12 * mjup epochs = np.linspace(0, 300, 100) + tau_ref_epoch # nearly the full period, MJD ra_model, dec_model, vz_model = kepler.calc_orbit( epochs, sma, ecc, inc, aop, pan, tau, plx, mtot, tau_ref_epoch=tau_ref_epoch) # generate some fake measurements just to feed into system.py to test bookkeeping t = table.Table([ epochs, np.ones(epochs.shape, dtype=int), ra_model, np.zeros(ra_model.shape), dec_model, np.zeros(dec_model.shape) ], names=[ "epoch", "object", "raoff", "raoff_err", "decoff", "decoff_err" ]) filename = os.path.join(orbitize.DATADIR, "rebound_1planet.csv") t.write(filename) # create the orbitize system and generate model predictions using the ground truth astrom_dat = read_input.read_file(filename) sys = system.System(1, astrom_dat, mtot, plx, tau_ref_epoch=tau_ref_epoch) params = np.array([sma, ecc, inc, aop, pan, tau, plx, mtot]) radec_orbitize, _ = sys.compute_model(params) ra_orb = radec_orbitize[:, 0] dec_orb = radec_orbitize[:, 1] # now project the orbit with rebound manom = basis.tau_to_manom(epochs[0], sma, mtot, tau, tau_ref_epoch) sim = rebound.Simulation() sim.units = ('yr', 'AU', 'Msun') # add star sim.add(m=mtot - mass_b) # add one planet sim.add(m=mass_b, a=sma, e=ecc, M=manom, omega=aop, Omega=pan + np.pi / 2, inc=inc) ps = sim.particles sim.move_to_com() # Use Wisdom Holman integrator (fast), with the timestep being < 5% of inner planet's orbital period sim.integrator = "ias15" sim.dt = ps[1].P / 1000. # integrate and measure star/planet separation ra_reb = [] dec_reb = [] for t in epochs: sim.integrate(t / 365.25) ra_reb.append(-(ps[1].x - ps[0].x)) # ra is negative x dec_reb.append(ps[1].y - ps[0].y) ra_reb = np.array(ra_reb) dec_reb = np.array(dec_reb) diff_ra = ra_reb - ra_orb / plx diff_dec = dec_reb - dec_orb / plx assert np.all(np.abs(diff_ra) < 1e-9) assert np.all(np.abs(diff_dec) < 1e-9)
def test_2planet_massive(): """ Compare multiplanet to rebound for planets with mass. """ # generate a planet orbit mjup = u.Mjup.to(u.Msun) mass_b = 12 * mjup mass_c = 9 * mjup params = np.array([ 10, 0.1, np.radians(45), np.radians(45), np.radians(45), 0.5, 3, 0.1, np.radians(45), np.radians(190), np.radians(45), 0.2, 50, mass_b, mass_c, 1.5 - mass_b - mass_c ]) params_noc = np.array([ 10, 0.1, np.radians(45), np.radians(45), np.radians(45), 0.5, 3, 0.1, np.radians(45), np.radians(190), np.radians(45), 0.2, 50, mass_b, 0, 1.5 - mass_b ]) tau_ref_epoch = 0 epochs = np.linspace(0, 365.25 * 10, 100) + tau_ref_epoch # nearly the full period, MJD # doesn't matter that this is right, just needs to be the same size. below doesn't include effect of c # just want to generate some measurements of plaent b to test compute model b_ra_model, b_dec_model, b_vz_model = kepler.calc_orbit( epochs, params[0], params[1], params[2], params[3], params[4], params[5], params[-2], params[-1], tau_ref_epoch=tau_ref_epoch) # generate some fake measurements of planet b, just to feed into system.py to test bookkeeping t = table.Table([ epochs, np.ones(epochs.shape, dtype=int), b_ra_model, np.zeros(b_ra_model.shape), b_dec_model, np.zeros(b_dec_model.shape) ], names=[ "epoch", "object", "raoff", "raoff_err", "decoff", "decoff_err" ]) filename = os.path.join(orbitize.DATADIR, "rebound_2planet_outer.csv") t.write(filename) #### TEST THE OUTER PLANET #### # create the orbitize system and generate model predictions using the ground truth astrom_dat = read_input.read_file(filename) sys = system.System(2, astrom_dat, params[-1], params[-4], tau_ref_epoch=tau_ref_epoch, fit_secondary_mass=True) # generate measurement radec_orbitize, _ = sys.compute_model(params) b_ra_orb = radec_orbitize[:, 0] b_dec_orb = radec_orbitize[:, 1] # debug, generate measurement without c having any mass radec_orb_noc, _ = sys.compute_model(params_noc) b_ra_orb_noc = radec_orb_noc[:, 0] b_dec_orb_noc = radec_orb_noc[:, 1] # check that planet c's perturbation is imprinted (nonzero)) #assert np.all(b_ra_orb_noc != b_ra_orb) # now project the orbit with rebound b_manom = basis.tau_to_manom(epochs[0], params[0], params[-1] + params[-3], params[5], tau_ref_epoch) c_manom = basis.tau_to_manom(epochs[0], params[0 + 6], params[-1] + params[-2], params[5 + 6], tau_ref_epoch) sim = rebound.Simulation() sim.units = ('yr', 'AU', 'Msun') # add star sim.add(m=params[-1]) # add two planets sim.add(m=mass_c, a=params[0 + 6], e=params[1 + 6], M=c_manom, omega=params[3 + 6], Omega=params[4 + 6] + np.pi / 2, inc=params[2 + 6]) sim.add(m=mass_b, a=params[0], e=params[1], M=b_manom, omega=params[3], Omega=params[4] + np.pi / 2, inc=params[2]) ps = sim.particles sim.move_to_com() # Use Wisdom Holman integrator (fast), with the timestep being < 5% of inner planet's orbital period sim.integrator = "ias15" sim.dt = ps[1].P / 1000. # integrate and measure star/planet separation b_ra_reb = [] b_dec_reb = [] for t in epochs: sim.integrate(t / 365.25) b_ra_reb.append(-(ps[2].x - ps[0].x)) # ra is negative x b_dec_reb.append(ps[2].y - ps[0].y) b_ra_reb = np.array(b_ra_reb) b_dec_reb = np.array(b_dec_reb) diff_ra = b_ra_reb - b_ra_orb / params[6 * 2] diff_dec = b_dec_reb - b_dec_orb / params[6 * 2] # we placed the planets far apart to minimize secular interactions but there are still some, so relax precision assert np.all(np.abs(diff_ra) / (params[0]) < 1e-3) assert np.all(np.abs(diff_dec) / (params[0]) < 1e-3) ###### NOW TEST THE INNER PLANET ####### # generate some fake measurements of planet c, just to feed into system.py to test bookkeeping t = table.Table([ epochs, np.ones(epochs.shape, dtype=int) * 2, b_ra_model, np.zeros(b_ra_model.shape), b_dec_model, np.zeros(b_dec_model.shape) ], names=[ "epoch", "object", "raoff", "raoff_err", "decoff", "decoff_err" ]) filename = os.path.join(orbitize.DATADIR, "rebound_2planet_inner.csv") t.write(filename) # create the orbitize system and generate model predictions using the ground truth astrom_dat = read_input.read_file(filename) sys = system.System(2, astrom_dat, params[-1], params[-2], tau_ref_epoch=tau_ref_epoch, fit_secondary_mass=True) # generate measurement radec_orbitize, _ = sys.compute_model(params) c_ra_orb = radec_orbitize[:, 0] c_dec_orb = radec_orbitize[:, 1] # start the REBOUND sim again sim = rebound.Simulation() sim.units = ('yr', 'AU', 'Msun') # add star sim.add(m=params[-1]) # add two planets sim.add(m=mass_c, a=params[0 + 6], e=params[1 + 6], M=c_manom, omega=params[3 + 6], Omega=params[4 + 6] + np.pi / 2, inc=params[2 + 6]) sim.add(m=mass_b, a=params[0], e=params[1], M=b_manom, omega=params[3], Omega=params[4] + np.pi / 2, inc=params[2]) ps = sim.particles sim.move_to_com() # Use Wisdom Holman integrator (fast), with the timestep being < 5% of inner planet's orbital period sim.integrator = "ias15" sim.dt = ps[1].P / 1000. # integrate and measure star/planet separation c_ra_reb = [] c_dec_reb = [] for t in epochs: sim.integrate(t / 365.25) c_ra_reb.append(-(ps[1].x - ps[0].x)) # ra is negative x c_dec_reb.append(ps[1].y - ps[0].y) c_ra_reb = np.array(c_ra_reb) c_dec_reb = np.array(c_dec_reb) diff_ra = c_ra_reb - c_ra_orb / params[6 * 2] diff_dec = c_dec_reb - c_dec_orb / params[6 * 2] # planet is 3 times closer, so roughly 3 times larger secular errors. assert np.all(np.abs(diff_ra) / (params[0]) < 3e-3) assert np.all(np.abs(diff_dec) / (params[0]) < 3e-3)
def print_prediction(planet_name, date_mjd, chains, tau_ref_epoch, num_samples=None): """ Prints out a prediction for the prediction of a planet given a set of posterior draws Args: planet_name (str): name of the planet, already checked to be in the list date_mjd (float): MJD of date for which we want a prediction chains (np.array): Nx8 array of N orbital elements. Orbital elements are ordered as: sma, ecc, inc, aop, pan, tau, plx, mtot tau_ref_epoch (float): MJD for reference epoch of tau (see orbitize for details on tau) num_samples (int): number of random samples for prediction. If None, will use all of them Returns: ra_args (tuple): a two-element tuple of the median RA offset, and stddev of RA offset dec_args (tuple): a two-element tuple of the median Dec offset, and stddev of Dec offset sep_args (tuple): a two-element tuple of the median separation offset, and stddev of sep offset pa_args (tuple): a two-element tuple of the median PA offset, and stddev of PA offset """ if num_samples is None: num_samples = chains.shape[0] rand_draws = np.arange(num_samples) # don't need to randomize else: if num_samples > chains.shape[0]: print("Requested too many samples. Maximum is {0}.".format( chains.shape[0])) return # randomly draw values rand_draws = np.random.randint(0, chains.shape[0], num_samples) rand_orbits = chains[rand_draws] if planet_name not in multi_dict: # single Keplerian orbit fits sma = rand_orbits[:, 0] ecc = rand_orbits[:, 1] inc = rand_orbits[:, 2] aop = rand_orbits[:, 3] pan = rand_orbits[:, 4] tau = rand_orbits[:, 5] plx = rand_orbits[:, 6] mtot = rand_orbits[:, 7] rand_ras, rand_decs, rand_vzs = kepler.calc_orbit( date_mjd, sma, ecc, inc, aop, pan, tau, plx, mtot, tau_ref_epoch=tau_ref_epoch) else: # massive multi-planet orbit fits planet_num, tot_planets = multi_dict[planet_name] orb_indices = np.arange(6 * planet_num, 6 * planet_num + 6, 1) sma, ecc, inc, aop, pan, tau = rand_orbits[:, orb_indices].T plx = rand_orbits[:, 6 * tot_planets] mass_planets = rand_orbits[:, -1 - tot_planets:-1] mass_star = rand_orbits[:, -1] all_pl_smas = rand_orbits[0, 0:6 * tot_planets:6] within_orbit = np.where(all_pl_smas <= all_pl_smas[planet_num]) mtot = mass_star + np.sum(mass_planets[:, within_orbit[0]], axis=1) rand_ras, rand_decs, rand_vzs = kepler.calc_orbit( date_mjd, sma, ecc, inc, aop, pan, tau, plx, mtot, tau_ref_epoch=tau_ref_epoch) # add perturbation from other planets for inner_pl in within_orbit[0]: if inner_pl == planet_num: continue within_inner_orbit = np.where(all_pl_smas < all_pl_smas[inner_pl]) inner_mtot = mass_star + np.sum( mass_planets[:, within_inner_orbit[0]], axis=1) mass_inner = mass_planets[:, inner_pl] inner_orb_indices = np.arange(6 * inner_pl, 6 * inner_pl + 6, 1) in_sma, in_ecc, in_inc, in_aop, in_pan, in_tau = rand_orbits[:, inner_orb_indices].T inner_rand_ras, inner_rand_decs, inner_rand_vzs = kepler.calc_orbit( date_mjd, in_sma, in_ecc, in_inc, in_aop, in_pan, in_tau, plx, inner_mtot, tau_ref_epoch=tau_ref_epoch) rand_ras += mass_inner / inner_mtot * inner_rand_ras rand_decs += mass_inner / inner_mtot * inner_rand_decs rand_vzs += mass_inner / inner_mtot * inner_rand_vzs rand_seps = np.sqrt(rand_ras**2 + rand_decs**2) rand_pas = np.degrees(np.arctan2(rand_ras, rand_decs)) % 360 ra_args = np.median(rand_ras), np.std(rand_ras) dec_args = np.median(rand_decs), np.std(rand_decs) sep_args = np.median(rand_seps), np.std(rand_seps) pa_args = np.median(rand_pas), np.std(rand_pas) rv_args = np.median(rand_vzs), np.std(rand_vzs) print("RA Offset = {0:.3f} +/- {1:.3f} mas".format(ra_args[0], ra_args[1])) print("Dec Offset = {0:.3f} +/- {1:.3f} mas".format( dec_args[0], dec_args[1])) print("Separation = {0:.3f} +/- {1:.3f} mas".format( sep_args[0], sep_args[1])) print("PA = {0:.3f} +/- {1:.3f} deg".format(pa_args[0], pa_args[1])) print("Planetary RV = {0:.3f} +/- {1:.3f} km/s".format( rv_args[0], rv_args[1])) return ra_args, dec_args, sep_args, pa_args
def test_1planet(): """ Check that for the 2-body case, the primary orbit around the barycenter is equal to -m2/(m1 + m2) times the secondary orbit around the primary. """ # generate a planet orbit sma = 1 ecc = 0.1 inc = np.radians(45) aop = np.radians(45) pan = np.radians(45) tau = 0.5 plx = 1 m0 = 1 tau_ref_epoch = 0 mjup = u.Mjup.to(u.Msun) mass_b = 100 * mjup mtot = mass_b + m0 epochs = np.linspace(0, 300, 100) + tau_ref_epoch # nearly the full period, MJD ra_model, dec_model, _ = kepler.calc_orbit(epochs, sma, ecc, inc, aop, pan, tau, plx, mtot, tau_ref_epoch=tau_ref_epoch) # generate some fake measurements to feed into system.py to test bookkeeping t = table.Table([ epochs, np.ones(epochs.shape, dtype=int), ra_model, np.zeros(ra_model.shape), dec_model, np.zeros(dec_model.shape) ], names=[ "epoch", "object", "raoff", "raoff_err", "decoff", "decoff_err" ]) filename = os.path.join(orbitize.DATADIR, "rebound_1planet.csv") t.write(filename) # create the orbitize system and generate model predictions using ground truth astrom_dat = read_input.read_file(filename) sys = system.System(1, astrom_dat, m0, plx, tau_ref_epoch=tau_ref_epoch, fit_secondary_mass=True) sys.track_planet_perturbs = True params = np.array([sma, ecc, inc, aop, pan, tau, plx, mass_b, m0]) ra, dec, _ = sys.compute_all_orbits(params) # the planet and stellar orbit should just be scaled versions of one another planet_ra = ra[:, 1, :] planet_dec = dec[:, 1, :] star_ra = ra[:, 0, :] star_dec = dec[:, 0, :] assert np.all(np.abs(star_ra + (mass_b / mtot) * planet_ra) < 1e-16) assert np.all(np.abs(star_dec + (mass_b / mtot) * planet_dec) < 1e-16) # remove the created csv file to clean up os.system('rm {}'.format(filename))
def plot_orbits(self, object_to_plot=1, start_mjd=51544., num_orbits_to_plot=100, num_epochs_to_plot=100, square_plot=True, show_colorbar=True, cmap=cmap, sep_pa_color='lightgrey', sep_pa_end_year=2025.0, cbar_param='Epoch [year]', mod180=False, rv_time_series=False, plot_astrometry=True, fig=None): """ Plots one orbital period for a select number of fitted orbits for a given object, with line segments colored according to time Args: object_to_plot (int): which object to plot (default: 1) start_mjd (float): MJD in which to start plotting orbits (default: 51544, the year 2000) num_orbits_to_plot (int): number of orbits to plot (default: 100) num_epochs_to_plot (int): number of points to plot per orbit (default: 100) square_plot (Boolean): Aspect ratio is always equal, but if square_plot is True (default), then the axes will be square, otherwise, white space padding is used show_colorbar (Boolean): Displays colorbar to the right of the plot [True] cmap (matplotlib.cm.ColorMap): color map to use for making orbit tracks (default: modified Purples_r) sep_pa_color (string): any valid matplotlib color string, used to set the color of the orbit tracks in the Sep/PA panels (default: 'lightgrey'). sep_pa_end_year (float): decimal year specifying when to stop plotting orbit tracks in the Sep/PA panels (default: 2025.0). cbar_param (string): options are the following: epochs, sma1, ecc1, inc1, aop1, pan1, tau1. Number can be switched out. Default is epochs. mod180 (Bool): if True, PA will be plotted in range [180, 540]. Useful for plotting short arcs with PAs that cross 360 deg during observations (default: False) rv_time_series (Boolean): if fitting for secondary mass using MCMC for rv fitting and want to display time series, set to True. astrometry (Boolean): set to True by default. Plots the astrometric data. fig (matplotlib.pyplot.Figure): optionally include a predefined Figure object to plot the orbit on. Most users will not need this keyword. Return: ``matplotlib.pyplot.Figure``: the orbit plot if input is valid, ``None`` otherwise (written): Henry Ngo, Sarah Blunt, 2018 Additions by Malena Rice, 2019 """ if Time(start_mjd, format='mjd').decimalyear >= sep_pa_end_year: raise ValueError( 'start_mjd keyword date must be less than sep_pa_end_year keyword date.' ) if object_to_plot > self.num_secondary_bodies: raise ValueError( "Only {0} secondary bodies being fit. Requested to plot body {1} which is out of range" .format(self.num_secondary_bodies, object_to_plot)) if object_to_plot == 0: raise ValueError( "Plotting the primary's orbit is currently unsupported. Stay tuned.." ) with warnings.catch_warnings(): warnings.simplefilter('ignore', ErfaWarning) dict_of_indices = { 'sma': 0, 'ecc': 1, 'inc': 2, 'aop': 3, 'pan': 4, 'tau': 5, 'plx': 6 * self.num_secondary_bodies, } if cbar_param == 'Epoch [year]': pass elif cbar_param[0:3] in dict_of_indices: try: object_id = np.int(cbar_param[3:]) except ValueError: object_id = 1 index = dict_of_indices[cbar_param[0:3]] + 6 * (object_id - 1) else: raise Exception( 'Invalid input; acceptable inputs include epochs, sma1, ecc1, inc1, aop1, pan1, tau1, sma2, ecc2, ...' ) start_index = (object_to_plot - 1) * 6 sma = self.post[:, start_index + dict_of_indices['sma']] ecc = self.post[:, start_index + dict_of_indices['ecc']] inc = self.post[:, start_index + dict_of_indices['inc']] aop = self.post[:, start_index + dict_of_indices['aop']] pan = self.post[:, start_index + dict_of_indices['pan']] tau = self.post[:, start_index + dict_of_indices['tau']] plx = self.post[:, dict_of_indices['plx']] # Then, get the other parameters if 'mtot' in self.labels: mtot = self.post[:, -1] elif 'm0' in self.labels: m0 = self.post[:, -1] m1 = self.post[:, -(self.num_secondary_bodies + 1) + (object_to_plot - 1)] mtot = m0 + m1 # Select random indices for plotted orbit if num_orbits_to_plot > len(sma): num_orbits_to_plot = len(sma) choose = np.random.randint(0, high=len(sma), size=num_orbits_to_plot) raoff = np.zeros((num_orbits_to_plot, num_epochs_to_plot)) deoff = np.zeros((num_orbits_to_plot, num_epochs_to_plot)) vz_star = np.zeros((num_orbits_to_plot, num_epochs_to_plot)) epochs = np.zeros((num_orbits_to_plot, num_epochs_to_plot)) # Loop through each orbit to plot and calcualte ra/dec offsets for all points in orbit # Need this loops since epochs[] vary for each orbit, unless we want to just plot the same time period for all orbits for i in np.arange(num_orbits_to_plot): orb_ind = choose[i] # Compute period (from Kepler's third law) period = np.sqrt(4 * np.pi**2.0 * (sma * u.AU)**3 / (consts.G * (mtot * u.Msun))) period = period.to(u.day).value # Create an epochs array to plot num_epochs_to_plot points over one orbital period epochs[i, :] = np.linspace(start_mjd, float(start_mjd + period[orb_ind]), num_epochs_to_plot) # Calculate ra/dec offsets for all epochs of this orbit raoff0, deoff0, _ = kepler.calc_orbit( epochs[i, :], sma[orb_ind], ecc[orb_ind], inc[orb_ind], aop[orb_ind], pan[orb_ind], tau[orb_ind], plx[orb_ind], mtot[orb_ind], tau_ref_epoch=self.tau_ref_epoch, tau_warning=False) raoff[i, :] = raoff0 deoff[i, :] = deoff0 # Create a linearly increasing colormap for our range of epochs if cbar_param != 'Epoch [year]': cbar_param_arr = self.post[:, index] norm = mpl.colors.Normalize(vmin=np.min(cbar_param_arr), vmax=np.max(cbar_param_arr)) norm_yr = mpl.colors.Normalize(vmin=np.min(cbar_param_arr), vmax=np.max(cbar_param_arr)) elif cbar_param == 'Epoch [year]': norm = mpl.colors.Normalize(vmin=np.min(epochs), vmax=np.max(epochs[-1, :])) norm_yr = mpl.colors.Normalize( vmin=np.min(Time(epochs, format='mjd').decimalyear), vmax=np.max(Time(epochs, format='mjd').decimalyear)) # Create figure for orbit plots if fig is None: fig = plt.figure(figsize=(14, 6)) if rv_time_series: fig = plt.figure(figsize=(14, 9)) ax = plt.subplot2grid((3, 14), (0, 0), rowspan=2, colspan=6) else: fig = plt.figure(figsize=(14, 6)) ax = plt.subplot2grid((2, 14), (0, 0), rowspan=2, colspan=6) else: plt.set_current_figure(fig) if rv_time_series: ax = plt.subplot2grid((3, 14), (0, 0), rowspan=2, colspan=6) else: ax = plt.subplot2grid((2, 14), (0, 0), rowspan=2, colspan=6) data = self.data astr_inds = np.where((~np.isnan(data['quant1'])) & (~np.isnan(data['quant2']))) astr_epochs = data['epoch'][astr_inds] sep_data, sep_err = data['quant1'][astr_inds], data['quant1_err'][ astr_inds] pa_data, pa_err = data['quant2'][astr_inds], data['quant2_err'][ astr_inds] # Plot each orbit (each segment between two points coloured using colormap) for i in np.arange(num_orbits_to_plot): points = np.array([raoff[i, :], deoff[i, :]]).T.reshape(-1, 1, 2) segments = np.concatenate([points[:-1], points[1:]], axis=1) lc = LineCollection(segments, cmap=cmap, norm=norm, linewidth=1.0) if cbar_param != 'Epoch [year]': lc.set_array(np.ones(len(epochs[0])) * cbar_param_arr[i]) elif cbar_param == 'Epoch [year]': lc.set_array(epochs[i, :]) ax.add_collection(lc) if plot_astrometry: ra_data, dec_data = orbitize.system.seppa2radec( sep_data, pa_data) ax.scatter(ra_data, dec_data, marker='*', c='#FF7F11', zorder=10, s=60) # modify the axes if square_plot: adjustable_param = 'datalim' else: adjustable_param = 'box' ax.set_aspect('equal', adjustable=adjustable_param) ax.set_xlabel('$\\Delta$RA [mas]') ax.set_ylabel('$\\Delta$Dec [mas]') ax.locator_params(axis='x', nbins=6) ax.locator_params(axis='y', nbins=6) ax.invert_xaxis() # To go to a left-handed coordinate system # Rob: Moved colorbar size to the bottom after tight_layout() because the cbar scaling was not compatible with tight_layout() # plot sep/PA and/or rv zoom-in panels if rv_time_series: ax1 = plt.subplot2grid((3, 14), (0, 8), colspan=6) ax2 = plt.subplot2grid((3, 14), (1, 8), colspan=6) ax3 = plt.subplot2grid((3, 14), (2, 0), colspan=14, rowspan=1) ax2.set_ylabel('PA [$^{{\\circ}}$]') ax1.set_ylabel('$\\rho$ [mas]') ax3.set_ylabel('RV [km/s]') ax3.set_xlabel('Epoch') ax2.set_xlabel('Epoch') plt.subplots_adjust(hspace=0.3) else: ax1 = plt.subplot2grid((2, 14), (0, 9), colspan=6) ax2 = plt.subplot2grid((2, 14), (1, 9), colspan=6) ax2.set_ylabel('PA [$^{{\\circ}}$]') ax1.set_ylabel('$\\rho$ [mas]') ax2.set_xlabel('Epoch') epochs_seppa = np.zeros((num_orbits_to_plot, num_epochs_to_plot)) for i in np.arange(num_orbits_to_plot): orb_ind = choose[i] epochs_seppa[i, :] = np.linspace( start_mjd, Time(sep_pa_end_year, format='decimalyear').mjd, num_epochs_to_plot) # Calculate ra/dec offsets for all epochs of this orbit if rv_time_series: raoff0, deoff0, vzoff0 = kepler.calc_orbit( epochs_seppa[i, :], sma[orb_ind], ecc[orb_ind], inc[orb_ind], aop[orb_ind], pan[orb_ind], tau[orb_ind], plx[orb_ind], mtot[orb_ind], tau_ref_epoch=self.tau_ref_epoch, mass_for_Kamp=m0[orb_ind], tau_warning=False) raoff[i, :] = raoff0 deoff[i, :] = deoff0 else: raoff0, deoff0, _ = kepler.calc_orbit( epochs_seppa[i, :], sma[orb_ind], ecc[orb_ind], inc[orb_ind], aop[orb_ind], pan[orb_ind], tau[orb_ind], plx[orb_ind], mtot[orb_ind], tau_ref_epoch=self.tau_ref_epoch, tau_warning=False) raoff[i, :] = raoff0 deoff[i, :] = deoff0 yr_epochs = Time(epochs_seppa[i, :], format='mjd').decimalyear seps, pas = orbitize.system.radec2seppa(raoff[i, :], deoff[i, :], mod180=mod180) plt.sca(ax1) plt.plot(yr_epochs, seps, color=sep_pa_color) # plot separations from data points plt.scatter(Time(astr_epochs, format='mjd').decimalyear, sep_data, s=10, marker='*', c='purple', zorder=10) plt.sca(ax2) plt.plot(yr_epochs, pas, color=sep_pa_color) plt.scatter(Time(astr_epochs, format='mjd').decimalyear, pa_data, s=10, marker='*', c='purple', zorder=10) if rv_time_series: # switch current axis to rv panel plt.sca(ax3) # get list of instruments insts = np.unique(data['instrument']) insts = [ i if isinstance(i, str) else i.decode() for i in insts ] insts = [i for i in insts if 'def' not in i] # get gamma/sigma labels and corresponding positions in the posterior gams = ['gamma_' + inst for inst in insts] if isinstance(self.labels, list): labels = np.array(self.labels) else: labels = self.labels # get the indices corresponding to each gamma within self.labels gam_idx = [ np.where(labels == inst_gamma)[0][0] for inst_gamma in gams ] # indices corresponding to each instrument in the datafile inds = {} for i in range(len(insts)): inds[insts[i]] = np.where( data['instrument'] == insts[i].encode())[0] # choose the orbit with the best log probability best_like = np.where(self.lnlike == np.amin(self.lnlike))[0][0] med_ga = [self.post[best_like, i] for i in gam_idx] # colour/shape scheme scheme for rv data points clrs = ['0496FF', '372554', 'FF1053', '3A7CA5', '143109'] symbols = ['o', '^', 'v', 's'] # get rvs and plot them for i, name in enumerate(inds.keys()): rv_inds = np.where((np.isnan(data['quant2']))) inst_data = data[inds[name]] rvs = inst_data['quant1'] epochs = inst_data['epoch'] epochs = Time(epochs, format='mjd').decimalyear rvs -= med_ga[i] plt.scatter(epochs, rvs, marker=symbols[i], s=5, label=name, c=f'#{clrs[i]}', zorder=5) inds[insts[i]] = np.where(data['instrument'] == insts[i])[0] plt.legend() # calculate the predicted rv trend using the best orbit raa, decc, vz = kepler.calc_orbit( epochs_seppa[i, :], sma[best_like], ecc[best_like], inc[best_like], aop[best_like], pan[best_like], tau[best_like], plx[best_like], mtot[best_like], tau_ref_epoch=self.tau_ref_epoch, mass_for_Kamp=m0[best_like]) vz = vz * -(m1[best_like]) / np.median(m0[best_like]) # plot rv trend plt.plot(Time(epochs_seppa[i, :], format='mjd').decimalyear, vz, color=sep_pa_color) # add colorbar if show_colorbar: if rv_time_series: # Create an axes for colorbar. The position of the axes is calculated based on the position of ax. # You can change x1.0.05 to adjust the distance between the main image and the colorbar. # You can change 0.02 to adjust the width of the colorbar. cbar_ax = fig.add_axes([ ax.get_position().x1 + 0.005, ax.get_position().y0, 0.02, ax.get_position().height ]) cbar = mpl.colorbar.ColorbarBase(cbar_ax, cmap=cmap, norm=norm_yr, orientation='vertical', label=cbar_param) else: # xpos, ypos, width, height, in fraction of figure size cbar_ax = fig.add_axes([0.47, 0.15, 0.015, 0.7]) cbar = mpl.colorbar.ColorbarBase(cbar_ax, cmap=cmap, norm=norm_yr, orientation='vertical', label=cbar_param) ax1.locator_params(axis='x', nbins=6) ax1.locator_params(axis='y', nbins=6) ax2.locator_params(axis='x', nbins=6) ax2.locator_params(axis='y', nbins=6) return fig
def test_compute_model(): """ Tests that the perturbation of a second planet using compute model gives roughly the amplitude we expect. """ # generate planet b orbital parameters b_params = [1, 0, 0, 0, 0, 0] 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] mass_c = 0.002 # Msun mtot = m0 + mass_b + mass_c period_c = np.sqrt(c_params[0]**3 / mtot) period_b = np.sqrt(b_params[0]**3 / mtot) epochs = np.linspace(0, period_c * 365.25, 100) + tau_ref_epoch # the full period of c, MJD ra_model, dec_model, 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, tau_ref_epoch=tau_ref_epoch) # generate some fake measurements just to feed into system.py to test bookkeeping # just make a 1 planet fit for now t = table.Table([ epochs, np.ones(epochs.shape, dtype=int), ra_model, np.zeros(ra_model.shape), dec_model, np.zeros(dec_model.shape) ], names=[ "epoch", "object", "raoff", "raoff_err", "decoff", "decoff_err" ]) filename = os.path.join(orbitize.DATADIR, "multiplanet_fake_1planettest.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_1body = system.System(1, astrom_dat, m0, plx, tau_ref_epoch=tau_ref_epoch, fit_secondary_mass=True) sys_2body = system.System(2, astrom_dat, m0, plx, tau_ref_epoch=tau_ref_epoch, fit_secondary_mass=True) # model predictions for the 1 body case # we had put one measurements of planet b in the data table, so compute_model only does planet b measurements params = np.append(b_params, [plx, mass_b, m0]) radec_1body, _ = sys_1body.compute_model(params) ra_1body = radec_1body[:, 0] dec_1body = radec_1body[:, 1] # model predictions for the 2 body case # still only generates predictions of b's location, but including the perturbation for c params = np.append(b_params, np.append(c_params, [plx, mass_b, mass_c, m0])) radec_2body, _ = sys_2body.compute_model(params) ra_2body = radec_2body[:, 0] dec_2body = radec_2body[:, 1] ra_diff = ra_2body - ra_1body dec_diff = dec_2body - dec_1body total_diff = np.sqrt(ra_diff**2 + dec_diff**2) # the expected influence of c is mass_c/m0 * sma_c * plx in amplitude # just test the first value, because of the face on orbit, we should see it immediately. assert total_diff[0] == pytest.approx(mass_c / m0 * c_params[0] * plx, abs=0.01 * mass_c / m0 * b_params[0] * plx)
def plot_orbits(self, parallax=None, total_mass=None, object_mass=0, object_to_plot=1, start_mjd=51544., num_orbits_to_plot=100, num_epochs_to_plot=100, square_plot=True, show_colorbar=True, cmap=cmap): """ Plots one orbital period for a select number of fitted orbits for a given object, with line segments colored according to time Args: parallax (float): parallax (in mas), however, if plx_err was passed to system, then this is ignored and the posterior samples for plx will be used instead (default: None) total_mass (float): total mass of system in solar masses, however, if mass_err was passed to system, then this is ignored and the posterior samples for mtot will be used instead (default: None) object_mass (float): mass of the object, in solar masses (default: 0) Note: this input has no effect at this time object_to_plot (int): which object to plot (default: 1) start_mjd (float): MJD in which to start plotting orbits (default: 51544, the year 2000) num_orbits_to_plot (int): number of orbits to plot (default: 100) num_epochs_to_plot (int): number of points to plot per orbit (default: 100) square_plot (Boolean): Aspect ratio is always equal, but if square_plot is True (default), then the axes will be square, otherwise, white space padding is used show_colorbar (Boolean): Displays colorbar to the right of the plot [True] cmap (matplotlib.cm.ColorMap): color map to use for making orbit tracks (default: modified Purples_r) Return: ``matplotlib.pyplot.Figure``: the orbit plot if input is valid, None otherwise (written): Henry Ngo, Sarah Blunt, 2018 """ # Split the 2-D post array into series of 1-D arrays for each orbital parameter num_objects, remainder = np.divmod(self.post.shape[1],6) if object_to_plot > num_objects: return None first_index = 0 + 6*(object_to_plot-1) sma = self.post[:,first_index+0] ecc = self.post[:,first_index+1] inc = self.post[:,first_index+2] aop = self.post[:,first_index+3] pan = self.post[:,first_index+4] tau = self.post[:,first_index+5] # Then, get the other parameters if remainder == 2: # have samples for parallax and mtot plx = self.post[:,-2] mtot = self.post[:,-1] else: # otherwise make arrays out of user provided value if total_mass is not None: mtot = np.ones(len(sma))*total_mass else: raise Exception('results.Results.plot_orbits(): total mass must be provided if not part of samples') if parallax is not None: plx = np.ones(len(sma))*parallax else: raise Exception('results.Results.plot_orbits(): parallax must be provided if not part of samples') mplanet = np.ones(len(sma))*object_mass # Select random indices for plotted orbit if num_orbits_to_plot > len(sma): num_orbits_to_plot = len(sma) choose = np.random.randint(0, high=len(sma), size=num_orbits_to_plot) raoff = np.zeros((num_orbits_to_plot, num_epochs_to_plot)) deoff = np.zeros((num_orbits_to_plot, num_epochs_to_plot)) epochs = np.zeros((num_orbits_to_plot, num_epochs_to_plot)) # Loop through each orbit to plot and calcualte ra/dec offsets for all points in orbit # Need this loops since epochs[] vary for each orbit, unless we want to just plot the same time period for all orbits for i in np.arange(num_orbits_to_plot): orb_ind = choose[i] # Compute period (from Kepler's third law) period = np.sqrt(4*np.pi**2.0*(sma*u.AU)**3/(consts.G*(mtot*u.Msun))) period = period.to(u.day).value # Create an epochs array to plot num_epochs_to_plot points over one orbital period epochs[i,:] = np.linspace(start_mjd, float(start_mjd+period[orb_ind]), num_epochs_to_plot) # Calculate ra/dec offsets for all epochs of this orbit raoff0, deoff0, _ = kepler.calc_orbit( epochs[i,:], sma[orb_ind], ecc[orb_ind], inc[orb_ind], aop[orb_ind], pan[orb_ind], tau[orb_ind], plx[orb_ind], mtot[orb_ind], mass=mplanet[orb_ind] ) raoff[i,:] = raoff0 deoff[i,:] = deoff0 # Create a linearly increasing colormap for our range of epochs norm = mpl.colors.Normalize(vmin=np.min(epochs), vmax=np.max(epochs[-1,:])) norm_yr = mpl.colors.Normalize(vmin=np.min(epochs/365.25), vmax=np.max(epochs[-1,:]/365.25)) # Create figure for orbit plots fig, ax = plt.subplots() # Plot each orbit (each segment between two points coloured using colormap) for i in np.arange(num_orbits_to_plot): points = np.array([raoff[i,:], deoff[i,:]]).T.reshape(-1,1,2) segments = np.concatenate([points[:-1], points[1:]], axis=1) lc = LineCollection( segments, cmap=cmap, norm=norm, linewidth=1.0 ) lc.set_array(epochs[i,:]) ax.add_collection(lc) # Modify the axes if square_plot: adjustable_param='datalim' else: adjustable_param='box' ax.set_aspect('equal', adjustable=adjustable_param) ax.set_xlabel('$\Delta$RA [mas]') ax.set_ylabel('$\Delta$Dec [mas]') ax.locator_params(axis='x', nbins=6) ax.locator_params(axis='y', nbins=6) # Add colorbar if show_colorbar: sm = mpl.cm.ScalarMappable(cmap=cmap, norm=norm_yr) sm.set_array([]) # magic? (just needs to *not* be None) cbar = fig.colorbar(sm, format='%g') # Alternative implementation example for right-hand colorbar # fig.subplots_adjust(right=0.8) # cbar_ax = fig.add_axes([0.825, 0.15, 0.05, 0.7]) # xpos, ypos, width, height, in fraction of figure size # cbar = mpl.colorbar.ColorbarBase(cbar_ax, cmap=cmap, norm=norm_yr, orientation='vertical') return fig