Merge pull request #85 from sblunt/fix-negative-st-mass

Do not include negative values in the Gaussian prior

orbitize/priors.py

@@ -38,19 +38,23 @@ class GaussianPrior(Prior):
     Args:
         mu (float): mean of the distribution
         sigma (float): standard deviation of the distribution
+        no_negatives (bool): if True, only positive values will be drawn from
+            this prior, and the probability of negative values will be 0 (default:True).
 
     (written) Sarah Blunt, 2018
     """
-    def __init__(self, mu, sigma):
+    def __init__(self, mu, sigma, no_negatives=True):
         self.mu = mu
         self.sigma = sigma
+        self.no_negatives = no_negatives
 
     def __repr__(self):
         return "Gaussian"
 
     def draw_samples(self, num_samples):
         """
-        Draw samples from a Gaussian distribution.
+        Draw positive samples from a Gaussian distribution.
+        Negative samples will not be returned. 
 
         Args:
             num_samples (float): the number of samples to generate
@@ -59,14 +63,29 @@ def draw_samples(self, num_samples):
             numpy array of float: samples drawn from the appropriate
             Gaussian distribution. Array has length `num_samples`. 
         """
+
         samples = np.random.normal(
             loc=self.mu, scale=self.sigma, size=num_samples
-            )
+        )        
+        bad = np.inf
+
+        if self.no_negatives:
+
+            while bad != 0:
+
+                bad_samples = np.where(samples < 0)[0]
+                bad = len(bad_samples)
+
+                samples[bad_samples] = np.random.normal(
+                    loc=self.mu, scale=self.sigma, size=bad
+                )   
+
         return samples
 
     def compute_lnprob(self, element_array):
         """
         Compute log(probability) of an array of numbers wrt a Gaussian distibution.
+        Negative numbers return a probability of -inf.
 
         Args:
             element_array (float or np.array of float): array of numbers. We want the 
@@ -80,6 +99,11 @@ def compute_lnprob(self, element_array):
         """
         lnprob = -0.5*np.log(2.*np.pi*self.sigma) - 0.5*((element_array - self.mu) / self.sigma)**2
 
+        if self.no_negatives:
+
+            bad_samples = np.where(element_array < 0)[0]
+            lnprob[bad_samples] = -np.inf
+
         return lnprob
 
 class JeffreysPrior(Prior):

orbitize/results.py

@@ -360,6 +360,7 @@ def plot_orbits(self, parallax=None, total_mass=None, object_mass=0,
 
         # Create a linearly increasing colormap for our range of epochs
         norm = mpl.colors.Normalize(vmin=np.min(epochs), vmax=np.max(epochs[-1,:]))
+
         norm_yr = mpl.colors.Normalize(
             vmin=np.min(Time(epochs,format='mjd').decimalyear), 
             vmax=np.max(Time(epochs,format='mjd').decimalyear)

tests/test_priors.py

@@ -7,31 +7,31 @@
 threshold = 1e-1
 
 initialization_inputs = {
-	priors.GaussianPrior : [0., 1.], 
+	priors.GaussianPrior : [1000., 1.], 
 	priors.JeffreysPrior : [1., 2.], 
 	priors.UniformPrior : [0., 1.], 
 	priors.SinPrior : [], 
 	priors.LinearPrior : [-2., 2.]
 }
 
 expected_means_mins_maxes = {
-	priors.GaussianPrior : (0.,-np.inf,np.inf), 
+	priors.GaussianPrior : (1000.,0.,np.inf), 
 	priors.JeffreysPrior : (1/np.log(2),1., 2.), 
 	priors.UniformPrior : (0.5, 0., 1.), 
 	priors.SinPrior : (np.pi/2., 0., np.pi), 
 	priors.LinearPrior : (1./3.,0.,1.0)
 }
 
 lnprob_inputs = {
-	priors.GaussianPrior : np.array([0., -np.inf, np.inf, 3., 3.]),
+	priors.GaussianPrior : np.array([-3.0, np.inf, 1000., 999.]),
 	priors.JeffreysPrior : np.array([-1., 0., 1., 1.5, 2., 2.5]),
 	priors.UniformPrior : np.array([0., 0.5, 1., -1., 2.]),
 	priors.SinPrior : np.array([0., np.pi/2., np.pi, 10., -1.]),
 	priors.LinearPrior : np.array([0., 0.5, 1., 2., -1.])
 }
 
 expected_probs = {
-	priors.GaussianPrior : np.array([nm(0,1).pdf(0), 0., 0., nm(0,1).pdf(3), nm(0,1).pdf(-3)]),
+	priors.GaussianPrior : np.array([0., 0., nm(1000.,1.).pdf(1000.), nm(1000.,1.).pdf(999.)]),
 	priors.JeffreysPrior : np.array([0., 0., 1., 2./3., 0.5, 0.])/np.log(2),
 	priors.UniformPrior : np.array([1., 1., 1., 0., 0.]),
 	priors.SinPrior : np.array([0., 0.5, 0., 0., 0.]),
@@ -65,6 +65,7 @@ def test_compute_lnprob():
 		values2test = lnprob_inputs[Prior]
 
 		lnprobs = TestPrior.compute_lnprob(values2test)
+
 		assert np.log(expected_probs[Prior]) == pytest.approx(lnprobs, abs=threshold)