arvi 0.1.23__tar.gz → 0.1.24__tar.gz

{arvi-0.1.23 → arvi-0.1.24}/.github/workflows/install.yml

@@ -16,7 +16,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.8", "3.9", "3.10", "3.11"]
+        python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
 
     steps:
     - uses: actions/checkout@v4

{arvi-0.1.23 → arvi-0.1.24}/.gitignore

@@ -2,7 +2,7 @@
 *.pkl
 *.rdb
 arvi/data/KOBE-translate.csv
-
+*_downloads
 
 
 # Byte-compiled / optimized / DLL files

{arvi-0.1.23 → arvi-0.1.24}/PKG-INFO

@@ -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 → arvi-0.1.24}/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-0.1.23 → arvi-0.1.24}/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-0.1.24/arvi/programs.py

@@ -0,0 +1,247 @@
+import os
+import multiprocessing
+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
+
+from .setup_logger import logger
+from .timeseries import RV
+
+__all__ = ['ESPRESSO_GTO']
+
+path = os.path.join(os.path.dirname(__file__), 'data')
+
+
+def get_star(star, instrument=None, verbose=False, **kwargs):
+    return RV(star, instrument=instrument,
+              _raise_on_error=False, verbose=verbose, **kwargs)
+
+
+class LazyRV:
+    def __init__(self, stars: list, instrument: str = None,
+                 _parallel_limit=10):
+        self.stars = stars
+        if isinstance(self.stars, str):
+            self.stars = [self.stars]
+        self.instrument = instrument
+        self._saved = None
+        self._parallel_limit = _parallel_limit
+
+    @property
+    def N(self):
+        return len(self.stars)
+
+    def __repr__(self):
+        return f"RV({self.N} stars)"
+
+    def _get(self, **kwargs):
+        if self.N > self._parallel_limit:
+            # logger.info('Querying DACE...')
+            _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 (can take a while)'))
+                print('')
+        else:
+            result = []
+            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, **kwargs))
+
+        return result
+
+        # # use a with statement to ensure threads are cleaned up promptly
+        # with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool:
+        #     star_to_RV = {
+        #         pool.submit(get_star, star, self.instrument): star 
+        #         for star in self.stars
+        #     }
+        #     logger.info('Querying DACE...')
+        #     pbar = tqdm(concurrent.futures.as_completed(star_to_RV),
+        #                 total=self.N, unit='star')
+        #     for future in pbar:
+        #         star = star_to_RV[future]
+        #         pbar.set_description(star)
+        #         try:
+        #             result.append(future.result())
+        #         except ValueError:
+        #             print(f'{star} generated an exception')
+        #             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, **kwargs):
+        if not self._saved:
+            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 = [
+    "HIP11533",  # F2
+    "HD63077",  # F9
+    "HD102365",  # G2
+    "HD160691",  # G3
+    "HD20794",  # G6
+    "HD115617",  # G6
+    "HD10700",  # G8
+    "HD69830",  # G8
+    "HD26965",  # K0
+    "HD100623",  # K0
+    "HD154088",  # K0
+    "HD72673",  # K1
+    "HD4628",  # K2
+    "HD191408",  # K2
+    "HD192310",  # K2
+    "HD16160",  # K3
+    "HD32147",  # K3
+    "HD22496",  # K5
+    "HIP93069",  # K5
+    "HD209100",  # K5
+    "HIP23708",  # K7
+    "HD152606",  # K8
+    "HD260655",  # M0
+    "HIP40239",  # M0
+    "HD304636",  # M0
+    "HIP85647",  # M0
+    "HD165222",  # M0
+    "HD191849",  # M0
+    "GJ191",  # M1
+    "HD42581",  # M1
+    "HIP42748",  # M1
+    "HIP65859",  # M1
+    "HIP86287",  # M1
+    "HD176029",  # M1
+    "GJ825",  # M1
+    "HD225213",  # M2
+    "GJ54",  # M2
+    "HIP22627",  # M2
+    "HIP51317",  # M2
+    "HD119850",  # M2
+    "GJ832",  # M2
+    "HD217987",  # M2
+    "GJ273",  # M3
+    "GJ388",  # M3
+    "HIP62452",  # M3
+    "HIP67164",  # M3
+    "HIP71253",  # M3
+    "GJ628",  # M3
+    "HIP85523",  # M3
+    "HIP86214",  # M3
+    "HIP92403",  # M3
+    "HIP113020",  # M3
+    "HIP1242",  # M4
+    "GJ83.1",  # M4
+    "HIP53020",  # M4
+    "Ross128",  # M4
+    "GJ699",  # M4
+    "GJ1002",  # M5
+    "GJ1061",  # M5
+    "GJ3618",  # M5
+    "Proxima",  # M5
+    "GJ406",  # M6
+]
+
+ESPRESSO_GTO_nt = namedtuple('ESPRESSO_GTO', ['WG1', 'WG2', 'WG3'])
+ESPRESSO_GTO = ESPRESSO_GTO_nt(
+    WG1=LazyRV(WG1_stars, instrument='ESPRESSO'),
+    WG2=LazyRV([], instrument='ESPRESSO'), # TODO
+    WG3=LazyRV([], instrument='ESPRESSO'), # TODO
+)
+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-0.1.23 → arvi-0.1.24}/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-0.1.23 → arvi-0.1.24}/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-0.1.23 → arvi-0.1.24}/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 → arvi-0.1.24}/arvi.egg-info/PKG-INFO

@@ -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/arvi/programs.py

@@ -1,159 +0,0 @@
-import os
-import multiprocessing
-from functools import partial
-from itertools import chain
-from collections import namedtuple
-from tqdm import tqdm
-# import numpy as np
-
-from .setup_logger import logger
-from .timeseries import RV
-
-__all__ = ['ESPRESSO_GTO']
-
-path = os.path.join(os.path.dirname(__file__), 'data')
-
-
-def get_star(star, instrument=None):
-    return RV(star, instrument=instrument,
-              _raise_on_error=False, verbose=False, load_extra_data=False)
-
-
-class LazyRV:
-    def __init__(self, stars: list, instrument: str = None,
-                 _parallel_limit=10):
-        self.stars = stars
-        if isinstance(self.stars, str):
-            self.stars = [self.stars]
-        self.instrument = instrument
-        self._saved = None
-        self._parallel_limit = _parallel_limit
-
-    @property
-    def N(self):
-        return len(self.stars)
-
-    def __repr__(self):
-        return f"RV({self.N} stars)"
-
-    def _get(self):
-        if self.N > self._parallel_limit:
-            # logger.info('Querying DACE...')
-            _get_star = partial(get_star, instrument=self.instrument)
-            with multiprocessing.Pool() 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)
-        else:
-            result = []
-            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))
-
-        return result
-
-        # # use a with statement to ensure threads are cleaned up promptly
-        # with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool:
-        #     star_to_RV = {
-        #         pool.submit(get_star, star, self.instrument): star 
-        #         for star in self.stars
-        #     }
-        #     logger.info('Querying DACE...')
-        #     pbar = tqdm(concurrent.futures.as_completed(star_to_RV),
-        #                 total=self.N, unit='star')
-        #     for future in pbar:
-        #         star = star_to_RV[future]
-        #         pbar.set_description(star)
-        #         try:
-        #             result.append(future.result())
-        #         except ValueError:
-        #             print(f'{star} generated an exception')
-        #             result.append(None)
-        # return result
-
-    def __iter__(self):
-        return self._get()
-
-    def __call__(self):
-        if not self._saved:
-            self._saved = self._get()
-        return self._saved
-
-
-# sorted by spectral type
-WG1_stars = [
-    "HIP11533",  # F2
-    "HD63077",  # F9
-    "HD102365",  # G2
-    "HD160691",  # G3
-    "HD20794",  # G6
-    "HD115617",  # G6
-    "HD10700",  # G8
-    "HD69830",  # G8
-    "HD26965",  # K0
-    "HD100623",  # K0
-    "HD154088",  # K0
-    "HD72673",  # K1
-    "HD4628",  # K2
-    "HD191408",  # K2
-    "HD192310",  # K2
-    "HD16160",  # K3
-    "HD32147",  # K3
-    "HD22496",  # K5
-    "HIP93069",  # K5
-    "HD209100",  # K5
-    "HIP23708",  # K7
-    "HD152606",  # K8
-    "HD260655",  # M0
-    "HIP40239",  # M0
-    "HD304636",  # M0
-    "HIP85647",  # M0
-    "HD165222",  # M0
-    "HD191849",  # M0
-    "GJ191",  # M1
-    "HD42581",  # M1
-    "HIP42748",  # M1
-    "HIP65859",  # M1
-    "HIP86287",  # M1
-    "HD176029",  # M1
-    "GJ825",  # M1
-    "HD225213",  # M2
-    "GJ54",  # M2
-    "HIP22627",  # M2
-    "HIP51317",  # M2
-    "HD119850",  # M2
-    "GJ832",  # M2
-    "HD217987",  # M2
-    "GJ273",  # M3
-    "GJ388",  # M3
-    "HIP62452",  # M3
-    "HIP67164",  # M3
-    "HIP71253",  # M3
-    "GJ628",  # M3
-    "HIP85523",  # M3
-    "HIP86214",  # M3
-    "HIP92403",  # M3
-    "HIP113020",  # M3
-    "HIP1242",  # M4
-    "GJ83.1",  # M4
-    "HIP53020",  # M4
-    "Ross128",  # M4
-    "GJ699",  # M4
-    "GJ1002",  # M5
-    "GJ1061",  # M5
-    "GJ3618",  # M5
-    "Proxima",  # M5
-    "GJ406",  # M6
-]
-
-ESPRESSO_GTO_nt = namedtuple('ESPRESSO_GTO', ['WG1', 'WG2', 'WG3'])
-ESPRESSO_GTO = ESPRESSO_GTO_nt(
-    WG1=LazyRV(WG1_stars, instrument='ESPRESSO'),
-    WG2=LazyRV([], instrument='ESPRESSO'), # TODO
-    WG3=LazyRV([], instrument='ESPRESSO'), # TODO
-)
-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.'