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

arvi/config.py

@@ -1 +1,2 @@
 return_self = False
+check_internet = False

arvi/dace_wrapper.py

@@ -5,7 +5,7 @@ import numpy as np
 from dace_query import DaceClass
 from dace_query.spectroscopy import SpectroscopyClass, Spectroscopy as default_Spectroscopy
 from .setup_logger import logger
-from .utils import create_directory, all_logging_disabled, stdout_disabled
+from .utils import create_directory, all_logging_disabled, stdout_disabled, tqdm
 
 
 def load_spectroscopy() -> SpectroscopyClass:
@@ -156,6 +156,13 @@ def check_existing(output_directory, files, type):
         f.partition('.fits')[0] for f in os.listdir(output_directory)
         if type in f
     ]
+
+    # also check for lowercase type
+    existing += [
+        f.partition('.fits')[0] for f in os.listdir(output_directory)
+        if type.lower() in f
+    ]
+
     if os.name == 'nt':  # on Windows, be careful with ':' in filename
         import re
         existing = [re.sub(r'T(\d+)_(\d+)_(\d+)', r'T\1:\2:\3', f) for f in existing]
@@ -175,7 +182,7 @@ def download(files, type, output_directory):
     """ Download files from DACE """
     Spectroscopy = load_spectroscopy()
     # with stdout_disabled(), all_logging_disabled():
-    Spectroscopy.download_files(files, file_type=type,
+    Spectroscopy.download_files(files, file_type=type.lower(),
                                 output_directory=output_directory)
 
 def extract_fits(output_directory):
@@ -194,15 +201,16 @@ def extract_fits(output_directory):
     return files
 
 
-def do_download_ccf(raw_files, output_directory, clobber=False, verbose=True):
-    """ Download CCFs from DACE """
+def do_download_filetype(type, raw_files, output_directory, clobber=False,
+                         verbose=True, chunk_size=20):
+    """ Download CCFs / S1Ds / S2Ds from DACE """
     raw_files = np.atleast_1d(raw_files)
 
     create_directory(output_directory)
 
     # check existing files to avoid re-downloading
     if not clobber:
-        raw_files = check_existing(output_directory, raw_files, 'CCF')
+        raw_files = check_existing(output_directory, raw_files, type)
 
     # any file left to download?
     if raw_files.size == 0:
@@ -212,68 +220,71 @@ def do_download_ccf(raw_files, output_directory, clobber=False, verbose=True):
 
     if verbose:
         n = raw_files.size
-        logger.info(f"Downloading {n} CCFs into '{output_directory}'...")
+        logger.info(f"Downloading {n} {type}s into '{output_directory}'...")
 
-    download(raw_files, 'ccf', output_directory)
+    # avoid an empty chunk
+    if chunk_size > n:
+        chunk_size = n
 
-    if verbose:
-        logger.info('Extracting .fits files')
+    for files in tqdm(zip(*(iter(raw_files),) * chunk_size), total=n // chunk_size):
+        download(files, type, output_directory)
+        extract_fits(output_directory)
 
-    extract_fits(output_directory)
+    logger.info('Extracted .fits files')
 
 
-def do_download_s1d(raw_files, output_directory, clobber=False, verbose=True):
-    """ Download S1Ds from DACE """
-    raw_files = np.atleast_1d(raw_files)
+# def do_download_s1d(raw_files, output_directory, clobber=False, verbose=True):
+#     """ Download S1Ds from DACE """
+#     raw_files = np.atleast_1d(raw_files)
 
-    create_directory(output_directory)
+#     create_directory(output_directory)
 
-    # check existing files to avoid re-downloading
-    if not clobber:
-        raw_files = check_existing(output_directory, raw_files, 'S1D')
+#     # check existing files to avoid re-downloading
+#     if not clobber:
+#         raw_files = check_existing(output_directory, raw_files, 'S1D')
 
-    # any file left to download?
-    if raw_files.size == 0:
-        if verbose:
-            logger.info('no files to download')
-        return
+#     # any file left to download?
+#     if raw_files.size == 0:
+#         if verbose:
+#             logger.info('no files to download')
+#         return
 
-    if verbose:
-        n = raw_files.size
-        logger.info(f"Downloading {n} S1Ds into '{output_directory}'...")
+#     if verbose:
+#         n = raw_files.size
+#         logger.info(f"Downloading {n} S1Ds into '{output_directory}'...")
 
-    download(raw_files, 's1d', output_directory)
+#     download(raw_files, 's1d', output_directory)
 
-    if verbose:
-        logger.info('Extracting .fits files')
+#     if verbose:
+#         logger.info('Extracting .fits files')
 
-    extract_fits(output_directory)
+#     extract_fits(output_directory)
 
 
-def do_download_s2d(raw_files, output_directory, clobber=False, verbose=True):
-    """ Download S2Ds from DACE """
-    raw_files = np.atleast_1d(raw_files)
+# def do_download_s2d(raw_files, output_directory, clobber=False, verbose=True):
+#     """ Download S2Ds from DACE """
+#     raw_files = np.atleast_1d(raw_files)
 
-    create_directory(output_directory)
+#     create_directory(output_directory)
 
-    # check existing files to avoid re-downloading
-    if not clobber:
-        raw_files = check_existing(output_directory, raw_files, 'S2D')
+#     # check existing files to avoid re-downloading
+#     if not clobber:
+#         raw_files = check_existing(output_directory, raw_files, 'S2D')
 
-    # any file left to download?
-    if raw_files.size == 0:
-        if verbose:
-            logger.info('no files to download')
-        return
+#     # any file left to download?
+#     if raw_files.size == 0:
+#         if verbose:
+#             logger.info('no files to download')
+#         return
 
-    if verbose:
-        n = raw_files.size
-        logger.info(f"Downloading {n} S2Ds into '{output_directory}'...")
+#     if verbose:
+#         n = raw_files.size
+#         logger.info(f"Downloading {n} S2Ds into '{output_directory}'...")
 
-    download(raw_files, 's2d', output_directory)
+#     download(raw_files, 's2d', output_directory)
 
-    if verbose:
-        logger.info('Extracting .fits files')
+#     if verbose:
+#         logger.info('Extracting .fits files')
 
-    extracted_files = extract_fits(output_directory)
-    return extracted_files
+#     extracted_files = extract_fits(output_directory)
+#     return extracted_files

arvi/extra_data.py

@@ -19,15 +19,24 @@ def get_extra_data(star, instrument=None, path=None, verbose=True):
     # print(metadata)
 
     files = glob(os.path.join(path, star + '*'))
+    files = [f for f in files if os.path.isfile(f)]
+    files = [f for f in files if not os.path.basename(f).endswith('.zip')]
+
     if len(files) == 0:
         raise FileNotFoundError
 
-    instruments = [os.path.basename(f).split('.')[0] for f in files]
-    instruments = [i.split('_', maxsplit=1)[1] for i in instruments]
+    def get_instruments(files):
+        instruments = [os.path.basename(f).split('.')[0] for f in files]
+        instruments = [i.split('_', maxsplit=1)[1] for i in instruments]
+        return instruments
+    
+    instruments = get_instruments(files)
 
     if instrument is not None:
         if not any([instrument in i for i in instruments]):
             raise FileNotFoundError
+        files = [f for f in files if instrument in f]
+        instruments = get_instruments(files)
 
     if verbose:
         logger.info(f'loading extra data for {star}')

arvi/plots.py

@@ -1,6 +1,8 @@
 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
@@ -15,7 +17,7 @@ from . import config
 
 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, **kwargs):
+         versus_n=False, show_histogram=False, bw=False, **kwargs):
     """ Plot the RVs
 
     Args:
@@ -38,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
@@ -58,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)
@@ -66,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
@@ -81,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:
@@ -147,9 +168,43 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
                 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()
@@ -340,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

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,
@@ -45,6 +48,8 @@ def run_query(query):
         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):
@@ -94,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])
 

arvi/stats.py

@@ -19,15 +19,19 @@ def wmean(a, e):
         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.
@@ -36,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))
 

arvi/timeseries.py

@@ -11,16 +11,15 @@ import numpy as np
 from astropy import units
 
 from .setup_logger import logger
-from .config import return_self
+from .config import return_self, check_internet
 from .translations import translate
-from .dace_wrapper import get_observations, get_arrays
-from .dace_wrapper import do_download_ccf, do_download_s1d, do_download_s2d
+from .dace_wrapper import do_download_filetype, get_observations, get_arrays
 from .simbad_wrapper import simbad
 from .extra_data import get_extra_data
 from .stats import wmean, wrms
 from .binning import bin_ccf_mask, binRV
 from .HZ import getHZ_period
-from .utils import strtobool
+from .utils import strtobool, there_is_internet
 
 
 @dataclass
@@ -71,6 +70,8 @@ class RV:
         self.__star__ = translate(self.star)
 
         if not self._child:
+            if check_internet and not there_is_internet():
+                raise ConnectionError('There is no internet connection?')
             # complicated way to query Simbad with self.__star__ or, if that
             # fails, try after removing a trailing 'A'
             for target in (self.__star__, self.__star__.replace('A', '')):
@@ -226,6 +227,8 @@ class RV:
     @property
     def N_nights(self) -> int:
         """ Number of individual nights """
+        if self.mtime.size == 0:
+            return 0
         return binRV(self.mtime, None, None, binning_bins=True).size - 1
 
     @property
@@ -402,7 +405,7 @@ class RV:
             files = [files]
 
         if star is None:
-            star_ = np.unique([os.path.splitext(f)[0].split('_')[0] for f in files])
+            star_ = np.unique([os.path.splitext(os.path.basename(f))[0].split('_')[0] for f in files])
             if star_.size == 1:
                 logger.info(f'assuming star is {star_[0]}')
                 star = star_[0]
@@ -439,12 +442,12 @@ class RV:
                     names = header.split()
 
             if len(names) > 3:
-                kw = dict(skip_header=2, dtype=None, encoding=None)
-                try:
-                    data = np.genfromtxt(f, **kw)
-                except ValueError:
+                kw = dict(skip_header=0, comments='--', names=True, dtype=None, encoding=None)
+                if '\t' in header:
                     data = np.genfromtxt(f, **kw, delimiter='\t')
-                data.dtype.names = names
+                else:
+                    data = np.genfromtxt(f, **kw)
+                # data.dtype.names = names
             else:
                 data = np.array([], dtype=np.dtype([]))
 
@@ -463,8 +466,10 @@ class RV:
 
             if 'rhk' in data.dtype.fields:
                 _s.rhk = data['rhk']
-                if 'srhk' in data.dtype.fields:
-                    _s.rhk_err = data['srhk']
+                _s.rhk_err = np.full_like(time, np.nan)
+                for possible_name in ['srhk', 'rhk_err']:
+                    if possible_name in data.dtype.fields:
+                        _s.rhk_err = data[possible_name]
             else:
                 _s.rhk = np.zeros_like(time)
                 _s.rhk_err = np.full_like(time, np.nan)
@@ -624,11 +629,12 @@ class RV:
                 setattr(self, q, arr)
 
 
-    def download_ccf(self, instrument=None, limit=None, directory=None):
+    def download_ccf(self, instrument=None, index=None, limit=None, directory=None, **kwargs):
         """ Download CCFs from DACE
 
         Args:
             instrument (str): Specific instrument for which to download data
+            index (int): Specific index of point for which to download data (0-based)
             limit (int): Maximum number of files to download.
             directory (str): Directory where to store data.
         """
@@ -638,18 +644,27 @@ class RV:
         if instrument is None:
             files = [file for file in self.raw_file if file.endswith('.fits')]
         else:
-            instrument = self._check_instrument(instrument)
+            strict = kwargs.pop('strict', False)
+            instrument = self._check_instrument(instrument, strict=strict)
             files = []
             for inst in instrument:
                 files += list(getattr(self, inst).raw_file)
 
-        do_download_ccf(files[:limit], directory)
+        if index is not None:
+            index = np.atleast_1d(index)
+            files = list(np.array(files)[index])
+
+        # remove empty strings
+        files = list(filter(None, files))
+
+        do_download_filetype('CCF', files[:limit], directory, **kwargs)
 
-    def download_s1d(self, instrument=None, limit=None, directory=None):
+    def download_s1d(self, instrument=None, index=None, limit=None, directory=None, **kwargs):
         """ Download S1Ds from DACE
 
         Args:
             instrument (str): Specific instrument for which to download data
+            index (int): Specific index of point for which to download data (0-based)
             limit (int): Maximum number of files to download.
             directory (str): Directory where to store data.
         """
@@ -659,18 +674,27 @@ class RV:
         if instrument is None:
             files = [file for file in self.raw_file if file.endswith('.fits')]
         else:
-            instrument = self._check_instrument(instrument)
+            strict = kwargs.pop('strict', False)
+            instrument = self._check_instrument(instrument, strict=strict)
             files = []
             for inst in instrument:
                 files += list(getattr(self, inst).raw_file)
 
-        do_download_s1d(files[:limit], directory)
+        if index is not None:
+            index = np.atleast_1d(index)
+            files = list(np.array(files)[index])
 
-    def download_s2d(self, instrument=None, limit=None, directory=None):
+        # remove empty strings
+        files = list(filter(None, files))
+
+        do_download_filetype('S1D', files[:limit], directory, **kwargs)
+
+    def download_s2d(self, instrument=None, index=None, limit=None, directory=None, **kwargs):
         """ Download S2Ds from DACE
 
         Args:
             instrument (str): Specific instrument for which to download data
+            index (int): Specific index of point for which to download data (0-based)
             limit (int): Maximum number of files to download.
             directory (str): Directory where to store data.
         """
@@ -680,12 +704,20 @@ class RV:
         if instrument is None:
             files = [file for file in self.raw_file if file.endswith('.fits')]
         else:
-            instrument = self._check_instrument(instrument)
+            strict = kwargs.pop('strict', False)
+            instrument = self._check_instrument(instrument, strict=strict)
             files = []
             for inst in instrument:
                 files += list(getattr(self, inst).raw_file)
 
-        extracted_files = do_download_s2d(files[:limit], directory)
+        if index is not None:
+            index = np.atleast_1d(index)
+            files = list(np.array(files)[index])
+
+        # remove empty strings
+        files = list(filter(None, files))
+
+        do_download_filetype('S2D', files[:limit], directory, **kwargs)
 
 
     from .plots import plot, plot_fwhm, plot_bis, plot_rhk, plot_quantity
@@ -757,7 +789,9 @@ class RV:
             return self
 
     def remove_point(self, index):
-        """ Remove individual observations at a given index (or indices)
+        """
+        Remove individual observations at a given index (or indices).
+        NOTE: Like Python, the index is 0-based.
         
         Args:
             index (int, list, ndarray):
@@ -1203,7 +1237,8 @@ class RV:
             self._build_arrays()
 
 
-    def save(self, directory=None, instrument=None, full=False, save_nans=True):
+    def save(self, directory=None, instrument=None, full=False,
+             save_masked=False, save_nans=True):
         """ Save the observations in .rdb files.
 
         Args:
@@ -1238,11 +1273,18 @@ class RV:
                 continue
 
             if full:
-                d = np.c_[
-                    _s.mtime, _s.mvrad, _s.msvrad,
-                    _s.fwhm[_s.mask], _s.fwhm_err[_s.mask],
-                    _s.rhk[_s.mask], _s.rhk_err[_s.mask],
-                ]
+                if save_masked:
+                    d = np.c_[
+                        _s.time, _s.vrad, _s.svrad,
+                        _s.fwhm, _s.fwhm_err,
+                        _s.rhk, _s.rhk_err,
+                    ]
+                else:
+                    d = np.c_[
+                        _s.mtime, _s.mvrad, _s.msvrad,
+                        _s.fwhm[_s.mask], _s.fwhm_err[_s.mask],
+                        _s.rhk[_s.mask], _s.rhk_err[_s.mask],
+                    ]
                 if not save_nans:
                     if np.isnan(d).any():
                         # remove observations where any of the indicators are # NaN
@@ -1254,7 +1296,10 @@ class RV:
                 header =  'bjd\tvrad\tsvrad\tfwhm\tsfwhm\trhk\tsrhk\n'
                 header += '---\t----\t-----\t----\t-----\t---\t----'
             else:
-                d = np.c_[_s.mtime, _s.mvrad, _s.msvrad]
+                if save_masked:
+                    d = np.c_[_s.time, _s.vrad, _s.svrad]
+                else:
+                    d = np.c_[_s.mtime, _s.mvrad, _s.msvrad]
                 header = 'bjd\tvrad\tsvrad\n---\t----\t-----'
             
             file = f'{star_name}_{inst}.rdb'

arvi/utils.py

@@ -13,6 +13,12 @@ import logging
 from glob import glob
 import numpy as np
 
+try:
+    from tqdm import tqdm, trange
+except ImportError:
+    tqdm = lambda x, *args, **kwargs: x
+    trange = lambda *args, **kwargs: range(*args, **kwargs)
+
 
 def create_directory(directory):
     """ Create a directory if it does not exist """
@@ -70,6 +76,14 @@ def strtobool(val):
     else:
         raise ValueError("invalid truth value {!r}".format(val))
 
+def there_is_internet(timeout=1):
+    from socket import create_connection
+    try:
+        create_connection(('8.8.8.8', 53), timeout=timeout)
+        return True
+    except OSError:
+        pass
+    return False
 
 def find_data_file(file):
     here = os.path.dirname(os.path.abspath(__file__))

{arvi-0.1.10.dist-info → arvi-0.1.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: arvi
-Version: 0.1.10
+Version: 0.1.11
 Summary: The Automated RV Inspector
 Author-email: João Faria <joao.faria@unige.ch>
 License: MIT

{arvi-0.1.10.dist-info → arvi-0.1.11.dist-info}/RECORD

@@ -1,29 +1,29 @@
 arvi/HZ.py,sha256=btqDFTxHeQlmvXOQKrjOpP8YFtmFjGhSZwsYYehLOaM,2885
 arvi/__init__.py,sha256=XjBsNGlwuLftPK50u7SOmYKpmOpt9X5AIqA9LdZ8GyU,344
 arvi/binning.py,sha256=G1nZvH2ev02quDPZCReygLkH7tX6B3hFAOoOYTfhLec,15249
-arvi/config.py,sha256=VM23kaAuyOWp9clfCIztJEnT5BWbG7o8see0UEqDeK0,20
-arvi/dace_wrapper.py,sha256=kDKcMratEDCfWT0ckexKjeuKLaOoULaaRcThzQQYT3s,9665
-arvi/extra_data.py,sha256=4NTYc6btntuuiFwQtRddxkjDqV9fAv-9P2d1JlWujhw,2211
+arvi/config.py,sha256=XqLmIsPl9dJvco1ngANG5W868BexfuJJEaRkuKZsAK8,43
+arvi/dace_wrapper.py,sha256=bfNzHKQUryCkZsJdEJImJY3oLIwVjflJOFQRRJ4RoAE,10125
+arvi/extra_data.py,sha256=WEEaYeLh52Zdv0uyHO72Ys5MWS3naTAP4wJV2BJ1mbk,2551
 arvi/instrument_specific.py,sha256=gDNu6xM0nL3AItAxipTdg3dIGaGGi7tNWE5nnad5MxM,2873
 arvi/lbl_wrapper.py,sha256=_ViGVkpakvuBR_xhu9XJRV5EKHpj5Go6jBZGJZMIS2Y,11850
 arvi/nasaexo_wrapper.py,sha256=xrdF4KKXmlZAqu4IQtdcqfOeUXJXoui5M-eZpnVXNlk,7311
-arvi/plots.py,sha256=7nHlXAkioJi4dab70NzOZUE7BZsA_moXm3vZTSHpatU,13968
+arvi/plots.py,sha256=J1FR-7NgWL0zdjtD0Sumj6ZHLXSRC_EtnpXkTguxQ-s,17549
 arvi/programs.py,sha256=OHmOW1BJUUg7xmzpTDzyAunxiGZ77QpVIN4Z9HxmE4M,3556
 arvi/reports.py,sha256=yrdajC-zz5_kH1Ucc6cU70DK-5dpG0Xyeru-EITKpNo,3355
 arvi/setup_logger.py,sha256=pBzaRTn0hntozjbaRVx0JIbWGuENkvYUApa6uB-FsRo,279
-arvi/simbad_wrapper.py,sha256=Yfo8qjx8C7DPg37D6GJwWm7PH9up_5Wxgmxx-BF4p_c,4669
-arvi/stats.py,sha256=5Itcsaqg6AVdDxluWks_36AHdr8xIXPuXxKtYe5YXCI,2076
-arvi/timeseries.py,sha256=QC-QA3zyl_y0XcFRy3g0urivFVvUS9KUJ8lgL6-C8eI,50318
+arvi/simbad_wrapper.py,sha256=F_VOHBI-KqT9q06sO2n2F1Ov6Un_1gAMP3ItBMqx-Yw,5284
+arvi/stats.py,sha256=MQiyLvdiAFxIC29uajTy8kuxD-b2Y6mraL4AfWkRJkM,2576
+arvi/timeseries.py,sha256=IvpjJqRL97bOwm4RpAGqvevkySj0Odjzx6vPtvlKzSI,52294
 arvi/translations.py,sha256=eyUJei8wlsyBecTz8E_ntpP35-Mre2Tzm_mUMMGaZWY,150
-arvi/utils.py,sha256=pSXyhFTw-0tzjLBzILeyDE7PjEFDvchn-oLoF-Ww7mI,3302
+arvi/utils.py,sha256=ay3GouZYqn66zF43KgSs8Ngsfb4YInKZEmk70YqVDoA,3676
 arvi/data/info.svg,sha256=0IMI6W-eFoTD8acnury79WJJakpBwLa4qKS4JWpsXiI,489
 arvi/data/obs_affected_ADC_issues.dat,sha256=pbwzvRBjS8KsLrKi1Cdgfihvo3hK3FUN8Csdysw-vAw,10653
 arvi/data/obs_affected_blue_cryostat_issues.dat,sha256=z4AK17xfz8tGTDv1FjRvQFnio4XA6PNNfDXuicewHk4,1771
 arvi/data/extra/HD86226_PFS1.rdb,sha256=vfAozbrKHM_j8dYkCBJsuHyD01KEM1asghe2KInwVao,3475
 arvi/data/extra/HD86226_PFS2.rdb,sha256=F2P7dB6gVyzCglUjNheB0hIHVClC5RmARrGwbrY1cfo,4114
 arvi/data/extra/metadata.json,sha256=C69hIw6CohyES6BI9vDWjxwSz7N4VOYX0PCgjXtYFmU,178
-arvi-0.1.10.dist-info/LICENSE,sha256=6JfQgl7SpM55t0EHMFNMnNh-AdkpGW25MwMiTnhdWQg,1068
-arvi-0.1.10.dist-info/METADATA,sha256=RzxIuuUzpwW2Y9W_-AId3iGwulGfLEaDL5dVvw-9Lls,1306
-arvi-0.1.10.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-arvi-0.1.10.dist-info/top_level.txt,sha256=4EeiKDVLD45ztuflTGfQ3TU8GVjJg5Y95xS5XjI-utU,5
-arvi-0.1.10.dist-info/RECORD,,
+arvi-0.1.11.dist-info/LICENSE,sha256=6JfQgl7SpM55t0EHMFNMnNh-AdkpGW25MwMiTnhdWQg,1068
+arvi-0.1.11.dist-info/METADATA,sha256=cb2feqRY6mVlakZeoO86IYmPQzFLTWP8IboRNaBKVrw,1306
+arvi-0.1.11.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+arvi-0.1.11.dist-info/top_level.txt,sha256=4EeiKDVLD45ztuflTGfQ3TU8GVjJg5Y95xS5XjI-utU,5
+arvi-0.1.11.dist-info/RECORD,,