arvi 0.2.7__tar.gz → 0.2.9__tar.gz

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

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

{arvi-0.2.7 → 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)
@@ -95,7 +96,7 @@ def get_arrays(result, latest_pipeline=True, ESPRESSO_mode='HR11', NIRPS_mode='H
             npipe = len(pipelines)
             if 'NIRPS' in inst and any(['LBL' in p for p in pipelines]):
                 # TODO: correctly load both CCF and LBL
-                pipelines = [pipelines[1]]
+                pipelines = [pipelines[0]]
             if 'HARPS' in inst and npipe > 1 and pipelines[1] == pipelines[0] + '-EGGS':
                 pipelines = pipelines[:2]
             else:
@@ -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():
@@ -544,6 +545,7 @@ def do_download_filetype(type, raw_files, output_directory, clobber=False, user=
     """ Download CCFs / S1Ds / S2Ds from DACE """
     logger = setup_logger()
     raw_files = np.atleast_1d(raw_files)
+    raw_files_original = raw_files.copy()
 
     create_directory(output_directory)
 
@@ -557,7 +559,7 @@ def do_download_filetype(type, raw_files, output_directory, clobber=False, user=
     if n == 0:
         if verbose:
             logger.info('no files to download')
-        return
+        return list(map(os.path.basename, raw_files_original))
 
     # avoid an empty chunk
     if chunk_size > n:
@@ -575,7 +577,9 @@ def do_download_filetype(type, raw_files, output_directory, clobber=False, user=
 
     if n < parallel_limit:
         iterator = [raw_files[i:i + chunk_size] for i in range(0, n, chunk_size)]
-        for files in tqdm(iterator, total=len(iterator)):
+        if len(iterator) > 1:
+            iterator = tqdm(iterator, total=len(iterator))
+        for files in iterator:
             download(files, type, output_directory, quiet=False, user=user)
             extract_fits(output_directory)
 
@@ -604,6 +608,7 @@ def do_download_filetype(type, raw_files, output_directory, clobber=False, user=
 
     sys.stdout.flush()
     logger.info('extracted .fits files')
+    return list(map(os.path.basename, raw_files_original))
 
 
 # def do_download_s1d(raw_files, output_directory, clobber=False, verbose=True):

{arvi-0.2.7 → arvi-0.2.9}/arvi/instrument_specific.py

@@ -127,9 +127,9 @@ def HARPS_commissioning(self, mask=True, plot=True):
     total_affected = affected.sum()
 
     if self.verbose:
-        n = total_affected
-        logger.info(f"there {'are'[:n^1]}{'is'[n^1:]} {n} frame{'s'[:n^1]} "
-                     "during HARPS commissioning")
+        n, i = total_affected, int(total_affected != 1)
+        logger.info(f"there {['is', 'are'][i]} {n} frame{['', 's'][i]} "
+                    "during HARPS commissioning")
 
     if mask:
         self.mask[affected] = False
@@ -159,9 +159,9 @@ def HARPS_fiber_commissioning(self, mask=True, plot=True):
     total_affected = affected.sum()
 
     if self.verbose:
-        n = total_affected
-        logger.info(f"there {'are'[:n^1]}{'is'[n^1:]} {n} frame{'s'[:n^1]} "
-                     "during the HARPS fiber commissioning period")
+        n, i = total_affected, int(total_affected != 1)
+        logger.info(f"there {['is', 'are'][i]} {n} frame{['', 's'][i]} "
+                    "during the HARPS fiber commissioning period")
 
     if mask:
         self.mask[affected] = False
@@ -191,15 +191,15 @@ def ESPRESSO_commissioning(self, mask=True, plot=True):
     total_affected = affected.sum()
 
     if self.verbose:
-        n = total_affected
-        logger.info(f"there {'are'[:n^1]}{'is'[n^1:]} {n} frame{'s'[:n^1]} "
-                     "during ESPRESSO commissioning")
+        n, i = total_affected, int(total_affected != 1)
+        logger.info(f"there {['is', 'are'][i]} {n} frame{['', 's'][i]} "
+                    "during ESPRESSO commissioning")
 
     if mask:
         self.mask[affected] = False
         self._propagate_mask_changes()
 
-        if plot:
+        if plot and total_affected > 0:
             self.plot(show_masked=True)
 
     return affected
@@ -240,9 +240,9 @@ def ADC_issues(self, mask=True, plot=True, check_headers=False):
     total_affected = intersect.sum()
 
     if self.verbose:
-        n = total_affected
-        logger.info(f"there {'are'[:n^1]}{'is'[n^1:]} {n} frame{'s'[:n^1]} "
-                     "affected by ADC issues")
+        n, i = total_affected, int(total_affected != 1)
+        logger.info(f"there {['is', 'are'][i]} {n} frame{['', 's'][i]} "
+                    "affected by ADC issues")
 
     if mask:
         self.mask[intersect] = False
@@ -276,9 +276,9 @@ def blue_cryostat_issues(self, mask=True, plot=True):
     total_affected = intersect.sum()
 
     if self.verbose:
-        n = total_affected
-        logger.info(f"there {'are'[:n^1]}{'is'[n^1:]} {n} frame{'s'[:n^1]} "
-                     "affected by blue cryostat issues")
+        n, i = total_affected, int(total_affected != 1)
+        logger.info(f"there {['is', 'are'][i]} {n} frame{['', 's'][i]} "
+                    "affected by blue cryostat issues")
 
     if mask:
         self.mask[intersect] = False
@@ -322,8 +322,9 @@ def qc_scired_issues(self, plot=False, **kwargs):
     n = affected.sum()
 
     if self.verbose:
-        logger.info(f"there {'are'[:n^1]}{'is'[n^1:]} {n} frame{'s'[:n^1]} "
-                     "where QC SCIRED CHECK is 0")
+        i = int(n != 1)
+        logger.info(f"there {['is', 'are'][i]} {n} frame{['', 's'][i]} "
+                    "where QC SCIRED CHECK is 0")
 
     if n == 0:
         return

{arvi-0.2.7 → arvi-0.2.9}/arvi/kima_wrapper.py

@@ -1,13 +1,14 @@
 import os
 import numpy as np
 
-from .setup_logger import logger
+from .setup_logger import setup_logger
 
 try:
     import kima
     from kima.pykima.utils import chdir
     from kima import distributions
-    from kima import RVData, RVmodel
+    from kima import RVData, HGPMdata
+    from kima import RVmodel, GPmodel, RVHGPMmodel
     kima_available = True
 except ImportError:
     kima_available = False
@@ -21,9 +22,12 @@ def try_to_guess_prior(model, prior):
     return None
 
 
-def run_kima(self, run=False, load=False, run_directory=None, priors={}, **kwargs):
+def run_kima(self, run=False, load=False, run_directory=None,
+             model=RVmodel, priors={}, **kwargs):
     if not kima_available:
         raise ImportError('kima not available, please install with `pip install kima`')
+    
+    logger = setup_logger()
 
     instruments = [inst for inst in self.instruments if self.NN[inst] > 1]
     time = [getattr(self, inst).mtime for inst in instruments]
@@ -33,12 +37,34 @@ def run_kima(self, run=False, load=False, run_directory=None, priors={}, **kwarg
 
     fix = kwargs.pop('fix', False)
     npmax = kwargs.pop('npmax', 1)
-    model = RVmodel(fix=fix, npmax=npmax, data=data)
+
+    if isinstance(model, str):
+        try:
+            model = {
+                'RVmodel': RVmodel,
+                'GPmodel': GPmodel,
+                'RVHGPMmodel': RVHGPMmodel
+            }[model]
+        except KeyError:
+            raise ValueError(f'unknown model: {model}')
+
+    if model is RVHGPMmodel:
+        pm_data = HGPMdata(self.simbad.gaia_id)
+        model = model(fix=fix, npmax=npmax, data=data, pm_data=pm_data)
+    else:
+        model = model(fix=fix, npmax=npmax, data=data)
 
     model.trend = kwargs.pop('trend', False)
     model.degree = kwargs.pop('degree', 0)
 
-    model.studentt = kwargs.pop('studentt', False)
+    if isinstance(model, RVmodel):
+        model.studentt = kwargs.pop('studentt', False)
+
+    if isinstance(model, GPmodel):
+        if 'kernel' in kwargs:
+            model.kernel = kwargs.pop('kernel')
+
+
     model.enforce_stability = kwargs.pop('enforce_stability', False)
     model.star_mass = kwargs.pop('star_mass', 1.0)
 
@@ -49,6 +75,13 @@ def run_kima(self, run=False, load=False, run_directory=None, priors={}, **kwarg
     if kwargs.pop('kuma', False):
         model.conditional.eprior = distributions.Kumaraswamy(0.867, 3.03)
 
+    if isinstance(model, RVHGPMmodel):
+        model.pm_ra_bary_prior = priors.pop('pm_ra_bary_prior', 
+                                            distributions.Gaussian(pm_data.pm_ra_hg, pm_data.sig_hg_ra))
+        model.pm_dec_bary_prior = priors.pop('pm_dec_bary_prior',
+                                             distributions.Gaussian(pm_data.pm_dec_hg, pm_data.sig_hg_dec))
+
+
     for k, v in priors.items():
         try:
             if 'conditional' in k:
@@ -67,7 +100,6 @@ def run_kima(self, run=False, load=False, run_directory=None, priors={}, **kwarg
         run_directory = os.getcwd()
 
     if run:
-        
         # TODO: use signature of kima.run to pop the correct kwargs
         # model_name = model.__class__.__name__
         # model_name = f'kima.{model_name}.{model_name}'
@@ -75,7 +107,10 @@ def run_kima(self, run=False, load=False, run_directory=None, priors={}, **kwarg
 
         with chdir(run_directory):
             kima.run(model, **kwargs)
-        
+
+    if isinstance(model, RVHGPMmodel):
+        data = (data, pm_data)
+
     if load:
         with chdir(run_directory):
             res = kima.load_results(model)

{arvi-0.2.7 → 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.7 → arvi-0.2.9}/arvi/programs.py

@@ -3,11 +3,13 @@ import multiprocessing
 from functools import partial, lru_cache
 from itertools import chain
 from collections import namedtuple
-from multiprocessing.pool import ThreadPool
+from multiprocessing.pool import ThreadPool, Pool
+import concurrent.futures
+import sys
 from tqdm import tqdm
 # import numpy as np
 
-from .setup_logger import logger
+from .setup_logger import setup_logger
 from .timeseries import RV
 
 __all__ = ['ESPRESSO_GTO']
@@ -22,13 +24,14 @@ def get_star(star, instrument=None, verbose=False, **kwargs):
 
 class LazyRV:
     def __init__(self, stars: list, instrument: str = None,
-                 _parallel_limit=10):
+                 _parallel_limit=10, _parallel_workers=8):
         self.stars = stars
         if isinstance(self.stars, str):
             self.stars = [self.stars]
         self.instrument = instrument
         self._saved = None
         self._parallel_limit = _parallel_limit
+        self._parallel_workers = _parallel_workers
 
     @property
     def N(self):
@@ -38,10 +41,11 @@ class LazyRV:
         return f"RV({self.N} stars)"
 
     def _get(self, **kwargs):
+        logger = setup_logger()
         if self.N > self._parallel_limit:
             # logger.info('Querying DACE...')
             _get_star = partial(get_star, instrument=self.instrument, **kwargs)
-            with ThreadPool(8) as pool:
+            with Pool(self._parallel_workers) as pool:
                 result = list(tqdm(pool.imap(_get_star, self.stars), 
                                    total=self.N, unit='star',
                                    desc='Querying DACE (can take a while)'))

{arvi-0.2.7 → arvi-0.2.9}/arvi/simbad_wrapper.py

@@ -70,6 +70,17 @@ FROM ident AS id1 JOIN ident AS id2 USING(oidref)
 WHERE id1.id = '{star}' AND id2.id LIKE '{name}%';
 """
 
+RA_DEC_QUERY = """
+SELECT DISTINCT basic.OID,
+       RA,
+       DEC,
+       main_id,
+       DISTANCE(POINT('ICRS', RA, DEC), POINT('ICRS', {ra}, {dec})) as "dist"
+FROM basic JOIN flux ON oidref = oid
+WHERE CONTAINS(POINT('ICRS', RA, DEC), CIRCLE('ICRS', {ra}, {dec}, 0.02)) = 1
+ORDER BY "dist";
+"""
+
 def find_identifier(identifier, star):
     response = run_query(HD_GJ_HIP_QUERY.format(name=identifier, star=star))
     if identifier in response:
@@ -288,6 +299,15 @@ class simbad:
     def bmv(self):
         return self.B - self.V
 
+    @classmethod
+    def from_ra_dec(cls, ra, dec, **kwargs):
+        table1 = run_query(query=RA_DEC_QUERY.format(ra=ra, dec=dec))
+        cols, values = parse_tablen(table1)
+        if len(values) == 0:
+            raise ValueError(f'no Simbad results for ra={ra}, dec={dec}')
+        assert cols == ['oid', 'ra', 'dec', 'main_id', 'dist']
+        star = values[0][cols.index('main_id')].replace('"', '')
+        return cls(star, **kwargs)
 
 def argsort_by_spectral_type(sptypes):
     STs = [f'{letter}{n}' for letter in ('F', 'G', 'K', 'M') for n in range(10)]

{arvi-0.2.7 → 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
@@ -408,16 +411,41 @@ class RV(ISSUES, REPORTS):
         self._did_correct_berv = False
         self.__post_init__()
 
-    def snapshot(self):
+    def snapshot(self, directory=None, delete_others=False):
         import pickle
         from datetime import datetime
         ts = datetime.now().timestamp()
         star_name = self.star.replace(' ', '')
         file = f'{star_name}_{ts}.pkl'
-        pickle.dump(self, open(file, 'wb'), protocol=0)
+
+        if directory is None:
+            directory = '.'
+        else:
+            os.makedirs(directory, exist_ok=True)
+
+        file = os.path.join(directory, file)
+
+        if delete_others:
+            import re
+            other_pkls = [
+                f for f in os.listdir(directory)
+                if re.search(fr'{star_name}_\d+.\d+.pkl', f)
+            ]
+            for pkl in other_pkls:
+                os.remove(os.path.join(directory, pkl))
+
+        metadata = {
+            'star': self.star,
+            'timestamp': ts,
+            'description': 'arvi snapshot'
+        }
+        pickle.dump((self, metadata), open(file, 'wb'), protocol=0)
+
         if self.verbose:
             logger.info(f'saved snapshot to {file}')
 
+        return file
+
     @property
     def N(self) -> int:
         """Total number of observations"""
@@ -552,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)
 
@@ -618,6 +647,8 @@ class RV(ISSUES, REPORTS):
             logger.info(f'reading snapshot of {star} from {dt}')
 
         s = pickle.load(open(file, 'rb'))
+        if isinstance(s, tuple) and len(s) == 2:
+            s, _metadata = s
         s._snapshot = file
         return s
 
@@ -659,10 +690,11 @@ class RV(ISSUES, REPORTS):
             file_object = hasattr(files, 'read')
             files = [files]
 
-        # if len(files) == 0:
-        #     if verbose:
-        #         logger.error('no files found')
-        #     return
+        if len(files) == 0:
+            if verbose:
+                logger.error('from_rdb: no files found')
+            return
+
         def get_star_name(file):
             return splitext(basename(file))[0].split('_')[0].replace('-', '_')
         
@@ -1260,8 +1292,9 @@ class RV(ISSUES, REPORTS):
                 logger.warning('may need to provide `top_level` in kwargs to find file')
             do_symlink_filetype('CCF', files[:limit], directory, **kwargs)
         else:
-            do_download_filetype('CCF', files[:limit], directory, clobber=clobber,
-                                 verbose=self.verbose, user=self.user, **kwargs)
+            downloaded = do_download_filetype('CCF', files[:limit], directory, 
+                                              clobber=clobber, verbose=self.verbose, 
+                                              user=self.user, **kwargs)
 
         if load:
             try:
@@ -1273,17 +1306,26 @@ class RV(ISSUES, REPORTS):
                     for f in files[:limit]
                 ]
                 downloaded = [
-                    skysub 
+                    skysub
                     if exists(skysub := f.replace('CCF_A.fits', 'CCF_SKYSUB_A.fits')) else f
                     for f in downloaded
                 ]
                 if self.verbose:
                     logger.info('loading the CCF(s) into `.CCF` attribute')
 
-                self.CCF = iCCF.from_file(downloaded)
+                self.CCF = iCCF.from_file(downloaded, verbose=False)
+                if len(self.CCF) == 1:
+                    self.CCF = [self.CCF]
 
-            except (ImportError, ValueError):
-                pass
+                if self.simbad is None:
+                    if self.verbose:
+                        logger.info('querying Simbad with RA/DEC from CCF header')
+                    ra = self.CCF[0].HDU[0].header['RA']
+                    dec = self.CCF[0].HDU[0].header['DEC']
+                    self._simbad = simbad.from_ra_dec(ra, dec)
+
+            except (ImportError, ValueError, FileNotFoundError):
+                logger.error('could not load CCF(s) into `.CCF` attribute')
 
     def download_s1d(self, instrument=None, index=None, limit=None,
                      directory=None, clobber=False, apply_mask=True, symlink=False, **kwargs):
@@ -1633,6 +1675,9 @@ class RV(ISSUES, REPORTS):
             inst = self.instruments[self.obs[m] - 1]
             n_before = (self.obs < self.obs[m]).sum()
             getattr(self, inst).mask[m - n_before] = False
+        for inst in self.instruments:
+            if getattr(self, inst).mtime.size == 0:
+                self.remove_instrument(inst, strict=True)
 
     def secular_acceleration(self, epoch=None, just_compute=False, force_simbad=False):
         """ 
@@ -1806,9 +1851,11 @@ class RV(ISSUES, REPORTS):
         for inst in instruments:
             m = self.instrument_array == inst
             result = dosigmaclip(self.vrad[m], low=sigma, high=sigma)
-            n = self.vrad[m].size - result.clipped.size
+            # n = self.vrad[m].size - result.clipped.size
 
-            ind = m & ((self.vrad < result.lower) | (self.vrad > result.upper))
+            ind = m & self.mask & \
+                  ((self.vrad < result.lower) | (self.vrad > result.upper))
+            n = ind.sum()
 
             if self.verbose and n > 0:
                 s = 's' if (n == 0 or n > 1) else ''
@@ -1869,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 = []
 
@@ -2058,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)))
 
@@ -2084,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)
@@ -2098,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
-        
+
 
 
 
@@ -2258,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()
 
@@ -2298,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
@@ -2334,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',
-                                    ])
+            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:
+                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)
+
+            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:
-                if save_masked:
-                    arrays = [_s.time, _s.vrad, _s.svrad]
-                else:
-                    arrays = [_s.mtime, _s.mvrad, _s.msvrad]
-
-                # d = np.stack(arrays, axis=1)
-                header = 'rjd\tvrad\tsvrad\n---\t----\t-----'
+                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.7 → 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.7 → arvi-0.2.9/arvi.egg-info}/PKG-INFO

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