arvi 0.2.7__tar.gz → 0.2.8__tar.gz

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

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

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

@@ -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:
@@ -544,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)
 
@@ -557,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:
@@ -575,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)
 
@@ -604,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.7 → 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.7 → 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.7 → 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.7 → 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.7 → arvi-0.2.8}/arvi/timeseries.py

@@ -408,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"""
@@ -618,6 +643,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 +686,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 +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:
@@ -1273,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):
@@ -1633,6 +1671,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 +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 ''

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

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