Merge branch 'master' of https://github.com/megbedell/wobble

README.md

@@ -1,4 +1,4 @@
-## wobble *(/wäb.lā/)* <a href='http://astro.uchicago.edu/~bmontet/wobble.mp3'>:sound:</a>
+## wobble *(/wäb.lā/)* <a href='https://benmontet.github.io/assets/wobble.mp3'>:sound:</a>
 
 [![Documentation Status](https://readthedocs.org/projects/wobble/badge/?version=latest)](https://wobble.readthedocs.io/en/latest/?badge=latest)[![arXiv](https://img.shields.io/badge/arXiv-1901.00503-orange.svg)](https://arxiv.org/abs/1901.00503)
 

docs/index.rst

@@ -6,7 +6,7 @@
 .. image:: wobble.png
 	:align: center
 
-*wobble* (pronounced `*(/wäb.lā/)* <http://astro.uchicago.edu/~bmontet/wobble.mp3>`_) is an open-source python implementation of a data-driven method for deriving stellar spectra, telluric spectra, and extremely precise radial velocities simultaneously without reliance on spectral models.
+*wobble* (pronounced `*(/wäb.lā/)* <https://benmontet.github.io/assets/wobble.mp3>`_) is an open-source python implementation of a data-driven method for deriving stellar spectra, telluric spectra, and extremely precise radial velocities simultaneously without reliance on spectral models.
 
 Read the paper `on arXiv <https://arxiv.org/abs/1901.00503>`_.
 
@@ -28,4 +28,4 @@ Bug Reports & Questions
 *wobble* is an open source project under the MIT license. The source code is available on `GitHub`_. In case of any questions or problems, please contact us via the `Git Issues`_.
 
 .. _GitHub: http://github.com/megbedell/wobble
-.. _Git Issues: http://github.com/megbedell/wobble/issues
+.. _Git Issues: http://github.com/megbedell/wobble/issues

wobble/regularization/regularization.py

@@ -10,7 +10,7 @@
 from wobble.results import Results
 from wobble.utils import get_session
 
-def generate_regularization_file(filename, R, type='star'):
+def generate_regularization_file(filename, R, type='star',defaults=None): 
     """
     Create a regularization parameter file with default values.
     
@@ -23,18 +23,23 @@ def generate_regularization_file(filename, R, type='star'):
     type : str, optional
         Type of object; sets which default values to use. 
         Acceptable values are 'star' (default) or 'telluric'.
+    defaults : list, optional
+    	list with 5 numbers, corresponding to the 5 regularization
+    	parameters
     """
     regularization_par = ['L1_template', 'L2_template', 
                               'L1_basis_vectors', 'L2_basis_vectors', 'L2_basis_weights']
     star_defaults = [1.e-2, 1.e2, 1.e5, 1.e6, 1.]
     #telluric_defaults = [1.e4, 1.e6, 1.e3, 1.e8, 1.]
     telluric_defaults = [1.e2, 1.e3, 1.e3, 1.e6, 1.]
-    if type=='star':
-        defaults = star_defaults
-    elif type=='telluric':
-        defaults = telluric_defaults
-    else:
-        assert False, "ERROR: type not recognized."
+    if defaults==None:
+        if type=='star':
+            defaults = star_defaults
+        elif type=='telluric':
+            defaults = telluric_defaults
+        else:
+            assert False, "ERROR: type not recognized."
+
     with h5py.File(filename,'w') as f:
         for par,val in zip(regularization_par, defaults):
             f.create_dataset(par, data=np.zeros(R)+val)

wobble/results.py

@@ -390,4 +390,91 @@ def plot_spectrum(self, r, n, data, filename, xlim=None, ylim=[0., 1.3], ylim_re
         fig.tight_layout()
         fig.subplots_adjust(hspace=0.05)
         plt.savefig(filename)
-        plt.close(fig)
+        plt.close(fig)
+
+    def plot_chromatic_rvs(self, min_order=None, max_order=None, percentiles=(16,84), wavelengths=None, scale='log',  ylim=None, center=True, filename=None):
+        """Output a representative percentile plot showing the chromaticity apparent in the rv signal.
+        
+        Parameters
+        ----------
+        min_order : 'int'
+                    Minimum order to plot.
+        max_order : 'int'
+                    Maximum order to plot. 
+        percentiles : 'tuple'
+                    Optional upper and lower percentile to plot.
+        wavelengths : 'tuple'
+                    Optional wavelength range (in Angstroms) to use instead of orders. 
+        scale : 'str'
+                    Optional scale; passed to matplotlib.
+        ylim : 'tuple'
+                    Optional ylim; passed to matplotlib.
+        center : 'boolean' 
+                    Determines whether the epochs are median centered before percentiles are calculated. 
+        filename : 'str'
+                    Saves plot if given. Optional filename; passed to matplotlib e.g. 'filename.png'
+   	"""
+        if min_order == None:
+            min_order = 0
+        if max_order == None:
+            max_order = len(self.orders)
+        orders = np.arange(min_order, max_order)
+        x = np.array([np.exp(np.mean(order)) for order in self.star_template_xs[min_order:max_order]])
+        if wavelengths != None:
+            orders = ((x > int(wavelengths[0])) & (x < int(wavelengths[1])))
+            x = x[orders]
+        if center == True:
+            median = np.median(np.array(self.star_rvs)[orders], axis=1).T
+        else:
+            median = 0
+        upper = np.percentile(np.array(self.star_rvs)[orders].T - median, percentiles[1], axis=0).T
+        upper_sigma = np.sqrt(1/np.array(self.star_ivars_rvs))[orders, np.argmin(abs(np.array(self.star_rvs)[orders].T - upper.T).T, axis=1)]
+        lower = np.percentile(np.array(self.star_rvs)[orders].T - median, percentiles[0], axis=0).T
+        lower_sigma = np.sqrt(1/np.array(self.star_ivars_rvs))[orders, np.argmin(abs(np.array(self.star_rvs)[orders].T - lower.T).T, axis=1)]
+        m, b = np.polyfit(x, upper, 1)
+        m2, b2 = np.polyfit(x, lower, 1)
+        y = m*x + b
+        y2 = m2*x + b2
+        plt.errorbar(x, upper, upper_sigma, fmt='o')
+        plt.errorbar(x, lower, lower_sigma, fmt='o')
+        plt.plot(x, y, color='tab:blue')
+        plt.plot(x, y2, color='tab:orange')
+        plt.fill_between(x, y, y2, color='lightgray')
+        plt.ylabel('Radial Velocity [m/s]')
+        plt.xlabel('Wavelength λ [Å]')
+        plt.ylim(ylim)
+        plt.xscale(scale)
+        if filename != None:
+            plt.savefig(filename)
+        plt.show()
+
+
+
+    def chromatic_index(self, min_order=None, max_order=None, wavelengths=None):
+        """ Returns a 2 by n array representing the chromatic index along with uncertainty for each epoch (calculated as slope of the linear least squares fit).
+        
+        Parameters
+        ----------
+        min_order : 'int'
+                    Minimum order to use in the calculation of the chromatic indices.
+        max_order : 'int'
+                    Maximum order to use in the calculation of the chromatic indices.
+        wavelengths : 'tuple'
+                    Optional wavelength range (in Angstroms) to use instead of orders. 
+        """
+        if min_order == None:
+            min_order = 0
+        if max_order == None:
+            max_order = len(self.orders)
+        orders = np.arange(min_order, max_order)
+        x = np.array([np.mean(order) for order in self.star_template_xs[min_order:max_order]])
+        if wavelengths != None:
+            orders = ((x > int(wavelengths[0])) & (x < int(wavelengths[1])))
+            x = x[orders]
+        chromatic_indices = []
+        sigmas = []
+        for epoch in range(len(self.epochs)): 
+           coefs = np.polyfit(x, np.array(self.star_rvs)[orders, epoch], 1, full=True)
+           chromatic_indices.append(coefs[0][0])
+           sigmas.append(np.sqrt(coefs[1][0]))
+        return [chromatic_indices, sigmas]