Add brigtness calc to comp all orbits

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,,

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))
 

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
 

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
             )
 

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

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))
     )