arvi 0.1.13__py3-none-any.whl → 0.1.15__py3-none-any.whl

arvi/plots.py

@@ -1,12 +1,7 @@
-import os
-from functools import partial, partialmethod
+from functools import partialmethod, wraps
 from itertools import cycle
 
-import matplotlib.collections
 import numpy as np
-import matplotlib
-import matplotlib.pyplot as plt
-import mplcursors
 
 from astropy.timeseries import LombScargle
 
@@ -14,11 +9,24 @@ from .setup_logger import logger
 from . import config
 from .stats import wmean
 
+from .utils import lazy_import
+plt = lazy_import('matplotlib.pyplot')
+
+
+def plot_fast(func):
+    @wraps(func)
+    def wrapper(*args, **kwargs):
+        with plt.style.context('fast'):
+            return func(*args, **kwargs)
+    return wrapper
+
+
 
 class BlittedCursor:
     """ A cross-hair cursor using blitting for faster redraw. """
     def __init__(self, axes, vertical=True, horizontal=True, show_text=None,
                  transforms_x=None, transforms_y=None):
+        import matplotlib # delay import
         if isinstance(axes, matplotlib.axes.Axes):
             axes = [axes]
         self.axes = axes
@@ -102,10 +110,31 @@ class BlittedCursor:
                     ax.draw_artist(self.text)
                 ax.figure.canvas.blit(ax.bbox)
 
+def clickable_legend(fig, ax, leg):
+    from matplotlib.text import Text
+    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, 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
+    return on_pick_legend
 
+# @plot_fast
 def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
-         remove_50000=False, tooltips=False, label=None, N_in_label=False,
-         versus_n=False, show_histogram=False, bw=False, **kwargs):
+         remove_50000=False, tooltips=False, show_legend=True, label=None,
+         N_in_label=False, versus_n=False, show_histogram=False, bw=False, **kwargs):
     """ Plot the RVs
 
     Args:
@@ -121,6 +150,8 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
             Whether to subtract 50000 from time. Defaults to False.
         tooltips (bool, optional):
             Show information upon clicking a point. Defaults to True.
+        show_legend (bool, optional):
+            Show legend. Defaults to True.
         N_in_label (bool, optional):
             Show number of observations in legend. Defaults to False.
         versus_n (bool, optional):
@@ -145,12 +176,14 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
             fig, (ax, axh) = plt.subplots(1, 2, constrained_layout=True, gridspec_kw={'width_ratios': [3, 1]})
         else:
             fig, ax = plt.subplots(1, 1, constrained_layout=True)
+    elif ax == -1:
+        ax = plt.gca()
+        fig = ax.figure
     else:
         if show_histogram:
             ax, axh = ax
         fig = ax.figure
 
-
     kwargs.setdefault('ls', '')
     kwargs.setdefault('capsize', 0)
     kwargs.setdefault('ms', 4)
@@ -172,7 +205,6 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
     except AttributeError:
         zorders = cycle([1] * len(instruments))
 
-    cursors = {}
     containers = {}
     for _i, inst in enumerate(instruments):
         s = self if self._child else getattr(self, inst)
@@ -207,22 +239,6 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
             hlabel = f'{s.mvrad.std():.2f} {self.units}'
             axh.hist(s.mvrad, **kw, label=hlabel)
 
-                # 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:
                 pass
@@ -239,62 +255,104 @@ 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)
 
-    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 show_legend:
+        leg = ax.legend()
+        on_pick_legend = clickable_legend(fig, ax, leg)
+        plt.connect('pick_event', on_pick_legend)
 
     if tooltips:
+        from matplotlib.lines import Line2D
+        from matplotlib.collections import LineCollection
         annotations = []
-        def on_pick_point(event):
-            print('annotations:', annotations)
+        axes_artists = []
+        selected_inds = []
+        selected_insts = []
+
+        def _cleanup():
             for text in annotations:
                 text.remove()
                 annotations.remove(text)
+            for art in axes_artists:
+                art.remove()
+                axes_artists.remove(art)
+            for ind in selected_inds:
+                selected_inds.remove(ind)
+
+        def on_press(event):
+            # print('press', event.key)
+            if event.key in ('r',):
+                for i, inst in zip(selected_inds, selected_insts):
+                    i = self._index_from_instrument_index(i, inst)
+                    self.remove_point(i)
+                _cleanup()
+                fig.canvas.draw_idle()
+            if event.key == 'escape':
+                _cleanup()
+                kp_cid = None
+                for k, _f in fig.canvas.callbacks.callbacks['key_press_event'].items():
+                    if _f._obj == on_press:
+                        kp_cid = k
+                if kp_cid is not None:
+                    fig.canvas.callbacks.disconnect(k)
+                fig.canvas.draw_idle()
 
+        def on_pick_point(event):
+            _cleanup()            
             artist = event.artist
-            if isinstance(artist, (matplotlib.lines.Line2D, matplotlib.collections.LineCollection)):
-                print(event.ind, artist)
-                if isinstance(artist, matplotlib.lines.Line2D):
+            if isinstance(artist, (Line2D, LineCollection)):
+                # print(event.ind, artist)
+                if isinstance(artist, Line2D):
+                    ind = event.ind
                     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="-"))
-                    )
+
+                    if event.ind.size > 1:
+                        mint, maxt = _s.mtime[ind].min(), _s.mtime[ind].max()
+                        miny, maxy = _s.mvrad[ind].min(), _s.mvrad[ind].max()
+                        axins = ax.inset_axes([0.1, 0.5, 0.5, 0.4],
+                            xlim=(mint - 0.1 * (maxt - mint), maxt + 0.1 * (maxt - mint)), 
+                            ylim=(miny - 0.1 * (maxy - miny), maxy + 0.1 * (maxy - miny)), 
+                            xticklabels=[], yticklabels=[]
+                            )
+                        axins.errorbar(_s.mtime[ind], _s.mvrad[ind], _s.msvrad[ind], fmt='o', ms=3,
+                                       color=artist.get_color())
+                        axins.margins(x=0.5)
+                        axins.autoscale_view()
+                        rectangle_patch, connector_lines = ax.indicate_inset_zoom(axins, edgecolor="black")
+                        axes_artists.append(axins)
+                        axes_artists.append(rectangle_patch)
+                        for line in connector_lines:
+                            axes_artists.append(line)
+                    else:
+                        ind = event.ind[0]
+                        selected_inds.append(ind)
+                        selected_insts.append(inst)
+                        # print(_s.mtime[ind], _s.mvrad[ind], _s.msvrad[ind])
+
+                        text = f'{inst} ({ind})\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())
+
+                    _ = fig.canvas.mpl_connect('key_press_event', on_press)
+
         plt.connect('pick_event', on_pick_point)
 
 
+
+
     if show_histogram:
         axh.legend()
 
@@ -309,25 +367,23 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
         else:
             ax.set_xlabel('BJD - 2400000 [days]')
 
-    from matplotlib.backend_tools import ToolBase, ToolToggleBase
-    tm = fig.canvas.manager.toolmanager
-
-    class InfoTool(ToolToggleBase):
-        description = "Show extra information about each observation"
-        image = os.path.abspath(os.path.join(os.path.dirname(__file__), 'data', 'info.svg'))
-        # def enable(self, *args, **kwargs):
-        #     self.figure.add_axes([1, 0, 0.3, 1])
-        #     self.figure.canvas.draw_idle()
-
-    from PIL import UnidentifiedImageError
-    try:
-        tm.add_tool("infotool", InfoTool)
-        fig.canvas.manager.toolbar.add_tool(tm.get_tool("infotool"), "toolgroup")
-        raise UnidentifiedImageError
-    except AttributeError:
-        pass
-    except UnidentifiedImageError:
-        pass
+    # from matplotlib.backend_tools import ToolBase, ToolToggleBase
+    # tm = fig.canvas.manager.toolmanager
+    # class InfoTool(ToolToggleBase):
+    #     description = "Show extra information about each observation"
+    #     image = os.path.abspath(os.path.join(os.path.dirname(__file__), 'data', 'info.svg'))
+    #     # def enable(self, *args, **kwargs):
+    #     #     self.figure.add_axes([1, 0, 0.3, 1])
+    #     #     self.figure.canvas.draw_idle()
+    # from PIL import UnidentifiedImageError
+    # try:
+    #     tm.add_tool("infotool", InfoTool)
+    #     fig.canvas.manager.toolbar.add_tool(tm.get_tool("infotool"), "toolgroup")
+    #     raise UnidentifiedImageError
+    # except AttributeError:
+    #     pass
+    # except UnidentifiedImageError:
+    #     pass
 
 
     if config.return_self:
@@ -339,8 +395,9 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
     return fig, ax
 
 
+@plot_fast
 def plot_quantity(self, quantity, ax=None, show_masked=False, instrument=None,
-                  time_offset=0, remove_50000=False, tooltips=False,
+                  time_offset=0, remove_50000=False, tooltips=False, show_legend=True,
                   N_in_label=False, **kwargs):
     if self.N == 0:
         if self.verbose:
@@ -399,16 +456,21 @@ def plot_quantity(self, quantity, ax=None, show_masked=False, instrument=None,
                     getattr(self, quantity + '_err')[~self.mask],
                     label='masked', fmt='x', ms=10, color='k', zorder=-2)
 
-    ax.legend()
+    if show_legend:
+        leg = ax.legend()
+        on_pick_legend = clickable_legend(fig, ax, leg)
+        plt.connect('pick_event', on_pick_legend)
+
     ax.minorticks_on()
 
     ylabel = {
+        quantity: quantity,
         'fwhm': f'FWHM [{self.units}]',
         'bispan': f'BIS [{self.units}]',
         'rhk': r"$\log$ R'$_{HK}$",
-        'berv': f'BERV [km/s]',
+        'berv': 'BERV [km/s]',
     }
-    ax.set_ylabel(ylabel[quantity])
+    ax.set_ylabel(ylabel[quantity.lower()])
 
     if remove_50000:
         ax.set_xlabel('BJD - 2450000 [days]')
@@ -427,6 +489,7 @@ plot_rhk = partialmethod(plot_quantity, quantity='rhk')
 plot_berv = partialmethod(plot_quantity, quantity='berv')
 
 
+@plot_fast
 def gls(self, ax=None, label=None, fap=True, instrument=None, adjust_means=config.adjust_means_gls,
         picker=True, **kwargs):
     """
@@ -470,7 +533,7 @@ def gls(self, ax=None, label=None, fap=True, instrument=None, adjust_means=confi
         if self.verbose:
             logger.info(f'calculating periodogram for instrument {instrument}')
 
-        if adjust_means:
+        if adjust_means and not self._child:
             if self.verbose:
                 logger.info('adjusting instrument means before gls')
             means = np.empty_like(y)
@@ -485,7 +548,7 @@ def gls(self, ax=None, label=None, fap=True, instrument=None, adjust_means=confi
         y = self.vrad[self.mask].copy()
         e = self.svrad[self.mask].copy()
 
-        if adjust_means:
+        if adjust_means and not self._child:
             if self.verbose:
                 logger.info('adjusting instrument means before gls')
             means = np.empty_like(y)
@@ -528,6 +591,7 @@ def gls(self, ax=None, label=None, fap=True, instrument=None, adjust_means=confi
         ax.legend()
 
     if ax.get_legend() is not None:
+        from matplotlib.text import Text
         leg = ax.get_legend()
         for text in leg.get_texts():
             text.set_picker(True)
@@ -535,7 +599,7 @@ def gls(self, ax=None, label=None, fap=True, instrument=None, adjust_means=confi
         def on_pick_legend(event):
             handles, labels = ax.get_legend_handles_labels()
             artist = event.artist
-            if isinstance(artist, matplotlib.text.Text):
+            if isinstance(artist, Text):
                 # print('handles:', handles)
                 # print('labels:', labels)
                 # print(artist.get_text())
@@ -559,6 +623,7 @@ def gls(self, ax=None, label=None, fap=True, instrument=None, adjust_means=confi
         return fig, ax
 
 
+@plot_fast
 def gls_quantity(self, quantity, ax=None, fap=True, picker=True):
     if not hasattr(self, quantity):
         logger.error(f"cannot find '{quantity}' attribute")
@@ -571,7 +636,10 @@ def gls_quantity(self, quantity, ax=None, fap=True, picker=True):
 
     t = self.mtime
     y = getattr(self, quantity)[self.mask]
-    ye = getattr(self, quantity + '_err')[self.mask]
+    try:
+        ye = getattr(self, quantity + '_err')[self.mask]
+    except AttributeError:
+        ye = None
 
     if np.isnan(y).any():
         if self.verbose:

arvi/programs.py

@@ -20,12 +20,14 @@ def get_star(star, instrument=None):
 
 
 class LazyRV:
-    def __init__(self, stars: list, instrument: str = None):
+    def __init__(self, stars: list, instrument: str = None,
+                 _parallel_limit=10):
         self.stars = stars
         if isinstance(self.stars, str):
             self.stars = [self.stars]
         self.instrument = instrument
         self._saved = None
+        self._parallel_limit = _parallel_limit
 
     @property
     def N(self):
@@ -35,7 +37,7 @@ class LazyRV:
         return f"RV({self.N} stars)"
 
     def _get(self):
-        if self.N > 10:
+        if self.N > self._parallel_limit:
             # logger.info('Querying DACE...')
             _get_star = partial(get_star, instrument=self.instrument)
             with multiprocessing.Pool() as pool:

arvi/reports.py

@@ -1,9 +1,6 @@
 from functools import partial
 import numpy as np
 from astropy.timeseries import LombScargle
-import matplotlib.pyplot as plt
-import matplotlib.gridspec as gridspec
-from matplotlib.backends.backend_pdf import PdfPages
 
 from .setup_logger import logger
 
@@ -31,6 +28,10 @@ def sine_picker(event, self, fig, ax, ax1):
 
 
 def report(self, save=None):
+    import matplotlib.pyplot as plt
+    import matplotlib.gridspec as gridspec
+    from matplotlib.backends.backend_pdf import PdfPages
+
     # size = A4
     size = 8.27, 11.69
     fig = plt.figure(figsize=size, constrained_layout=True)

arvi/simbad_wrapper.py

@@ -1,7 +1,6 @@
 import os
 import requests
 
-from astropy.coordinates import SkyCoord
 import pysweetcat
 
 DATA_PATH = os.path.dirname(__file__)
@@ -44,7 +43,7 @@ def run_query(query):
     url = 'http://simbad.u-strasbg.fr/simbad/sim-tap/sync'
     data = dict(query=query, request='doQuery', lang='ADQL', format='text/plain', phase='run')
     try:
-        response = requests.post(url, data=data, timeout=10)
+        response = requests.post(url, data=data, timeout=5)
     except requests.ReadTimeout as err:
         raise IndexError(err)
     except requests.ConnectionError as err:
@@ -97,6 +96,8 @@ class simbad:
         Args:
             star (str): The name of the star to query simbad
         """
+        from astropy.coordinates import SkyCoord
+
         self.star = star
 
         if 'kobe' in self.star.lower():

arvi/stats.py

@@ -1,7 +1,4 @@
 import numpy as np
-from scipy.stats import median_abs_deviation
-from scipy.stats._stats_py import SigmaclipResult
-
 
 def wmean(a, e):
     """Weighted mean of array `a`, with uncertainties given by `e`.
@@ -69,6 +66,8 @@ def sigmaclip_median(a, low=4.0, high=4.0):
             - `lower`: Lower clipping limit
             - `upper`: Upper clipping limit
     """
+    from scipy.stats import median_abs_deviation
+    from scipy.stats._stats_py import SigmaclipResult
     c = np.asarray(a).ravel()
     delta = 1
     while delta:

arvi/stellar.py

@@ -0,0 +1,89 @@
+import numpy as np
+
+class prot_age_result:
+    prot_n84: float
+    prot_n84_err: float
+    prot_m08: float
+    prot_m08_err: float
+    age_m08: float
+    age_m08_err: float
+    def __init__(self):
+        pass
+    def __repr__(self):
+        s = f'{self.prot_n84=:.2f} ± {self.prot_n84_err:.2f}, '
+        s += f'{self.prot_m08=:.2f} ± {self.prot_m08_err:.2f}, '
+        s += f'{self.age_m08=:.2f} ± {self.age_m08_err:.2f}'
+        return s.replace('self.', '')
+
+
+def calc_prot_age(self, bv=None):
+    """
+    Calculate rotation period and age from logR'HK activity level, based on the
+    empirical relations of Noyes et al. (1984) and Mamajek & Hillenbrand (2008).
+
+    Args:
+        self (`arvi.RV`):
+            RV object
+        bv (float, optional):
+            B-V colour. If None, use value from Simbad
+
+    Returns:
+        An object with the following attributes:
+
+        prot_n84 (float, array):
+            Chromospheric rotational period via Noyes et al. (1984)
+        prot_n84_err (float, array):
+            Error on 'prot_n84'
+        prot_m08 (float, array):
+            Chromospheric rotational period via Mamajek & Hillenbrand (2008)
+        prot_m08_err (float, array):
+            Error on 'prot_m08'
+        age_m08 (float, array):
+            Gyrochronology age via Mamajek & Hillenbrand (2008)
+        age_m08_err (float, array):
+            Error on 'age_m08'
+
+    Range of logR'HK-Prot relation: -5.5 < logR'HK < -4.3
+    Range of Mamajek & Hillenbrand (2008) relation for ages: 0.5 < B-V < 0.9
+    """
+    log_rhk = np.nanmean(self.rhk[self.mask])
+    if bv is None:
+        bv = self.simbad.B - self.simbad.V
+
+    # Calculate chromospheric Prot:
+    if np.any(log_rhk < -4.3) & np.any(log_rhk > -5.5):
+        if bv < 1:
+            tau = 1.362 - 0.166*(1-bv) + 0.025*(1-bv)**2 - 5.323*(1-bv)**3
+        else:
+            tau = 1.362 - 0.14*(1-bv)
+
+        prot_n84 = 0.324 - 0.400*(5 + log_rhk) - 0.283*(5 + log_rhk)**2 - 1.325*(5 + log_rhk)**3 + tau
+        prot_n84 = 10**prot_n84
+        prot_n84_err = np.log(10)*0.08*prot_n84
+
+        prot_m08 = (0.808 - 2.966*(log_rhk + 4.52))*10**tau
+        prot_m08_err = 4.4*bv*1.7 - 1.7
+    else:
+        prot_n84 = np.nan
+        prot_n84_err = np.nan
+        prot_m08 = np.nan
+        prot_m08_err = np.nan
+
+    # Calculate gyrochronology age:
+    if np.any(prot_m08 > 0.0) & (bv > 0.50) & (bv < 0.9):
+        age_m08 = 1e-3*(prot_m08/0.407/(bv - 0.495)**0.325)**(1./0.566)
+        #age_m08_err = 0.05*np.log(10)*age_m08
+        age_m08_err  = 0.2 * age_m08 * np.log(10) # using 0.2 dex typical error from paper
+    else:
+        age_m08 = np.nan
+        age_m08_err = np.nan
+
+    r = prot_age_result()
+    r.prot_n84 = prot_n84
+    r.prot_n84_err = prot_n84_err
+    r.prot_m08 = prot_m08
+    r.prot_m08_err = prot_m08_err
+    r.age_m08 = age_m08
+    r.age_m08_err = age_m08_err
+    return r
+    # return prot_n84, prot_n84_err, prot_m08, prot_m08_err, age_m08, age_m08_err