diff --git a/orbitize/example_data/reflected_light_example.csv b/orbitize/example_data/reflected_light_example.csv new file mode 100644 index 00000000..5ae4a029 --- /dev/null +++ b/orbitize/example_data/reflected_light_example.csv @@ -0,0 +1,6 @@ +epoch,object,sep,sep_err,pa,pa_err,brightness,brightness_err +54781,1,210.0,27.0,211.49,1.9,0.9,0.1 +52953,1,413.0,22.0,34,4,0.6,0.2 +55129,1,299.0,14.0,211,3,, +55194,1,306.0,9.0,212.1,1.7,, +55296,1,346.0,7.0,209.9,1.2,, \ No newline at end of file diff --git a/orbitize/kepler.py b/orbitize/kepler.py index 447bbcc7..c7506686 100644 --- a/orbitize/kepler.py +++ b/orbitize/kepler.py @@ -44,6 +44,62 @@ def tau_to_manom(date, sma, mtot, tau, tau_ref_epoch): return mean_anom +def times2trueanom_and_eccanom( + sma, + epochs, + mtot, + ecc, + tau, + tau_ref_epoch=58849, + tolerance=1e-9, + max_iter=100, + use_c=True, + use_gpu=False, +): + """ + Convert times to true anomaly and eccentric anomaly by solving Kepler's Equation. + + Args: + sma (np.array): semi-major axis of orbit [au] + epochs (np.array): MJD times for which we want the positions of the planet + mtot (np.array): total mass of the two-body orbit (M_* + M_planet) [Solar masses] + ecc (np.array): eccentricity of the orbit [0,1] + tau (np.array): epoch of periastron passage in fraction of orbital period past MJD=0 [0,1] + tau_ref_epoch (float, optional): reference date that tau is defined with respect to (default: 58849) + tolerance (float, optional): absolute tolerance of iterative computation. Defaults to 1e-9. + max_iter (int, optional): maximum number of iterations before switching. Defaults to 100. + use_c (bool, optional): Use the C solver if configured. Defaults to True + use_gpu (bool, optional): Use the GPU solver if configured. Defaults to False + + Returns: + 2-tuple: + + np.array: true anomalies (shape n_epochs) + + np.array: eccentric anomalies (shape n_epochs) + """ + + n_orbs = np.size(sma) # num sets of input orbital parameters + n_dates = np.size(epochs) # number of dates to compute offsets and vz + + + # Necessary for _calc_ecc_anom, for now + if np.isscalar(epochs): # just in case epochs is given as a scalar + epochs = np.array([epochs]) + ecc_arr = np.tile(ecc, (n_dates, 1)) + + # # compute mean anomaly (size: n_orbs x n_dates) + manom = tau_to_manom(epochs[:, None], sma, mtot, tau, tau_ref_epoch) + # compute eccentric anomalies (size: n_orbs x n_dates) + eanom = _calc_ecc_anom(manom, ecc_arr, tolerance=tolerance, max_iter=max_iter, use_c=use_c, use_gpu=use_gpu) + + # compute the true anomalies (size: n_orbs x n_dates) + # Note: matrix multiplication makes the shapes work out here and below + tanom = 2.*np.arctan(np.sqrt((1.0 + ecc)/(1.0 - ecc))*np.tan(0.5*eanom)) + + return tanom, eanom + + def calc_orbit( epochs, sma, ecc, inc, aop, pan, tau, plx, mtot, mass_for_Kamp=None, tau_ref_epoch=58849, tolerance=1e-9, @@ -70,7 +126,7 @@ def calc_orbit( For example, if you want to return the stellar RV, this is the planet mass. If you want to return the planetary RV, this is the stellar mass. [Solar masses]. For planet mass ~ 0, mass_for_Kamp ~ M_tot, and function returns planetary RV (default). - tau_ref_epoch (float, optional): reference date that tau is defined with respect to (i.e., tau=0) + tau_ref_epoch (float, optional): reference date that tau is defined with respect to (default: 58849) tolerance (float, optional): absolute tolerance of iterative computation. Defaults to 1e-9. max_iter (int, optional): maximum number of iterations before switching. Defaults to 100. use_c (bool, optional): Use the C solver if configured. Defaults to True @@ -89,26 +145,14 @@ def calc_orbit( Written: Jason Wang, Henry Ngo, 2018 """ - n_orbs = np.size(sma) # num sets of input orbital parameters - n_dates = np.size(epochs) # number of dates to compute offsets and vz - # return planetary RV if `mass_for_Kamp` is not defined + # return planetary RV if `mass_for_Kamp` is not defined if mass_for_Kamp is None: mass_for_Kamp = mtot + ecc - # Necessary for _calc_ecc_anom, for now - if np.isscalar(epochs): # just in case epochs is given as a scalar - epochs = np.array([epochs]) - ecc_arr = np.tile(ecc, (n_dates, 1)) + tanom, eanom = times2trueanom_and_eccanom(sma, epochs, mtot, ecc, tau, tau_ref_epoch=tau_ref_epoch, tolerance=tolerance, max_iter=max_iter, use_c=use_c, use_gpu=use_gpu) - # # compute mean anomaly (size: n_orbs x n_dates) - manom = tau_to_manom(epochs[:, None], sma, mtot, tau, tau_ref_epoch) - # compute eccentric anomalies (size: n_orbs x n_dates) - eanom = _calc_ecc_anom(manom, ecc_arr, tolerance=tolerance, max_iter=max_iter, use_c=use_c, use_gpu=use_gpu) - - # compute the true anomalies (size: n_orbs x n_dates) - # Note: matrix multiplication makes the shapes work out here and below - tanom = 2.*np.arctan(np.sqrt((1.0 + ecc)/(1.0 - ecc))*np.tan(0.5*eanom)) # compute 3-D orbital radius of second body (size: n_orbs x n_dates) radius = sma * (1.0 - ecc * np.cos(eanom)) diff --git a/orbitize/read_input.py b/orbitize/read_input.py index a17966ba..3b094982 100644 --- a/orbitize/read_input.py +++ b/orbitize/read_input.py @@ -179,6 +179,10 @@ def read_file(filename): have_seppacorr = np.zeros( num_measurements, dtype=bool ) # zeros are False + if "brightness" in input_table.columns: + have_brightness = ~input_table["brightness"].mask + else: + have_brightness = np.zeros(num_measurements, dtype=bool) if "rv" in input_table.columns: have_rv = ~input_table["rv"].mask else: @@ -231,11 +235,14 @@ def read_file(filename): else: have_rv = np.zeros(num_measurements, dtype=bool) # zeros are False - # Rob: not sure if we need this but adding just in case if "instrument" in input_table.columns: have_inst = np.ones(num_measurements, dtype=bool) else: have_inst = np.zeros(num_measurements, dtype=bool) + if "brightness" in input_table.columns: + have_brightness = np.ones(num_measurements, dtype=bool) + else: + have_brightness = np.zeros(num_measurements, dtype=bool) # orbitize! backwards compatability since we added new columns, some old data formats may not have them # fill in with default values @@ -280,9 +287,9 @@ def read_file(filename): if not isinstance(row["object"], (int, np.int32, np.int64)): raise Exception("Invalid object ID. Object IDs must be integers.") - # determine input quantity type (RA/DEC, SEP/PA, or RV) + # determine input quantity type (RA/DEC, SEP/PA, RV, or BRIGHTNESS) if orbitize_style: - if row["quant_type"] == "rv": # special format for rv rows + if row["quant_type"] == "rv" or row["quant_type"] == 'brightness': # special format for rv rows output_table.add_row( [ MJD, @@ -332,6 +339,7 @@ def read_file(filename): ) else: # When not in orbitize style + if have_ra[index] and have_dec[index]: # check if there's a covariance term if have_radeccorr[index]: @@ -428,6 +436,20 @@ def read_file(filename): "defrv", ] ) + if have_brightness[index]: + output_table.add_row( + [ + MJD, + row["object"], + row["brightness"], + row["brightness_err"], + None, + None, + None, + "brightness", + "defbr", + ] + ) return output_table diff --git a/orbitize/sampler.py b/orbitize/sampler.py index 5ee50f21..20968ff0 100644 --- a/orbitize/sampler.py +++ b/orbitize/sampler.py @@ -101,14 +101,14 @@ def _logl(self, params): if self.system.hipparcos_IAD is not None: # compute Ra/Dec predictions at the Hipparcos IAD epochs - raoff_model, deoff_model, _ = self.system.compute_all_orbits( + raoff_model, deoff_model, _, _ = self.system.compute_all_orbits( params, epochs=self.system.hipparcos_IAD.epochs_mjd ) ( raoff_model_hip_epoch, deoff_model_hip_epoch, - _, + _, _ ) = self.system.compute_all_orbits( params, epochs=Time([1991.25], format="decimalyear").mjd ) @@ -134,7 +134,7 @@ def _logl(self, params): ).mjd # compute Ra/Dec predictions at the Gaia epoch - raoff_model, deoff_model, _ = self.system.compute_all_orbits( + raoff_model, deoff_model, _, _ = self.system.compute_all_orbits( params, epochs=gaiahip_epochs ) diff --git a/orbitize/system.py b/orbitize/system.py index a78fa586..33850c39 100644 --- a/orbitize/system.py +++ b/orbitize/system.py @@ -98,10 +98,15 @@ def __init__( # List of index arrays corresponding to each rv for each body self.rv = [] + # index arrays corresponding to brightness for each body + self.brightness = [] + self.fit_astrometry = True radec_indices = np.where(self.data_table["quant_type"] == "radec") seppa_indices = np.where(self.data_table["quant_type"] == "seppa") + brightness_indices = np.where(self.data_table["quant_type"] == "brightness") + 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") @@ -140,6 +145,7 @@ def __init__( np.intersect1d(self.body_indices[body_num], seppa_indices) ) self.rv.append(np.intersect1d(self.body_indices[body_num], rv_indices)) + self.brightness.append(np.intersect1d(self.body_indices[body_num], brightness_indices)) # we should track the influence of the planet(s) on each other/the star if: # we are not fitting massless planets and @@ -302,6 +308,7 @@ def __init__( self.param_idx = self.basis.param_idx + def save(self, hf): """ Saves the current object to an hdf5 file @@ -365,6 +372,11 @@ def compute_all_orbits(self, params_arr, epochs=None, comp_rebound=False): vz (np.array of float): N_epochs x N_bodies x N_orbits array of radial velocities at each epoch. + brightness (np.array of float): N_epochs x N_bodies x N_orbits of + photometric brightness predictions, assuming a Lambertian disk + reflection law, at each epoch. Normalized so that brightness=1 + at maximum. + """ if epochs is None: @@ -386,6 +398,7 @@ def compute_all_orbits(self, params_arr, epochs=None, comp_rebound=False): 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)) + brightness_out = 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: @@ -489,16 +502,33 @@ def compute_all_orbits(self, params_arr, epochs=None, comp_rebound=False): tau_ref_epoch=self.tau_ref_epoch, ) + tanom, eanom = kepler.times2trueanom_and_eccanom(sma, epochs, mtot, ecc, tau, tau_ref_epoch=self.tau_ref_epoch) + + + R = (sma*(1-ecc**2))/(1+ecc*np.cos(tanom)) + + z = (R)*(-np.cos(argp)*np.sin(inc)*np.sin(tanom)-np.cos(tanom)*np.sin(inc)*np.sin(argp)) + + B = np.arctan2(-R, z)+ np.pi + + alpha = (1/np.pi)*(np.sin(B)+(np.pi-B)*np.cos(B)) + + albedo = 0.5 # NOTE: we're only fitting relative changes in brightness, so the actual value of albedo doesn't matter + brightness = albedo*alpha/R**2 + + # 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]) + brightness = np.array([brightness]) # 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)) + brightness_out[:, body_num, :] = np.reshape(brightness, (n_epochs, n_orbits)) # vz_i is the ith companion radial velocity if self.fit_secondary_mass: @@ -572,11 +602,12 @@ def compute_all_orbits(self, params_arr, epochs=None, comp_rebound=False): raoff[:, :, bad_orbits] = np.inf deoff[:, :, bad_orbits] = np.inf vz[:, :, bad_orbits] = np.inf - return raoff, deoff, vz + brightness_out[:, :, bad_orbits] = np.inf + return raoff, deoff, vz, brightness_out else: - return raoff, deoff, vz + return raoff, deoff, vz, brightness_out else: - return raoff, deoff, vz + return raoff, deoff, vz, brightness_out def compute_model(self, params_arr, use_rebound=False): """ @@ -611,7 +642,7 @@ def compute_model(self, params_arr, use_rebound=False): standard_params_arr, comp_rebound=True ) else: - raoff, decoff, vz = self.compute_all_orbits(standard_params_arr) + raoff, decoff, vz, brightness = self.compute_all_orbits(standard_params_arr) if len(standard_params_arr.shape) == 1: n_orbits = 1 @@ -663,6 +694,11 @@ def compute_model(self, params_arr, use_rebound=False): model[self.rv[body_num], 0] = vz[self.rv[body_num], body_num, :] model[self.rv[body_num], 1] = np.nan + # Brightness + if len(self.brightness[body_num]) > 0: + model[self.brightness[body_num], 0] = brightness[self.brightness[body_num], body_num, :] + model[self.brightness[body_num], 1] = np.nan + # if we have abs astrometry measurements in the input file (i.e. not # from Hipparcos or Gaia), add the parallactic & proper motion here by # calling AbsAstrom compute_model diff --git a/tests/test_abs_astrometry.py b/tests/test_abs_astrometry.py index ca6c3a8a..84f9dc0c 100644 --- a/tests/test_abs_astrometry.py +++ b/tests/test_abs_astrometry.py @@ -58,7 +58,7 @@ def test_1planet(): 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) + ra, dec, _, _ = sys.compute_all_orbits(params) # the planet and stellar orbit should just be scaled versions of one another planet_ra = ra[:, 1, :] @@ -143,7 +143,7 @@ def test_arbitrary_abs_astrom(): ] ) - plxonly_fullorbit_ra, plxonly_fullorbit_dec, _ = mySystem.compute_all_orbits( + plxonly_fullorbit_ra, plxonly_fullorbit_dec, _, _ = mySystem.compute_all_orbits( plx_only_params, epochs=np.linspace(epochs[0], epochs[0] + 365.25 / 2, int(1e6)) ) diff --git a/tests/test_brightness.py b/tests/test_brightness.py new file mode 100644 index 00000000..809fc25f --- /dev/null +++ b/tests/test_brightness.py @@ -0,0 +1,94 @@ +""" +Tests for the reflected-light calculation in system.compute_all_orbits +""" + +from orbitize import system, sampler, plot +from orbitize import DATADIR, read_input +import os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +from astropy.time import Time +import pytest + + +def test_brightness_calculation(): + + num_secondary_bodies = 1 + + input_file = os.path.join(DATADIR, "reflected_light_example.csv") + data_table = read_input.read_file(input_file) + + times = data_table["epoch"].value + + system_mass = 1.47 + plx = 24.30 + + test_system = system.System(num_secondary_bodies, data_table, system_mass, plx) + + params = np.array( + [ + 10.0, + 0.3, + np.radians(89), + np.radians(21), + np.radians(70), + 0.0, + 51.5, + 1.75, + ] + ) + + ra, dec, vz, brightness = test_system.compute_all_orbits(params) + + expected_brightness = np.zeros((len(times), num_secondary_bodies + 1, 1)) + expected_brightness[:,1,:] = np.array([ + .00040277, .00040277, .00164703, .00164703, .00032383, .00031209, .00029531 + ]).reshape((7,1)) + + assert expected_brightness == pytest.approx(brightness, abs=1e-8) + + +def test_read_input_with_brightness(): + + num_secondary_bodies = 1 + + input_file = os.path.join(DATADIR, "reflected_light_example.csv") + + data_table = read_input.read_file(input_file) + + assert len(data_table[data_table['quant_type'] == 'brightness']) == 2 + assert len(data_table[data_table['quant_type'] == 'seppa']) == 5 + + brightness_data = data_table[data_table['quant_type'] == 'brightness'] + assert np.all(brightness_data['quant1'].value == [0.9, 0.6]) + assert np.all(brightness_data['quant1_err'].value == [0.1, 0.2]) + assert np.all(np.isnan(brightness_data['quant2'].value)) + assert np.all(np.isnan(brightness_data['quant2_err'].value)) + + +def test_compute_posterior(): + """ + Test that a short mcmc runs to completion when we include photometric + data in the csv. orbitize should automatically detect that we are + including photometry in the fit and incorporate it into the likelihood. + """ + + num_secondary_bodies = 1 + + input_file = os.path.join(DATADIR, "reflected_light_example.csv") + data_table = read_input.read_file(input_file) + + system_mass = 1.47 + plx = 24.30 + + test_system = system.System(num_secondary_bodies, data_table, system_mass, plx) + test_mcmc = sampler.MCMC(test_system, num_temps=1, num_walkers=30, num_threads=1) + test_mcmc.run_sampler(1) + + + +if __name__ == "__main__": + test_brightness_calculation() + # test_read_input_with_brightness() + # test_compute_posterior() diff --git a/tests/test_rebound.py b/tests/test_rebound.py index 9f06afb8..0a53ae4e 100644 --- a/tests/test_rebound.py +++ b/tests/test_rebound.py @@ -119,10 +119,10 @@ def test_8799_rebound_vs_kepler(plotname=None): assert hr8799_sys.track_planet_perturbs - rra, rde, _ = hr8799_sys.compute_all_orbits( + rra, rde, _, _ = hr8799_sys.compute_all_orbits( params_arr, epochs=epochs, comp_rebound=True ) - kra, kde, _ = hr8799_sys.compute_all_orbits( + kra, kde, _, _ = hr8799_sys.compute_all_orbits( params_arr, epochs=epochs, comp_rebound=False ) @@ -224,5 +224,5 @@ def test_rebound_mcmc(): if __name__ == "__main__": - # test_8799_rebound_vs_kepler(plotname="hr8799_diffs") - test_rebound_mcmc() + test_8799_rebound_vs_kepler() + # test_rebound_mcmc() diff --git a/tests/test_results.py b/tests/test_results.py index 0840a993..5dbd91a9 100644 --- a/tests/test_results.py +++ b/tests/test_results.py @@ -339,29 +339,29 @@ def test_save_and_load_gaia_and_hipparcos(): test_save_and_load_hipparcos_only() test_save_and_load_gaia_and_hipparcos() - test_results = test_init_and_add_samples() - - test_results_printing(test_results) - test_plot_long_periods(test_results) - test_results_radec = test_init_and_add_samples(radec_input=True) - - test_save_and_load_results(test_results, has_lnlike=True) - test_save_and_load_results(test_results, has_lnlike=True) - test_save_and_load_results(test_results, has_lnlike=False) - test_save_and_load_results(test_results, has_lnlike=False) - test_corner_fig1, test_corner_fig2, test_corner_fig3 = test_plot_corner( - test_results - ) - test_orbit_figs = test_plot_orbits(test_results) - test_orbit_figs = test_plot_orbits(test_results_radec) - test_corner_fig1.savefig("test_corner1.png") - test_corner_fig2.savefig("test_corner2.png") - test_corner_fig3.savefig("test_corner3.png") - test_orbit_figs[0].savefig("test_orbit1.png") - test_orbit_figs[1].savefig("test_orbit2.png") - test_orbit_figs[2].savefig("test_orbit3.png") - test_orbit_figs[3].savefig("test_orbit4.png") - test_orbit_figs[4].savefig("test_orbit5.png") - - # clean up - os.system("rm test_*.png") + # test_results = test_init_and_add_samples() + + # test_results_printing(test_results) + # test_plot_long_periods(test_results) + # test_results_radec = test_init_and_add_samples(radec_input=True) + + # test_save_and_load_results(test_results, has_lnlike=True) + # test_save_and_load_results(test_results, has_lnlike=True) + # test_save_and_load_results(test_results, has_lnlike=False) + # test_save_and_load_results(test_results, has_lnlike=False) + # test_corner_fig1, test_corner_fig2, test_corner_fig3 = test_plot_corner( + # test_results + # ) + # test_orbit_figs = test_plot_orbits(test_results) + # test_orbit_figs = test_plot_orbits(test_results_radec) + # test_corner_fig1.savefig("test_corner1.png") + # test_corner_fig2.savefig("test_corner2.png") + # test_corner_fig3.savefig("test_corner3.png") + # test_orbit_figs[0].savefig("test_orbit1.png") + # test_orbit_figs[1].savefig("test_orbit2.png") + # test_orbit_figs[2].savefig("test_orbit3.png") + # test_orbit_figs[3].savefig("test_orbit4.png") + # test_orbit_figs[4].savefig("test_orbit5.png") + + # # clean up + # os.system("rm test_*.png") diff --git a/tests/test_secondary_rvs.py b/tests/test_secondary_rvs.py index f5064734..fa8e5800 100644 --- a/tests/test_secondary_rvs.py +++ b/tests/test_secondary_rvs.py @@ -58,7 +58,7 @@ def test_secondary_rv_lnlike_calc(): os.system("rm tmp.csv") # assert that the secondary orbit is the primary orbit scaled - _, _, rv = mySys.compute_all_orbits(orbitize_params_list) + _, _, rv, _ = mySys.compute_all_orbits(orbitize_params_list) rv0 = rv[:, 0] rv1 = rv[:, 1] @@ -75,5 +75,5 @@ def test_read_input(): mySystem = system.System(1, input_data, 1, 1, fit_secondary_mass=False) if __name__ == "__main__": - # test_secondary_rv_lnlike_calc() + test_secondary_rv_lnlike_calc() test_read_input()