arvi 0.2.6__tar.gz → 0.2.8__tar.gz

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

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

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

@@ -2,7 +2,7 @@ import os
 import sys
 import tarfile
 import collections
-from functools import lru_cache
+from functools import lru_cache, partial
 from itertools import islice
 import numpy as np
 
@@ -95,7 +95,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:
@@ -170,7 +170,7 @@ def get_observations_from_instrument(star, instrument, user=None, main_id=None,
     """
     Spectroscopy = load_spectroscopy(user)
     found_dace_id = False
-    with timer('simbad query'):
+    with timer('dace_id query'):
         try:
             dace_id = get_dace_id(star, verbose=verbose, raise_error=True)
             found_dace_id = True
@@ -212,17 +212,22 @@ def get_observations_from_instrument(star, instrument, user=None, main_id=None,
 
     for inst in np.unique(result['ins_name']):
         mask1 = result['ins_name'] == inst
-        r[inst] = {}
+        r[str(inst)] = {}
+
+        key2 = 'ins_drs_version'
+        n_key2 = len(np.unique(result[key2][mask1]))
+        if len(np.unique(result['pub_bibcode'][mask1])) >= n_key2:
+            key2 = 'pub_bibcode'
 
-        for pipe in np.unique(result['ins_drs_version'][mask1]):
-            mask2 = mask1 & (result['ins_drs_version'] == pipe)
-            r[inst][pipe] = {}
+        for pipe in np.unique(result[key2][mask1]):
+            mask2 = mask1 & (result[key2] == pipe)
+            r[str(inst)][str(pipe)] = {}
 
             for ins_mode in np.unique(result['ins_mode'][mask2]):
                 mask3 = mask2 & (result['ins_mode'] == ins_mode)
                 _nan = np.full(mask3.sum(), np.nan)
 
-                r[inst][pipe][ins_mode] = {
+                r[str(inst)][str(pipe)][str(ins_mode)] = {
                     'texp': result['texp'][mask3],
                     'bispan': result['spectro_ccf_bispan'][mask3],
                     'bispan_err': result['spectro_ccf_bispan_err'][mask3],
@@ -372,9 +377,11 @@ def get_observations(star, instrument=None, user=None, main_id=None, verbose=Tru
     #     # For all other strings, sort alphabetically
     #     return (2, s)
 
-    def custom_key(val):
+    def custom_key(val, strip_EGGS=False):
+        if strip_EGGS:
+            val = val.replace('-EGGS', '').replace(' EGGS', '')
         key = 0
-        key -= 2 if val == '3.5' else 0
+        key -= 1 if '3.5' in val else 0
         key -= 1 if 'EGGS' in val else 0
         key -= 1 if ('UHR' in val or 'MR' in val) else 0
         key -= 1 if 'LBL' in val else 0
@@ -385,6 +392,8 @@ def get_observations(star, instrument=None, user=None, main_id=None, verbose=Tru
         # new_result[inst] = dict(
         #     sorted(result[inst].items(), key=custom_sort_key, reverse=True)
         # )
+        if all(['EGGS' in k for k in result[inst].keys()]):
+            custom_key = partial(custom_key, strip_EGGS=True)
         # WARNING: not the same as reverse=True (not sure why)
         sorted_keys = sorted(result[inst].keys(), key=custom_key)[::-1]
         new_result[inst] = {}
@@ -535,6 +544,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)
 
@@ -548,7 +558,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:
@@ -566,7 +576,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)
 
@@ -595,6 +607,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.6 → arvi-0.2.8}/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.6 → arvi-0.2.8}/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.6 → arvi-0.2.8}/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.6 → arvi-0.2.8}/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.6 → arvi-0.2.8}/arvi/timeseries.py

@@ -40,24 +40,53 @@ class RV(ISSUES, REPORTS):
     """
     A class holding RV observations
 
+    Args:
+        star (str):
+            Name of the star
+        instrument (str, list):
+            Name of the instrument or list of instruments
+        verbose (bool):
+            Print logging messages
+        do_maxerror:
+            Mask points based on a maximum RV uncertainty
+        do_secular_acceleration:
+            Apply secular acceleration correction. This only applies
+            to certain instruments.
+        do_sigma_clip (bool):
+            Apply sigma clipping on the RVs
+        do_adjust_means (bool):
+            Subtract individual weighted mean RV from each instrument
+        only_latest_pipeline (bool):
+            Select only the latest pipeline from each instrument
+        load_extra_data (bool):
+        check_drs_qc (bool):
+            Mask points based on DRS quality control flags
+        user (str):
+            User name for DACE queries (should be a section in `~/.dacerc` file)
+
     Examples:
         >>> s = RV('Proxima')
+        >>> s = RV('HD10180', instrument='HARPS')
 
-    Attributes:
-        star (str):
-            The name of the star
-        N (int):
-            Total number of observations
-        instruments (list):
-            List of instruments for which there are RVs. Each instrument is also
-            stored as an attribute (e.g. `self.CORALIE98` or `self.HARPS`)
-        simbad (simbad):
-            Information on the target from Simbad
     """
+    # Attributes:
+    #     star (str):
+    #         The name of the star
+    #     N (int):
+    #         Total number of observations
+    #     NN (dict):
+    #         Number of observations per instrument
+    #     instruments (list):
+    #         List of instruments for which there are RVs. Each instrument is also
+    #         stored as an attribute (e.g. `self.CORALIE98` or `self.HARPS`)
+    #     simbad (simbad):
+    #         Information on the target from Simbad
+    #     gaia (gaia):
+    #         Information on the target from Gaia DR3
     star: str
     instrument: Union[str, list] = field(init=True, repr=False, default=None)
     verbose: bool = field(init=True, repr=False, default=True)
-    do_maxerror: Union[bool, float] = field(init=True, repr=False, default=False)
+    do_maxerror: float = field(init=True, repr=False, default=None)
     do_secular_acceleration: bool = field(init=True, repr=False, default=True)
     do_sigma_clip: bool = field(init=True, repr=False, default=False)
     do_adjust_means: bool = field(init=True, repr=False, default=True)
@@ -217,7 +246,7 @@ class RV(ISSUES, REPORTS):
                     else:
                         mid = None
 
-                    with timer():
+                    with timer('dace query'):
                         self.dace_result = get_observations(self.__star__, self.instrument,
                                                             user=self.user, main_id=mid, verbose=self.verbose)
                 except ValueError as e:
@@ -318,7 +347,7 @@ class RV(ISSUES, REPORTS):
 
             if self.do_adjust_means:
                 self.adjust_means()
-        
+
         _star_no_space = self.star.replace(' ', '')
         _directory = sanitize_path(_star_no_space)
         self._download_directory = f'{_directory}_downloads'
@@ -379,16 +408,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"""
@@ -543,7 +597,7 @@ class RV(ISSUES, REPORTS):
         return s
 
     @classmethod
-    def from_arrays(cls, star, time, vrad, svrad, inst, *args, **kwargs):
+    def from_arrays(cls, star, time, vrad, svrad, inst, **kwargs):
         s = cls(star, _child=True)
         time, vrad, svrad = map(np.atleast_1d, (time, vrad, svrad))
 
@@ -558,11 +612,15 @@ class RV(ISSUES, REPORTS):
         s.time = time
         s.vrad = vrad
         s.svrad = svrad
-        # mask
-        s.mask = kwargs.get('mask', np.full_like(s.time, True, dtype=bool))
+
+        s.mask = kwargs.pop('mask', np.full_like(s.time, True, dtype=bool))
+        s.units = kwargs.pop('units', 'm/s')
+
+        for k, v in kwargs.items():
+            setattr(s, k, np.atleast_1d(v))
 
         s.instruments = [inst]
-        s._quantities = np.array([])
+        s._quantities = np.array(list(kwargs.keys()))
 
         return s
 
@@ -583,8 +641,10 @@ class RV(ISSUES, REPORTS):
         dt = datetime.fromtimestamp(float(timestamp))
         if verbose:
             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
 
@@ -595,7 +655,7 @@ class RV(ISSUES, REPORTS):
 
         Args:
             files (str, list):
-                File name or list of file names
+                File name, file object, or list of file names
             star (str, optional):
                 Name of the star. If None, try to infer it from file name
             instrument (str, list, optional):
@@ -606,37 +666,69 @@ class RV(ISSUES, REPORTS):
                 Number of lines to skip in the header. Defaults to 2.
 
         Examples:
-            s = RV.from_rdb('star_HARPS.rdb')
+            >>> s = RV.from_rdb('star_HARPS.rdb')
         """
         from glob import glob
         from os.path import splitext, basename
 
         verbose = kwargs.pop('verbose', True)
 
+        file_object = False
+
         if isinstance(files, str):
             if '*' in files:
                 files = glob(files)
             else:
                 files = [files]
+        elif isinstance(files, list):
+            pass
+        else:
+            file_object = hasattr(files, 'read')
+            files = [files]
 
         if len(files) == 0:
             if verbose:
-                logger.error('no files found')
+                logger.error('from_rdb: no files found')
             return
 
-        if star is None:
-            star_ = np.unique([splitext(basename(f))[0].split('_')[0] for f in files])
-            if star_.size == 1:
-                star = star_[0].replace('-', '_')
+        def get_star_name(file):
+            return splitext(basename(file))[0].split('_')[0].replace('-', '_')
+        
+        def get_instrument(file):
+            return splitext(basename(file))[0].split('_')[1]
+
+        if file_object:
+            if star is None:
+                try:
+                    star = get_star_name(files[0].name)
+                except Exception:
+                    star ='unknown'
                 if verbose:
                     logger.info(f'assuming star is {star}')
 
-        if instrument is None:
-            instruments = np.array([splitext(basename(f))[0].split('_')[1] for f in files])
-            if verbose:
-                logger.info(f'assuming instruments: {instruments}')
+            if instrument is None:
+                try:
+                    instrument = get_instrument(files[0].name)
+                except Exception:
+                    instrument = 'unknown'
+                if verbose:
+                    logger.info(f'assuming instrument is {instrument}')
+
+            instruments = np.array([instrument])
         else:
-            instruments = np.atleast_1d(instrument)
+            if star is None:
+                star = np.unique([get_star_name(f) for f in files])[0]
+                if verbose:
+                    logger.info(f'assuming star is {star}')
+            else:
+                star = 'unknown'
+
+            if instrument is None:
+                instruments = np.array([splitext(basename(f))[0].split('_')[1] for f in files])
+                if verbose:
+                    logger.info(f'assuming instruments: {instruments}')
+            else:
+                instruments = np.atleast_1d(instrument)
 
         if instruments.size == 1 and len(files) > 1:
             instruments = np.repeat(instruments, len(files))
@@ -669,12 +761,16 @@ class RV(ISSUES, REPORTS):
             _quantities = []
 
             #! hack
-            with open(f) as ff:
-                header = ff.readline().strip()
-                if '\t' in header:
-                    names = header.split('\t')
-                else:
-                    names = header.split()
+            if file_object:
+                header = f.readline().strip()
+            else:
+                with open(f) as ff:
+                    header = ff.readline().strip()
+
+            if '\t' in header:
+                names = header.split('\t')
+            else:
+                names = header.split()
 
             if len(names) > 3:
                 # if f.endswith('.rdb'):
@@ -687,7 +783,7 @@ class RV(ISSUES, REPORTS):
                     data = np.genfromtxt(f, **kw, delimiter='\t')
                 else:
                     data = np.genfromtxt(f, **kw)
-                
+
                 # if data.ndim in (0, 1):
                 #     data = data.reshape(-1, 1)
 
@@ -831,7 +927,22 @@ class RV(ISSUES, REPORTS):
 
     @classmethod
     def from_ccf(cls, files, star=None, instrument=None, **kwargs):
-        """ Create an RV object from a CCF file or a list of CCF files """
+        """ Create an RV object from a CCF file or a list of CCF files
+
+        !!! Note
+            This function relies on the `iCCF` package
+
+        Args:
+            files (str or list):
+                CCF file or list of CCF files
+            star (str):
+                Star name. If not provided, it will be inferred from the header
+                of the CCF file
+            instrument (str):
+                Instrument name. If not provided, it will be inferred from the
+                header of the CCF file
+        
+        """
         try:
             import iCCF
         except ImportError:
@@ -856,7 +967,7 @@ class RV(ISSUES, REPORTS):
             objects = np.unique([i.HDU[0].header['OBJECT'].replace(' ', '') for i in CCFs])
 
         if len(objects) != 1:
-            logger.warning(f'found {objects.size} different stars in the CCF files, '
+            logger.warning(f'found {objects.size} different stars in the CCF files ({objects}), '
                            'choosing the first one')
         star = objects[0]
 
@@ -882,9 +993,22 @@ class RV(ISSUES, REPORTS):
             _quantities.append('contrast')
             _quantities.append('contrast_err')
 
+            _s.bispan = np.array([i.BIS*1e3 for i in CCFs])
+            _s.bispan_err = np.array([i.BISerror*1e3 for i in CCFs])
+            _quantities.append('bispan')
+            _quantities.append('bispan_err')
+
+            _s.rhk = np.full_like(time, np.nan)
+            _s.rhk_err = np.full_like(time, np.nan)
+            _quantities.append('rhk')
+            _quantities.append('rhk_err')
+
             _s.texp = np.array([i.HDU[0].header['EXPTIME'] for i in CCFs])
             _quantities.append('texp')
 
+            _s.berv = np.array([i.HDU[0].header['HIERARCH ESO QC BERV'] for i in CCFs])
+            _quantities.append('berv')
+
             _s.date_night = np.array([
                 i.HDU[0].header['DATE-OBS'].split('T')[0] for i in CCFs
             ])
@@ -955,17 +1079,17 @@ class RV(ISSUES, REPORTS):
             if fits_file not in tar.getnames():
                 logger.error(f'KOBE file not found for {star}')
                 return
-            
+
             hdul = fits.open(tar.extractfile(fits_file))
 
         else:
             resp = requests.get(f'https://kobe.caha.es/internal/fitsfiles/{fits_file}',
                                 auth=HTTPBasicAuth('kobeteam', config.kobe_password))
-            
+
             if resp.status_code != 200:
                 # something went wrong, try to extract the file by downloading the
                 # full tar.gz archive
-                
+
                 logger.warning(f'could not find "{fits_file}" on server, trying to download full archive')
                 resp = requests.get('https://kobe.caha.es/internal/fitsfiles.tar.gz',
                                     auth=HTTPBasicAuth('kobeteam', config.kobe_password))
@@ -983,7 +1107,7 @@ class RV(ISSUES, REPORTS):
                 if fits_file not in tar.getnames():
                     logger.error(f'KOBE file not found for {star}')
                     return
-                
+
                 hdul = fits.open(tar.extractfile(fits_file))
 
             else:
@@ -1164,8 +1288,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:
@@ -1177,17 +1302,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):
@@ -1382,7 +1516,9 @@ class RV(ISSUES, REPORTS):
     def remove_point(self, index):
         """
         Remove individual observations at a given index (or indices).
-        NOTE: Like Python, the index is 0-based.
+        
+        !!! Note
+            Like Python, the index is 0-based.
 
         Args:
             index (int, list, ndarray):
@@ -1414,11 +1550,14 @@ class RV(ISSUES, REPORTS):
     def restore_point(self, index):
         """
         Restore previously deleted individual observations at a given index (or
-        indices). NOTE: Like Python, the index is 0-based.
+        indices). 
+        
+        !!! Note
+            Like Python, the index is 0-based
 
         Args:
             index (int, list, ndarray):
-                Single index, list, or array of indices to restore.
+                Single index, list, or array of indices to restore
         """
         index = np.atleast_1d(index)
         try:
@@ -1444,6 +1583,14 @@ class RV(ISSUES, REPORTS):
         self.mask = self.mask & self.public
         self._propagate_mask_changes()
 
+    def remove_public(self):
+        """ Remove public observations """
+        if self.verbose:
+            n = self.public.sum()
+            logger.info(f'masking public observations ({n})')
+        self.mask = self.mask & (~self.public)
+        self._propagate_mask_changes()
+
     def remove_single_observations(self):
         """ Remove instruments for which there is a single observation """
         singles = [i for i in self.instruments if getattr(self, i).mtime.size == 1]
@@ -1469,26 +1616,26 @@ class RV(ISSUES, REPORTS):
                 if self.verbose:
                     logger.warning(f'no observations for prog_id "{prog_id}"')
 
-    def remove_after_bjd(self, bjd):
+    def remove_after_bjd(self, bjd: float):
         """ Remove observations after a given BJD """
         if (self.time > bjd).any():
             ind = np.where(self.time > bjd)[0]
             self.remove_point(ind)
 
-    def remove_before_bjd(self, bjd):
+    def remove_before_bjd(self, bjd: float):
         """ Remove observations before a given BJD """
         if (self.time < bjd).any():
             ind = np.where(self.time < bjd)[0]
             self.remove_point(ind)
-    
-    def remove_between_bjds(self, bjd1, bjd2):
+
+    def remove_between_bjds(self, bjd1: float, bjd2: float):
         """ Remove observations between two BJDs """
         to_remove = (self.time > bjd1) & (self.time < bjd2)
         if to_remove.any():
             ind = np.where(to_remove)[0]
             self.remove_point(ind)
 
-    def choose_n_points(self, n, seed=None, instrument=None):
+    def choose_n_points(self, n: int, seed=None, instrument=None):
         """ Randomly choose `n` observations and mask out the remaining ones
 
         Args:
@@ -1524,17 +1671,25 @@ 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):
-        """
-        Remove secular acceleration from RVs
+        """ 
+        Remove secular acceleration from RVs. This uses the proper motions from
+        Gaia (in `self.gaia`) if available, otherwise from Simbad (in
+        `self.simbad`), unless `force_simbad=True`.
+
 
         Args:
             epoch (float, optional):
                 The reference epoch (DACE uses 55500, 31/10/2010)
-            instruments (bool or collection of str):
-                Only remove secular acceleration for some instruments, or for all
-                if `instruments=True`
+            just_compute (bool, optional):
+                Just compute the secular acceleration and return, without
+                changing the RVs
+            force_simbad (bool, optional):
+                Use Simbad proper motions even if Gaia is available
         """
         if self._did_secular_acceleration and not just_compute:  # don't do it twice
             return
@@ -1643,7 +1798,7 @@ class RV(ISSUES, REPORTS):
 
         if config.return_self:
             return self
-    
+
     def _undo_secular_acceleration(self):
         if self._did_secular_acceleration:
             _old_verbose = self.verbose
@@ -1671,7 +1826,15 @@ class RV(ISSUES, REPORTS):
             self._did_secular_acceleration = False
 
     def sigmaclip(self, sigma=5, instrument=None, strict=True):
-        """ Sigma-clip RVs (per instrument!) """
+        """
+        Sigma-clip RVs (per instrument!), by MAD away from the median.
+
+        Args:
+            sigma (float):
+                Number of MADs to clip
+            instrument (str, list):
+                Instrument(s) to sigma-clip
+        """
         #from scipy.stats import sigmaclip as dosigmaclip
         from .stats import sigmaclip_median as dosigmaclip
 
@@ -1684,9 +1847,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 ''
@@ -1708,7 +1873,7 @@ class RV(ISSUES, REPORTS):
 
         self._propagate_mask_changes()
 
-        if self._did_adjust_means:
+        if len(changed_instruments) > 0 and self._did_adjust_means:
             self._did_adjust_means = False
             self.adjust_means(instrument=changed_instruments)
 
@@ -1741,7 +1906,8 @@ class RV(ISSUES, REPORTS):
         """
         Nightly bin the observations.
 
-        WARNING: This creates and returns a new object and does not modify self.
+        !!! Warning
+            This creates and returns a new object and does not modify self.
         """
 
         # create copy of self to be returned
@@ -1845,7 +2011,7 @@ class RV(ISSUES, REPORTS):
         return np.nanmean(z, axis=0)
 
     def subtract_mean(self):
-        """ Subtract (single) mean RV from all instruments """
+        """ Subtract (a single) non-weighted mean RV from all instruments """
         self._meanRV = meanRV = self.mvrad.mean()
         for inst in self.instruments:
             s = getattr(self, inst)
@@ -1881,7 +2047,7 @@ class RV(ISSUES, REPORTS):
         #     row = []
         #     if print_as_table:
         #         logger.info('subtracted weighted average from each instrument:')
-    
+
         others = ('fwhm', 'bispan', )
 
         instruments = self._check_instrument(instrument, strict=kwargs.get('strict', False))
@@ -2160,10 +2326,10 @@ class RV(ISSUES, REPORTS):
         """ Sort instruments by first or last observation date.
 
         Args:
-            by_first_observation (bool, optional, default=True):
-                Sort by first observation date.
-            by_last_observation (bool, optional, default=False):
-                Sort by last observation date.
+            by_first_observation (bool, optional):
+                Sort by first observation date
+            by_last_observation (bool, optional):
+                Sort by last observation date
         """
         if by_last_observation:
             by_first_observation = False
@@ -2239,7 +2405,7 @@ class RV(ISSUES, REPORTS):
                     #     if self.verbose:
                     #         logger.warning(f'masking {nan_mask.sum()} observations with NaN in indicators')
 
-                header = '\t'.join(['rjd', 'vrad', 'svrad', 
+                header = '\t'.join(['rjd', 'vrad', 'svrad',
                                     'fwhm', 'sig_fwhm',
                                     'bispan', 'sig_bispan',
                                     'contrast', 'sig_contrast',

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

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