arvi 0.1.23__py3-none-any.whl → 0.1.24__py3-none-any.whl

arvi/dace_wrapper.py

@@ -1,7 +1,9 @@
 import os
+import sys
 import tarfile
 import collections
 from functools import lru_cache
+from itertools import islice
 import numpy as np
 from dace_query import DaceClass
 from dace_query.spectroscopy import SpectroscopyClass, Spectroscopy as default_Spectroscopy
@@ -22,14 +24,15 @@ def load_spectroscopy() -> SpectroscopyClass:
     return default_Spectroscopy
 
 @lru_cache()
-def get_dace_id(star):
+def get_dace_id(star, verbose=True):
     filters = {"obj_id_catname": {"equal": [star]}}
     try:
-        with stdout_disabled(), all_logging_disabled():
+        with all_logging_disabled():
             r = load_spectroscopy().query_database(filters=filters, limit=1)
         return r['obj_id_daceid'][0]
     except KeyError:
-        logger.error(f"Could not find DACE ID for {star}")
+        if verbose:
+            logger.error(f"Could not find DACE ID for {star}")
         raise ValueError from None
 
 def get_arrays(result, latest_pipeline=True, ESPRESSO_mode='HR11', NIRPS_mode='HE', verbose=True):
@@ -55,15 +58,19 @@ def get_arrays(result, latest_pipeline=True, ESPRESSO_mode='HR11', NIRPS_mode='H
                 pipelines = [pipelines[i]]
 
         if latest_pipeline:
+            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]]
+            if 'HARPS' in inst and npipe > 1 and pipelines[1] == pipelines[0] + '-EGGS':
+                pipelines = pipelines[:2]
             else:
                 pipelines = [pipelines[0]]
 
-            if verbose and len(pipelines) > 1:
+            if verbose and npipe > 1:
                 logger.info(f'selecting latest pipeline ({pipelines[0]}) for {inst}')
 
+
         for pipe in pipelines:
             modes = [m for m in result[inst][pipe].keys()]
 
@@ -80,15 +87,25 @@ def get_arrays(result, latest_pipeline=True, ESPRESSO_mode='HR11', NIRPS_mode='H
 
             # HARPS15 observations should not be separated by 'mode' if some are
             # done together with NIRPS
-            if 'HARPS15' in inst and 'HARPS+NIRPS' in modes:
-                m0 = modes[0]
-                data = {
-                    k: np.concatenate([result[inst][pipe][m][k] for m in modes])
-                    for k in result[inst][pipe][m0].keys()
-                }
-                arrays.append(
-                    ((inst, pipe, m0), data)
-                )
+            if 'HARPS15' in inst:
+                if 'HARPS+NIRPS' in modes:
+                    m0 = modes[0]
+                    data = {
+                        k: np.concatenate([result[inst][pipe][m][k] for m in modes])
+                        for k in result[inst][pipe][m0].keys()
+                    }
+                    arrays.append(
+                        ((str(inst), str(pipe), str(m0)), data)
+                    )
+                if 'EGGS+NIRPS' in modes:
+                    m0 = modes[0]
+                    data = {
+                        k: np.concatenate([result[inst][pipe][m][k] for m in modes])
+                        for k in result[inst][pipe][m0].keys()
+                    }
+                    arrays.append(
+                        ((str(inst + '_EGGS'), str(pipe), str(m0)), data)
+                    )
                 continue
 
             for mode in modes:
@@ -97,39 +114,61 @@ def get_arrays(result, latest_pipeline=True, ESPRESSO_mode='HR11', NIRPS_mode='H
                     raise ValueError
 
                 arrays.append(
-                    ((inst, pipe, mode), result[inst][pipe][mode])
+                    ((str(inst), str(pipe), str(mode)), result[inst][pipe][mode])
                 )
 
     return arrays
 
-def get_observations_from_instrument(star, instrument, main_id=None):
+def get_observations_from_instrument(star, instrument, main_id=None, verbose=True):
     """ Query DACE for all observations of a given star and instrument
 
     Args:
-        star (str): name of the star
-        instrument (str): instrument name
-        main_id (str, optional): Simbad main id of target to query DACE id. Defaults to None.
+        star (str):
+            name of the star
+        instrument (str):
+            instrument name
+        main_id (str, optional):
+            Simbad main id of target to query DACE id. Defaults to None.
+        verbose (bool, optional):
+            whether to print warnings. Defaults to True.
 
     Raises:
-        ValueError: If query for DACE id fails
+        ValueError:
+            If query for DACE id fails
 
     Returns:
-        dict: dictionary with data from DACE
+        dict:
+            dictionary with data from DACE
     """
+    Spectroscopy = load_spectroscopy()
+    found_dace_id = False
     try:
-        dace_id = get_dace_id(star)
+        dace_id = get_dace_id(star, verbose=verbose)
+        found_dace_id = True
     except ValueError as e:
         if main_id is not None:
-            dace_id = get_dace_id(main_id)
-        else:
-            raise e
+            try:
+                dace_id = get_dace_id(main_id, verbose=verbose)
+                found_dace_id = True
+            except ValueError:
+                pass
+    
+    if not found_dace_id:
+        try:
+            with all_logging_disabled():
+                result = Spectroscopy.get_timeseries(target=star,
+                                                     sorted_by_instrument=True,
+                                                     output_format='numpy')
+                return result
+        except TypeError:
+            msg = f'no {instrument} observations for {star}'
+            raise ValueError(msg) from None
 
-    Spectroscopy = load_spectroscopy()
     filters = {
         "ins_name": {"contains": [instrument]},
         "obj_id_daceid": {"contains": [dace_id]}
     }
-    with stdout_disabled(), all_logging_disabled():
+    with all_logging_disabled():
         result = Spectroscopy.query_database(filters=filters)
     
     if len(result) == 0:
@@ -218,7 +257,7 @@ def get_observations(star, instrument=None, main_id=None, verbose=True):
             raise ValueError(msg) from None
     else:
         try:
-            result = get_observations_from_instrument(star, instrument, main_id)
+            result = get_observations_from_instrument(star, instrument, main_id, verbose)
         except ValueError:
             msg = f'no {instrument} observations for {star}'
             raise ValueError(msg) from None
@@ -323,16 +362,29 @@ def check_existing(output_directory, files, type):
 
     return np.array(missing)
 
-def download(files, type, output_directory):
+def download(files, type, output_directory, output_filename=None, quiet=True, pbar=None):
     """ Download files from DACE """
     Spectroscopy = load_spectroscopy()
-    with stdout_disabled(), all_logging_disabled():
+    if isinstance(files, str):
+        files = [files]
+    if quiet:
+        with all_logging_disabled():
+            Spectroscopy.download_files(files, file_type=type.lower(),
+                                        output_directory=output_directory,
+                                        output_filename=output_filename)
+    else:
         Spectroscopy.download_files(files, file_type=type.lower(),
-                                    output_directory=output_directory)
+                                    output_directory=output_directory,
+                                    output_filename=output_filename)
+    if pbar is not None:
+        pbar.update()
 
-def extract_fits(output_directory):
+
+def extract_fits(output_directory, filename=None):
     """ Extract fits files from tar.gz file """
-    file = os.path.join(output_directory, 'spectroscopy_download.tar.gz')
+    if filename is None:
+        filename = 'spectroscopy_download.tar.gz'
+    file = os.path.join(output_directory, filename)
     with tarfile.open(file, "r") as tar:
         files = []
         for member in tar.getmembers():
@@ -390,7 +442,7 @@ def do_symlink_filetype(type, raw_files, output_directory, clobber=False, top_le
 
 
 def do_download_filetype(type, raw_files, output_directory, clobber=False,
-                         verbose=True, chunk_size=20):
+                         verbose=True, chunk_size=20, parallel_limit=30):
     """ Download CCFs / S1Ds / S2Ds from DACE """
     raw_files = np.atleast_1d(raw_files)
 
@@ -410,7 +462,7 @@ def do_download_filetype(type, raw_files, output_directory, clobber=False,
 
     # avoid an empty chunk
     if chunk_size > n:
-        chunk_size = n
+        chunk_size = n  
 
     if verbose:
         if chunk_size < n:
@@ -422,11 +474,36 @@ def do_download_filetype(type, raw_files, output_directory, clobber=False,
             msg = f"downloading {n} {type}s into '{output_directory}'..."
             logger.info(msg)
 
-    iterator = [raw_files[i:i + chunk_size] for i in range(0, n, chunk_size)]
-    for files in tqdm(iterator, total=len(iterator)):
-        download(files, type, output_directory)
-        extract_fits(output_directory)
+    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)):
+            download(files, type, output_directory, quiet=False)
+            extract_fits(output_directory)
 
+    else:
+        def chunker(it, size):
+            iterator = iter(it)
+            while chunk := list(islice(iterator, size)):
+                yield chunk
+
+        chunks = list(chunker(raw_files, chunk_size))
+        pbar = tqdm(total=len(chunks))
+        it1 = [
+            (files, type, output_directory, f'spectroscopy_download{i+1}.tar.gz', True, pbar)
+            for i, files in enumerate(chunks)
+        ]
+        it2 = [(output_directory, f'spectroscopy_download{i+1}.tar.gz') for i in range(len(chunks))]
+
+        # import multiprocessing as mp
+        # with mp.Pool(4) as pool:
+        from multiprocessing.pool import ThreadPool
+
+        with ThreadPool(4) as pool:
+            pool.starmap(download, it1)
+            pool.starmap(extract_fits, it2)
+            print('')
+
+    sys.stdout.flush()
     logger.info('extracted .fits files')
 
 

arvi/kima_wrapper.py

@@ -25,10 +25,11 @@ def run_kima(self, run=False, load=False, run_directory=None, priors={}, **kwarg
     if not kima_available:
         raise ImportError('kima not available, please install with `pip install kima`')
 
-    time = [getattr(self, inst).mtime for inst in self.instruments]
-    vrad = [getattr(self, inst).mvrad for inst in self.instruments]
-    err = [getattr(self, inst).msvrad for inst in self.instruments]
-    data = RVData(time, vrad, err, instruments=self.instruments)
+    instruments = [inst for inst in self.instruments if self.NN[inst] > 1]
+    time = [getattr(self, inst).mtime for inst in instruments]
+    vrad = [getattr(self, inst).mvrad for inst in instruments]
+    err = [getattr(self, inst).msvrad for inst in instruments]
+    data = RVData(time, vrad, err, instruments=instruments)
 
     fix = kwargs.pop('fix', False)
     npmax = kwargs.pop('npmax', 1)
@@ -41,6 +42,13 @@ def run_kima(self, run=False, load=False, run_directory=None, priors={}, **kwarg
     model.enforce_stability = kwargs.pop('enforce_stability', False)
     model.star_mass = kwargs.pop('star_mass', 1.0)
 
+    if kwargs.pop('gaussian_priors_individual_offsets', False):
+        from kima.pykima.utils import get_gaussian_priors_individual_offsets
+        model.individual_offset_prior = get_gaussian_priors_individual_offsets(data, use_std=True)
+
+    if kwargs.pop('kuma', False):
+        model.conditional.eprior = distributions.Kumaraswamy(0.867, 3.03)
+
     for k, v in priors.items():
         try:
             if 'conditional' in k:
@@ -55,9 +63,10 @@ def run_kima(self, run=False, load=False, run_directory=None, priors={}, **kwarg
             logger.warning(msg)
             return
 
+    if run_directory is None:
+        run_directory = os.getcwd()
+
     if run:
-        if run_directory is None:
-            run_directory = os.getcwd()
         
         # TODO: use signature of kima.run to pop the correct kwargs
         # model_name = model.__class__.__name__
@@ -67,8 +76,9 @@ def run_kima(self, run=False, load=False, run_directory=None, priors={}, **kwarg
         with chdir(run_directory):
             kima.run(model, **kwargs)
         
-        if load:
+    if load:
+        with chdir(run_directory):
             res = kima.load_results(model)
-            return data, model, res
+        return data, model, res
 
     return data, model

arvi/programs.py

@@ -1,8 +1,9 @@
 import os
 import multiprocessing
-from functools import partial
+from functools import partial, lru_cache
 from itertools import chain
 from collections import namedtuple
+from multiprocessing.pool import ThreadPool
 from tqdm import tqdm
 # import numpy as np
 
@@ -14,9 +15,9 @@ __all__ = ['ESPRESSO_GTO']
 path = os.path.join(os.path.dirname(__file__), 'data')
 
 
-def get_star(star, instrument=None):
+def get_star(star, instrument=None, verbose=False, **kwargs):
     return RV(star, instrument=instrument,
-              _raise_on_error=False, verbose=False, load_extra_data=False)
+              _raise_on_error=False, verbose=verbose, **kwargs)
 
 
 class LazyRV:
@@ -36,21 +37,22 @@ class LazyRV:
     def __repr__(self):
         return f"RV({self.N} stars)"
 
-    def _get(self):
+    def _get(self, **kwargs):
         if self.N > self._parallel_limit:
             # logger.info('Querying DACE...')
-            _get_star = partial(get_star, instrument=self.instrument)
-            with multiprocessing.Pool() as pool:
+            _get_star = partial(get_star, instrument=self.instrument, **kwargs)
+            with ThreadPool(8) as pool:
                 result = list(tqdm(pool.imap(_get_star, self.stars), 
-                                   total=self.N, unit='star', desc='Querying DACE'))
-                # result = pool.map(get_star, self.stars)
+                                   total=self.N, unit='star',
+                                   desc='Querying DACE (can take a while)'))
+                print('')
         else:
             result = []
-            logger.info('Querying DACE...')
+            logger.info('querying DACE...')
             pbar = tqdm(self.stars, total=self.N, unit='star')
             for star in pbar:
                 pbar.set_description(star)
-                result.append(get_star(star, self.instrument))
+                result.append(get_star(star, self.instrument, **kwargs))
 
         return result
 
@@ -73,14 +75,25 @@ class LazyRV:
         #             result.append(None)
         # return result
 
+    def reload(self, **kwargs):
+        self._saved = self._get(**kwargs)
+        return self._saved
+
     def __iter__(self):
         return self._get()
 
-    def __call__(self):
+    def __call__(self, **kwargs):
         if not self._saved:
-            self._saved = self._get()
+            self._saved = self._get(**kwargs)
         return self._saved
 
+    @lru_cache(maxsize=10)
+    def __getitem__(self, index):
+        star = self.stars[index]
+        if self._saved is not None:
+            return self._saved[index]
+        return get_star(star, self.instrument, verbose=True)
+
 
 # sorted by spectral type
 WG1_stars = [
@@ -157,3 +170,78 @@ ESPRESSO_GTO = ESPRESSO_GTO_nt(
 ESPRESSO_GTO.WG1.__doc__ = 'RV observations for all WG1 targets. Call ESPRESSO_GTO.WG1() to load them.'
 ESPRESSO_GTO.WG2.__doc__ = 'RV observations for all WG2 targets. Call ESPRESSO_GTO.WG2() to load them.'
 ESPRESSO_GTO.WG3.__doc__ = 'RV observations for all WG3 targets. Call ESPRESSO_GTO.WG3() to load them.'
+
+
+import requests
+
+def _get_NIRPS_GTO_stars(WP=1):
+    from io import StringIO
+    import numpy as np
+
+    url = 'https://www.eso.org/sci/observing/teles-alloc/gto/113/NIRPS/P113_NIRPS-consortium.csv'
+    file = StringIO(requests.get(url).content.decode())
+    stars_P113 = np.loadtxt(file, delimiter=',', usecols=(0,), dtype=str, skiprows=3)
+    
+    url = 'https://www.eso.org/sci/observing/teles-alloc/gto/114/NIRPS/P114_NIRPS-consortium.csv'
+    file = StringIO(requests.get(url).content.decode())
+    stars_P114 = np.loadtxt(file, delimiter=',', usecols=(0,), dtype=str, skiprows=3)
+
+    url = 'https://www.eso.org/sci/observing/teles-alloc/gto/115/NIRPS/P115_NIRPS-consortium.csv'
+    file = StringIO(requests.get(url).content.decode())
+    stars_P115 = np.loadtxt(file, delimiter=',', usecols=(0,), dtype=str, skiprows=3)
+
+    def _get_stars_period(stars, WP):
+        stars = np.delete(stars, stars=='')
+
+        stars = np.char.replace(stars, '_', ' ')
+        stars = np.char.replace(stars, "Proxima Centauri", "Proxima")
+        stars = np.char.replace(stars, "Barnard's star", "GJ699")
+        stars = np.char.replace(stars, "Teegarden's Star", 'Teegarden')
+
+        if WP in (1, 'WP1'):
+            wp1_indices = slice(np.where(stars == 'WP1')[0][0] + 1, np.where(stars == 'WP2')[0][0])
+            return stars[wp1_indices]
+        elif WP == 2:
+            wp2_indices = slice(np.where(stars == 'WP2')[0][0] + 1, np.where(stars == 'WP3')[0][0])
+            return stars[wp2_indices]
+        elif WP == 3:
+            wp3_indices = slice(np.where(stars == 'WP3')[0][0] + 1, np.where(stars == 'Other Science 1')[0][0])
+            return stars[wp3_indices]
+        elif WP == 'OS1':
+            os1_indices = slice(np.where(stars == 'Other Science 1')[0][0] + 1, np.where(stars == 'Other Science 2')[0][0])
+            return stars[os1_indices]
+        elif WP == 'OS2':
+            os2_indices = slice(np.where(stars == 'Other Science 2')[0][0] + 1, None)
+            stars = np.char.replace(stars, 'MMU', 'No')
+            stars = np.char.replace(stars, 'Cl*', '')
+            return stars[os2_indices]
+    
+    stars_P113 = _get_stars_period(stars_P113, WP)
+    stars_P114 = _get_stars_period(stars_P114, WP)
+    stars_P115 = _get_stars_period(stars_P115, WP)
+    return np.union1d(np.union1d(stars_P113, stars_P114), stars_P115)
+
+try:
+    NIRPS_GTO_WP1_stars = _get_NIRPS_GTO_stars(WP=1)
+    NIRPS_GTO_WP2_stars = _get_NIRPS_GTO_stars(WP=2)
+    NIRPS_GTO_WP3_stars = _get_NIRPS_GTO_stars(WP=3)
+    NIRPS_GTO_OS1_stars = _get_NIRPS_GTO_stars(WP='OS1')
+    NIRPS_GTO_OS2_stars = _get_NIRPS_GTO_stars(WP='OS2')
+except requests.ConnectionError:
+    from .setup_logger import logger
+    logger.error('Cannot download NIRPS GTO protected target list')
+else:
+    NIRPS_GTO_nt = namedtuple('NIRPS_GTO', ['WP1', 'WP2', 'WP3', 'OS1', 'OS2'])
+    NIRPS_GTO_nt.__doc__ = 'RV observations for all NIRPS GTO targets. See NIRPS_GTO.WP1, NIRPS_GTO.WP2, ...'
+    NIRPS_GTO = NIRPS_GTO_nt(
+        WP1=LazyRV(NIRPS_GTO_WP1_stars, instrument='NIRPS'),
+        WP2=LazyRV(NIRPS_GTO_WP2_stars, instrument='NIRPS'),
+        WP3=LazyRV(NIRPS_GTO_WP3_stars, instrument='NIRPS'),
+        OS1=LazyRV(NIRPS_GTO_OS1_stars, instrument='NIRPS'),
+        OS2=LazyRV(NIRPS_GTO_OS2_stars, instrument='NIRPS'),
+    )
+    NIRPS_GTO.WP1.__doc__ = 'RV observations for all WP1 targets. Call NIRPS_GTO.WP1() to load them.'
+    NIRPS_GTO.WP2.__doc__ = 'RV observations for all WP2 targets. Call NIRPS_GTO.WP2() to load them.'
+    NIRPS_GTO.WP3.__doc__ = 'RV observations for all WP3 targets. Call NIRPS_GTO.WP3() to load them.'
+    NIRPS_GTO.OS1.__doc__ = 'RV observations for all OS1 targets. Call NIRPS_GTO.OS1() to load them.'
+    NIRPS_GTO.OS2.__doc__ = 'RV observations for all OS2 targets. Call NIRPS_GTO.OS2() to load them.'

arvi/spectra.py

@@ -72,7 +72,7 @@ def fit_gaussian_to_line(wave, flux, center_wavelength, around=0.15 * u.angstrom
     ]).T
 
     try:
-        popt, pcov = curve_fit(gaussian, w, f, p0=[-f.ptp(), center_wavelength.value, 0.1, f.max()],
+        popt, pcov = curve_fit(gaussian, w, f, p0=[-np.ptp(f), center_wavelength.value, 0.1, f.max()],
                                bounds=(lower, upper))
     except RuntimeError as e:
         logger.warning(f'fit_gaussian_to_line: {e}')
@@ -115,7 +115,7 @@ def detrend(w, f):
 def build_master(self, limit=None, plot=True):
     files = sorted(glob(f'{self.star}_downloads/*S1D_A.fits'))
     if self.verbose:
-        logger.info(f'Found {len(files)} S1D files')
+        logger.info(f'found {len(files)} S1D files')
 
     files = files[:limit]
 
@@ -168,8 +168,8 @@ def determine_stellar_parameters(self, linelist: str, plot=True, **kwargs):
     ]
 
     if self.verbose:
-        logger.info(f'Found {len(lines)} lines in linelist')
-        logger.info('Measuring EWs...')
+        logger.info(f'found {len(lines)} lines in linelist')
+        logger.info('measuring EWs...')
 
     EW = []
     pbar = tqdm(linelist)
@@ -183,14 +183,14 @@ def determine_stellar_parameters(self, linelist: str, plot=True, **kwargs):
     EW = np.array(EW)[~np.isnan(EW)]
 
     if self.verbose:
-        logger.info('Determining stellar parameters (can take a few minutes)...')
+        logger.info('determining stellar parameters (can take a few minutes)...')
 
     callback = lambda p, r, A: print('current parameters:', p)
     result = Korg.Fit.ews_to_stellar_parameters(lines, EW, callback=callback)
     par, stat_err, sys_err = result
 
     if self.verbose:
-        logger.info(f'Best fit stellar parameters:')
+        logger.info(f'best-fit stellar parameters:')
         logger.info(f'  Teff: {par[0]:.0f} ± {sys_err[0]:.0f} K')
         logger.info(f'  logg: {par[1]:.2f} ± {sys_err[1]:.2f} dex')
         logger.info(f'  m/H :  {par[3]:.2f} ± {sys_err[3]:.2f} dex')

arvi/timeseries.py

@@ -193,7 +193,7 @@ class RV:
                 time_stamp = datetime.now(timezone.utc)  #.isoformat().split('.')[0]
                 self._last_dace_query = time_stamp
 
-        _replacements = (('-', '_'), ('.', '_'), ('__', '_'))
+        _replacements = (('-', '_'), ('.', '_'), (' ', '_'), ('__', '_'))
         def do_replacements(s):
             for a, b in _replacements:
                 s = s.replace(a, b)
@@ -444,7 +444,7 @@ class RV:
 
         s.instruments = [inst]
         s.pipelines = [pipe]
-        s.modes = [mode]
+        s.modes = [str(mode)]
 
         return s
 
@@ -598,29 +598,47 @@ class RV:
             _quantities.append('rhk')
             _quantities.append('rhk_err')
 
-            _s.bispan = np.zeros_like(time)
-            _s.bispan_err = np.full_like(time, np.nan)
+            # try to find BISPAN and uncertainty
+            if (v := find_column(data, ['bispan'])) is not False:
+                _s.bispan = v
+                _s.bispan_err = np.full_like(time, np.nan)
+                if (sv := find_column(data, ['sbispan'])) is not False:
+                    _s.bispan_err = sv
+            else:
+                _s.bispan = np.full_like(time, np.nan)
+                _s.bispan_err = np.full_like(time, np.nan)
+
+            _quantities.append('bispan')
+            _quantities.append('bispan_err')
+
+            # try to find BERV
+            if (v := find_column(data, ['berv', 'HIERARCH ESO QC BERV'])) is not False:
+                _s.berv = v
+            else:
+                _s.berv = np.full_like(time, np.nan)
+            _quantities.append('berv')
 
             # other quantities, but all NaNs
-            for q in ['bispan', 'caindex', 'ccf_asym', 'contrast', 'haindex', 'naindex', 'sindex']:
+            for q in ['caindex', 'ccf_asym', 'contrast', 'haindex', 'naindex', 'sindex']:
                 setattr(_s, q, np.full_like(time, np.nan))
                 setattr(_s, q + '_err', np.full_like(time, np.nan))
                 _quantities.append(q)
                 _quantities.append(q + '_err')
-            for q in ['berv', 'texp']:
+            for q in ['texp', ]:
                 setattr(_s, q, np.full_like(time, np.nan))
                 _quantities.append(q)
             for q in ['ccf_mask', 'date_night', 'prog_id', 'raw_file', 'pub_reference']:
                 setattr(_s, q, np.full(time.size, ''))
                 _quantities.append(q)
-            for q in ['drs_qc']:
+            for q in ['drs_qc', ]:
                 setattr(_s, q, np.full(time.size, True))
                 _quantities.append(q)
 
             _s.extra_fields = ExtraFields()
-            for field in data.dtype.names:
-                if field not in _quantities:
-                    setattr(_s.extra_fields, field, data[field])
+            for name in data.dtype.names:
+                if name not in _quantities:
+                    name_ = name.replace(' ', '_')
+                    setattr(_s.extra_fields, name_, data[name])
                     # _quantities.append(field)
 
             #! end hack
@@ -628,12 +646,12 @@ class RV:
             _s.mask = np.ones_like(time, dtype=bool)
             _s.obs = np.full_like(time, i + 1)
 
-            _s.instruments = [instrument]
+            _s.instruments = [str(instrument)]
             _s._quantities = np.array(_quantities)
             setattr(s, instrument, _s)
 
         s._child = False
-        s.instruments = list(instruments)
+        s.instruments = list(map(str, instruments))
         s._build_arrays()
 
         if kwargs.get('do_adjust_means', False):
@@ -726,15 +744,17 @@ class RV:
         return s
 
     @classmethod
-    @lru_cache(maxsize=10)
+    @lru_cache(maxsize=60)
     def from_KOBE_file(cls, star, **kwargs):
         assert 'KOBE' in star, f'{star} is not a KOBE star?'
         import requests
         from requests.auth import HTTPBasicAuth
         from io import BytesIO
         import tarfile
+        from time import time as pytime
         from astropy.io import fits
         from .config import config
+        from .utils import get_data_path
 
         try:
             config.kobe_password
@@ -743,23 +763,11 @@ class RV:
             return
 
         tar = None
+        local_targz_file = os.path.join(get_data_path(), 'KOBE_fitsfiles.tar.gz')
         fits_file = f'{star}_RVs.fits'
-        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))
-
-            if resp.status_code != 200:
-                logger.error(f'KOBE file not found for {star}')
-                return
 
-            tar = tarfile.open(fileobj=BytesIO(resp.content))
+        if os.path.exists(local_targz_file) and os.path.getmtime(local_targz_file) > pytime() - 60*60*2:
+            tar = tarfile.open(local_targz_file)
 
             if fits_file not in tar.getnames():
                 logger.error(f'KOBE file not found for {star}')
@@ -768,8 +776,36 @@ class RV:
             hdul = fits.open(tar.extractfile(fits_file))
 
         else:
-            # found the file on the server, read it directly
-            hdul = fits.open(BytesIO(resp.content))
+            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))
+
+                if resp.status_code != 200:
+                    logger.error(f'KOBE file not found for {star}')
+                    return
+
+                # save tar.gz file for later
+                with open(local_targz_file, 'wb') as tg:
+                    tg.write(resp.content)
+
+                tar = tarfile.open(fileobj=BytesIO(resp.content))
+
+                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:
+                # found the file on the server, read it directly
+                hdul = fits.open(BytesIO(resp.content))
 
         s = cls(star, _child=True)
 
@@ -780,26 +816,47 @@ class RV:
         s.vrad_preNZP = hdul[1].data['RVd']
         s.vrad_preNZP_err = hdul[1].data['eRVd']
 
+        s.fwhm = hdul[1].data['FWHM']
+        s.fwhm_err = hdul[1].data['eFWHM']
+
+        s.crx = hdul[1].data['CRX']
+        s.crx_err = hdul[1].data['eCRX']
+        s.dlw = hdul[1].data['DLW']
+        s.dlw_err = hdul[1].data['eDLW']
+        s.contrast = hdul[1].data['CONTRAST']
+        s.contrast_err = hdul[1].data['eCONTRAST']
+        s.bispan = hdul[1].data['BIS']
+        s.bispan_err = hdul[1].data['eBIS']
+
+
         s.drift = hdul[1].data['drift']
         s.drift_err = hdul[1].data['e_drift']
 
         s.nzp = hdul[1].data['NZP']
         s.nzp_err = hdul[1].data['eNZP']
 
+        s.texp = hdul[1].data['ExpTime']
         s.berv = hdul[1].data['BERV']
+        s.units = 'km/s'
 
+        s.obs = np.ones_like(s.time, dtype=int)
         s.mask = np.full_like(s.time, True, dtype=bool)
         s.instruments = ['CARMENES']
+        s._quantities = np.array(['berv', ])
 
         # so meta!
         setattr(s, 'CARMENES', s)
 
         s._kobe_result = hdul[1].data
 
+        s.mask = s._kobe_result['rvflag']
+        s._propagate_mask_changes()
+
         if tar is not None:
             tar.close()
         hdul.close()
 
+        s._child = False
         return s
 
 
@@ -1190,6 +1247,13 @@ class RV:
         if (self.time < bjd).any():
             ind = np.where(self.time < bjd)[0]
             self.remove_point(ind)
+    
+    def remove_between_bjds(self, bjd1, bjd2):
+        """ 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):
         """ Randomly choose `n` observations and mask out the remaining ones
@@ -1700,8 +1764,21 @@ class RV:
             s = getattr(self, inst)
             s.vrad *= factor
             s.svrad *= factor
-            s.fwhm *= factor
-            s.fwhm_err *= factor
+            try:
+                s.fwhm *= factor
+                s.fwhm_err *= factor
+            except AttributeError:
+                pass
+
+            for q in (
+                'bispan',
+                'nzp', 'vrad_preNZP',
+            ):
+                try:
+                    setattr(s, q, getattr(s, q) * factor)
+                    setattr(s, f'{q}_err', getattr(s, f'{q}_err') * factor)
+                except AttributeError:
+                    pass
 
         self._build_arrays()
         self.units = new_units

arvi/translations.py

@@ -6,6 +6,11 @@ STARS = {
     "Barnard's": 'GJ699',
     'Ross128': 'Ross  128',
     'Ross 128': 'Ross  128',
+    # 
+    'Teegarden': 'GAT1370',
+    "Teegarden's Star": 'GAT1370',
+    # 
+    "Smethells 20": 'TIC464410508',
 }
 
 

{arvi-0.1.23.dist-info → arvi-0.1.24.dist-info}/METADATA

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

{arvi-0.1.23.dist-info → arvi-0.1.24.dist-info}/RECORD

@@ -4,24 +4,24 @@ arvi/ariadne_wrapper.py,sha256=YvilopJa9T4NwPcj3Nah_U8smSeSAU5-HYZMb_GJ-BQ,2232
 arvi/berv.py,sha256=eKnpuPC1w45UrUEyFRbs9F9j3bXz3kxYzNXbnRgvFQM,17596
 arvi/binning.py,sha256=jbemJ-bM3aqoOsqMo_OhWt_co-JAQ0nhdG_GpTsrRsw,15403
 arvi/config.py,sha256=W-v8NNhRd_PROu0wCMilXmOhYcju4xbUalugd5u7SRU,1881
-arvi/dace_wrapper.py,sha256=U3N5iO_nRHU3FlxrjsMxFzl9O32UKpKuAjhkAM1dQfs,18432
+arvi/dace_wrapper.py,sha256=CUHKN5m7KULM1sES5x0GXx8HQgJE7XdwwWv6_zAYhb4,21324
 arvi/extra_data.py,sha256=WEEaYeLh52Zdv0uyHO72Ys5MWS3naTAP4wJV2BJ1mbk,2551
 arvi/gaia_wrapper.py,sha256=icm3LJjG9pjP47_bM30NFyocUQO3X3SHS5yQ-Dwcr5w,4653
 arvi/headers.py,sha256=uvdJebw1M5YkGjE3vJJwYBOnLikib75uuZE9FXB5JJM,1673
 arvi/instrument_specific.py,sha256=-pbm2Vk3iK_1K7nDa1avlJOKHBcXllwILI4lQn-Ze-A,7761
-arvi/kima_wrapper.py,sha256=y_Z0Hl2ECbs2-B6ZR9retrjId7-QxcRylG7b5aDsiFk,2306
+arvi/kima_wrapper.py,sha256=BvNTVqzM4lMNhLCyBFVh3T84hHfGKAFpgiYiOi4lh0g,2731
 arvi/lbl_wrapper.py,sha256=_ViGVkpakvuBR_xhu9XJRV5EKHpj5Go6jBZGJZMIS2Y,11850
 arvi/nasaexo_wrapper.py,sha256=mWt7eHgSZe4MBKCmUvMPTyUPGuiwGTqKugNBvmjOg9s,7306
 arvi/plots.py,sha256=WUm-sqN0aZTNXvE1kYpvmHTW9QPWqSCpKhNjwaqxjEk,29628
-arvi/programs.py,sha256=C0Fbldjf-QEZYYJp5wBKP3h7zraD0O2mJC7Su967STg,4607
+arvi/programs.py,sha256=BW7xBNKLei7NVLLW3_lsVskwzkaIoNRiHK2jn9Tn2ZM,8879
 arvi/reports.py,sha256=ayPdZ4HZO9iCDdnADQ18gQPJh79o-1UYG7TYkvm9Lrc,4051
 arvi/setup_logger.py,sha256=pBzaRTn0hntozjbaRVx0JIbWGuENkvYUApa6uB-FsRo,279
 arvi/simbad_wrapper.py,sha256=iAAwEMcr1Hgu6lnDctmaCC1TLPCB8yAfHG0wxh9K9C8,5791
-arvi/spectra.py,sha256=pTAWSW4vk96DWRQ-6l5mNJHUhiAyaPR-QDjZdOT6Ak0,7489
+arvi/spectra.py,sha256=ebF1ocodTastLx0CyqLSpE8EZNDXBF8riyfxMr3L6H0,7491
 arvi/stats.py,sha256=ilzzGL9ew-SyVa9eEdrYCpD3DliOAwhoNUg9LIlHjzU,2583
 arvi/stellar.py,sha256=veuL_y9kJvvApU_jqYQqP3EkcRnQffTc8Us6iT5UrFI,3790
-arvi/timeseries.py,sha256=qkBE6-NPBTXZyGLVHXRVUt5atrdcH3fEuaZ4SaHF7Bk,74443
-arvi/translations.py,sha256=FBF_2OrpMQBG4GtV4_UOspiaxetiGCY7TQFcwZMMVuQ,838
+arvi/timeseries.py,sha256=NdmSSYeDdS-cYBXPt8NCeKS1jLdv8LP6Zh561KRGfZc,77328
+arvi/translations.py,sha256=PUSrn4zvYO2MqGzUxlFGwev_tBkgJaJrIYs6NKHzbWo,951
 arvi/utils.py,sha256=LImV8iPjG8ZKjPCT9lp25_pDb-51ZZk42Hc8bzZt7M0,6568
 arvi/data/info.svg,sha256=0IMI6W-eFoTD8acnury79WJJakpBwLa4qKS4JWpsXiI,489
 arvi/data/obs_affected_ADC_issues.dat,sha256=tn93uOL0eCTYhireqp1wG-_c3CbxPA7C-Rf-pejVY8M,10853
@@ -29,8 +29,8 @@ arvi/data/obs_affected_blue_cryostat_issues.dat,sha256=z4AK17xfz8tGTDv1FjRvQFnio
 arvi/data/extra/HD86226_PFS1.rdb,sha256=vfAozbrKHM_j8dYkCBJsuHyD01KEM1asghe2KInwVao,3475
 arvi/data/extra/HD86226_PFS2.rdb,sha256=F2P7dB6gVyzCglUjNheB0hIHVClC5RmARrGwbrY1cfo,4114
 arvi/data/extra/metadata.json,sha256=C69hIw6CohyES6BI9vDWjxwSz7N4VOYX0PCgjXtYFmU,178
-arvi-0.1.23.dist-info/LICENSE,sha256=6JfQgl7SpM55t0EHMFNMnNh-AdkpGW25MwMiTnhdWQg,1068
-arvi-0.1.23.dist-info/METADATA,sha256=hI6gLlWx_4i0WSnXTaV6XGowzOOkWDfPW13t8Jx92Os,1852
-arvi-0.1.23.dist-info/WHEEL,sha256=a7TGlA-5DaHMRrarXjVbQagU3Man_dCnGIWMJr5kRWo,91
-arvi-0.1.23.dist-info/top_level.txt,sha256=4EeiKDVLD45ztuflTGfQ3TU8GVjJg5Y95xS5XjI-utU,5
-arvi-0.1.23.dist-info/RECORD,,
+arvi-0.1.24.dist-info/LICENSE,sha256=6JfQgl7SpM55t0EHMFNMnNh-AdkpGW25MwMiTnhdWQg,1068
+arvi-0.1.24.dist-info/METADATA,sha256=ijGk4XpvyZZ1IxV0_6iFuJkjZI4wtWW6A4kuCZE-gSA,1852
+arvi-0.1.24.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
+arvi-0.1.24.dist-info/top_level.txt,sha256=4EeiKDVLD45ztuflTGfQ3TU8GVjJg5Y95xS5XjI-utU,5
+arvi-0.1.24.dist-info/RECORD,,

{arvi-0.1.23.dist-info → arvi-0.1.24.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.4.0)
+Generator: setuptools (75.5.0)
 Root-Is-Purelib: true
 Tag: py3-none-any