arvi 0.2.8__tar.gz → 0.2.9__tar.gz

{arvi-0.2.8/arvi.egg-info → arvi-0.2.9}/PKG-INFO

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

{arvi-0.2.8 → arvi-0.2.9}/arvi/dace_wrapper.py

@@ -10,7 +10,7 @@ from .setup_logger import setup_logger
 from .utils import create_directory, all_logging_disabled, stdout_disabled, timer, tqdm
 
 
-def load_spectroscopy(user=None):
+def load_spectroscopy(user=None, verbose=True):
     logger = setup_logger()
     with all_logging_disabled():
         from dace_query.spectroscopy import SpectroscopyClass, Spectroscopy as default_Spectroscopy
@@ -19,7 +19,8 @@ def load_spectroscopy(user=None):
     from .config import config
     # requesting as public
     if config.request_as_public:
-        logger.warning('requesting DACE data as public')
+        if verbose:
+            logger.warning('requesting DACE data as public')
         with all_logging_disabled():
             dace = DaceClass(dace_rc_config_path='none')
         return SpectroscopyClass(dace_instance=dace)
@@ -168,7 +169,7 @@ def get_observations_from_instrument(star, instrument, user=None, main_id=None,
         dict:
             dictionary with data from DACE
     """
-    Spectroscopy = load_spectroscopy(user)
+    Spectroscopy = load_spectroscopy(user, verbose)
     found_dace_id = False
     with timer('dace_id query'):
         try:
@@ -283,7 +284,7 @@ def get_observations_from_instrument(star, instrument, user=None, main_id=None,
 def get_observations(star, instrument=None, user=None, main_id=None, verbose=True):
     logger = setup_logger()
     if instrument is None:
-        Spectroscopy = load_spectroscopy(user)
+        Spectroscopy = load_spectroscopy(user, verbose)
 
         try:
             with stdout_disabled(), all_logging_disabled():

{arvi-0.2.8 → arvi-0.2.9}/arvi/plots.py

@@ -1,6 +1,7 @@
 from functools import partialmethod, wraps
 from itertools import cycle
 
+import matplotlib.pyplot as plt
 import numpy as np
 
 from astropy.timeseries import LombScargle
@@ -9,9 +10,6 @@ from .setup_logger import setup_logger
 from .config import config
 from .stats import wmean
 
-from .utils import lazy_import
-plt = lazy_import('matplotlib.pyplot')
-
 
 def plot_settings(func):
     @wraps(func)

{arvi-0.2.8 → arvi-0.2.9}/arvi/timeseries.py

@@ -6,6 +6,8 @@ from glob import glob
 import warnings
 from copy import deepcopy
 from datetime import datetime, timezone
+
+# import lazy_loader as lazy
 import numpy as np
 
 from .setup_logger import setup_logger
@@ -24,10 +26,11 @@ from .HZ import getHZ_period
 from .instrument_specific import ISSUES
 from .reports import REPORTS
 from .utils import sanitize_path, strtobool, there_is_internet, timer, chdir
-from .utils import lazy_import
+# from .utils import lazy_import
 
-units = lazy_import('astropy.units')
-# from astropy import units
+# units = lazy_import('astropy.units')
+# units = lazy.load('astropy.units')
+from astropy import units
 
 class ExtraFields:
     @property
@@ -577,7 +580,8 @@ class RV(ISSUES, REPORTS):
                     # --> not just in rhk and rhk_err...
                     if data[arr].dtype == float and (bad := data[arr] == -99999).any():
                         data[arr][bad] = np.nan
-
+                    if data[arr].dtype == float and (bad := data[arr] == -99).any():
+                        data[arr][bad] = np.nan
                     setattr(s, arr, data[arr][ind])
                     s._quantities.append(arr)
 
@@ -1912,6 +1916,8 @@ class RV(ISSUES, REPORTS):
 
         # create copy of self to be returned
         snew = deepcopy(self)
+        # store original object
+        snew._unbinned = deepcopy(self)
 
         all_bad_quantities = []
 
@@ -2101,23 +2107,32 @@ class RV(ISSUES, REPORTS):
         if config.return_self:
             return self
 
-    def detrend(self, degree=1):
-        """ Detrend the RVs of all instruments """
+    def detrend(self, degree: int=1):
+        """
+        Detrend the RVs of all instruments using a polynomial of degree `degree`
+        """
         instrument_indices = np.unique_inverse(self.instrument_array).inverse_indices
-        def fun(p, t, degree, ninstruments, just_model=False, index=None):
+        instrument_indices_masked = np.unique_inverse(self.instrument_array[self.mask]).inverse_indices
+
+        def fun(p, t, degree, ninstruments, just_model=False, index=None, masked=True):
             polyp, offsets = p[:degree], p[-ninstruments:]
             polyp = np.r_[polyp, 0.0]
             if index is None:
-                model = offsets[instrument_indices] + np.polyval(polyp, t)
+                if masked:
+                    model = offsets[instrument_indices_masked] + np.polyval(polyp, t)
+                else:
+                    model = offsets[instrument_indices] + np.polyval(polyp, t)
             else:
                 model = offsets[index] + np.polyval(polyp, t)
             if just_model:
                 return model
             return self.mvrad - model
+
         coef = np.polyfit(self.mtime, self.mvrad, degree)
         x0 = np.append(coef, [0.0] * (len(self.instruments) - 1))
-        print(x0)
+        # print(x0)
         fun(x0, self.mtime, degree, len(self.instruments))
+
         from scipy.optimize import leastsq
         xbest, _ = leastsq(fun, x0, args=(self.mtime, degree, len(self.instruments)))
 
@@ -2127,12 +2142,13 @@ class RV(ISSUES, REPORTS):
         self.plot(ax=ax)
         for i, inst in enumerate(self.instruments):
             s = getattr(self, inst)
-            ax.plot(s.time, fun(xbest, s.time, degree, len(self.instruments), just_model=True, index=i),
+            ax.plot(s.time,
+                    fun(xbest, s.time, degree, len(self.instruments), just_model=True, index=i, masked=False),
                     color=f'C{i}')
         ax.set_title('original', loc='left', fontsize=10)
         ax.set_title(f'coefficients: {xbest[:degree]}', loc='right', fontsize=10)
 
-        self.add_to_vrad(-fun(xbest, self.time, degree, len(self.instruments), just_model=True))
+        self.add_to_vrad(-fun(xbest, self.time, degree, len(self.instruments), just_model=True, masked=False))
         ax = fig.add_subplot(2, 1, 2)
         self.plot(ax=ax)
         ax.set_title('detrended', loc='left', fontsize=10)
@@ -2141,7 +2157,7 @@ class RV(ISSUES, REPORTS):
         # axs[1].errorbar(self.mtime, fun(xbest, self.mtime, degree, len(self.instruments)), self.msvrad, fmt='o')
 
         return
-        
+
 
 
 
@@ -2301,24 +2317,28 @@ class RV(ISSUES, REPORTS):
         self.units = new_units
 
 
-    def put_at_systemic_velocity(self):
+    def put_at_systemic_velocity(self, factor=1.0):
         """
-        For instruments in which mean(RV) < ptp(RV), "move" RVs to the systemic
-        velocity from simbad. This is useful if some instruments are centered
-        at zero while others are not, and instead of calling `.adjust_means()`,
-        but it only works when the systemic velocity is smaller than ptp(RV).
+        For instruments in which mean(RV) < `factor` * ptp(RV), "move" RVs to
+        the systemic velocity from simbad. This is useful if some instruments
+        are centered at zero while others are not, and instead of calling
+        `.adjust_means()`, but it only works when the systemic velocity is
+        smaller than `factor` * ptp(RV).
         """
         changed = False
         for inst in self.instruments:
+            changed_inst = False
             s = getattr(self, inst)
             if s.mask.any():
-                if np.abs(s.mvrad.mean()) < np.ptp(s.mvrad):
+                if np.abs(s.mvrad.mean()) < factor * np.ptp(s.mvrad):
                     s.vrad += self.simbad.rvz_radvel * 1e3
-                    changed = True
+                    changed = changed_inst = True
             else:  # all observations are masked, use non-masked arrays
-                if np.abs(s.vrad.mean()) < np.ptp(s.vrad):
+                if np.abs(s.vrad.mean()) < factor * np.ptp(s.vrad):
                     s.vrad += self.simbad.rvz_radvel * 1e3
-                    changed = True
+                    changed = changed_inst = True
+            if changed_inst and self.verbose:
+                logger.info(f"putting {inst} RVs at systemic velocity")
         if changed:
             self._build_arrays()
 
@@ -2341,33 +2361,66 @@ class RV(ISSUES, REPORTS):
             self._build_arrays()
 
 
-    def save(self, directory=None, instrument=None, full=False, postfix=None,
-             save_masked=False, save_nans=True):
-        """ Save the observations in .rdb files.
+    def save(self, directory=None, instrument=None, format='rdb',
+             indicators=False, join_instruments=False, postfix=None,
+             save_masked=False, save_nans=True, **kwargs):
+        """ Save the observations in .rdb or .csv files.
 
         Args:
             directory (str, optional):
                 Directory where to save the .rdb files.
             instrument (str, optional):
                 Instrument for which to save observations.
-            full (bool, optional):
-                Save just RVs and errors (False) or more indicators (True).
+            format (str, optional):
+                Format to use ('rdb' or 'csv').
+            indicators (bool, str, list[str], optional):
+                Save only RVs and errors (False) or more indicators. If True,
+                use a default list, if `str`, use an existing list, if list[str]
+                provide a sequence of specific indicators.
+            join_instruments (bool, optional):
+                Join all instruments in a single file.
             postfix (str, optional):
                 Postfix to add to the filenames ([star]_[instrument]_[postfix].rdb).
+            save_masked (bool, optional)
+                If True, also save masked observations (those for which
+                self.mask == True)
             save_nans (bool, optional)
                 Whether to save NaN values in the indicators, if they exist. If
                 False, the full observation which contains NaN values is not saved.
         """
+        if format not in ('rdb', 'csv'):
+            logger.error(f"format must be 'rdb' or 'csv', got '{format}'")
+            return
+
         star_name = self.star.replace(' ', '')
 
-        if directory is None:
-            directory = '.'
-        else:
+        if directory is not None:
             os.makedirs(directory, exist_ok=True)
 
+        indicator_sets = {
+            "default": [
+                "fwhm", "fwhm_err",
+                "bispan", "bispan_err",
+                "contrast", "contrast_err",
+                "rhk", "rhk_err",
+                "berv",
+            ],
+            "CORALIE": [
+                "fwhm", "fwhm_err",
+                "bispan", "bispan_err",
+                "contrast", "contrast_err",
+                "haindex", "haindex_err",
+                "berv",
+            ],
+        } 
+
+        if 'full' in kwargs:
+            logger.warning('argument `full` is deprecated, use `indicators` instead')
+            indicators = kwargs['full']
+
         files = []
 
-        for inst in self.instruments:
+        for _i, inst in enumerate(self.instruments):
             if instrument is not None:
                 if instrument not in inst:
                     continue
@@ -2377,75 +2430,95 @@ class RV(ISSUES, REPORTS):
             if not _s.mask.any():  # all observations are masked, don't save
                 continue
 
-            if full:
-                if save_masked:
-                    arrays = [
-                        _s.time, _s.vrad, _s.svrad,
-                        _s.fwhm, _s.fwhm_err,
-                        _s.bispan, _s.bispan_err,
-                        _s.contrast, _s.contrast_err,
-                        _s.rhk, _s.rhk_err,
-                        _s.berv,
-                    ]
-                else:
-                    arrays = [
-                        _s.mtime, _s.mvrad, _s.msvrad,
-                        _s.fwhm[_s.mask], _s.fwhm_err[_s.mask],
-                        _s.bispan[_s.mask], _s.bispan_err[_s.mask],
-                        _s.contrast[_s.mask], _s.contrast_err[_s.mask],
-                        _s.rhk[_s.mask], _s.rhk_err[_s.mask],
-                        _s.berv[_s.mask],
-                    ]
-                if not save_nans:
-                    raise NotImplementedError
-                    # if np.isnan(d).any():
-                    #     # remove observations where any of the indicators are # NaN
-                    #     nan_mask = np.isnan(d[:, 3:]).any(axis=1)
-                    #     d = d[~nan_mask]
-                    #     if self.verbose:
-                    #         logger.warning(f'masking {nan_mask.sum()} observations with NaN in indicators')
-
-                header = '\t'.join(['rjd', 'vrad', 'svrad',
-                                    'fwhm', 'sig_fwhm',
-                                    'bispan', 'sig_bispan',
-                                    'contrast', 'sig_contrast',
-                                    'rhk', 'sig_rhk',
-                                    'berv',
-                                    ])
-                header += '\n'
-                header += '\t'.join(['-' * len(c) for c in header.strip().split('\t')])
+            if save_masked:
+                arrays = [_s.time, _s.vrad, _s.svrad]
+                if join_instruments:
+                    arrays += [_s.instrument_array]
+            else:
+                arrays = [_s.mtime, _s.mvrad, _s.msvrad]
+                if join_instruments:
+                    arrays += [_s.instrument_array[_s.mask]]
+            
+            if indicators in (False, None):
+                indicator_names = []
+            else:
+                if indicators is True:
+                    indicator_names = indicator_sets["default"]
+                elif isinstance(indicators, str):
+                    try:
+                        indicator_names = indicator_sets[indicators]
+                    except KeyError:
+                        logger.error(f"unknown indicator set '{indicators}'")
+                        logger.error(f"available: {list(indicator_sets.keys())}")
+                        return
+                elif isinstance(indicators, list) and all(isinstance(i, str) for i in indicators):
+                    indicator_names = indicators
 
+            if save_masked:
+                arrays += [getattr(_s, ind) for ind in indicator_names]
             else:
-                if save_masked:
-                    arrays = [_s.time, _s.vrad, _s.svrad]
-                else:
-                    arrays = [_s.mtime, _s.mvrad, _s.msvrad]
+                arrays += [getattr(_s, ind)[_s.mask] for ind in indicator_names]
+
+            d = np.stack(arrays, axis=1)
+            if not save_nans:
+                # raise NotImplementedError
+                if np.isnan(d).any():
+                    # remove observations where any of the indicators are # NaN
+                    nan_mask = np.isnan(d[:, 3:]).any(axis=1)
+                    d = d[~nan_mask]
+                    if self.verbose:
+                        msg = f'{inst}: masking {nan_mask.sum()} observations with NaN in indicators'
+                        logger.warning(msg)
 
-                # d = np.stack(arrays, axis=1)
-                header = 'rjd\tvrad\tsvrad\n---\t----\t-----'
+            cols = ['rjd', 'vrad', 'svrad']
+            cols += ['inst'] if join_instruments else []
+            cols += indicator_names
+
+            if format == 'rdb':
+                header = '\t'.join(cols)
+                header += '\n'
+                header += '\t'.join(['-' * len(c) for c in header.strip().split('\t')])
+            else:
+                header = ','.join(cols)
 
-            file = f'{star_name}_{inst}.rdb'
-            if postfix is not None:
-                file = f'{star_name}_{inst}_{postfix}.rdb'
+            if join_instruments:
+                file = f'{star_name}.{format}'
+                if postfix is not None:
+                    file = f'{star_name}_{postfix}.{format}'
+            else:
+                file = f'{star_name}_{inst}.{format}'
+                if postfix is not None:
+                    file = f'{star_name}_{inst}_{postfix}.{format}'
 
+            if directory is not None:
+                file = os.path.join(directory, file)
             files.append(file)
-            file = os.path.join(directory, file)
 
             N = len(arrays[0])
-            with open(file, 'w') as f:
-                f.write(header + '\n')
+            with open(file, 'a' if join_instruments and _i != 0 else 'w') as f:
+                if join_instruments and _i != 0:
+                    pass
+                else:
+                    f.write(header + '\n')
+
                 for i in range(N):
                     for j, a in enumerate(arrays):
                         f.write(str(a[i]))
                         if j < len(arrays) - 1:
-                            f.write('\t')
+                            f.write('\t' if format == 'rdb' else ',')
                     f.write('\n')
 
             # np.savetxt(file, d, header=header, delimiter='\t', comments='', fmt='%f')
 
-            if self.verbose:
+            if self.verbose and not join_instruments:
                 logger.info(f'saving to {file}')
 
+        if self.verbose and join_instruments:
+            logger.info(f'saving to {files[0]}')
+
+        if join_instruments:
+            files = [files[0]]
+
         return files
 
     def checksum(self, write_to=None):

{arvi-0.2.8 → arvi-0.2.9}/arvi/utils.py

@@ -2,6 +2,7 @@ import os
 import sys
 import time
 from contextlib import contextmanager
+
 try:
     from unittest.mock import patch
 except ImportError:
@@ -98,8 +99,9 @@ def sanitize_path(path):
     path = path.replace('*', '_')
     return path
 
-def pretty_print_table(rows, line_between_rows=True, logger=None):
-    """
+def pretty_print_table(rows, line_between_rows=True, string=False, 
+                       markdown=False, latex=False, logger=None):
+    r"""
     Example Output
     ┌──────┬─────────────┬────┬───────┐
     │ True │ short       │ 77 │ catty │
@@ -110,25 +112,76 @@ def pretty_print_table(rows, line_between_rows=True, logger=None):
     └──────┴─────────────┴────┴───────┘
     """
     _print = logger.info if logger else print
+    if string:
+        def _print(x, s):
+            s += x + '\n'
+            return s
+    else:
+        if logger:
+            def _print(x, _):
+                logger.info(x)
+        else:
+            def _print(x, _):
+                print(x)
+
+    if latex or markdown:
+        line_between_rows = False
+
+    s = ''
 
     # find the max length of each column
     max_col_lens = list(map(max, zip(*[(len(str(cell)) for cell in row) for row in rows])))
 
+    if markdown:
+        bar_char = '|'
+    else:
+        bar_char = r'│'
+
     # print the table's top border
-    _print('┌' + '┬'.join('─' * (n + 2) for n in max_col_lens) + '┐')
+    if markdown:
+        pass
+    elif latex:
+        s = _print(r'\begin{table*}', s)
+        # s = _print(r'\centering', s)
+        s = _print(r'\begin{tabular}' + '{' + ' c ' * len(rows[0]) + '}', s)
+    else:
+        s = _print(r'┌' + r'┬'.join(r'─' * (n + 2) for n in max_col_lens) + r'┐', s)
 
-    rows_separator = '├' + '┼'.join('─' * (n + 2) for n in max_col_lens) + '┤'
+    if markdown:
+        header_separator = bar_char + bar_char.join('-' * (n + 2) for n in max_col_lens) + bar_char
 
-    row_fstring = ' │ '.join("{: <%s}" % n for n in max_col_lens)
+    rows_separator = r'├' + r'┼'.join(r'─' * (n + 2) for n in max_col_lens) + r'┤'
+
+    if latex:
+        row_fstring = ' & '.join("{: <%s}" % n for n in max_col_lens)
+    else:
+        row_fstring = bar_char.center(3).join("{: <%s}" % n for n in max_col_lens)
 
     for i, row in enumerate(rows):
-        _print('│ ' + row_fstring.format(*map(str, row)) + ' │')
+        if markdown and i == 1:
+            s = _print(header_separator, s)
+
+        if latex:
+            s = _print(row_fstring.format(*map(str, row)) + r' \\', s)
+        else:
+            s = _print(bar_char + ' ' + row_fstring.format(*map(str, row)) + ' ' + bar_char, s)
         
+
         if line_between_rows and i < len(rows) - 1:
-            _print(rows_separator)
+            s = _print(rows_separator, s)
+
 
     # print the table's bottom border
-    _print('└' + '┴'.join('─' * (n + 2) for n in max_col_lens) + '┘')
+    if markdown:
+        pass
+    elif latex:
+        s = _print(r'\end{tabular}', s)
+        s = _print(r'\end{table*}', s)
+    else:
+        s = _print(r'└' + r'┴'.join(r'─' * (n + 2) for n in max_col_lens) + r'┘', s)
+
+    if string:
+        return s
 
 
 def strtobool(val):
@@ -238,3 +291,28 @@ def get_object_fast(file):
         value = f.read(20)
     return value.decode().split("'")[1].strip()
 
+
+def get_simbad_oid(self):
+    import requests
+    if isinstance(self, str):
+        star = self
+    else:
+        star = self.star
+    oid = requests.post('https://simbad.cds.unistra.fr/simbad/sim-tap/sync', 
+                        data=dict(format='text', request='doQuery', lang='adql', phase='run', 
+                                  query=f"SELECT basic.OID FROM basic JOIN ident ON oidref = oid WHERE id = '{star}';"))
+    oid = oid.text.split()[-1]
+    return oid
+
+
+
+# from https://stackoverflow.com/questions/37765197/darken-or-lighten-a-color-in-matplotlib
+def adjust_lightness(color, amount=0.5):
+    import matplotlib.colors as mc
+    import colorsys
+    try:
+        c = mc.cnames[color]
+    except KeyError:
+        c = color
+    c = colorsys.rgb_to_hls(*mc.to_rgb(c))
+    return colorsys.hls_to_rgb(c[0], max(0, min(1, amount * c[1])), c[2])

{arvi-0.2.8 → arvi-0.2.9/arvi.egg-info}/PKG-INFO

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