orbitize for morgan

orbitize/basis.py

@@ -491,6 +491,7 @@ def __init__(
         hipparcos_IAD=None,
         rv=False,
         rv_instruments=None,
+        astrometric_jitter=False
     ):
 
         super(Period, self).__init__(
@@ -505,6 +506,7 @@ def __init__(
             rv,
             rv_instruments,
         )
+        self.astrometric_jitter = astrometric_jitter
 
     def construct_priors(self):
         """
@@ -548,6 +550,11 @@ def construct_priors(self):
         # Add mass priors
         self.set_default_mass_priors(basis_priors, basis_labels)
 
+        # Add astrometric jitter priors
+        if self.astrometric_jitter:
+            basis_labels.append("sigma_ast")
+            basis_priors.append(priors.UniformPrior(0, 10))
+
         # 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))))

orbitize/hipparcos.py

@@ -164,6 +164,10 @@ class HipparcosLogProb(object):
             should be False, but it's helpful for testing. Check out
             `orbitize.hipparcos.nielsen_iad_refitting_test()` for an example
             using this renormalization.
+        include_hip_iad_in_likelihood (bool): if False, then don't add the Hipparcos
+            log(likelihood) to the overall log(likelihood computed in sampler.py)
+        brandt_correction (bool): if True, add delta_r = 0.140 mas and sigma_jit = 2.25 mas
+            to the residuals, following Brandt+ 2023.
 
     Written: Sarah Blunt & Rob de Rosa, 2021
     """
@@ -175,9 +179,12 @@ def __init__(
         num_secondary_bodies,
         alphadec0_epoch=1991.25,
         renormalize_errors=False,
+        include_hip_iad_in_likelihood=True,
+        brandt_correction=False,
     ):
         self.path_to_iad_file = path_to_iad_file
         self.renormalize_errors = renormalize_errors
+        self.include_hip_iad_in_likelihood = include_hip_iad_in_likelihood
 
         # infer if the IAD file is an older DVD file or a new file
         with open(path_to_iad_file, "r") as f:
@@ -259,7 +266,9 @@ def __init__(
             self.solution_type = solution_details["isol_n"].values[0]
 
             if self.solution_type == 1:
-                self.var = astrometric_solution["var"].values[0]
+                self.var = (
+                    10 * astrometric_solution["var"].values[0]
+                ) ** 2  # N.B. input is different units than var from Vizier catalog!!
             else:
                 self.var = 0
 
@@ -282,14 +291,16 @@ def __init__(
         else:
             iad = np.transpose(np.loadtxt(path_to_iad_file))
 
-        n_lines = len(iad)
-
         times = iad[1] + 1991.25
         self.cos_phi = iad[3]  # scan direction
         self.sin_phi = iad[4]
         self.R = iad[5]  # abscissa residual [mas]
+        if brandt_correction:
+            self.R += 0.140  # [mas]
         self.eps = iad[6]  # error on abscissa residual [mas]
 
+        n_lines = len(self.eps)
+
         # reject negative errors (scans that were rejected by Hipparcos team)
         good_scans = np.where(self.eps > 0)[0]
 
@@ -304,6 +315,9 @@ def __init__(
         # if the star has a type 1 (stochastic) solution, we need to undo the addition of a jitter term in quadrature
         self.eps = np.sqrt(self.eps**2 - self.var**2)
 
+        if brandt_correction:
+            self.eps = np.sqrt(self.eps**2 + 2.25**2)
+
         epochs = Time(times, format="decimalyear")
         self.epochs = epochs.decimalyear
         self.epochs_mjd = epochs.mjd
@@ -425,11 +439,15 @@ def compute_lnlike(self, raoff_model, deoff_model, samples, param_idx):
             )
 
         # compute chi2 (Nielsen+ 2020 Eq 7)
+        if "sigma_ast" in param_idx:
+            eps = np.sqrt(self.eps**2 + samples[param_idx["sigma_ast"]] ** 2)
+        else:
+            eps = self.eps
         chi2 = np.sum(
-            [(dist[:, i] / self.eps) ** 2 for i in np.arange(n_samples)],
+            [(dist[:, i] / eps) ** 2 for i in np.arange(n_samples)],
             axis=1,
         )
-        lnlike = -0.5 * chi2
+        lnlike = -0.5 * chi2 - np.sum(np.log(np.sqrt(2 * np.pi * eps)))
 
         return lnlike
 
@@ -513,51 +531,54 @@ def log_prob(model_pars):
             myHipLogProb.plx0 + 3 * myHipLogProb.plx0_err,
             1000,
         )
-        axes[0].hist(sampler.flatchain[:, 0], bins=50, density=True, color="r")
+        axes[0].hist(sampler.flatchain[:, 0], bins=50, density=True, color="grey")
         axes[0].plot(
             xs, norm(myHipLogProb.plx0, myHipLogProb.plx0_err).pdf(xs), color="k"
         )
-        axes[0].set_xlabel("plx [mas]")
+        axes[0].set_xlabel("$\pi$ [mas]")
 
         # PM RA
         xs = np.linspace(
             myHipLogProb.pm_ra0 - 3 * myHipLogProb.pm_ra0_err,
             myHipLogProb.pm_ra0 + 3 * myHipLogProb.pm_ra0_err,
             1000,
         )
-        axes[1].hist(sampler.flatchain[:, 1], bins=50, density=True, color="r")
+        axes[1].hist(sampler.flatchain[:, 1], bins=50, density=True, color="grey")
         axes[1].plot(
             xs, norm(myHipLogProb.pm_ra0, myHipLogProb.pm_ra0_err).pdf(xs), color="k"
         )
-        axes[1].set_xlabel("PM RA [mas/yr]")
+        axes[1].set_xlabel("$\mu_{{\\alpha_0^*}}$ [mas/yr]")
 
         # PM Dec
         xs = np.linspace(
             myHipLogProb.pm_dec0 - 3 * myHipLogProb.pm_dec0_err,
             myHipLogProb.pm_dec0 + 3 * myHipLogProb.pm_dec0_err,
             1000,
         )
-        axes[2].hist(sampler.flatchain[:, 2], bins=50, density=True, color="r")
+        axes[2].hist(sampler.flatchain[:, 2], bins=50, density=True, color="grey")
         axes[2].plot(
             xs, norm(myHipLogProb.pm_dec0, myHipLogProb.pm_dec0_err).pdf(xs), color="k"
         )
-        axes[2].set_xlabel("PM Dec [mas/yr]")
+        axes[2].set_xlabel("$\mu_{{\\delta_0}}$ [mas/yr]")
 
         # RA offset
-        axes[3].hist(sampler.flatchain[:, 3], bins=50, density=True, color="r")
+        axes[3].hist(sampler.flatchain[:, 3], bins=50, density=True, color="grey")
         xs = np.linspace(
             -3 * myHipLogProb.alpha0_err, 3 * myHipLogProb.alpha0_err, 1000
         )
         axes[3].plot(xs, norm(0, myHipLogProb.alpha0_err).pdf(xs), color="k")
-        axes[3].set_xlabel("RA Offset [mas]")
+        axes[3].set_xlabel("$\\alpha_0^*$ [mas]")
 
         # Dec offset
         xs = np.linspace(
             -3 * myHipLogProb.delta0_err, 3 * myHipLogProb.delta0_err, 1000
         )
-        axes[4].hist(sampler.flatchain[:, 4], bins=50, density=True, color="r")
+        axes[4].hist(sampler.flatchain[:, 4], bins=50, density=True, color="grey")
         axes[4].plot(xs, norm(0, myHipLogProb.delta0_err).pdf(xs), color="k")
-        axes[4].set_xlabel("Dec Offset [mas]")
+        axes[4].set_xlabel("$\delta_0$ [mas]")
+
+        for a in axes:
+            a.set_ylabel("relative prob.")
 
         plt.tight_layout()
         plt.savefig(saveplot, dpi=250)

orbitize/lnlike.py

@@ -44,9 +44,8 @@ def chi2_lnlike(
         in log scale, and defines pa chi-sq using complex phase representation.
             log sep chisq = (log sep - log sep_true)^2 / (sep_sigma / sep_true)^2
             pa chisq = 2 * (1 - cos(pa-pa_true))/pa_sigma^2
-i
-    """
 
+    """
     if np.ndim(model) == 3:
         # move M dimension to the primary axis, so that numpy knows to iterate over it
         model = np.rollaxis(model, 2, 0)  # now MxNobsx2 in dimensions

orbitize/plot.py

@@ -73,25 +73,25 @@ def plot_corner(results, param_list=None, **corner_kwargs):
     # Define array of default axis labels (overwritten if user specifies list)
     default_labels = {
 
-        "sma": "$a_{0}$ [au]",
-        "ecc": "$ecc_{0}$",
-        "inc": "$inc_{0}$ [$^\\circ$]",
-        "aop": "$\\omega_{0}$ [$^\\circ$]",
-        "pan": "$\\Omega_{0}$ [$^\\circ$]",
-        "tau": "$\\tau_{0}$",
+        "sma": "$a_{{B}}$ [au]",
+        "ecc": "$ecc$",
+        "inc": "$i$ [$^\\circ$]",
+        "aop": "$\\omega_{{B}}$ [$^\\circ$]",
+        "pan": "$\\Omega$ [$^\\circ$]",
+        "tau": "$\\tau$",
         "tp": "$T_{{\\mathrm{{P}}}}$",
         "plx": "$\\pi$ [mas]",
         "gam": "$\\gamma$ [km/s]",
         "sig": "$\\sigma$ [km/s]",
         "mtot": "$M_T$ [M$_{{\\odot}}$]",
-        "m0": "$M_0$ [M$_{{\\odot}}$]",
-        "m": "$M_{0}$ [M$_{{\\rm Jup}}$]",
-        "pm_ra": "$\\mu_{{\\alpha}}$ [mas/yr]",
+        "m0": "$M_{{a}}$ [M$_{{\\odot}}$]",
+        "m": "$M_{{b}}$ [M$_{{\\odot}}$]",
+        "pm_ra": "$\\mu_{{\\alpha^{{*}}}}$ [mas/yr]",
         "pm_dec": "$\\mu_{{\\delta}}$ [mas/yr]",
         "alpha0": "$\\alpha^{{*}}_{{0}}$ [mas]",
         "delta0": "$\\delta_0$ [mas]",
-        "m": "$M_{0}$ [M$_{{\\rm Jup}}$]",
-        "per": "$P_{0}$ [yr]",
+        # "m": "$M_{0}$ [M$_{{\\rm Jup}}$]",
+        "per": "$P$ [yr]",
         "K": "$K_{0}$ [km/s]",
         "x": "$X_{0}$ [AU]",
         "y": "$Y_{0}$ [AU]",
@@ -129,9 +129,9 @@ def plot_corner(results, param_list=None, **corner_kwargs):
     samples[:, angle_indices] = np.degrees(
         samples[:, angle_indices]
     )  # convert angles from rad to deg
-    samples[:, secondary_mass_indices] *= u.solMass.to(
-        u.jupiterMass
-    )  # convert to Jupiter masses for companions
+    # samples[:, secondary_mass_indices] *= u.solMass.to(
+    #     u.jupiterMass
+    # )  # convert to Jupiter masses for companions
 
     if (
         "labels" not in corner_kwargs

orbitize/sampler.py

@@ -89,6 +89,10 @@ def _logl(self, params):
         seppa_indices = self.system.all_seppa
 
         # compute lnlike
+        # NOTE: this won't work if we're fitting more than just astrometry. This is a hack for betelgeuse.
+        if "sigma_ast" in self.system.param_idx:
+            jitter = np.ones_like(model) * params[self.system.param_idx["sigma_ast"]]
+
         lnlikes = self.lnlike(
             data, errs, corrs, model, jitter, seppa_indices, chi2_type=self.chi2_type
         )
@@ -100,30 +104,31 @@ def _logl(self, params):
             lnlikes_sum += self.custom_lnlike(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(
-                params, epochs=self.system.hipparcos_IAD.epochs_mjd
-            )
+            if self.system.hipparcos_IAD.include_hip_iad_in_likelihood:
+                # compute Ra/Dec predictions at the Hipparcos IAD epochs
+                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
-            )
+                (
+                    raoff_model_hip_epoch,
+                    deoff_model_hip_epoch,
+                    _,
+                ) = self.system.compute_all_orbits(
+                    params, epochs=Time([1991.25], format="decimalyear").mjd
+                )
 
-            # subtract off position of star at reference Hipparcos epoch
-            raoff_model[:, 0, :] -= raoff_model_hip_epoch[:, 0, :]
-            deoff_model[:, 0, :] -= deoff_model_hip_epoch[:, 0, :]
+                # subtract off position of star at reference Hipparcos epoch
+                raoff_model[:, 0, :] -= raoff_model_hip_epoch[:, 0, :]
+                deoff_model[:, 0, :] -= deoff_model_hip_epoch[:, 0, :]
 
-            # select body 0 raoff/deoff predictions & feed into Hip IAD lnlike fn
-            lnlikes_sum += self.system.hipparcos_IAD.compute_lnlike(
-                raoff_model[:, 0, :],
-                deoff_model[:, 0, :],
-                params,
-                self.system.param_idx,
-            )
+                # select body 0 raoff/deoff predictions & feed into Hip IAD lnlike fn
+                lnlikes_sum += self.system.hipparcos_IAD.compute_lnlike(
+                    raoff_model[:, 0, :],
+                    deoff_model[:, 0, :],
+                    params,
+                    self.system.param_idx,
+                )
 
         if self.system.gaia is not None:
             gaiahip_epochs = Time(

orbitize/system.py

@@ -42,6 +42,9 @@ class System(object):
             basis to be used. See basis.py for a list of implemented fitting bases.
         use_rebound (bool): if True, use an n-body backend solver instead
             of a Keplerian solver.
+        astrometric_jitter (bool): if True, include a fitted jitter term in the 
+            model (will be applied to Hipparcos data and arbitrary absolute astrometry
+            data)
 
     Priors are initialized as a list of orbitize.priors.Prior objects and stored
     in the variable ``System.sys_priors``. You should initialize this class,
@@ -67,6 +70,7 @@ def __init__(
         gaia=None,
         fitting_basis="Standard",
         use_rebound=False,
+        astrometric_jitter=False,
     ):
         self.num_secondary_bodies = num_secondary_bodies
         self.data_table = data_table
@@ -81,6 +85,7 @@ def __init__(
         self.gaia = gaia
         self.fitting_basis = fitting_basis
         self.use_rebound = use_rebound
+        self.astrometric_jitter=astrometric_jitter
 
         self.best_epochs = []
         self.input_table = self.data_table.copy()
@@ -235,6 +240,7 @@ def __init__(
             hipparcos_IAD=self.hipparcos_IAD,
             rv=contains_rv,
             rv_instruments=self.rv_instruments,
+            astrometric_jitter=self.astrometric_jitter,
             **self.extra_basis_kwargs
         )
 
@@ -454,8 +460,11 @@ def compute_all_orbits(self, params_arr, epochs=None, comp_rebound=False):
                     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
+                    inside_masses = all_pl_masses[within_orbit].reshape((-1, n_orbits))
+                    if n_orbits == 1:
+                        mtot = np.sum(inside_masses) + m0
+                    else:
+                        mtot = np.sum(inside_masses, axis=0) + m0
 
                 else:
                     m_pl = np.zeros(self.num_secondary_bodies)