Merge pull request #245 from sblunt/ofti-pan-priors-correctbase

account for user-defined priors on PAN

README.md

@@ -7,10 +7,8 @@ Orbit-fitting for directly imaged objects. For installation instructions, tutori
 [![Build Status](https://travis-ci.com/sblunt/orbitize.svg?branch=main)](https://travis-ci.com/sblunt/orbitize)
 [![Coverage Status](https://coveralls.io/repos/github/sblunt/orbitize/badge.svg?branch=master&service=github)](https://coveralls.io/github/sblunt/orbitize?branch=master&service=github)
 [![Documentation Status](https://readthedocs.org/projects/orbitize/badge/?version=latest)](http://orbitize.readthedocs.io/en/latest/?badge=latest)
-[![Requirements Status](https://requires.io/github/sblunt/orbitize/requirements.svg?branch=main)](https://requires.io/github/sblunt/orbitize/requirements/?branch=main)
 [![PyPI version](https://badge.fury.io/py/orbitize.svg)](https://badge.fury.io/py/orbitize)
 [![PyPI downloads](https://img.shields.io/pypi/dm/orbitize.svg)](https://pypistats.org/packages/orbitize)
 
-
 [![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.3337378.svg)](https://zenodo.org/record/3337378#.XUHT3ZNKjUJ)
 [![License](https://img.shields.io/badge/License-BSD%203--Clause-blue.svg)](https://opensource.org/licenses/BSD-3-Clause)

orbitize/sampler.py

@@ -314,6 +314,27 @@ def reject(self, samples):
             corrs = None
         lnp_scaled = lnp - orbitize.lnlike.chi2_norm_term(errs, corrs)
 
+        # account for user-set priors on PAN that were destroyed by scale-and-rotate
+        pan_prior = self.system.sys_priors[
+            self.system.param_idx['pan1']
+        ]
+        if pan_prior is not orbitize.priors.UniformPrior:
+
+            # apply PAN prior
+            lnp_scaled += pan_prior.compute_lnprob(samples[4,:])
+
+        # prior is uniform but with different bounds that OFTI expects
+        elif (pan_prior.minval != 0) or (
+            (pan_prior.maxval != np.pi) or (pan_prior.maxval != 2*np.pi)
+        ):
+
+            samples_outside_pan_prior = np.where(
+                (samples[4,:] < pan_prior.minval) or 
+                (samples[4,:] > pan_prior.maxval)
+            )[0]
+
+            lnp_scaled[samples_outside_pan_prior] = -np.inf
+
         # reject orbits with probability less than a uniform random number
         random_samples = np.log(np.random.random(len(lnp)))
         saved_orbit_idx = np.where(lnp_scaled > random_samples)[0]

tests/test_OFTI.py

@@ -250,10 +250,42 @@ def test_OFTI_covariances():
     # test against seppa fits to see they are similar
     assert sma_seppa == pytest.approx(sma, abs=0.2 * sma_seppa)
 
+def test_OFTI_pan_priors():
+
+    # initialize sampler
+    myDriver = orbitize.driver.Driver(
+        input_file, 'OFTI', 1, 1.22, 56.95, mass_err=0.08, plx_err=0.26)
+
+    s = myDriver.sampler
+
+    # change PAN prior
+    new_min = 0.05
+    new_max = np.pi - 0.05
+    myDriver.system.sys_priors[4] = priors.UniformPrior(new_min, new_max)
+
+    # run sampler
+    orbits = s.run_sampler(100)
+
+    # check that bounds were applied correctly
+    assert np.max(orbits[:,4]) < new_max
+    assert np.min(orbits[:,4]) > new_min
+
+    # change PAN prior again
+    mu = np.pi / 2
+    sigma = 0.05
+    myDriver.system.sys_priors[4] = priors.GaussianPrior(mu, sigma = sigma)
+
+    # run sampler again
+    orbits = s.run_sampler(100)
+
+    # check that bounds were applied correctly
+    assert mu == pytest.approx(np.mean(orbits[:,4]), abs=0.01) 
+    assert sigma == pytest.approx(np.std(orbits[:,4]), abs=0.01)
 
 if __name__ == "__main__":
-    #test_scale_and_rotate()
-    test_run_sampler()
-    test_OFTI_covariances()
+
+    # test_scale_and_rotate()
+    # test_run_sampler()
     # test_OFTI_multiplanet()
-    # print("Done!")
+    test_OFTI_pan_priors()
+    print("Done!")