Merge pull request #161 from sblunt/multiplanet

Multiplanet support

MANIFEST.in

@@ -1,3 +1,4 @@
 global-include orbitize/*.pyx
 global-include orbitize/*.c
-include requirements.txt
+include requirements.txt
+include orbitize/example_data/*

docs/formatting_inputs.rst

@@ -24,8 +24,6 @@ You may mix and match these three valid measurement formats in the same input fi
 If you have, for example, one RV measurement of a star and three astrometric
 measurements of an orbiting planet, you should put ``0`` in the ``object`` column for the RV point, and ``1`` in the columns for the astrometric measurements.
 
-.. Warning:: For now, ``orbitize`` only accepts astrometric measurements    for one secondary body. In a future release, it will also handle            astrometric measurements for multiple secondaries, RV measurements of the primary and secondar(ies), and astrometric measurements of the primary. Stay tuned!
-
 This method will look for columns with the above labels in whatever file format you choose so if you encounter errors, be sure to double check the column labels in your input file.
 
 Putting it all together, here an example of a valid .csv input file::

docs/tutorials.rst

@@ -40,7 +40,8 @@ us if you are still confused).
 .. toctree::
    :maxdepth: 1
 
-   tutorials/Orbitize_RV_Tutorial.ipynb
+   tutorials/RV_MCMC_Tutorial.ipynb
+   tutorials/Multiplanet_Tutorial.ipynb
    tutorials/Modifying_Priors.ipynb
    tutorials/Plotting_tutorial.ipynb
    tutorials/MCMC_vs_OFTI.ipynb

docs/tutorials/RV_MCMC_Tutorial.ipynb

@@ -174,7 +174,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# RV Priors\n",
+    "## RV Priors\n",
     "\n",
     "The priors for the two RV parameters, the radial velocity offset (gamma), and jitter (sigma), have default uniform prior and log uniform prior respectively. The gamma uniform prior is set between $(-5,5)$ $\\mathrm{km/s}$, and the jitter log uniform prior is set for ($10^{-4},0.05$) $\\mathrm{km/s}$. The prior for `m1` is a log uniform prior and is set for ($10^{-3},2.0$)$M_\\odot$. The current version of orbitize addressed in this tutorial returns the stellar radial velocity only. \n",
     "\n",
@@ -331,7 +331,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Saving Results over Extended MCMC Run\n",
+    "## Saving Results over Extended MCMC Run\n",
     "\n",
     "Sometimes our MCMC run will need to run for an extended period of time to let the walkers converge. To observe the convergence, we often need to see the walkers' progress along parameter space. We can save the sampler results periodically and keep running the sampler until convergence. To run for a greater number of steps and periodically save the results, we can create a for-loop and run for as many iterations as we'd like.  "
    ]
@@ -455,7 +455,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Plotting and Accesing Saved Results"
+    "## Plotting and Accesing Saved Results"
    ]
   },
   {
@@ -644,4 +644,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
+}

orbitize/__init__.py

@@ -5,7 +5,8 @@
 __version__ = '1.10.0'
 
 # set Python env variable to keep track of example data dir
-DATADIR = os.path.join(sys.prefix, 'orbitize_example_data')
+orbitize_dir = os.path.dirname(__file__)
+DATADIR = os.path.join(orbitize_dir, 'example_data/')
 
 # define functions for pickling class methods
 def _pickle_method(method):

orbitize/basis.py

@@ -62,4 +62,28 @@ def switch_tau_epoch(old_tau, old_epoch, new_epoch, period):
     t0 = tau_to_t0(old_tau, old_epoch, period)
     new_tau = t0_to_tau(t0, new_epoch, period)
 
-    return new_tau
+    return new_tau
+
+def tau_to_manom(date, sma, mtot, tau, tau_ref_epoch):
+    """
+    Gets the mean anomlay
+    
+    Args:
+        date (float or np.array): MJD
+        sma (float): semi major axis (AU)
+        mtot (float): total mass (M_sun)
+        tau (float): epoch of periastron, in units of the orbital period
+        tau_ref_epoch (float): reference epoch for tau
+        
+    Returns:
+        mean_anom (float or np.array): mean anomaly on that date [0, 2pi)
+    """
+    period = sma**(1.5)/np.sqrt(mtot) # years
+
+    frac_date = (date - tau_ref_epoch)/365.25/period
+    frac_date %= 1
+
+    mean_anom = (frac_date - tau) * 2 * np.pi
+    mean_anom %= 2 * np.pi
+
+    return mean_anom

tests/GJ504.csv → orbitize/example_data/GJ504.csv

@@ -1,10 +1,10 @@
-# Table 12 of Blunt et al 2017
-# Previous Observations of GJ 504 b
-epoch,object,sep,sep_err,pa,pa_err
-55645.95,1,2479,16,327.94,0.39
-55702.89,1,2483,8,327.45,0.19
-55785.015,1,2481,33,326.84,0.94
-55787.935,1,2448,24,325.82,0.66
-55985.19400184,1,2483,15,326.46,0.36
-56029.11400323,1,2487,8,326.54,0.18
-56072.30200459,1,2499,26,326.14,0.61
+# Table 12 of Blunt et al 2017
+# Previous Observations of GJ 504 b
+epoch,object,sep,sep_err,pa,pa_err
+55645.95,1,2479,16,327.94,0.39
+55702.89,1,2483,8,327.45,0.19
+55785.015,1,2481,33,326.84,0.94
+55787.935,1,2448,24,325.82,0.66
+55985.19400184,1,2483,15,326.46,0.36
+56029.11400323,1,2487,8,326.54,0.18
+56072.30200459,1,2499,26,326.14,0.61

tests/GJ504_1epoch.csv → orbitize/example_data/GJ504_1epoch.csv

@@ -1,4 +1,4 @@
-# Table 12 of Blunt et al 2017
-# One-epoch Observation of GJ 504 b
-epoch,object,sep,sep_err,pa,pa_err
-55645.95,1,2479,16,327.94,0.39
+# Table 12 of Blunt et al 2017
+# One-epoch Observation of GJ 504 b
+epoch,object,sep,sep_err,pa,pa_err
+55645.95,1,2479,16,327.94,0.39

tests/test_val.csv → orbitize/example_data/test_val.csv

@@ -1,5 +1,5 @@
-epoch,object,raoff,raoff_err,decoff,decoff_err,sep,sep_err,pa,pa_err,rv,rv_err
-1234,1,0.01,0.005,0.5,0.05,,,,,,
-1235,1,,,,,1,0.005,89,0.1,,
-1236,1,,,,,1,0.005,89.3,0.3,,
+epoch,object,raoff,raoff_err,decoff,decoff_err,sep,sep_err,pa,pa_err,rv,rv_err
+1234,1,0.01,0.005,0.5,0.05,,,,,,
+1235,1,,,,,1,0.005,89,0.1,,
+1236,1,,,,,1,0.005,89.3,0.3,,
 1237,0,,,,,,,,,10,0.1

orbitize/example_data/test_val_multi.csv

@@ -0,0 +1,5 @@
+epoch,object,raoff,raoff_err,decoff,decoff_err,sep,sep_err,pa,pa_err,rv,rv_err
+53200,1,1471,6,887,6,,,,,,
+56226,1,1549,4,743,4,,,,,,
+56130,2,-578,5,761,5,,,,,,
+56226,2,-572,3,768,3,,,,,,

tests/test_val_rv.csv → orbitize/example_data/test_val_rv.csv

@@ -1,8 +1,8 @@
-epoch,object,raoff,raoff_err,decoff,decoff_err,sep,sep_err,pa,pa_err,rv,rv_err
-10,1,0.01,0.005,0.5,0.05,,,,,,
-65,1,,,,,1,0.005,89,0.1,,
-186,1,,,,,1,0.005,89.3,0.3,,
-423,0,,,,,,,,,10,0.1
-500,0,,,,,,,,,100,1
-754,0,,,,,,,,,15,1.5
+epoch,object,raoff,raoff_err,decoff,decoff_err,sep,sep_err,pa,pa_err,rv,rv_err
+10,1,0.01,0.005,0.5,0.05,,,,,,
+65,1,,,,,1,0.005,89,0.1,,
+186,1,,,,,1,0.005,89.3,0.3,,
+423,0,,,,,,,,,10,0.1
+500,0,,,,,,,,,100,1
+754,0,,,,,,,,,15,1.5
 1235,0,,,,,,,,,12,0.015

orbitize/read_input.py

@@ -45,11 +45,6 @@ def read_file(filename):
         taken at the same epoch), ``read_file()`` will generate a separate output row for
         each valid set.
 
-    .. Warning:: For now, ``orbitize`` only accepts astrometric measurements for one
-        secondary body. In a future release, it will also handle astrometric measurements for
-        multiple secondaries, RV measurements of the primary and secondar(ies), and astrometric
-        measurements of the primary. Stay tuned!
-
     Alternatively, you can also supply a data file with the columns already corresponding to
     the orbitize format (see the example in description of what this method returns). This may
     be useful if you are wanting to use the output of the `write_orbitize_input` method.

orbitize/results.py

@@ -41,6 +41,7 @@ class Results(object):
             the orbits described in ``post`` (default: None).
         tau_ref_epoch (float): date (in days, typically MJD) that tau is defined relative to
         labels (list of str): parameter labels in same order as `post`
+        num_secondary_bodies (int): number of companions fit 
 
     The ``post`` array is in the following order::
 
@@ -56,13 +57,15 @@ class Results(object):
     Written: Henry Ngo, Sarah Blunt, 2018
     """
 
-    def __init__(self, sampler_name=None, post=None, lnlike=None, tau_ref_epoch=None, labels=None):
+    def __init__(self, sampler_name=None, post=None, lnlike=None, tau_ref_epoch=None, labels=None,
+                 num_secondary_bodies=None):
 
         self.sampler_name = sampler_name
         self.post = post
         self.lnlike = lnlike
         self.tau_ref_epoch = tau_ref_epoch
         self.labels = labels
+        self.num_secondary_bodies=num_secondary_bodies
 
     def add_samples(self, orbital_params, lnlikes, labels):
         """
@@ -117,6 +120,9 @@ def save_results(self, filename):
         if self.labels is not None:
             hf['col_names'] = np.array(self.labels).astype('S')
         hf.attrs['parameter_labels'] = self.labels  # Rob: added this to account for the RV labels
+        if self.num_secondary_bodies is not None:
+            hf.attrs['num_secondary_bodies'] = self.num_secondary_bodies
+
         hf.close()  # Closes file object, which writes file to disk
 
     def load_results(self, filename, append=False):
@@ -149,7 +155,13 @@ def load_results(self, filename, append=False):
             labels = hf.attrs['parameter_labels']
         except KeyError:
             # again, probably an old file without saved parameter labels
+            # old files only fit single planets
             labels = ['sma1', 'ecc1', 'inc1', 'aop1', 'pan1', 'tau1', 'plx', 'mtot']
+        try:
+            num_secondary_bodies = int(hf.attrs['num_secondary_bodies'])
+        except KeyError:
+            # old, has to be single planet fit
+            num_secondary_bodies = 1
 
         hf.close()  # Closes file object
 
@@ -174,6 +186,11 @@ def load_results(self, filename, append=False):
             elif self.labels.any() != labels.any():
                 raise ValueError("Loaded data has parameter labels {} while Results object has already been initialized to {}.".format(
                     labels, self.labels))
+            if self.num_secondary_bodies == 0:
+                self.num_secondary_bodies = num_secondary_bodies
+            elif self.num_secondary_bodies != num_secondary_bodies:
+                raise ValueError("Loaded data has {} number of secondary bodies while Results object has already been initialized to {}.".format(
+                    num_secondary_bodies, self.num_secondary_bodies))
 
             # Now append post and lnlike
             self.add_samples(post, lnlike, self.labels)
@@ -184,6 +201,7 @@ def load_results(self, filename, append=False):
                 self.add_samples(post, lnlike, self.labels)
                 self.tau_ref_epoch = tau_ref_epoch
                 self.labels = labels
+                self.num_secondary_bodies = num_secondary_bodies
             else:
                 raise Exception(
                     'Unable to load file {} to Results object. append is set to False but object is not empty'.format(filename))
@@ -227,42 +245,51 @@ def plot_corner(self, param_list=None, **corner_kwargs):
         default_labels = {
             'sma': 'a [au]',
             'ecc': 'ecc',
-            'inc': 'inc [$^{\\circ}$]',
-            'aop': '$\\omega$ [$^{\\circ}$]',
-            'pan': '$\\Omega$ [$^{\\circ}$]',
+            'inc': 'inc [$^\\circ$]',
+            'aop': '$\\omega$ [$^\\circ$]',
+            'pan': '$\\Omega$ [$^\\circ$]',
             'tau': '$\\tau$',
             'plx': '$\\pi$ [mas]',
             'gam': '$\\gamma$ [m/s]',
             'sig': '$\\sigma$ [m/s]',
-            'mtot': '$M_T$ [M$_{\\odot}$]',
-            'm0': '$M_0$ [M$_{\\odot}$]',
-            'm1': '$M_1$ [M$_{\\odot}$]',
+            'mtot': '$M_T$ [M$_\\odot$]',
+            'm0': '$M_0$ [M$_\\odot$]',
+            'm': '$M_{0}$ [M$_\{Jup\}$]',
         }
 
         if param_list is None:
             param_list = self.labels
         param_indices = []
         angle_indices = []
+        secondary_mass_indices = []
         for i, param in enumerate(param_list):
             index_num = np.where(np.array(self.labels) == param)[0][0]
             param_indices.append(index_num)
-            label_key = param_list[i]
+            label_key = param
             if label_key.startswith('aop') or label_key.startswith('pan') or label_key.startswith('inc'):
-                angle_indices.append(index_num)
+                angle_indices.append(i)
+            if label_key.startswith('m') and label_key != 'm0' and label_key != 'mtot':
+                secondary_mass_indices.append(i)
 
         samples = copy.copy(self.post[:, param_indices])  # keep only chains for selected parameters
         samples[:, angle_indices] = np.degrees(
             self.post[:, angle_indices])  # convert angles from rad to deg
+        samples[:, secondary_mass_indices] *= u.solMass.to(u.jupiterMass) # convert to Jupiter masses for companions
 
         if 'labels' not in corner_kwargs:  # use default labels if user didn't already supply them
             reduced_labels_list = []
             for i in np.arange(len(param_indices)):
                 label_key = param_list[i]
-                if label_key.startswith('m') or label_key.startswith('plx'):
-                    pass
+                if label_key.startswith("m") and label_key != 'm0' and label_key != 'mtot':
+                    body_num = label_key[1]
+                    label_key = "m"
+                elif label_key == 'm0' or label_key == 'mtot' or label_key.startswith('plx'):
+                    body_num = ""
+                    # maintain original label key
                 else:
+                    body_num = label_key[3]
                     label_key = label_key[0:3]
-                reduced_labels_list.append(default_labels[label_key])
+                reduced_labels_list.append(default_labels[label_key].format(body_num))
             corner_kwargs['labels'] = reduced_labels_list
 
         figure = corner.corner(samples, **corner_kwargs)
@@ -310,7 +337,13 @@ def plot_orbits(self, object_to_plot=1, start_mjd=51544.,
         """
 
         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.')
+            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)
@@ -322,7 +355,7 @@ def plot_orbits(self, object_to_plot=1, start_mjd=51544.,
                 'aop': 3,
                 'pan': 4,
                 'tau': 5,
-                'plx': 6,
+                'plx': 6 * self.num_secondary_bodies,
             }
 
             if cbar_param == 'epochs':
@@ -338,29 +371,27 @@ def plot_orbits(self, object_to_plot=1, start_mjd=51544.,
                 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']]
+
+            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[:, -2]
+                m1 = self.post[:, -(self.num_secondary_bodies+1) + (object_to_plot-1)]
                 mtot = m0 + m1
             if 'gamma' in self.labels:
-                dict_of_indices['gamma'] = 7
-                dict_of_indices['sigma'] = 8
+                dict_of_indices['gamma'] = 6 * self.num_secondary_bodies + 1
+                dict_of_indices['sigma'] = 6 * self.num_secondary_bodies + 2
                 gamma = self.post[:, dict_of_indices['gamma']]
 
             # Select random indices for plotted orbit