arvi 0.2.11__py3-none-any.whl → 0.3.2__py3-none-any.whl

arvi/__init__.py

@@ -17,11 +17,15 @@ def __getattr__(name: str):
     if not config.fancy_import:
         raise AttributeError
     
+    if name.startswith('__'):
+        return
     if name in (
         '_ipython_canary_method_should_not_exist_',
         '_ipython_display_',
         '_repr_mimebundle_',
-        '__wrapped__'
+        # '__custom_documentations__',
+        # '__wrapped__',
+        # '__dataframe__'
     ):
         return
 

arvi/config.py

@@ -36,6 +36,8 @@ class config:
         'fancy_import': True,
         # use the 'dark_background' matplotlib theme
         'dark_plots': False,
+        # specific colors per instrument
+        'colors_per_instrument': False,
         # debug
         'debug': False,
     }
@@ -43,7 +45,7 @@ class config:
 
     __user_config = get_config()
 
-    def __getattr__(self, name):
+    def __getattr__(self, name: str):
         if name in ('__custom_documentations__', ):
             # return {'return_self': 'help!'}
             return {}
@@ -54,7 +56,6 @@ class config:
                 value = True if value == 'True' else value
                 value = False if value == 'False' else value
                 self.__conf[name] = value
-
             return self.__conf[name]
         except KeyError:
             raise KeyError(f"unknown config option '{name}'")

arvi/dace_wrapper.py

@@ -229,61 +229,95 @@ def get_observations_from_instrument(star, instrument, user=None, main_id=None,
                 mask3 = mask2 & (result['ins_mode'] == ins_mode)
                 _nan = np.full(mask3.sum(), np.nan)
 
-                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],
-                    'drift_noise': result['spectro_cal_drift_noise'][mask3],
-                    'rjd': result['obj_date_bjd'][mask3],
-                    'cal_therror': _nan,
-                    'fwhm': result['spectro_ccf_fwhm'][mask3],
-                    'fwhm_err': result['spectro_ccf_fwhm_err'][mask3],
-                    'rv': result['spectro_ccf_rv'][mask3],
-                    'rv_err': result['spectro_ccf_rv_err'][mask3],
-                    'berv': result['spectro_cal_berv'][mask3],
-                    '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],
-                    'contrast_err': result['spectro_ccf_contrast_err'][mask3],
-                    'cal_thfile': result['spectro_cal_thfile'][mask3],
-                    'spectroFluxSn50': result['spectro_flux_sn50'][mask3],
-                    'protm08': result['spectro_analysis_protm08'][mask3],
-                    'protm08_err': result['spectro_analysis_protm08_err'][mask3],
-                    'caindex': result['spectro_analysis_ca'][mask3],
-                    'caindex_err': result['spectro_analysis_ca_err'][mask3],
-                    'pub_reference': result['pub_ref'][mask3],
-                    'drs_qc': result['spectro_drs_qc'][mask3],
-                    'haindex': result['spectro_analysis_halpha'][mask3],
-                    'haindex_err': result['spectro_analysis_halpha_err'][mask3],
-                    'protn84': result['spectro_analysis_protn84'][mask3],
-                    'protn84_err': result['spectro_analysis_protn84_err'][mask3],
-                    'naindex': result['spectro_analysis_na'][mask3],
-                    'naindex_err': result['spectro_analysis_na_err'][mask3],
-                    'snca2': _nan,
-                    'mask': result['spectro_ccf_mask'][mask3],
-                    'public': result['public'][mask3],
-                    'spectroFluxSn20': result['spectro_flux_sn20'][mask3],
-                    'sindex': result['spectro_analysis_smw'][mask3],
-                    'sindex_err': result['spectro_analysis_smw_err'][mask3],
-                    'drift_used': _nan,
-                    'ccf_asym': result['spectro_ccf_asym'][mask3],
-                    'ccf_asym_err': result['spectro_ccf_asym_err'][mask3],
-                    'date_night': result['date_night'][mask3],
-                    'raw_file': result['file_rootpath'][mask3],
-                    'prog_id': result['prog_id'][mask3],
-                    'th_ar': result['th_ar'][mask3],
-                    'th_ar1': result['th_ar1'][mask3],
-                    'th_ar2': result['th_ar2'][mask3],
+                translations = {
+                    'obj_date_bjd': 'rjd',
+                    'spectro_drs_qc': 'drs_qc',
+                    'spectro_cal_berv_mx': 'bervmax',
+                    'pub_ref': 'pub_reference',
+                    'file_rootpath': 'raw_file',
+                    'spectro_ccf_asym': 'ccf_asym',
+                    'spectro_ccf_asym_err': 'ccf_asym_err',
                 }
+                new_result = {}
+                for key in result.keys():
+                    if key in translations:
+                        new_key = translations[key]
+                    else:
+                        new_key = key
+                        new_key = new_key.replace('spectro_ccf_', '')
+                        new_key = new_key.replace('spectro_cal_', '')
+                        new_key = new_key.replace('spectro_analysis_', '')
+                    new_result[new_key] = result[key][mask3]
+
+                new_result['ccf_noise'] = np.sqrt(
+                    np.square(result['spectro_ccf_rv_err'][mask3]) - np.square(result['spectro_cal_drift_noise'][mask3])
+                )
+
+                r[str(inst)][str(pipe)][str(ins_mode)] = new_result
+
+                # 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],
+                #     'drift_noise': result['spectro_cal_drift_noise'][mask3],
+                #     'rjd': result['obj_date_bjd'][mask3],
+                #     'cal_therror': _nan,
+                #     'fwhm': result['spectro_ccf_fwhm'][mask3],
+                #     'fwhm_err': result['spectro_ccf_fwhm_err'][mask3],
+                #     'rv': result['spectro_ccf_rv'][mask3],
+                #     'rv_err': result['spectro_ccf_rv_err'][mask3],
+                #     'berv': result['spectro_cal_berv'][mask3],
+                #     '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],
+                #     'contrast_err': result['spectro_ccf_contrast_err'][mask3],
+                #     'cal_thfile': result['spectro_cal_thfile'][mask3],
+                #     'spectroFluxSn50': result['spectro_flux_sn50'][mask3],
+                #     'protm08': result['spectro_analysis_protm08'][mask3],
+                #     'protm08_err': result['spectro_analysis_protm08_err'][mask3],
+                #     'caindex': result['spectro_analysis_ca'][mask3],
+                #     'caindex_err': result['spectro_analysis_ca_err'][mask3],
+                #     'pub_reference': result['pub_ref'][mask3],
+                #     'drs_qc': result['spectro_drs_qc'][mask3],
+                #     'haindex': result['spectro_analysis_halpha'][mask3],
+                #     'haindex_err': result['spectro_analysis_halpha_err'][mask3],
+                #     'protn84': result['spectro_analysis_protn84'][mask3],
+                #     'protn84_err': result['spectro_analysis_protn84_err'][mask3],
+                #     'naindex': result['spectro_analysis_na'][mask3],
+                #     'naindex_err': result['spectro_analysis_na_err'][mask3],
+                #     'snca2': _nan,
+                #     'mask': result['spectro_ccf_mask'][mask3],
+                #     'public': result['public'][mask3],
+                #     'spectroFluxSn20': result['spectro_flux_sn20'][mask3],
+                #     'sindex': result['spectro_analysis_smw'][mask3],
+                #     'sindex_err': result['spectro_analysis_smw_err'][mask3],
+                #     'drift_used': _nan,
+                #     'ccf_asym': result['spectro_ccf_asym'][mask3],
+                #     'ccf_asym_err': result['spectro_ccf_asym_err'][mask3],
+                #     'date_night': result['date_night'][mask3],
+                #     'raw_file': result['file_rootpath'][mask3],
+                #     'prog_id': result['prog_id'][mask3],
+                #     'th_ar': result['th_ar'][mask3],
+                #     'th_ar1': result['th_ar1'][mask3],
+                #     'th_ar2': result['th_ar2'][mask3],
+                # }
     
     # print(r.keys())    
     # print([r[k].keys() for k in r.keys()])
     # print([r[k1][k2].keys() for k1 in r.keys() for k2 in r[k1].keys()])
     return r
 
+def _warn_harpsn(instrument):
+    if 'HARPSN' in instrument or 'HARPS-N' in instrument:
+        logger = setup_logger()
+        logger.warning(f'Did you mean "HARPN" instead of "{instrument}"?')
+        return True
+    return False
+
+
 def get_observations(star, instrument=None, user=None, main_id=None, verbose=True):
     logger = setup_logger()
     if instrument is None:
@@ -305,6 +339,7 @@ def get_observations(star, instrument=None, user=None, main_id=None, verbose=Tru
             result = get_observations_from_instrument(star, instrument, user, main_id, verbose)
         except ValueError:
             msg = f'no {instrument} observations for {star}'
+            _warn_harpsn(instrument)
             raise ValueError(msg) from None
 
     # defaultdict --> dict
@@ -384,6 +419,8 @@ def get_observations(star, instrument=None, user=None, main_id=None, verbose=Tru
     def custom_key(val, strip_EGGS=False):
         if strip_EGGS:
             val = val.replace('-EGGS', '').replace(' EGGS', '')
+        if val in ('3.3', '3.4', '3.5', '3.6', '3.7', '3.8'):
+            val = '0.' + val
         key = 0
         key -= 1 if '3.5' in val else 0
         key -= 1 if 'EGGS' in val else 0
@@ -408,9 +445,10 @@ def get_observations(star, instrument=None, user=None, main_id=None, verbose=Tru
         logger.info('RVs available from')
         with logger.contextualize(indent='  '):
             _inst = ''
-            for inst in instruments:
+            max_len_inst = max([len(inst) for inst in instruments])
+            for inst in sorted(instruments):
                 pipelines = list(new_result[inst].keys())
-                max_len = max([len(pipe) for pipe in pipelines])
+                max_len_pipe = max([len(pipe) for pipe in pipelines])
                 for pipe in pipelines:
                     last_pipe = pipe == pipelines[-1]
                     modes = list(new_result[inst][pipe].keys())
@@ -418,11 +456,11 @@ def get_observations(star, instrument=None, user=None, main_id=None, verbose=Tru
                         N = len(new_result[inst][pipe][mode]['rjd'])
                         # LOG
                         if inst == _inst and last_pipe:
-                            logger.info(f'{" ":>12s} └ {pipe:{max_len}s} - {mode} ({N} observations)')
+                            logger.info(f'{" ":>{max_len_inst}s} └ {pipe:{max_len_pipe}s} - {mode} ({N} observations)')
                         elif inst == _inst:
-                            logger.info(f'{" ":>12s} ├ {pipe:{max_len}s} - {mode} ({N} observations)')
+                            logger.info(f'{" ":>{max_len_inst}s} ├ {pipe:{max_len_pipe}s} - {mode} ({N} observations)')
                         else:
-                            logger.info(f'{inst:>12s} ├ {pipe:{max_len}s} - {mode} ({N} observations)')
+                            logger.info(f'{inst:>{max_len_inst}s} ├ {pipe:{max_len_pipe}s} - {mode} ({N} observations)')
                             _inst = inst
 
     return new_result

arvi/gaia_wrapper.py

@@ -9,7 +9,8 @@ DATA_PATH = os.path.join(DATA_PATH, 'data')
 QUERY = """
 SELECT TOP 20 gaia_source.designation, gaia_source.source_id,
               gaia_source.ra, gaia_source.dec, 
-              gaia_source.parallax, gaia_source.pmra, gaia_source.pmdec,
+              gaia_source.parallax, gaia_source.parallax_error,
+              gaia_source.pmra, gaia_source.pmdec,
               gaia_source.ruwe, gaia_source.phot_g_mean_mag, gaia_source.bp_rp, 
               gaia_source.radial_velocity, gaia_source.radial_velocity_error
 FROM gaiadr3.gaia_source 
@@ -27,7 +28,8 @@ CONTAINS(
 QUERY_ID = """
 SELECT TOP 20 gaia_source.designation, gaia_source.source_id,
               gaia_source.ra, gaia_source.dec, 
-              gaia_source.parallax, gaia_source.pmra, gaia_source.pmdec,
+              gaia_source.parallax, gaia_source.parallax_error,
+              gaia_source.pmra, gaia_source.pmdec,
               gaia_source.ruwe, gaia_source.phot_g_mean_mag, gaia_source.bp_rp, 
               gaia_source.radial_velocity, gaia_source.radial_velocity_error
 FROM gaiadr3.gaia_source 
@@ -120,6 +122,7 @@ class gaia:
         self.pmdec = float(results['pmdec'])
         self.coords = SkyCoord(self.ra, self.dec, unit='deg')
         self.plx = float(results['parallax'])
+        self.plx_err = float(results['parallax_error'])
         try:
             self.radial_velocity = float(results['radial_velocity'])
         except ValueError:
@@ -131,5 +134,12 @@ class gaia:
 
         return
 
+    def distance(self):
+        """ Calculate the distance to the star as 1 / parallax [pc] """
+        from astropy import units as u
+        d = (self.plx * u.mas).to(u.parsec, 
+                                  equivalencies=u.equivalencies.parallax())
+        return d
+
     def __repr__(self):
         return f'{self.star} (DR3 id={self.dr3_id})'

arvi/kima_wrapper.py

@@ -1,19 +1,8 @@
 import os
-import numpy as np
+import ast
 
 from .setup_logger import setup_logger
 
-try:
-    import kima
-    from kima.pykima.utils import chdir
-    from kima import distributions
-    from kima import RVData, HGPMdata
-    from kima import RVmodel, GPmodel, RVHGPMmodel
-    kima_available = True
-except ImportError:
-    kima_available = False
-
-
 def try_to_guess_prior(model, prior):
     if 'jitter' in prior:
         return 'Jprior'
@@ -23,8 +12,14 @@ def try_to_guess_prior(model, prior):
 
 
 def run_kima(self, run=False, load=False, run_directory=None,
-             model=RVmodel, priors={}, **kwargs):
-    if not kima_available:
+             model='RVmodel', priors={}, **kwargs):
+    try:
+        import kima
+        from kima.pykima.utils import chdir
+        from kima import distributions
+        from kima import RVData, HGPMdata
+        from kima import RVmodel, GPmodel, RVHGPMmodel
+    except ImportError:
         raise ImportError('kima not available, please install with `pip install kima`')
     
     logger = setup_logger()
@@ -57,7 +52,7 @@ def run_kima(self, run=False, load=False, run_directory=None,
     model.trend = kwargs.pop('trend', False)
     model.degree = kwargs.pop('degree', 0)
 
-    if isinstance(model, RVmodel):
+    if isinstance(model, (RVmodel, RVHGPMmodel)):
         model.studentt = kwargs.pop('studentt', False)
 
     if isinstance(model, GPmodel):
@@ -81,6 +76,16 @@ def run_kima(self, run=False, load=False, run_directory=None,
         model.pm_dec_bary_prior = priors.pop('pm_dec_bary_prior',
                                              distributions.Gaussian(pm_data.pm_dec_hg, pm_data.sig_hg_dec))
 
+    KO = kwargs.pop('known_object', False)
+    if KO:
+        if isinstance(KO, int) and KO is not True:
+            model.set_known_object(KO)
+        else:
+            try:
+                model.set_known_object(kwargs.pop('n_known_object'))
+            except KeyError:
+                msg = 'if `known_object` is True, specify `n_known_object` or pass `known_object` as an integer'
+                raise ValueError(msg) from None
 
     for k, v in priors.items():
         try:
@@ -99,21 +104,27 @@ def run_kima(self, run=False, load=False, run_directory=None,
     if run_directory is None:
         run_directory = os.getcwd()
 
-    if run:
-        # TODO: use signature of kima.run to pop the correct kwargs
-        # model_name = model.__class__.__name__
-        # model_name = f'kima.{model_name}.{model_name}'
-        # signature, defaults = [sig for sig in kima.run.__nb_signature__ if model_name in sig[0]]
+    diagnostic = kwargs.pop('diagnostic', False)
 
+    if run:
+        model_name = model.__class__.__name__
+        model_name = f'kima.{model_name}.{model_name}'
+        signature, _, defaults = [sig for sig in kima._run_really.__nb_signature__ if model_name in sig[0]][0]
+        signature = signature.replace('\\', '')
+        args = ast.parse(signature + ':\n pass').body[0].args
+        defaults = {arg.arg: d for arg, d in zip(args.args[1:], defaults)}
+        defaults.update(kwargs)
         with chdir(run_directory):
-            kima.run(model, **kwargs)
+            kima.run(model, **defaults)
 
     if isinstance(model, RVHGPMmodel):
         data = (data, pm_data)
 
     if load:
         with chdir(run_directory):
-            res = kima.load_results(model)
+            res = kima.load_results(model, diagnostic=diagnostic)
+
+        res.star = self.star
         return data, model, res
 
     return data, model

arvi/plots.py

@@ -587,6 +587,32 @@ def plot_xy(self, x, y, ax=None, instrument=None, show_legend=True, **kwargs):
         return fig, ax
 
 
+def plot_several(obj, axs=None, **kwargs):
+    try:
+        from grid_strategy.strategies import SquareStrategy
+    except ImportError:
+        raise ImportError("Please install 'grid-strategy' to use this function")
+    if axs is None:
+        if isinstance(obj, (list, tuple)):
+            n = len(obj)
+            GS = SquareStrategy().get_grid(n)
+            fig = plt.gcf()
+            fig.set_size_inches(8, 6)
+            axs = [fig.add_subplot(gs) for gs in GS]
+        else:
+            raise NotImplementedError('`obj` should be a list or tuple')
+    else:
+        assert isinstance(axs, (list, tuple))
+        assert len(axs) == len(obj), 'Number of axes does not match number of objects'
+
+    kwargs.setdefault('show_legend', False)
+    kwargs.setdefault('show_title', True)
+
+    for ax, s in zip(axs, obj):
+        s.plot(ax=ax, **kwargs)
+        ax.set(xlabel='', ylabel='')
+
+
 # @plot_fast
 def gls(self, ax=None, label=None, instrument=None, 
         fap=True, fap_method='baluev', adjust_means=config.adjust_means_gls,

arvi/programs.py

@@ -23,6 +23,7 @@ def get_star(star, instrument=None, verbose=False, **kwargs):
 
 
 class LazyRV:
+    """ A lazy wrapper around `RV` """
     def __init__(self, stars: list, instrument: str = None,
                  _parallel_limit=10, _parallel_workers=8):
         self.stars = stars
@@ -60,25 +61,6 @@ class LazyRV:
 
         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
@@ -92,12 +74,24 @@ class LazyRV:
         return self._saved
 
     @lru_cache(maxsize=10)
-    def __getitem__(self, index):
-        star = self.stars[index]
+    def __getitem__(self, index_or_name):
+        if isinstance(index_or_name, str):
+            index = self.stars.index(index_or_name)
+        else:
+            index = index_or_name
+
         if self._saved is not None:
             return self._saved[index]
+
+        star = self.stars[index]
         return get_star(star, self.instrument, verbose=True)
 
+    def plot(self, split=20, **kwargs):
+        from .plots import plot_several
+        from itertools import batched
+        S = list(batched(self.__call__(), split))
+        for s in S:
+            plot_several(s, **kwargs)
 
 # sorted by spectral type
 WG1_stars = [

arvi/setup_logger.py

@@ -2,17 +2,17 @@ import sys
 
 def setup_logger():
     from loguru import logger
-    try:
-        import marimo as mo
-        if mo.running_in_notebook():
-            raise NotImplementedError
-    except (NotImplementedError, AttributeError):
-        pass
-    except (ImportError, ModuleNotFoundError):
-        logger.remove()
-    else:
-        logger.remove()
-
+    # try:
+    #     import marimo as mo
+    #     if mo.running_in_notebook():
+    #         return logger
+    #         raise NotImplementedError
+    # except (NotImplementedError, AttributeError):
+    #     pass
+    # except (ImportError, ModuleNotFoundError):
+    #     logger.remove()
+    # else:
+    logger.remove()
     logger.configure(extra={"indent": ""})
     logger.add(
         sys.stdout,

arvi/simbad_wrapper.py

@@ -1,4 +1,5 @@
 import os
+from urllib.error import URLError
 import requests
 from dataclasses import dataclass
 from functools import partial
@@ -64,6 +65,12 @@ JOIN ident ON oidref = oid
 WHERE id = '{star}';
 """
 
+OTYPE_QUERY = """
+SELECT otype FROM basic 
+JOIN ident ON oidref = oid 
+WHERE id = '{star}';
+"""
+
 HD_GJ_HIP_QUERY = """
 SELECT id2.id
 FROM ident AS id1 JOIN ident AS id2 USING(oidref)
@@ -204,11 +211,19 @@ class simbad:
         try:
             table2 = run_query(query=BV_QUERY.format(star=self.star))
             if _debug:
-                print('table2:', table2)
+                print('table2:\n', table2)
             cols, values = parse_table1(table2, cols, values)
         except IndexError:
             self.B = self.V = np.nan
 
+        try:
+            table21 = run_query(query=OTYPE_QUERY.format(star=self.star))
+            if _debug:
+                print('table21:\n', table21)
+            cols, values = parse_table1(table21, cols, values)
+        except IndexError:
+            self.otype = None
+
         try:
             table3 = run_query(query=IDS_QUERY.format(star=self.star))
             if _debug:
@@ -228,6 +243,14 @@ class simbad:
                 setattr(self, '_' + filter_name, ufloat(float(mag), float(mag_err)))
             except ValueError:
                 setattr(self, '_' + filter_name, float(mag))
+            
+            # substitute missing V magnitude
+            if filter_name == 'V' and values[cols.index('V')] == '':
+                values[cols.index('V')] = mag
+            # substitute missing B magnitude
+            if filter_name == 'B' and values[cols.index('B')] == '':
+                values[cols.index('B')] = mag
+
 
         # measurements table
         table5 = run_query(query=MEAS_QUERY.format(star=self.star))
@@ -258,6 +281,8 @@ class simbad:
             try:
                 setattr(self, col, float(val))
             except ValueError:
+                if val == '':
+                    val = None
                 setattr(self, col, val)
 
         self.coords = SkyCoord(self.ra, self.dec, unit='deg')
@@ -281,12 +306,17 @@ class simbad:
             else:
                 self.teff = data['teff']
                 self.sweetcat = data
-
+        except URLError:
+            pass
         except IndexError:
-            if self.sp_type == '':
+            if self.sp_type in (None, ''):
                 if len(self.measurements.teff) > 0:
-                    self.teff = int(np.mean(self.measurements.teff))
-                    self.sp_type = teff_to_sptype(self.teff)
+                    try:
+                        self.teff = int(np.mean(self.measurements.teff))
+                        self.sp_type = teff_to_sptype(self.teff)
+                    except ValueError:
+                        self.teff = np.nan
+                        self.sp_type = ''
             elif self.sp_type[:2] in EFFECTIVE_TEMPERATURES:
                 self.teff = EFFECTIVE_TEMPERATURES[self.sp_type[:2]]
 

arvi/sophie_wrapper.py

@@ -13,7 +13,7 @@ URL_HEADER = "http://atlas.obs-hp.fr/sophie/sophie.cgi?n=sophiecc&c=i&z=fd&a=t&o
 def extract_keyword(keyword, text, raise_error=True):
     for line in text.splitlines():
         if keyword in line:
-            value = re.findall(fr'{keyword}\s+([\'\w\d.]+)', line)[0]
+            value = re.findall(fr'{keyword}\s+([-\'\w\d.]+)', line)[0]
             value = value.replace("'", "")
             try:
                 return float(value)

arvi/stats.py

@@ -1,6 +1,6 @@
 from functools import partial
 import numpy as np
-from scipy.stats import norm
+
 
 def wmean(a, e):
     """Weighted mean of array `a`, with uncertainties given by `e`.
@@ -69,7 +69,7 @@ weighted_median = partial(weighted_quantiles_interpolate, quantiles=0.5)
 
 
 
-def sigmaclip_median(a, low=4.0, high=4.0, k=1/norm.ppf(3/4)):
+def sigmaclip_median(a, low=4.0, high=4.0, k=None):
     """
     Same as scipy.stats.sigmaclip but using the median and median absolute
     deviation instead of the mean and standard deviation.
@@ -91,8 +91,12 @@ def sigmaclip_median(a, low=4.0, high=4.0, k=1/norm.ppf(3/4)):
             - `lower`: Lower clipping limit
             - `upper`: Upper clipping limit
     """
-    from scipy.stats import median_abs_deviation
+    from scipy.stats import median_abs_deviation, norm
     from scipy.stats._stats_py import SigmaclipResult
+
+    if k is None:
+        k = 1 / norm.ppf(3 / 4)
+
     c = np.asarray(a).ravel()
     delta = 1
     while delta:

arvi/timeseries.py

@@ -145,6 +145,7 @@ class RV(ISSUES, REPORTS):
         if self.verbose:
             logger.info('querying Simbad...')
 
+        # TODO: removing the 'A' might not be a good idea 
         # complicated way to query Simbad with self.__star__ or, if that
         # fails, try after removing a trailing 'A'
         for target in set([self.__star__, self.__star__.replace('A', '')]):
@@ -590,9 +591,12 @@ class RV(ISSUES, REPORTS):
         """ Total time span of the (masked) observations """
         return np.ptp(self.mtime)
 
-    def _index_from_instrument_index(self, index, instrument):
+    def _index_from_instrument_index(self, index, instrument, masked=True):
         ind = np.where(self.instrument_array == instrument)[0]
-        return ind[getattr(self, instrument).mask][index]
+        if masked:
+            return ind[getattr(self, instrument).mask][index]
+        else:
+            return ind[index]
 
     # @property
     def _tt(self, f=20) -> np.ndarray:
@@ -651,7 +655,11 @@ class RV(ISSUES, REPORTS):
         return s
 
     @classmethod
-    def from_arrays(cls, star, time, vrad, svrad, inst, **kwargs):
+    def from_arrays(cls, star, time, vrad, svrad, instrument:str, **kwargs):
+        if 'inst' in kwargs:
+            logger.warning('`inst` is deprecated. Use `instrument` instead.')
+            instrument = kwargs.pop('inst')
+
         s = cls(star, _child=True)
         time, vrad, svrad = map(np.atleast_1d, (time, vrad, svrad))
 
@@ -673,8 +681,12 @@ class RV(ISSUES, REPORTS):
         for k, v in kwargs.items():
             setattr(s, k, np.atleast_1d(v))
 
-        s.instruments = [inst]
         s._quantities = np.array(list(kwargs.keys()))
+        _instrument = instrument.replace(' ', '_').replace('-', '_')
+        s.instruments = [_instrument]
+
+        setattr(s, _instrument, deepcopy(s))
+        s._child = False
 
         return s
 
@@ -1039,7 +1051,7 @@ class RV(ISSUES, REPORTS):
             time = np.array([i.bjd for i in CCFs])
             vrad = np.array([i.RV*1e3 for i in CCFs])
             svrad = np.array([i.RVerror*1e3 for i in CCFs])
-            _s = RV.from_arrays(star, time, vrad, svrad, inst=instrument)
+            _s = RV.from_arrays(star, time, vrad, svrad, instrument=instrument)
 
             _quantities = []
 
@@ -1066,9 +1078,13 @@ class RV(ISSUES, REPORTS):
             _s.texp = np.array([i.HDU[0].header['EXPTIME'] for i in CCFs])
             _quantities.append('texp')
 
-            _s.berv = np.array([i.HDU[0].header['HIERARCH ESO QC BERV'] for i in CCFs])
+            try:
+                _s.berv = np.array([i.HDU[0].header['HIERARCH ESO QC BERV'] for i in CCFs])
+            except KeyError:
+                _s.berv = np.full_like(time, np.nan)
             _quantities.append('berv')
 
+
             _s.date_night = np.array([
                 i.HDU[0].header['DATE-OBS'].split('T')[0] for i in CCFs
             ])
@@ -1567,6 +1583,11 @@ class RV(ISSUES, REPORTS):
             condition (ndarray):
                 Boolean array of the same length as the observations
         """
+        condition = np.asarray(condition, dtype=bool)
+        if not np.any(condition):
+            if self.verbose:
+                logger.info('no points to remove')
+            return
         if self.verbose:
             inst = np.unique(self.instrument_array[condition])
             logger.info(f"Removing {condition.sum()} points from instruments {inst}")
@@ -1657,6 +1678,20 @@ class RV(ISSUES, REPORTS):
         for inst in singles:
             self.remove_instrument(inst, strict=True)
 
+    def remove_more_than_n_per_night(self, n=2):
+        """ Remove whenever there are more than `n` observations per night """
+        ind = np.array([], dtype=int)
+        for s in self:
+            # how many observations per night
+            n_night = (np.abs(s.time[:, None] - s.time[None, :]) < 0.5).sum(axis=0)
+            # indices for this instrument
+            ind_s = np.where(n_night >= n)[0]
+            # translate to indices in self
+            ind_self = self._index_from_instrument_index(ind_s, s.instruments[0], masked=False)
+            ind = np.r_[ind, ind_self]
+        if len(ind) > 0:
+            self.remove_point(ind)
+
     def remove_prog_id(self, prog_id):
         """ Remove observations from a given program ID """
         from glob import has_magic
@@ -1695,7 +1730,7 @@ class RV(ISSUES, REPORTS):
             ind = np.where(to_remove)[0]
             self.remove_point(ind)
 
-    def choose_n_points(self, n: int, seed=None, instrument=None):
+    def choose_n_points(self, n: int, instrument=None, seed=None):
         """ Randomly choose `n` observations and mask out the remaining ones
 
         Args:
@@ -1706,21 +1741,21 @@ class RV(ISSUES, REPORTS):
             instrument (str or list, optional):
                 For which instrument to choose points (default is all).
         """
-        instruments = self._check_instrument(instrument)
+        if not self._check_instrument(instrument):
+            return
+        # instruments = self._check_instrument(instrument)
+        mask_for_this_inst = self._instrument_mask(instrument)
         rng = np.random.default_rng(seed=seed)
-        for inst in instruments:
-            # s = getattr(self, inst)
-            mask_for_this_inst = self.obs == self.instruments.index(inst) + 1
-            # only choose if there are more than n points
-            if self.mask[mask_for_this_inst].sum() > n:
-                if self.verbose:
-                    logger.info(f'selecting {n} points from {inst}')
-                # indices of points for this instrument which are not masked already
-                available = np.where(self.mask & mask_for_this_inst)[0]
-                # choose n randomly
-                i = rng.choice(available, size=n, replace=False)
-                # mask the others out
-                self.mask[np.setdiff1d(available, i)] = False
+        # only choose if there are more than n points
+        if self.mask[mask_for_this_inst].sum() > n:
+            if self.verbose:
+                logger.info(f'selecting {n} points from {instrument}')
+            # indices of points for this instrument which are not masked already
+            available = np.where(self.mask & mask_for_this_inst)[0]
+            # choose n randomly
+            i = rng.choice(available, size=n, replace=False)
+            # mask the others out
+            self.mask[np.setdiff1d(available, i)] = False
         self._propagate_mask_changes()
 
 
@@ -1924,10 +1959,14 @@ class RV(ISSUES, REPORTS):
             if np.isnan(d[m]).all():
                 continue
 
-            result = dosigmaclip(d[m], low=sigma, high=sigma)
+            result = dosigmaclip(d[m & self.mask], low=sigma, high=sigma)
             # n = self.vrad[m].size - result.clipped.size
 
-            ind = m & self.mask & ((d < result.lower) | (d > result.upper))
+            # NOTE: result.lower and result.upper are updated values, calculated
+            # *after* the last point has been removed. So the previous solution
+            # getting points outside the range [result.lower, result.upper] is
+            # not correct and we need to use result.clipped instead.
+            ind = m & self.mask & ~np.isin(d, result.clipped)
             n = ind.sum()
 
             if self.verbose and n > 0:
@@ -2020,9 +2059,9 @@ class RV(ISSUES, REPORTS):
 
 
 
-    def bin(self):
+    def bin(self, daily=False):
         """
-        Nightly bin the observations.
+        Bin the observations, nightly by default.
 
         !!! Warning
             This creates and returns a new object and does not modify self.
@@ -2033,6 +2072,8 @@ class RV(ISSUES, REPORTS):
         # store original object
         snew._unbinned = deepcopy(self)
 
+        time_offset = 0.5 if daily else 0.0
+
         all_bad_quantities = []
 
         for inst in snew.instruments:
@@ -2046,7 +2087,8 @@ class RV(ISSUES, REPORTS):
             if s.mtime.size == 0:
                 continue
 
-            tb, vb, svb = binRV(s.mtime, s.mvrad, s.msvrad)
+            tb, vb, svb = binRV(s.mtime + time_offset, s.mvrad, s.msvrad)
+            tb -= time_offset
             s.vrad = vb
             s.svrad = svb
 
@@ -2057,14 +2099,14 @@ class RV(ISSUES, REPORTS):
 
                 # treat date_night specially, basically doing a group-by
                 if q == 'date_night':
-                    inds = binRV(s.mtime, None, None, binning_indices=True)
+                    inds = binRV(s.mtime + time_offset, None, None, binning_indices=True)
                     setattr(s, q, Q[s.mask][inds])
                     continue
 
                 # treat ccf_mask specially, doing a 'unique' bin
                 if q == 'ccf_mask':
                     ccf_mask = getattr(s, q)[s.mask]
-                    setattr(s, q, bin_ccf_mask(s.mtime, ccf_mask))
+                    setattr(s, q, bin_ccf_mask(s.mtime + time_offset, ccf_mask))
                     continue
 
                 if Q.dtype != np.float64:
@@ -2079,12 +2121,14 @@ class RV(ISSUES, REPORTS):
                 elif q + '_err' in s._quantities:
                     Qerr = getattr(s, q + '_err')
                     if (Qerr == 0.0).all(): # if all errors are NaN, don't use them
-                        _, yb = binRV(s.mtime, Q[s.mask], stat='mean', tstat='mean')
+                        _, yb = binRV(s.mtime + time_offset, Q[s.mask],
+                                      stat='mean', tstat='mean')
                     else:
                         if (Qerr <= 0.0).any(): # if any error is <= 0, set it to NaN
                             Qerr[Qerr <= 0.0] = np.nan
 
-                        _, yb, eb = binRV(s.mtime, Q[s.mask], Qerr[s.mask], remove_nans=False)
+                        _, yb, eb = binRV(s.mtime + time_offset, Q[s.mask], Qerr[s.mask],
+                                          remove_nans=False)
                         setattr(s, q + '_err', eb)
 
                     setattr(s, q, yb)
@@ -2093,7 +2137,7 @@ class RV(ISSUES, REPORTS):
                     with warnings.catch_warnings():
                         warnings.filterwarnings('ignore', category=RuntimeWarning)
                         try:
-                            _, yb = binRV(s.mtime, Q[s.mask],
+                            _, yb = binRV(s.mtime + time_offset, Q[s.mask],
                                         stat=np.nanmean, tstat=np.nanmean)
                             setattr(s, q, yb)
                         except TypeError:
@@ -2395,7 +2439,7 @@ class RV(ISSUES, REPORTS):
         if new_units not in possible:
             msg = f"new_units must be one of 'm/s', 'km/s', 'ms', 'kms', got '{new_units}'"
             raise ValueError(msg)
-        
+
         new_units = possible[new_units]
         if new_units == self.units:
             return
@@ -2403,11 +2447,14 @@ class RV(ISSUES, REPORTS):
         if self.verbose:
             logger.info(f"changing units from {self.units} to {new_units}")
 
-        if new_units == 'm/s' and self.units == 'km/s':
+        if new_units == 'm/s' and self.units in ('km/s', 'kms'):
             factor = 1e3
-        elif new_units == 'km/s' and self.units == 'm/s':
+        elif new_units == 'km/s' and self.units in ('m/s', 'ms'):
             factor = 1e-3
-        
+        else:
+            logger.warning(f"no known conversion from {self.units} to {new_units}")
+            return
+
         for inst in self.instruments:
             s = getattr(self, inst)
             s.vrad *= factor
@@ -2469,6 +2516,10 @@ class RV(ISSUES, REPORTS):
             by_last_observation (bool, optional):
                 Sort by last observation date
         """
+        if self.N == 0:
+            if self.verbose:
+                logger.warning("no observations to sort")
+            return
         if by_last_observation:
             by_first_observation = False
         if by_first_observation:

{arvi-0.2.11.dist-info → arvi-0.3.2.dist-info}/METADATA

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

{arvi-0.2.11.dist-info → arvi-0.3.2.dist-info}/RECORD

@@ -1,29 +1,29 @@
 arvi/HZ.py,sha256=u7rguhlILRBW-LOczlY3dkIB4LM8p8W7Xfg4FnNaYG0,2850
-arvi/__init__.py,sha256=8IeHbu8LR1H3MxVi9aJHNpBU6SDbmzj9FZGKsiF0AKE,740
+arvi/__init__.py,sha256=XpHKf-hlJvZvjnhGia_0DAoL0OHagtR7zknIJwIx_UY,853
 arvi/ariadne_wrapper.py,sha256=YvilopJa9T4NwPcj3Nah_U8smSeSAU5-HYZMb_GJ-BQ,2232
 arvi/berv.py,sha256=eKnpuPC1w45UrUEyFRbs9F9j3bXz3kxYzNXbnRgvFQM,17596
 arvi/binning.py,sha256=NK9y9bUrdyWCbh79LkcRABHG-n5MtlETMHMvLj1z-OM,15437
-arvi/config.py,sha256=JkHSwF-EEqwwbcc8thGgbFc9udDZPjQH-9XFjqDepBY,2337
-arvi/dace_wrapper.py,sha256=i3aIRKSc2MylEzM5WbbU-BnuLbz8JlG1ez6sK_vvQH0,26125
+arvi/config.py,sha256=DbOhDKGbNoKVWMd965wY6gAfYf0TdbjR5_WIlHvilcA,2422
+arvi/dace_wrapper.py,sha256=-FHY9xc2Oyf7b_qzbZWy0RtQmLvm770yJeWVV8_8Zwc,27893
 arvi/exofop_wrapper.py,sha256=8S7UEcrBAgANIweMV0-CvaWaVTPgGVo8vQQk_KRa0nU,2414
 arvi/extra_data.py,sha256=Xi65pI5kkzqlMmHGl9xFoumtH699611pJJ5PV-a_IfU,3397
-arvi/gaia_wrapper.py,sha256=HTuigIduin3raWfSC7QYuQxDk2dEXYH_4egRkzzg7Xw,4379
+arvi/gaia_wrapper.py,sha256=JiD0BRix7wD3LXcnaAxWyRkC2M91pVJ2P9DhtzrBavM,4793
 arvi/headers.py,sha256=uvdJebw1M5YkGjE3vJJwYBOnLikib75uuZE9FXB5JJM,1673
 arvi/instrument_specific.py,sha256=94oMb6UeH6tp7H8YXnXHpxEhIz2evz0iYsT_HrNOCTo,12105
 arvi/kepmodel_wrapper.py,sha256=mmHudetAZ4cBxKDwzQzgUydzkjhomCWw5VVuyiKfXq8,10288
-arvi/kima_wrapper.py,sha256=GrAZWkDCg8ukhW41M1VTadSbab0GBa6BIzjtAtvjk58,3891
+arvi/kima_wrapper.py,sha256=T_NjUcbhwPh4gKW4uS2Bki1xKNvZ2LApcYfoJVrkbz8,4541
 arvi/lbl_wrapper.py,sha256=_ViGVkpakvuBR_xhu9XJRV5EKHpj5Go6jBZGJZMIS2Y,11850
 arvi/nasaexo_wrapper.py,sha256=ZKY3IUClqsJuysxDv0Gu51EnzMX7101zQb7UQy_urhI,7431
-arvi/plots.py,sha256=U4VUNyIx4h_rEFd7ZWgBcawUcIGcURES0A4VXIBKp3U,35240
-arvi/programs.py,sha256=M8o8hXr6W22dMiIX3Nxz4pgb8lsJXApDlq7HStyTfqs,9047
+arvi/plots.py,sha256=7u-EqZRYod1arakelhNwNZ_EJwxYDcplxn9KgcS33aQ,36131
+arvi/programs.py,sha256=2gTvHcJPEmmf4hXQcHvG487PXe__qsS-S-L_Ip6tgNM,8622
 arvi/reports.py,sha256=a38EZNhyGoSSzJh63wBQCAt3_xhqbpVGcDOXaZWTLXs,11127
-arvi/setup_logger.py,sha256=dHzO2gPjw6CaKWpYZd2f83z09tmxgi--qpp7k1jROjI,615
-arvi/simbad_wrapper.py,sha256=uZc8mcfNijXsQi29LReRTabZb2hRPhYdLsDLMgq1OEI,9927
-arvi/sophie_wrapper.py,sha256=KUeWccXud5_Lrx72S1HSemHIZRdjd2oLvqyofwsL0QQ,3440
+arvi/setup_logger.py,sha256=bRZ0CdYM406IoIKyaLurqGLxq-sQrHry8nSSiZGoMBw,656
+arvi/simbad_wrapper.py,sha256=ID43g5mlORx5pOsk3cMX2dGkKktrFBbzxhzEHeZ-0mw,10937
+arvi/sophie_wrapper.py,sha256=SUNGeOhS8FSlHsF9PTK3QIx8EUbEvmGlyIfogwjhtSc,3441
 arvi/spectra.py,sha256=ebF1ocodTastLx0CyqLSpE8EZNDXBF8riyfxMr3L6H0,7491
-arvi/stats.py,sha256=gvMkKzP83AV8_Oi71JHmA8QH8Y1z1viYykV9ELVDqZI,3547
+arvi/stats.py,sha256=oV1ouTV_c5bx-akTCXh6WKhoit9LpGHUw0xhR7RF0kY,3566
 arvi/stellar.py,sha256=GQ7yweuBRnfkJ0M5eWjvLd8uvGq_by81PbXfidBvWis,4918
-arvi/timeseries.py,sha256=suXPLmFlNbnFbfixXufNMwqwZWz-FwKdckJ1x2PtrXQ,104664
+arvi/timeseries.py,sha256=shQxIjbsWS3l2x5dh9VWk61bTbmZvQmMIvXjMHPDPDI,106826
 arvi/translations.py,sha256=PUSrn4zvYO2MqGzUxlFGwev_tBkgJaJrIYs6NKHzbWo,951
 arvi/utils.py,sha256=MuAgjyXr297Sm_T6QmB1riVUktyT9ud1qngGMgKlXMc,10863
 arvi/data/info.svg,sha256=0IMI6W-eFoTD8acnury79WJJakpBwLa4qKS4JWpsXiI,489
@@ -32,8 +32,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.2.11.dist-info/licenses/LICENSE,sha256=6JfQgl7SpM55t0EHMFNMnNh-AdkpGW25MwMiTnhdWQg,1068
-arvi-0.2.11.dist-info/METADATA,sha256=G2yojj2vhNVuEPtoY81O3aWPvgMMxNw5ax4_40zVhb8,1933
-arvi-0.2.11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-arvi-0.2.11.dist-info/top_level.txt,sha256=4EeiKDVLD45ztuflTGfQ3TU8GVjJg5Y95xS5XjI-utU,5
-arvi-0.2.11.dist-info/RECORD,,
+arvi-0.3.2.dist-info/licenses/LICENSE,sha256=6JfQgl7SpM55t0EHMFNMnNh-AdkpGW25MwMiTnhdWQg,1068
+arvi-0.3.2.dist-info/METADATA,sha256=FmP8zMlErd55jcJNCbH_KSG5ergHedbcpwQOBpCOwZw,1932
+arvi-0.3.2.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+arvi-0.3.2.dist-info/top_level.txt,sha256=4EeiKDVLD45ztuflTGfQ3TU8GVjJg5Y95xS5XjI-utU,5
+arvi-0.3.2.dist-info/RECORD,,

{arvi-0.2.11.dist-info → arvi-0.3.2.dist-info}/WHEEL

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