Naut combined

orbitize/plot.py

@@ -28,8 +28,7 @@
     cmap(np.linspace(0.0, 0.7, 1000)),
 )
 
-
-def plot_corner(results, param_list=None, **corner_kwargs):
+def plot_corner(results, param_list=None, downsample=None, **corner_kwargs):
     """
     Make a corner plot of posterior on orbit fit from any sampler
 
@@ -55,6 +54,9 @@ def plot_corner(results, param_list=None, **corner_kwargs):
                 sigma: rv jitter
                 mi: mass of individual body i, for i = 0, 1, 2, ... (only if fit_secondary_mass == True)
                 mtot: total mass (only if fit_secondary_mass == False)
+
+        downsample (int):
+            amount of samples to randomly draw from the posterior using ``results.downsample``
 
         **corner_kwargs: any remaining keyword args are sent to ``corner.corner``.
                             See `here <https://corner.readthedocs.io/>`_.
@@ -125,8 +127,15 @@ def plot_corner(results, param_list=None, **corner_kwargs):
         else:
             fixed_indices.append(i)
 
+    if downsample is not None:
+        post, _ = results.downsample(downsample)
+        weights = None
+    else:
+        post = results.weighted_post
+        weights = results.weights
+
     samples = np.copy(
-        results.post[:, param_indices]
+        post[:, param_indices]
     )  # keep only chains for selected parameters
     samples[:, angle_indices] = np.degrees(
         samples[:, angle_indices]
@@ -161,7 +170,7 @@ def plot_corner(results, param_list=None, **corner_kwargs):
 
         corner_kwargs["labels"] = reduced_labels_list
 
-    figure = corner.corner(samples, **corner_kwargs)
+    figure = corner.corner(samples, weights=weights, **corner_kwargs)
     return figure
 
 

orbitize/results.py

@@ -9,7 +9,6 @@
 import orbitize.plot
 import orbitize.gaia, orbitize.hipparcos
 
-
 class Results(object):
     """
     A class to store accepted orbital configurations from the sampler
@@ -28,6 +27,15 @@ class Results(object):
         data (astropy.table.Table): output from ``orbitize.read_input.read_file()``
         curr_pos (np.array of float): for MCMC only. A multi-D array of the 
             current walker positions that is used for restarting a MCMC sampler. 
+        weighted_post (np.array of float): RxN array of orbital parameters
+            (posterior output from weighted orbit-fitting process), where R is the
+            number of orbits generated, and N is the number of varying orbital
+            parameters in the fit (default: None).
+        weighted_lnlike (np.array of float): R array of log-likelihoods corresponding to
+            the orbits described in ``weighted_post`` (default: None).
+        lnweights (np.array of float): R array of log-weights corresponding to the orbits
+            and log-likelihoods described in ``weighted_post`` and ``weighted_lnlike`` (default: None)
+
 
     Written: Henry Ngo, Sarah Blunt, 2018
 
@@ -36,7 +44,7 @@ class Results(object):
 
     def __init__(
         self, system=None, sampler_name=None, post=None, lnlike=None,
-        version_number=None, curr_pos=None
+        version_number=None, curr_pos=None, weighted_post=None, weighted_lnlike=None, lnweight=None
     ):
 
         self.system = system
@@ -45,6 +53,9 @@ def __init__(
         self.lnlike = lnlike
         self.curr_pos = curr_pos
         self.version_number = version_number
+        self._weighted_post = weighted_post
+        self._weighted_lnlike = weighted_lnlike
+        self.lnweight = lnweight
         self.ln_evidence = None
         self.ln_evidence_err = None
 
@@ -58,7 +69,60 @@ def __init__(
             self.param_idx = self.system.param_idx
             self.standard_param_idx = self.system.basis.standard_basis_idx
 
-    def add_samples(self, orbital_params, lnlikes, curr_pos=None): 
+    @property
+    def weighted_post(self):
+        """
+        Returns the weighted posterior if it exists,
+        otherwise returns the unweighted posterior.
+        """
+        if self._weighted_post is not None:
+            return self._weighted_post
+        return self.post
+
+    @property
+    def weighted_lnlike(self):
+        """
+        Returns the weighted log-likelihoods if it exists,
+        otherwise returns the unweighted log-likelihoods.
+        """
+        if self._weighted_lnlike is not None:
+            return self._weighted_lnlike
+        return self.lnlike
+
+    @property
+    def weights(self):
+        """
+        Returns the weights of ``weighted_post`` and ``weighted_lnlike``
+        if it exists, otherwise returns None.
+        """
+        if self.lnweight is None:
+            return None
+        return np.exp(self.lnweight)
+
+    def downsample(self, amount, duplicates=True):
+        """
+        Samples from the posterior, or the weighted posterior if it exists.
+
+        Args:
+            amount (int): number of samples to draw from the posetrior
+            duplicates (bool): whether to replace sampled orbits from the posterior,
+                which may cause duplicate samples (default: True)
+
+        Returns:
+            tuple:
+
+                numpy.array of float: orbital parameters of the samples (``amount``xN, where N is number of varying orbital parameters)
+
+                numpy.array of float: log-likelihoods of the samples (length ``amount``)
+
+        """
+
+        indexes = np.random.choice(len(self.weighted_post), amount, duplicates, self.weights)
+        if self.weighted_lnlike is None:
+            return self.weighted_post[indexes], None
+        return self.weighted_post[indexes], self.weighted_lnlike[indexes]
+
+    def add_samples(self, orbital_params, lnlikes, curr_pos=None, weighted_post=None, weighted_lnlike=None, lnweight=None): 
         """
         Add accepted orbits, their likelihoods, and the orbitize version number 
         to the results
@@ -69,6 +133,12 @@ def add_samples(self, orbital_params, lnlikes, curr_pos=None):
             lnlike (np.array): add corresponding lnlike values to results
             curr_pos (np.array of float): for MCMC only. A multi-D array of the 
                 current walker positions
+            weighted_post (np.array): for Nautilus only. Sets of orbital params
+                associated with ``lnweight`` to add to results.
+            weighted_lnlike (np.array): for Nautilus only. Corresponding weighted lnlike
+                values for ``weighted_post`` to add to results.
+            weighted_post (np.array): for Nautilus only. Log of weights associated
+                with ``weighted_post`` and ``weighted_lnlike``.
 
         Written: Henry Ngo, 2018
 
@@ -88,6 +158,16 @@ def add_samples(self, orbital_params, lnlikes, curr_pos=None):
         else:
             self.post = np.vstack((self.post, orbital_params))
             self.lnlike = np.append(self.lnlike, lnlikes)
+
+        if weighted_post is not None:
+            if self._weighted_post is None:
+                self._weighted_post = weighted_post
+                self._weighted_lnlike = weighted_lnlike
+                self.lnweight = lnweight
+            else:
+                self._weighted_post = np.vstack((self._weighted_post, weighted_post))
+                self._weighted_lnlike = np.append(self._weighted_lnlike, weighted_lnlike)
+                self.lnweight = np.append(self.lnweight, lnweight)
 
         if curr_pos is not None:
             self.curr_pos = curr_pos
@@ -132,6 +212,12 @@ def save_results(self, filename):
         if self.curr_pos is not None:
             hf.create_dataset("curr_pos", data=self.curr_pos)
 
+        if self._weighted_post is not None and self._weighted_lnlike is not None and self.lnweight is not None:
+            hf.create_dataset("weighted_post", data=self._weighted_post)
+            hf.create_dataset("weighted_lnlike", data=self._weighted_lnlike)
+            hf.create_dataset("lnweight", data=self.lnweight)
+
+
         self.system.save(hf)
 
         hf.close()  # Closes file object, which writes file to disk
@@ -160,12 +246,12 @@ def load_results(self, filename, append=False):
             version_number = str(hf.attrs['version_number'])
         else:
             version_number = "<= 1.13"
-        post = hf.get('post')
-        if post is not None:
-            post = np.array(post)
-        lnlike = hf.get('lnlike')
-        if lnlike is not None:
-            lnlike = np.array(lnlike)
+
+        post = array_not_none(hf.get('post'))
+        lnlike = array_not_none(hf.get('lnlike'))
+        weighted_post = array_not_none(hf.get('weighted_post'))
+        weighted_lnlike = array_not_none(hf.get('weighted_lnlike'))
+        lnweight = array_not_none(hf.get("lnweight"))
 
         if 'num_secondary_bodies' in hf.attrs:
             num_secondary_bodies = int(hf.attrs['num_secondary_bodies'])
@@ -287,10 +373,7 @@ def load_results(self, filename, append=False):
         if 'ln_evidence_err' in hf.attrs:
             self.ln_evidence_err = hf.attrs['ln_evidence_err']
 
-        try:
-            curr_pos = np.array(hf.get('curr_pos'))
-        except KeyError:
-            curr_pos = None
+        curr_pos = array_not_none(hf.get('curr_pos'))
 
         hf.close()  # Closes file object
 
@@ -316,14 +399,14 @@ def load_results(self, filename, append=False):
                     'Unable to append file {} to Results object. version_number of object and file do not match'.format(filename))
 
             # Now append post and lnlike
-            self.add_samples(post, lnlike)#, self.labels)
+            self.add_samples(post, lnlike, weighted_post=weighted_post, weighted_lnlike=weighted_lnlike, lnweight=lnweight)#, self.labels)
         else:
 
             # Only proceed if object is completely empty
-            if self.sampler_name is None and self.post is None and self.lnlike is None and self.version_number is None:# and self.tau_ref_epoch is None :
+            if self.sampler_name is None and self.post is None and self.lnlike is None and self.version_number is None and self._weighted_lnlike is None and self._weighted_post is None and self.lnweight is None:# and self.tau_ref_epoch is None :
                 self.sampler_name = sampler_name
                 self.version_number = version_number
-                self.add_samples(post, lnlike)#, self.labels)
+                self.add_samples(post, lnlike, weighted_post=weighted_post, weighted_lnlike=weighted_lnlike, lnweight=lnweight)#, self.labels)
 
             else:
                 raise Exception(
@@ -390,11 +473,11 @@ def print_results(self):
         self.results_str += '-------------------\n'
         print(self.results_str)
 
-    def plot_corner(self, param_list=None, **corner_kwargs):
+    def plot_corner(self, param_list=None, downsample=None, **corner_kwargs):
         """
         Wrapper for orbitize.plot.plot_corner
         """
-        return orbitize.plot.plot_corner(self, param_list, **corner_kwargs)
+        return orbitize.plot.plot_corner(self, param_list, downsample, **corner_kwargs)
 
     def plot_orbits(self, object_to_plot=1, start_mjd=51544.,
         num_orbits_to_plot=100, num_epochs_to_plot=100,
@@ -444,4 +527,12 @@ def plot_propermotion(self,
                         cmap=cmap,
                         cbar_param=cbar_param,
                         # fig=fig
-                        )
+                        )
+
+def array_not_none(raw):
+    """
+    Returns a numpy.array of the input if it is not None, else returns None
+    """
+    if raw is not None:
+        return np.array(raw)
+    return raw