arvi 0.1.30__tar.gz → 0.2.11__tar.gz

{arvi-0.1.30 → arvi-0.2.11}/.github/workflows/docs-gh-pages.yml

@@ -26,13 +26,13 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
-        uses: actions/checkout@v4
+        uses: actions/checkout@v5
         
       - name: Setup Pages
         uses: actions/configure-pages@v5
 
       - name: Set up Python
-        uses: actions/setup-python@v5
+        uses: actions/setup-python@v6
         with:
           python-version: "3.10"
   
@@ -42,6 +42,11 @@ jobs:
           pip install mkdocs mkdocs-material 
           pip install "mkdocstrings[crystal,python]" mkdocs-autorefs mkdocstrings-python
           pip install mkdocs-git-revision-date-localized-plugin
+          pip install mkdocs-jupyter
+
+      - name: Install arvi itself
+        run: |
+          pip install arvi
 
       - name: Build with Mkdocs
         run: |

{arvi-0.1.30 → arvi-0.2.11}/.github/workflows/install.yml

@@ -16,9 +16,9 @@ jobs:
         python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
 
     steps:
-    - uses: actions/checkout@v4
+    - uses: actions/checkout@v5
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v5
+      uses: actions/setup-python@v6
       with:
         python-version: ${{ matrix.python-version }}
 

{arvi-0.1.30 → arvi-0.2.11}/.github/workflows/python-publish.yml

@@ -22,9 +22,9 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v4
+    - uses: actions/checkout@v5
     - name: Set up Python
-      uses: actions/setup-python@v5
+      uses: actions/setup-python@v6
       with:
         python-version: '3.x'
     - name: Install dependencies
@@ -34,7 +34,7 @@ jobs:
     - name: Build package
       run: python -m build
     - name: Publish package
-      uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
+      uses: pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e
       with:
         user: __token__
         password: ${{ secrets.PYPI_API_TOKEN }}

{arvi-0.1.30/arvi.egg-info → arvi-0.2.11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: arvi
-Version: 0.1.30
+Version: 0.2.11
 Summary: The Automated RV Inspector
 Author-email: João Faria <joao.faria@unige.ch>
 License: MIT
@@ -23,7 +23,7 @@ Requires-Dist: kepmodel
 Dynamic: license-file
 
 <p align="center">
-  <img width = "140" src="https://github.com/j-faria/arvi/blob/main/docs/logo/logo.png?raw=true"/>
+  <img width = "140" src="https://raw.githubusercontent.com/j-faria/arvi/refs/heads/main/docs/logo/logo.png"/>
 </p>
 
 This package sits alongside [DACE](https://dace.unige.ch/) to help with the

{arvi-0.1.30 → arvi-0.2.11}/README.md

@@ -1,5 +1,5 @@
 <p align="center">
-  <img width = "140" src="https://github.com/j-faria/arvi/blob/main/docs/logo/logo.png?raw=true"/>
+  <img width = "140" src="https://raw.githubusercontent.com/j-faria/arvi/refs/heads/main/docs/logo/logo.png"/>
 </p>
 
 This package sits alongside [DACE](https://dace.unige.ch/) to help with the

{arvi-0.1.30 → arvi-0.2.11}/arvi/__init__.py

@@ -1,22 +1,22 @@
-__all__ = ['RV']
+__all__ = ['RV', 'config', 'simbad', 'gaia']
 
 from importlib.metadata import version, PackageNotFoundError
-
-from .config import config
-from .timeseries import RV
-
 try:
     __version__ = version("arvi")
 except PackageNotFoundError:
     # package is not installed
     pass
 
-## OLD
-# # the __getattr__ function is always called twice, so we need this
-# # to only build and return the RV object on the second time
-# _ran_once = False
+from .config import config
+from .simbad_wrapper import simbad
+from .gaia_wrapper import gaia
+
+from .timeseries import RV
 
 def __getattr__(name: str):
+    if not config.fancy_import:
+        raise AttributeError
+    
     if name in (
         '_ipython_canary_method_should_not_exist_',
         '_ipython_display_',
@@ -29,15 +29,5 @@ def __getattr__(name: str):
         globals()[name] = RV(name)
         return globals()[name]
     except ValueError as e:
-        raise ImportError(e) from None
-        # raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
-
-    ## OLD
-    # # can't do it any other way :(
-    # global _ran_once
+        raise AttributeError(e)
 
-    # if _ran_once:
-    #     _ran_once = False
-    #     return RV(name)
-    # else:
-    #     _ran_once = True

{arvi-0.1.30 → arvi-0.2.11}/arvi/binning.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from .setup_logger import logger
+from .setup_logger import setup_logger
 
 ###############################################################################
 # the following is mostly a copy of the scipy implementation of
@@ -390,6 +390,7 @@ def binRV(time, rv, err=None, stat='wmean', tstat='wmean', estat='addquad',
 
 
 def bin_ccf_mask(time, ccf_mask):
+    logger = setup_logger()
     indices = binRV(time, None, binning_indices=True)
     indices = np.r_[indices, time.size]
     bmask = []

{arvi-0.1.30 → arvi-0.2.11}/arvi/config.py

@@ -8,6 +8,7 @@ def get_config_path():
 
 def get_config():
     config = configparser.ConfigParser()
+    config.add_section('config')
     if (path := get_config_path()).exists():
         config.read(path)
     return config
@@ -31,6 +32,10 @@ class config:
         'check_internet': False,
         # make all DACE requests without using a .dacerc file
         'request_as_public': False,
+        # enable from arvi import star_name
+        'fancy_import': True,
+        # use the 'dark_background' matplotlib theme
+        'dark_plots': False,
         # debug
         'debug': False,
     }
@@ -43,10 +48,16 @@ class config:
             # return {'return_self': 'help!'}
             return {}
 
-        if self.__user_config.has_option('config', name):
-            self.__conf[name] = self.__user_config.get('config', name)
+        try:
+            if self.__user_config.has_option('config', name):
+                value = self.__user_config.get('config', name)
+                value = True if value == 'True' else value
+                value = False if value == 'False' else value
+                self.__conf[name] = value
 
-        return self.__conf[name]
+            return self.__conf[name]
+        except KeyError:
+            raise KeyError(f"unknown config option '{name}'")
 
     def __setattr__(self, name, value):
         if name in config.__setters:
@@ -54,7 +65,7 @@ class config:
         else:
             if 'config' not in self.__user_config:
                 self.__user_config.add_section('config')
-            self.__user_config.set('config', name, value)
+            self.__user_config.set('config', name, str(value))
             save_config(self.__user_config)
             # raise NameError(f"unknown configuration name '{name}'")
 

{arvi-0.1.30 → arvi-0.2.11}/arvi/dace_wrapper.py

@@ -2,20 +2,25 @@ import os
 import sys
 import tarfile
 import collections
-from functools import lru_cache
+from functools import lru_cache, partial
 from itertools import islice
 import numpy as np
-from dace_query import DaceClass
-from dace_query.spectroscopy import SpectroscopyClass, Spectroscopy as default_Spectroscopy
-from .setup_logger import logger
-from .utils import create_directory, all_logging_disabled, stdout_disabled, tqdm
 
+from .setup_logger import setup_logger
+from .utils import create_directory, all_logging_disabled, stdout_disabled, timer, tqdm
+
+
+def load_spectroscopy(user=None, verbose=True):
+    logger = setup_logger()
+    with all_logging_disabled():
+        from dace_query.spectroscopy import SpectroscopyClass, Spectroscopy as default_Spectroscopy
+        from dace_query import DaceClass
 
-def load_spectroscopy(user=None) -> SpectroscopyClass:
     from .config import config
     # requesting as public
     if config.request_as_public:
-        logger.warning('requesting DACE data as public')
+        if verbose:
+            logger.warning('requesting DACE data as public')
         with all_logging_disabled():
             dace = DaceClass(dace_rc_config_path='none')
         return SpectroscopyClass(dace_instance=dace)
@@ -39,21 +44,29 @@ def load_spectroscopy(user=None) -> SpectroscopyClass:
         logger.info(f'using credentials for user {user} in ~/.dacerc')
         return SpectroscopyClass(dace_instance=dace)
     # default
+    if not os.path.exists(os.path.expanduser('~/.dacerc')):
+        logger.warning('requesting DACE data as public (no .dacerc file found)')
     return default_Spectroscopy
 
-@lru_cache()
-def get_dace_id(star, verbose=True):
+
+@lru_cache(maxsize=1024)
+def get_dace_id(star, verbose=True, raise_error=False):
+    logger = setup_logger()
     filters = {"obj_id_catname": {"equal": [star]}}
     try:
         with all_logging_disabled():
             r = load_spectroscopy().query_database(filters=filters, limit=1)
-        return r['obj_id_daceid'][0]
+        return str(r['obj_id_daceid'][0])
     except KeyError:
         if verbose:
             logger.error(f"Could not find DACE ID for {star}")
+        if not raise_error:
+            return None
         raise ValueError from None
 
+
 def get_arrays(result, latest_pipeline=True, ESPRESSO_mode='HR11', NIRPS_mode='HE', verbose=True):
+    logger = setup_logger()
     arrays = []
     instruments = [str(i) for i in result.keys()]
 
@@ -62,7 +75,6 @@ def get_arrays(result, latest_pipeline=True, ESPRESSO_mode='HR11', NIRPS_mode='H
 
         # select ESPRESSO mode, which is defined at the level of the pipeline
         if 'ESPRESSO' in inst:
-
             find_mode = [ESPRESSO_mode in pipe for pipe in pipelines]
             # the mode was not found
             if not any(find_mode):
@@ -75,11 +87,16 @@ def get_arrays(result, latest_pipeline=True, ESPRESSO_mode='HR11', NIRPS_mode='H
                 i = [i for i, pipe in enumerate(pipelines) if ESPRESSO_mode in pipe][0]
                 pipelines = [pipelines[i]]
 
+        # select NIRPS mode
+        if 'NIRPS' in inst:
+            if any(this_mode := [p for p in pipelines if NIRPS_mode in p]):
+                pipelines = this_mode
+
         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]]
+                pipelines = [pipelines[0]]
             if 'HARPS' in inst and npipe > 1 and pipelines[1] == pipelines[0] + '-EGGS':
                 pipelines = pipelines[:2]
             else:
@@ -152,19 +169,20 @@ def get_observations_from_instrument(star, instrument, user=None, main_id=None,
         dict:
             dictionary with data from DACE
     """
-    Spectroscopy = load_spectroscopy(user)
+    Spectroscopy = load_spectroscopy(user, verbose)
     found_dace_id = False
-    try:
-        dace_id = get_dace_id(star, verbose=verbose)
-        found_dace_id = True
-    except ValueError as e:
-        if main_id is not None:
-            try:
-                dace_id = get_dace_id(main_id, verbose=verbose)
-                found_dace_id = True
-            except ValueError:
-                pass
-    
+    with timer('dace_id query'):
+        try:
+            dace_id = get_dace_id(star, verbose=verbose, raise_error=True)
+            found_dace_id = True
+        except ValueError as e:
+            if main_id is not None:
+                try:
+                    dace_id = get_dace_id(main_id, verbose=verbose, raise_error=True)
+                    found_dace_id = True
+                except ValueError:
+                    pass
+
     if not found_dace_id:
         try:
             with all_logging_disabled():
@@ -176,10 +194,16 @@ def get_observations_from_instrument(star, instrument, user=None, main_id=None,
             msg = f'no {instrument} observations for {star}'
             raise ValueError(msg) from None
 
-    filters = {
-        "ins_name": {"contains": [instrument]},
-        "obj_id_daceid": {"contains": [dace_id]}
-    }
+    if (isinstance(instrument, str)):
+        filters = {
+            "ins_name": {"contains": [instrument]},
+            "obj_id_daceid": {"contains": [dace_id]}
+        }
+    elif (isinstance(instrument, (list, tuple, np.ndarray))):
+        filters = {
+            "ins_name": {"contains": instrument},
+            "obj_id_daceid": {"contains": [dace_id]}
+        }
     with all_logging_disabled():
         result = Spectroscopy.query_database(filters=filters)
     
@@ -190,17 +214,22 @@ def get_observations_from_instrument(star, instrument, user=None, main_id=None,
 
     for inst in np.unique(result['ins_name']):
         mask1 = result['ins_name'] == inst
-        r[inst] = {}
+        r[str(inst)] = {}
 
-        for pipe in np.unique(result['ins_drs_version'][mask1]):
-            mask2 = mask1 & (result['ins_drs_version'] == pipe)
-            r[inst][pipe] = {}
+        key2 = 'ins_drs_version'
+        n_key2 = len(np.unique(result[key2][mask1]))
+        if len(np.unique(result['pub_bibcode'][mask1])) >= n_key2:
+            key2 = 'pub_bibcode'
+
+        for pipe in np.unique(result[key2][mask1]):
+            mask2 = mask1 & (result[key2] == pipe)
+            r[str(inst)][str(pipe)] = {}
 
             for ins_mode in np.unique(result['ins_mode'][mask2]):
                 mask3 = mask2 & (result['ins_mode'] == ins_mode)
                 _nan = np.full(mask3.sum(), np.nan)
 
-                r[inst][pipe][ins_mode] = {
+                r[str(inst)][str(pipe)][str(ins_mode)] = {
                     'texp': result['texp'][mask3],
                     'bispan': result['spectro_ccf_bispan'][mask3],
                     'bispan_err': result['spectro_ccf_bispan_err'][mask3],
@@ -212,7 +241,9 @@ def get_observations_from_instrument(star, instrument, user=None, main_id=None,
                     'rv': result['spectro_ccf_rv'][mask3],
                     'rv_err': result['spectro_ccf_rv_err'][mask3],
                     'berv': result['spectro_cal_berv'][mask3],
-                    'ccf_noise': _nan,
+                    'ccf_noise': np.sqrt(
+                        np.square(result['spectro_ccf_rv_err'][mask3]) - np.square(result['spectro_cal_drift_noise'][mask3])
+                    ),
                     'rhk': result['spectro_analysis_rhk'][mask3],
                     'rhk_err': result['spectro_analysis_rhk_err'][mask3],
                     'contrast': result['spectro_ccf_contrast'][mask3],
@@ -254,8 +285,10 @@ def get_observations_from_instrument(star, instrument, user=None, main_id=None,
     return r
 
 def get_observations(star, instrument=None, user=None, main_id=None, verbose=True):
+    logger = setup_logger()
     if instrument is None:
-        Spectroscopy = load_spectroscopy(user)
+        Spectroscopy = load_spectroscopy(user, verbose)
+
         try:
             with stdout_disabled(), all_logging_disabled():
                 result = Spectroscopy.get_timeseries(target=star,
@@ -288,12 +321,14 @@ def get_observations(star, instrument=None, user=None, main_id=None, verbose=Tru
             result[inst] = dict(result[inst])
     #
 
-    instruments = list(result.keys())
+    instruments = list(map(str, result.keys()))
 
     if instrument is not None:
         # select only the provided instrument (if it's there)
-        instruments = [inst for inst in instruments if instrument in inst]
-
+        if (isinstance(instrument, str)):
+            instruments = [inst for inst in instruments if instrument in inst]
+        elif (isinstance(instrument, list)):
+            instruments = [inst for inst in instruments if any(i in inst for i in instrument)]
     if len(instruments) == 0:
         if instrument is None:
             msg = f'no observations for {star}'
@@ -317,38 +352,57 @@ def get_observations(star, instrument=None, user=None, main_id=None, verbose=Tru
     #     else:
     #         return -1
 
-    # sort pipelines, must be extra careful with HARPS/HARPN pipeline version numbers
-    # got here with the help of DeepSeek
-    from re import match
-    def custom_sort_key(s):
-        s = s[0]
-        print(s)
-        # Check for version number pattern (e.g., 3.2.5 or 3.2.5-EGGS)
-        version_match = match(r'^(\d+(?:\.\d+)*)(?:-(.*))?$', s)
-        if version_match:
-            version_parts = tuple(map(int, version_match.group(1).split('.')))
-            suffix = version_match.group(2)
-            
-            if suffix is not None:
-                # Suffixed versions: sort in ascending order (3.2.5-HR11 < 3.3.1-HR11)
-                return (0, 0, version_parts, suffix)
-            else:
-                # Unsuffixed versions: sort in descending order (3.5 > 3.2.5)
-                return (0, 1, tuple(-x for x in version_parts))
-        
-        # Check for scientific reference pattern (e.g., 2004A&A...)
-        year_match = match(r'^(\d{4})', s)
-        if year_match:
-            year = int(year_match.group(1))
-            return (1, year)
-        
-        # For all other strings, sort alphabetically
-        return (2, s)
-
-    # from functools import cmp_to_key
+    # # sort pipelines, must be extra careful with HARPS/HARPN pipeline version numbers
+    # # got here with the help of DeepSeek
+    # # from functools import cmp_to_key
+    # from re import match
+    # def custom_sort_key(s):
+    #     s = s[0]
+    #     # Check for version number pattern (e.g., 3.2.5 or 3.2.5-EGGS)
+    #     version_match = match(r'^(\d+(?:\.\d+)*)(?:[-\s](.*))?$', s)
+    #     if version_match:
+    #         version_parts = list(map(int, version_match.group(1).split('.')))
+    #         if len(version_parts) == 2:
+    #             version_parts.insert(1, -1)
+    #         # if version_match.group(2) and 'LBL' in version_match.group(2):
+    #         #     version_parts.append(-1)
+    #         # else:
+    #         #     version_parts.append(0)
+    #         if version_match.group(2) is None:
+    #             version_parts.append('')
+    #         else:
+    #             version_parts.append(version_match.group(2))
+    #         return (0, 1, version_parts)
+    #     # Check for scientific reference pattern (e.g., 2004A&A...)
+    #     year_match = match(r'^(\d{4})', s)
+    #     if year_match:
+    #         year = int(year_match.group(1))
+    #         return (1, year)
+    #     # For all other strings, sort alphabetically
+    #     return (2, s)
+
+    def custom_key(val, strip_EGGS=False):
+        if strip_EGGS:
+            val = val.replace('-EGGS', '').replace(' EGGS', '')
+        key = 0
+        key -= 1 if '3.5' in val else 0
+        key -= 1 if 'EGGS' in val else 0
+        key -= 1 if ('UHR' in val or 'MR' in val) else 0
+        key -= 1 if 'LBL' in val else 0
+        return str(key) if key != 0 else val
+
     new_result = {}
     for inst in instruments:
-        new_result[inst] = dict(sorted(result[inst].items(), key=custom_sort_key, reverse=True))
+        # new_result[inst] = dict(
+        #     sorted(result[inst].items(), key=custom_sort_key, reverse=True)
+        # )
+        if all(['EGGS' in k for k in result[inst].keys()]):
+            custom_key = partial(custom_key, strip_EGGS=True)
+        # WARNING: not the same as reverse=True (not sure why)
+        sorted_keys = sorted(result[inst].keys(), key=custom_key)[::-1]
+        new_result[inst] = {}
+        for key in sorted_keys:
+            new_result[inst][key] = result[inst][key]
 
     if verbose:
         logger.info('RVs available from')
@@ -382,7 +436,7 @@ def check_existing(output_directory, files, type):
     ]
 
     if type == 'S2D':
-        existing = [
+        existing += [
             f.partition('.fits')[0] for f in os.listdir(output_directory)
             if 'e2ds' in f
         ]   
@@ -446,6 +500,7 @@ def extract_fits(output_directory, filename=None):
 
 
 def do_symlink_filetype(type, raw_files, output_directory, clobber=False, top_level=None, verbose=True):
+    logger = setup_logger()
     terminations = {
         'CCF': '_CCF_A.fits',
         'S1D': '_S1D_A.fits',
@@ -491,21 +546,23 @@ def do_symlink_filetype(type, raw_files, output_directory, clobber=False, top_le
 def do_download_filetype(type, raw_files, output_directory, clobber=False, user=None,
                          verbose=True, chunk_size=20, parallel_limit=30):
     """ 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)
 
     # check existing files to avoid re-downloading
     if not clobber:
         raw_files = check_existing(output_directory, raw_files, type)
-
+    
     n = raw_files.size
 
     # any file left to download?
     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:
@@ -523,7 +580,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)
 
@@ -552,6 +611,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.11/arvi/exofop_wrapper.py

@@ -0,0 +1,62 @@
+import csv
+import requests
+import time
+import importlib.resources as resources
+import numpy as np
+
+from .setup_logger import setup_logger
+
+def get_toi_list(verbose=True):
+    logger = setup_logger()
+    toi_list = resources.files('arvi') / 'data' / 'exofop_toi_list.csv'
+    now = time.time()
+    download = not toi_list.exists() or toi_list.stat().st_mtime < now - 48 * 60 * 60
+    if download:
+        if verbose:
+            logger.info('Downloading exofop TOI list (can take a while)...')
+        r = requests.get('https://exofop.ipac.caltech.edu/tess/download_toi.php?sort=toi&output=csv')
+        with open(toi_list, 'wb') as f:
+            f.write(r.content)
+    return toi_list
+
+class exofop:
+    def __init__(self, star: str, verbose=True, _debug=False):
+        self.star = star
+        self.verbose = verbose
+        
+        toi_list = get_toi_list(verbose=verbose)
+        tsv = ('|'.join(i) for i in csv.reader(open(toi_list, encoding='utf8')))
+        self.data = np.genfromtxt(tsv, delimiter='|', 
+                                  names=True, encoding='utf8', dtype=None)
+
+
+        try:
+            if self.star.startswith('TIC'):
+                self.tic = self.star
+                w = self.data['TIC_ID'] == int(self.star[3:])
+                self.toi = 'TOI-' + str(int(self.data['TOI'][w][0]))
+            else:
+                toi = self.star.replace('TOI-', '')
+                toi = toi if toi.endswith('.01') else toi + '.01'
+                toi_float = float(toi)
+                if toi_float not in self.data['TOI']:
+                    raise ValueError
+                w = self.data['TOI'] == toi_float
+                self.tic = 'TIC' + str(int(self.data['TIC_ID'][w][0]))
+                self.toi = self.star
+        except ValueError:
+            raise ValueError(f'{self.star} not found in exofop TOI list')
+        else:
+            self.ra = str(self.data['RA'][w][0])
+            self.dec = str(self.data['Dec'][w][0])
+
+            self.epoch = float(self.data['Epoch_BJD'][w][0])
+            self.period = float(self.data['Period_days'][w][0])
+            if self.period == 0.0:
+                self.period = np.nan
+            self.duration = float(self.data['Duration_hours'][w][0])
+            self.depth = float(self.data['Depth_ppm'][w][0])
+
+        
+    def __repr__(self):
+        return f'{self.star} (TIC={self.tic}, epoch={self.epoch:.3f}, period={self.period:.3f})'

{arvi-0.1.30 → arvi-0.2.11}/arvi/extra_data.py

@@ -3,8 +3,7 @@ from glob import glob
 import json
 
 from numpy import full
-from .setup_logger import logger
-from . import timeseries
+from .setup_logger import setup_logger
 
 refs = {
     'HD86226': 'Teske et al. 2020 (AJ, 160, 2)'
@@ -12,16 +11,21 @@ refs = {
 
 def get_extra_data(star, instrument=None, path=None, verbose=True,
                    check_for_kms=True):
+    from . import timeseries
+    logger = setup_logger()
     if path is None:
         path = os.path.dirname(__file__)
         path = os.path.join(path, 'data', 'extra')
+        metadata = json.load(open(os.path.join(path, 'metadata.json'), 'r'))
+        # print(metadata)
+    else:
+        metadata = {}
 
-    metadata = json.load(open(os.path.join(path, 'metadata.json'), 'r'))
-    # print(metadata)
-
-    files = glob(os.path.join(path, star.replace(' ', '') + '*.rdb'))
+    files = glob(os.path.join(path, star + '*.rdb'))
+    files += glob(os.path.join(path, star.replace(' ', '') + '*.rdb'))
     files = [f for f in files if os.path.isfile(f)]
-    files = [f for f in files if not os.path.basename(f).endswith('.zip')]
+    files = [f for f in files if not f.endswith('_actin.rdb')]
+    files = list(set(files))
 
     if len(files) == 0:
         raise FileNotFoundError