arvi 0.1.8__py3-none-any.whl → 0.1.11__py3-none-any.whl

arvi/plots.py

@@ -1,7 +1,10 @@
 import os
 from functools import partial, partialmethod
+from itertools import cycle
 
+import matplotlib.collections
 import numpy as np
+import matplotlib
 import matplotlib.pyplot as plt
 from matplotlib.collections import LineCollection
 import mplcursors
@@ -13,8 +16,8 @@ from . import config
 
 
 def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
-         remove_50000=False, tooltips=True, label=None, N_in_label=False,
-         versus_n=False, show_histogram=False, **kwargs):
+         remove_50000=False, tooltips=False, label=None, N_in_label=False,
+         versus_n=False, show_histogram=False, bw=False, **kwargs):
     """ Plot the RVs
 
     Args:
@@ -37,6 +40,8 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
         show_histogram (bool, optional)
             Whether to show a panel with the RV histograms (per intrument).
             Defaults to False.
+        bw (bool, optional):
+            Adapt plot to black and white. Defaults to False.
 
     Returns:
         Figure: the figure
@@ -57,7 +62,7 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
             ax, axh = ax
         fig = ax.figure
 
-    kwargs.setdefault('marker', 'o')
+
     kwargs.setdefault('ls', '')
     kwargs.setdefault('capsize', 0)
     kwargs.setdefault('ms', 4)
@@ -65,10 +70,22 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
     if remove_50000:
         time_offset = 50000
 
-    instruments = self._check_instrument(instrument)
+    strict = kwargs.pop('strict', False)
+    instruments = self._check_instrument(instrument, strict=strict)
+
+    if bw:
+        markers = cycle(('o', 'P', 's', '^', '*'))
+    else:
+        markers = cycle(('o',) * len(instruments))
+
+    try:
+        zorders = cycle(-np.argsort([getattr(self, i).error for i in instruments])[::-1])
+    except AttributeError:
+        zorders = cycle([1] * len(instruments))
 
     cursors = {}
-    for inst in instruments:
+    containers = {}
+    for _i, inst in enumerate(instruments):
         s = self if self._child else getattr(self, inst)
         if s.mask.sum() == 0:
             continue
@@ -80,37 +97,42 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
                 p = p.replace('_', '.')
                 _label = f'{i}-{p}'
         else:
-            _label = label
+            if isinstance(label, list):
+                _label = label[_i]
+            else:
+                _label = label
 
         if versus_n:
-            container = ax.errorbar(np.arange(1, s.mtime.size + 1),
-                                    s.mvrad, s.msvrad, label=_label, picker=True, **kwargs)
+            container = ax.errorbar(np.arange(1, s.mtime.size + 1), s.mvrad, s.msvrad,
+                                    label=_label, picker=True, marker=next(markers), zorder=next(zorders),
+                                    **kwargs)
         else:
-            container = ax.errorbar(s.mtime - time_offset,
-                                    s.mvrad, s.msvrad, label=_label, picker=True, **kwargs)
+            container = ax.errorbar(s.mtime - time_offset, s.mvrad, s.msvrad,
+                                    label=_label, picker=True, marker=next(markers), zorder=next(zorders),
+                                    **kwargs)
+
+        containers[inst] = list(container)
 
         if show_histogram:
             kw = dict(histtype='step', bins='doane', orientation='horizontal')
             hlabel = f'{s.mvrad.std():.2f} {self.units}'
             axh.hist(s.mvrad, **kw, label=hlabel)
 
-        if tooltips:
-            cursors[inst] = crsr = mplcursors.cursor(container, multiple=False)
-
-            @crsr.connect("add")
-            def _(sel):
-                inst = sel.artist.get_label()
-                _s = getattr(self, inst)
-                vrad, svrad = _s.vrad[sel.index], _s.svrad[sel.index]
-                sel.annotation.get_bbox_patch().set(fc="white")
-                text = f'{inst}\n'
-                text += f'BJD: {sel.target[0]:9.5f}\n'
-                text += f'RV: {vrad:.3f} ± {svrad:.3f}'
-                if fig.canvas.manager.toolmanager.get_tool('infotool').toggled:
-                    text += '\n\n'
-                    text += f'date: {_s.date_night[sel.index]}\n'
-                    text += f'mask: {_s.ccf_mask[sel.index]}'
-                sel.annotation.set_text(text)
+                # cursors[inst] = crsr = mplcursors.cursor(container, multiple=False)
+                # @crsr.connect("add")
+                # def _(sel):
+                #     inst = sel.artist.get_label()
+                #     _s = getattr(self, inst)
+                #     vrad, svrad = _s.vrad[sel.index], _s.svrad[sel.index]
+                #     sel.annotation.get_bbox_patch().set(fc="white")
+                #     text = f'{inst}\n'
+                #     text += f'BJD: {sel.target[0]:9.5f}\n'
+                #     text += f'RV: {vrad:.3f} ± {svrad:.3f}'
+                #     # if fig.canvas.manager.toolmanager.get_tool('infotool').toggled:
+                #     #     text += '\n\n'
+                #     #     text += f'date: {_s.date_night[sel.index]}\n'
+                #     #     text += f'mask: {_s.ccf_mask[sel.index]}'
+                #     sel.annotation.set_text(text)
 
         if show_masked:
             if versus_n:
@@ -128,7 +150,62 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
             ax.errorbar(self.time[~self.mask] - time_offset, self.vrad[~self.mask], self.svrad[~self.mask],
                         label='masked', fmt='x', ms=10, color='k', zorder=-2)
 
-    ax.legend()
+    leg = ax.legend()
+    handles, labels = ax.get_legend_handles_labels()
+    for text in leg.get_texts():
+        text.set_picker(True)
+
+    def on_pick_legend(event):
+        artist = event.artist
+        if isinstance(artist, matplotlib.text.Text):
+            try:
+                h = handles[labels.index(artist.get_text())]
+                alpha_text = {None:0.2, 1.0: 0.2, 0.2:1.0}[artist.get_alpha()]
+                alpha_point = {None: 0.0, 1.0: 0.0, 0.2: 1.0}[artist.get_alpha()]
+                h[0].set_alpha(alpha_point)
+                h[2][0].set_alpha(alpha_point)
+                artist.set_alpha(alpha_text)
+                fig.canvas.draw()
+            except ValueError:
+                pass
+    plt.connect('pick_event', on_pick_legend)
+
+    if tooltips:
+        annotations = []
+        def on_pick_point(event):
+            print('annotations:', annotations)
+            for text in annotations:
+                text.remove()
+                annotations.remove(text)
+
+            artist = event.artist
+            if isinstance(artist, (matplotlib.lines.Line2D, matplotlib.collections.LineCollection)):
+                print(event.ind, artist)
+                if isinstance(artist, matplotlib.lines.Line2D):
+                    matching_instrument = [k for k, v in containers.items() if artist in v]
+                    print(matching_instrument)
+                    if len(matching_instrument) == 0:
+                        return
+                    inst = matching_instrument[0]
+                    _s = getattr(self, inst)
+                    ind = event.ind[0]
+                    # print(_s.mtime[ind], _s.mvrad[ind], _s.msvrad[ind])
+
+                    text = f'{inst}\n'
+                    text += f'{_s.mtime[ind]:9.5f}\n'
+                    text += f'RV: {_s.mvrad[ind]:.1f} ± {_s.msvrad[ind]:.1f}'
+
+                    annotations.append(
+                        ax.annotate(text, (_s.mtime[ind], _s.mvrad[ind]), xycoords='data',
+                                    xytext=(5, 10), textcoords='offset points', fontsize=9,
+                                    bbox={'boxstyle': 'round', 'fc': 'w'}, arrowprops=dict(arrowstyle="-"))
+                    )
+                    # ax.annotate(f'{inst}', (0.5, 0.5), xycoords=artist, ha='center', va='center')
+                    fig.canvas.draw()
+                    # print(event.ind, artist.get_label())
+        plt.connect('pick_event', on_pick_point)
+
+
     if show_histogram:
         axh.legend()
 
@@ -209,7 +286,11 @@ def plot_quantity(self, quantity, ax=None, show_masked=False, instrument=None,
             label = f'{i}-{p}'
 
         y = getattr(s, quantity)
-        ye = getattr(s, quantity + '_err')
+        try:
+            ye = getattr(s, quantity + '_err')
+        except AttributeError:
+            ye = np.zeros_like(y)
+
 
         if np.isnan(y).all() or np.isnan(ye).all():
             lines, *_ = ax.errorbar([], [], [],
@@ -256,6 +337,21 @@ plot_rhk = partialmethod(plot_quantity, quantity='rhk')
 
 
 def gls(self, ax=None, label=None, fap=True, picker=True, instrument=None, **kwargs):
+    """
+    Calculate and plot the Generalised Lomb-Scargle periodogram of the radial
+    velocities.
+
+    Args:
+        ax (matplotlib.axes.Axes):
+            The matplotlib axes to plot on. If None, a new figure will be
+            created.
+        label (str):
+            The label to use for the plot.
+        fap (bool):
+            Whether to show the false alarm probability.
+        instrument (str or list):
+            Which instruments' data to include in the periodogram.
+    """
     if self.N == 0:
         if self.verbose:
             logger.error('no data to compute gls')
@@ -267,7 +363,8 @@ def gls(self, ax=None, label=None, fap=True, picker=True, instrument=None, **kwa
         fig = ax.figure
 
     if instrument is not None:
-        instrument = self._check_instrument(instrument)
+        strict = kwargs.pop('strict', False)
+        instrument = self._check_instrument(instrument, strict=strict)
         if instrument is not None:
             instrument_mask = np.isin(self.instrument_array, instrument)
             t = self.time[instrument_mask & self.mask]
@@ -298,6 +395,31 @@ def gls(self, ax=None, label=None, fap=True, picker=True, instrument=None, **kwa
     if label is not None:
         ax.legend()
 
+    if ax.get_legend() is not None:
+        leg = ax.get_legend()
+        for text in leg.get_texts():
+            text.set_picker(True)
+
+        def on_pick_legend(event):
+            handles, labels = ax.get_legend_handles_labels()
+            artist = event.artist
+            if isinstance(artist, matplotlib.text.Text):
+                # print('handles:', handles)
+                # print('labels:', labels)
+                # print(artist.get_text())
+                try:
+                    h = handles[labels.index(artist.get_text())]
+                    alpha_text = {None:0.2, 1.0: 0.2, 0.2:1.0}[artist.get_alpha()]
+                    alpha_point = {None: 0.0, 1.0: 0.0, 0.2: 1.0}[artist.get_alpha()]
+                    h.set_alpha(alpha_point)
+                    artist.set_alpha(alpha_text)
+                    fig.canvas.draw()
+                except ValueError:
+                    pass
+
+        if 'pick_event' not in fig.canvas.callbacks.callbacks:
+            plt.connect('pick_event', on_pick_legend)
+
 
     if config.return_self:
         return self
@@ -347,3 +469,41 @@ def gls_quantity(self, quantity, ax=None, fap=True, picker=True):
 gls_fwhm = partialmethod(gls_quantity, quantity='fwhm')
 gls_bis = partialmethod(gls_quantity, quantity='bispan')
 gls_rhk = partialmethod(gls_quantity, quantity='rhk')
+
+
+def histogram_svrad(self, ax=None, instrument=None, label=None):
+    """ Plot an histogram of the radial velocity uncertainties. 
+    
+    Args:
+        ax (matplotlib.axes.Axes):
+            The matplotlib axes to plot on. If None, a new figure will be
+            created.
+        instrument (str or list):
+            Which instruments' data to include in the histogram.
+        label (str):
+            The label to use for the plot.
+    """
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, constrained_layout=True)
+    else:
+        fig = ax.figure
+
+    instruments = self._check_instrument(instrument)
+
+    for inst in instruments:
+        s = self if self._child else getattr(self, inst)
+
+        if label is None:
+            _label = inst
+            if not self.only_latest_pipeline:
+                i, p = _label.split('_', 1)
+                p = p.replace('_', '.')
+                _label = f'{i}-{p}'
+        else:
+            _label = label
+
+        kw = dict(bins=40, histtype='step', density=False, lw=2)
+        ax.hist(s.msvrad, label=_label, **kw)
+    ax.legend()
+    ax.set(xlabel=f'RV uncertainty [m/s]', ylabel='Number')
+

arvi/simbad_wrapper.py

@@ -1,9 +1,12 @@
+import os
 from dataclasses import dataclass, field
 import requests
 
 from astropy.coordinates import SkyCoord
 import pysweetcat
 
+DATA_PATH = os.path.dirname(__file__)
+DATA_PATH = os.path.join(DATA_PATH, 'data')
 
 QUERY = """
 SELECT basic.OID,
@@ -40,12 +43,13 @@ WHERE id = '{star}';
 
 def run_query(query):
     url = 'http://simbad.u-strasbg.fr/simbad/sim-tap/sync'
-    response = requests.post(url,
-                             data=dict(query=query,
-                                       request='doQuery',
-                                       lang='ADQL',
-                                       format='text/plain',
-                                       phase='run'))
+    data = dict(query=query, request='doQuery', lang='ADQL', format='text/plain', phase='run')
+    try:
+        response = requests.post(url, data=data, timeout=10)
+    except requests.ReadTimeout as err:
+        raise IndexError(err)
+    except requests.ConnectionError as err:
+        raise IndexError(err)
     return response.content.decode()
 
 def parse_table(table, cols=None, values=None):
@@ -95,6 +99,16 @@ class simbad:
         """
         self.star = star
 
+        if 'kobe' in self.star.lower():
+            fname = os.path.join(DATA_PATH, 'KOBE-translate.csv')
+            kobe_translate = {}
+            if os.path.exists(fname):
+                with open(fname) as f:
+                    for line in f.readlines():
+                        kobe_id, catname = line.strip().split(',')
+                        kobe_translate[kobe_id] = catname
+                self.star = star = kobe_translate[self.star]
+
         # oid = run_query(query=OID_QUERY.format(star=star))
         # self.oid = str(oid.split()[-1])
 
@@ -122,6 +136,9 @@ class simbad:
 
         self.coords = SkyCoord(self.ra, self.dec, unit='deg')
 
+        if self.plx_value == '':
+            self.plx_value = None
+
         try:
             swc_data = pysweetcat.get_data()
             data = swc_data.find(star)

arvi/stats.py

@@ -11,17 +11,27 @@ def wmean(a, e):
         a (array): Array containing data
         e (array): Uncertainties on `a`
     """
+    if (e == 0).any():
+        raise ZeroDivisionError
+    if (e < 0).any():
+        raise ValueError
+    if (a.shape != e.shape):
+        raise ValueError
     return np.average(a, weights=1 / e**2)
 
-def rms(a):
+def rms(a, ignore_nans=False):
     """ Root mean square of array `a`
 
     Args:
         a (array): Array containing data
     """
+    if ignore_nans:
+        a = a[~np.isnan(a)]
+    if len(a) == 0:
+        return np.nan
     return np.sqrt((a**2).mean())
 
-def wrms(a, e):
+def wrms(a, e, ignore_nans=False):
     """ Weighted root mean square of array `a`, with uncertanty given by `e`.
     The weighted rms is calculated using the weighted mean, where the weights
     are equal to 1/e**2.
@@ -30,6 +40,16 @@ def wrms(a, e):
         a (array): Array containing data 
         e (array): Uncertainties on `a`
     """
+    if ignore_nans:
+        nans = np.logical_or(np.isnan(a), np.isnan(e))
+        a = a[~nans]
+        e = e[~nans]
+    if (e == 0).any():
+        raise ZeroDivisionError('uncertainty cannot be zero')
+    if (e < 0).any():
+        raise ValueError('uncertainty cannot be negative')
+    if (a.shape != e.shape):
+        raise ValueError('arrays must have the same shape')
     w = 1 / e**2
     return np.sqrt(np.sum(w * (a - np.average(a, weights=w))**2) / sum(w))
 
@@ -38,6 +58,16 @@ def sigmaclip_median(a, low=4.0, high=4.0):
     """
     Same as scipy.stats.sigmaclip but using the median and median absolute
     deviation instead of the mean and standard deviation.
+
+    Args:
+        a (array): Array containing data
+        low (float): Number of MAD to use for the lower clipping limit
+        high (float): Number of MAD to use for the upper clipping limit
+    Returns:
+        SigmaclipResult: Object with the following attributes:
+            - `clipped`: Masked array of data
+            - `lower`: Lower clipping limit
+            - `upper`: Upper clipping limit
     """
     c = np.asarray(a).ravel()
     delta = 1