Merge pull request #365 from sblunt/obs-priors

Obs priors

.gitignore

@@ -10,6 +10,7 @@ orbitize/example_data/*test.hdf5
 !orbitize/example_data/v1_posterior.hdf5
 *.fits
 *.png
+!docs/tutorials/eift_hd206893.png
 tests/test_results.h5
 tests/multiplanet*test.csv
 

orbitize/basis.py

@@ -51,6 +51,7 @@ def __init__(
             "pan": priors.UniformPrior(0.0, angle_upperlim),
             "tau": priors.UniformPrior(0.0, 1.0),
             "K": priors.LogUniformPrior(1e-4, 10),
+            "tp": priors.UniformPrior(0.0, 10000.0),
         }
 
     @abc.abstractmethod
@@ -167,6 +168,186 @@ def set_default_mass_priors(self, priors_arr, labels_arr):
             priors_arr.append(self.stellar_or_system_mass)
 
 
+class ObsPriors(Basis):
+    """
+    Basis used in Kosmo O'Neil+ 2019, and implemented here for use with the
+    orbitize.priors.ObsPriors prior, following that paper. The basis is the same
+    as the orbitize.basis.Period basis, except tp (time of periastron passage) is
+    used in place of tau.
+
+    Args:
+        stellar_or_system_mass (float): mass of the primary star (if fitting for
+            dynamical masses of both components) or total system mass (if
+            fitting using relative astrometry only) [M_sol]
+        mass_err (float): uncertainty on 'stellar_or_system_mass', in M_sol
+        plx (float): mean parallax of the system, in mas
+        plx_err (float): uncertainty on 'plx', in mas
+        num_secondary_bodies (int): number of secondary bodies in the system, should be at least 1
+        fit_secondary_mass (bool): if True, include the dynamical mass of orbitting body as fitted parameter, if False,
+            'stellar_or_system_mass' is taken to be total mass
+        angle_upperlim (float): either pi or 2pi, to restrict the prior range for 'pan' parameter (default: 2pi)
+        tau_ref_epoch (float): reference epoch for defining tau. Default 58849,
+            the same as it is elsewhere in the code.
+
+    Limiations:
+        This basis cannot be used with RVs or absolute astrometry. It assumes
+        inputs are vanilla Ra/Dec for relative astrometry.
+    """
+
+    def __init__(
+        self,
+        stellar_or_system_mass,
+        mass_err,
+        plx,
+        plx_err,
+        num_secondary_bodies,
+        fit_secondary_mass,
+        angle_upperlim=2 * np.pi,
+        hipparcos_IAD=None,
+        rv=False,
+        rv_instruments=None,
+        tau_ref_epoch=58849,
+    ):
+
+        self.tau_ref_epoch = tau_ref_epoch
+
+        if hipparcos_IAD is not None or rv or len(rv_instruments) > 0:
+            raise ValueError(
+                "This basis cannot be used with RVs or absolute astrometry."
+            )
+
+        super(ObsPriors, self).__init__(
+            stellar_or_system_mass,
+            mass_err,
+            plx,
+            plx_err,
+            num_secondary_bodies,
+            fit_secondary_mass,
+            angle_upperlim,
+            None,
+            False,
+            None,
+        )
+
+    def construct_priors(self):
+        """
+        Generates the parameter label array and initializes the corresponding priors for each
+        parameter to be sampled. For this basis, the parameters common to each
+        companion are: per, ecc, inc, aop, pan, Tp. Parallax and mass priors both
+        assumed to be fixed, and are added at the end.
+
+        Returns:
+            tuple:
+
+                list: list of strings (labels) that indicate the names of each parameter to sample
+
+                list: list of orbitize.priors.Prior objects that indicate the prior distribution of each label
+        """
+        base_labels = ["per", "ecc", "inc", "aop", "pan", "tp"]
+        basis_priors = []
+        basis_labels = []
+
+        # Add the priors common to each companion
+        for body in np.arange(self.num_secondary_bodies):
+            for elem in base_labels:
+                basis_priors.append(self.default_priors[elem])
+                basis_labels.append(elem + str(body + 1))
+
+        # Add parallax prior
+        basis_labels.append("plx")
+        if self.plx_err > 0:
+            raise ValueError("Parallax must be fixed for this prior type.")
+        else:
+            basis_priors.append(self.plx)
+
+        # Add mass priors
+        basis_labels.append("mtot")
+        if self.mass_err > 0:
+            raise ValueError("Mtot must be fixed for this prior type.")
+        else:
+            basis_priors.append(self.stellar_or_system_mass)
+
+        # Define param label dictionary in current basis & standard basis
+        self.param_idx = dict(zip(basis_labels, np.arange(len(basis_labels))))
+        self.standard_basis_idx = dict(zip(basis_labels, np.arange(len(basis_labels))))
+
+        for body_num in np.arange(self.num_secondary_bodies) + 1:
+            self.standard_basis_idx["sma{}".format(body_num)] = self.param_idx[
+                "per{}".format(body_num)
+            ]
+            self.standard_basis_idx["tau{}".format(body_num)] = self.param_idx[
+                "tp{}".format(body_num)
+            ]
+
+        return basis_priors, basis_labels
+
+    def to_standard_basis(self, param_arr):
+        """
+        Convert parameter array from ObsPriors basis to Standard basis.
+
+        Args:
+            param_arr (np.array of float): RxM array of fitting parameters in the ObsPriors basis,
+                where R is the number of parameters being fit, and M is the number of orbits. If
+                M=1 (for MCMC), this can be a 1D array.
+
+        Returns:
+            np.array of float: modifies 'param_arr' to contain Standard basis elements.
+                Shape of 'param_arr' remains the same.
+        """
+        for body_num in np.arange(self.num_secondary_bodies) + 1:
+            per = param_arr[self.param_idx["per{}".format(body_num)]]
+
+            mtot = param_arr[self.param_idx["mtot"]]
+
+            # Compute semi-major axis using Kepler's Third Law and replace period
+            sma = np.cbrt(
+                (consts.G * (mtot * u.Msun) * ((per * u.year) ** 2)) / (4 * np.pi**2)
+            )
+            sma = sma.to(u.AU).value
+            param_arr[self.standard_basis_idx["sma{}".format(body_num)]] = sma
+
+            tp = param_arr[self.param_idx["tp{}".format(body_num)]]
+
+            tau = tp_to_tau(tp, self.tau_ref_epoch, per)
+
+            param_arr[self.standard_basis_idx["tau{}".format(body_num)]] = tau
+
+        return param_arr
+
+    def to_obspriors_basis(self, param_arr, after_date):
+        """
+        Convert parameter array from Standard basis to ObsPriors basis. This function
+        is used primarily for testing purposes.
+
+        Args:
+            param_arr (np.array of float): RxM array of fitting parameters in the Standard basis,
+                where R is the number of parameters being fit, and M is the number of orbits. If
+                M=1 (for MCMC), this can be a 1D array.
+            after_date (float or np.array): tp will be the first periastron after this date. If None, use ref_epoch.
+
+        Returns:
+            np.array of float: modifies 'param_arr' to contain ObsPriors elements.
+                Shape of 'param_arr' remains the same.
+        """
+        for body_num in np.arange(self.num_secondary_bodies) + 1:
+            sma = param_arr[self.standard_basis_idx["sma{}".format(body_num)]]
+            mtot = param_arr[self.standard_basis_idx["mtot"]]
+
+            per = np.sqrt(
+                (4 * (np.pi**2) * (sma * u.AU) ** 3) / (consts.G * (mtot * u.Msun))
+            )
+            per = per.to(u.year).value
+            param_arr[self.param_idx["per{}".format(body_num)]] = per
+
+            tau = param_arr[self.standard_basis_idx["tau{}".format(body_num)]]
+
+            tp = tau_to_tp(tau, self.tau_ref_epoch, per, after_date=after_date)
+
+            param_arr[self.standard_basis_idx["tp{}".format(body_num)]] = tp
+
+        return param_arr
+
+
 class Standard(Basis):
     """
     Standard basis set based upon the 6 standard Keplarian elements: (sma, ecc, inc, aop, pan, tau).

orbitize/example_data/hd206893b.csv

@@ -0,0 +1,10 @@
+epoch,object,raoff,decoff,raoff_err,decoff_err
+57298,1,253.72,92.35,2.98,2.85
+57606,1,236.63,127.94,9.77,9.18
+57645,1,234.52,123.39,1.79,1.03
+57946,1,210.76,152.09,1.94,1.88
+58276,1,167.49,180.87,1.61,16.97
+58287,1,177.67,174.6,1.67,1.67
+58365,1,165.7,185.33,3.28,3.66
+58368,1,170.38,185.94,2.52,2.74
+58414,1,161.64,176.21,13.6,14.31

orbitize/plot.py

@@ -78,6 +78,7 @@ def plot_corner(results, param_list=None, **corner_kwargs):
         "aop": "$\\omega_{0}$ [$^\\circ$]",
         "pan": "$\\Omega_{0}$ [$^\\circ$]",
         "tau": "$\\tau_{0}$",
+        "tp": "$T_{{\\mathrm{{P}}}}$",
         "plx": "$\\pi$ [mas]",
         "gam": "$\\gamma$ [km/s]",
         "sig": "$\\sigma$ [km/s]",
@@ -175,9 +176,10 @@ def plot_orbits(
     cbar_param="Epoch [year]",
     mod180=False,
     rv_time_series=False,
-    rv_time_series2=False,
     plot_astrometry=True,
     plot_astrometry_insts=False,
+    plot_errorbars=True,
+    rv_time_series2=False,
     primary_instrument_name=None,
     fontsize=20,
     fig=None,
@@ -240,6 +242,16 @@ def plot_orbits(
             "Plotting the primary's orbit is currently unsupported. Stay tuned."
         )
 
+    if rv_time_series and "m0" not in results.labels:
+        rv_time_series = False
+
+        warnings.warn(
+            "It seems that the stellar and companion mass "
+            "have not been fitted separately. Setting "
+            "rv_time_series=True is therefore not possible "
+            "so the argument is set to False instead."
+        )
+
     with warnings.catch_warnings():
         warnings.simplefilter("ignore", ErfaWarning)