arvi 0.1.11__py3-none-any.whl → 0.1.13__py3-none-any.whl

arvi/timeseries.py

@@ -11,15 +11,16 @@ import numpy as np
 from astropy import units
 
 from .setup_logger import logger
-from .config import return_self, check_internet
+from .config import return_self, check_internet, debug
 from .translations import translate
-from .dace_wrapper import do_download_filetype, get_observations, get_arrays
+from .dace_wrapper import do_download_filetype, do_symlink_filetype, get_observations, get_arrays
 from .simbad_wrapper import simbad
+from .gaia_wrapper import gaia
 from .extra_data import get_extra_data
 from .stats import wmean, wrms
 from .binning import bin_ccf_mask, binRV
 from .HZ import getHZ_period
-from .utils import strtobool, there_is_internet
+from .utils import strtobool, there_is_internet, timer
 
 
 @dataclass
@@ -49,7 +50,7 @@ class RV:
     do_sigma_clip: bool = field(init=True, repr=False, default=False)
     do_adjust_means: bool = field(init=True, repr=False, default=True)
     only_latest_pipeline: bool = field(init=True, repr=False, default=True)
-    load_extra_data: Union[bool, str] = field(init=True, repr=False, default=True)
+    load_extra_data: Union[bool, str] = field(init=True, repr=False, default=False)
     #
     _child: bool = field(init=True, repr=False, default=False)
     _did_secular_acceleration: bool = field(init=False, repr=False, default=False)
@@ -72,6 +73,7 @@ class RV:
         if not self._child:
             if check_internet and not there_is_internet():
                 raise ConnectionError('There is no internet connection?')
+
             # complicated way to query Simbad with self.__star__ or, if that
             # fails, try after removing a trailing 'A'
             for target in (self.__star__, self.__star__.replace('A', '')):
@@ -84,12 +86,26 @@ class RV:
                 if self.verbose:
                     logger.error(f'simbad query for {self.__star__} failed')
 
+            # complicated way to query Gaia with self.__star__ or, if that
+            # fails, try after removing a trailing 'A'
+            for target in (self.__star__, self.__star__.replace('A', '')):
+                try:
+                    self.gaia = gaia(target)
+                    break
+                except ValueError:
+                    continue
+            else:
+                if self.verbose:
+                    logger.error(f'Gaia query for {self.__star__} failed')
+
             # query DACE
             if self.verbose:
                 logger.info(f'querying DACE for {self.__star__}...')
             try:
-                self.dace_result = get_observations(self.__star__, self.instrument, 
-                                                    verbose=self.verbose)
+                with timer():
+                    self.dace_result = get_observations(self.__star__, self.instrument,
+                                                        main_id=self.simbad.main_id,
+                                                        verbose=self.verbose)
             except ValueError as e:
                 # querying DACE failed, should we raise an error?
                 if self._raise_on_error:
@@ -243,12 +259,16 @@ class RV:
         table += ' | '.join(map(str, self.NN.values())) + '\n'
         return table
 
+    @property
+    def point(self):
+        return [(t.round(4), v.round(4), sv.round(4)) for t, v, sv in zip(self.time, self.vrad, self.svrad)]
+
     @property
     def mtime(self) -> np.ndarray:
         """ Masked array of times """
         return self.time[self.mask]
 
-    @property  
+    @property
     def mvrad(self) -> np.ndarray:
         """ Masked array of radial velocities """
         return self.vrad[self.mask]
@@ -374,7 +394,7 @@ class RV:
             star, timestamp = file.replace('.pkl', '').split('_')
         else:
             try:
-                file = sorted(glob(f'{star}_*.pkl'))[-1]
+                file = sorted(glob(f'{star}_*.*.pkl'))[-1]
             except IndexError:
                 raise ValueError(f'cannot find any file matching {star}_*.pkl')
             star, timestamp = file.replace('.pkl', '').split('_')
@@ -409,7 +429,7 @@ class RV:
             if star_.size == 1:
                 logger.info(f'assuming star is {star_[0]}')
                 star = star_[0]
-        
+
         if instrument is None:
             instruments = np.array([os.path.splitext(f)[0].split('_')[1] for f in files])
             logger.info(f'assuming instruments: {instruments}')
@@ -495,7 +515,7 @@ class RV:
             for q in ['drs_qc']:
                 setattr(_s, q, np.full(time.size, True))
                 _quantities.append(q)
-            
+
             #! end hack
 
             _s.mask = np.ones_like(time, dtype=bool)
@@ -522,10 +542,10 @@ class RV:
         except ImportError:
             logger.error('iCCF is not installed. Please install it with `pip install iCCF`')
             return
-        
+
         if isinstance(files, str):
             files = [files]
-    
+
         I = iCCF.from_file(files)
 
         objects = np.unique([i.HDU[0].header['OBJECT'].replace(' ', '') for i in I])
@@ -552,7 +572,7 @@ class RV:
         return s
 
 
-    def _check_instrument(self, instrument, strict=False):# -> list | None:
+    def _check_instrument(self, instrument, strict=False, log=False):# -> list | None:
         """
         Check if there are observations from `instrument`.
 
@@ -585,6 +605,11 @@ class RV:
                 if any([instrument in inst for inst in self.instruments]):
                     return [inst for inst in self.instruments if instrument in inst]
 
+        if log:
+            logger.error(f"No data from instrument '{instrument}'")
+            logger.info(f'available: {self.instruments}')
+            return
+
 
     def _build_arrays(self):
         """ build all concatenated arrays of `self` from each of the `.inst`s """
@@ -629,7 +654,8 @@ class RV:
                 setattr(self, q, arr)
 
 
-    def download_ccf(self, instrument=None, index=None, limit=None, directory=None, **kwargs):
+    def download_ccf(self, instrument=None, index=None, limit=None,
+                     directory=None, symlink=False, **kwargs):
         """ Download CCFs from DACE
 
         Args:
@@ -657,9 +683,15 @@ class RV:
         # remove empty strings
         files = list(filter(None, files))
 
-        do_download_filetype('CCF', files[:limit], directory, **kwargs)
+        if symlink:
+            if 'top_level' not in kwargs:
+                logger.warning('may need to provide `top_level` in kwargs to find file')
+            do_symlink_filetype('CCF', files[:limit], directory, **kwargs)
+        else:
+            do_download_filetype('CCF', files[:limit], directory, verbose=self.verbose, **kwargs)
 
-    def download_s1d(self, instrument=None, index=None, limit=None, directory=None, **kwargs):
+    def download_s1d(self, instrument=None, index=None, limit=None,
+                     directory=None, symlink=False, **kwargs):
         """ Download S1Ds from DACE
 
         Args:
@@ -687,9 +719,15 @@ class RV:
         # remove empty strings
         files = list(filter(None, files))
 
-        do_download_filetype('S1D', files[:limit], directory, **kwargs)
+        if symlink:
+            if 'top_level' not in kwargs:
+                logger.warning('may need to provide `top_level` in kwargs to find file')
+            do_symlink_filetype('S1D', files[:limit], directory, **kwargs)
+        else:
+            do_download_filetype('S1D', files[:limit], directory, verbose=self.verbose, **kwargs)
 
-    def download_s2d(self, instrument=None, index=None, limit=None, directory=None, **kwargs):
+    def download_s2d(self, instrument=None, index=None, limit=None,
+                     directory=None, symlink=False, **kwargs):
         """ Download S2Ds from DACE
 
         Args:
@@ -717,11 +755,16 @@ class RV:
         # remove empty strings
         files = list(filter(None, files))
 
-        do_download_filetype('S2D', files[:limit], directory, **kwargs)
+        if symlink:
+            if 'top_level' not in kwargs:
+                logger.warning('may need to provide `top_level` in kwargs to find file')
+            do_symlink_filetype('S2D', files[:limit], directory, **kwargs)
+        else:
+            do_download_filetype('S2D', files[:limit], directory, verbose=self.verbose, **kwargs)
 
 
-    from .plots import plot, plot_fwhm, plot_bis, plot_rhk, plot_quantity
-    from .plots import gls, gls_fwhm, gls_bis, gls_rhk
+    from .plots import plot, plot_fwhm, plot_bis, plot_rhk, plot_berv, plot_quantity
+    from .plots import gls, gls_fwhm, gls_bis, gls_rhk, window_function
     from .reports import report
 
     from .instrument_specific import known_issues
@@ -729,13 +772,13 @@ class RV:
 
     def remove_instrument(self, instrument, strict=False):
         """ Remove all observations from one instrument
-        
+
         Args:
             instrument (str or list):
                 The instrument(s) for which to remove observations.
             strict (bool):
                 Whether to match (each) `instrument` exactly
-        
+
         Note:
             A common name can be used to remove observations for several subsets
             of a given instrument. For example
@@ -788,11 +831,24 @@ class RV:
         if return_self:
             return self
 
+    def remove_condition(self, condition):
+        """ Remove all observations that satisfy a condition
+
+        Args:
+            condition (np.ndarray):
+                Boolean array of the same length as the observations
+        """
+        if self.verbose:
+            inst = np.unique(self.instrument_array[condition])
+            logger.info(f"Removing {condition.sum()} points from instruments {inst}")
+        self.mask = self.mask & ~condition
+        self._propagate_mask_changes()
+
     def remove_point(self, index):
         """
         Remove individual observations at a given index (or indices).
         NOTE: Like Python, the index is 0-based.
-        
+
         Args:
             index (int, list, ndarray):
                 Single index, list, or array of indices to remove.
@@ -899,45 +955,76 @@ class RV:
             n_before = (self.obs < self.obs[m]).sum()
             getattr(self, inst).mask[m - n_before] = False
 
-    def secular_acceleration(self, epoch=55500, plot=False):
+    def secular_acceleration(self, epoch=None, just_compute=False, force_simbad=False):
         """
         Remove secular acceleration from RVs
 
         Args:
-            epoch (float):
+            epoch (float, optional):
                 The reference epoch (DACE uses 55500, 31/10/2010)
             instruments (bool or collection of str):
-                Only remove secular acceleration for some instruments, or for all 
+                Only remove secular acceleration for some instruments, or for all
                 if `instruments=True`
-            plot (bool):
-                Show a plot of the RVs with the secular acceleration
         """
-        if self._did_secular_acceleration:  # don't do it twice
+        if self._did_secular_acceleration and not just_compute:  # don't do it twice
             return
-        
+
+        #as_yr = units.arcsec / units.year
+        mas_yr = units.milliarcsecond / units.year
+        mas = units.milliarcsecond
+
         try:
-            self.simbad
-        except AttributeError:
+            if force_simbad:
+                raise AttributeError
+
+            self.gaia
+            self.gaia.plx
+
             if self.verbose:
-                logger.error('no information from simbad, cannot remove secular acceleration')
-            return
+                logger.info('using Gaia information to remove secular acceleration')
+
+            if epoch is None:
+                # Gaia DR3 epoch (astropy.time.Time('J2016.0', format='jyear_str').jd)
+                epoch = 57389.0
+
+            π = self.gaia.plx * mas
+            d = π.to(units.pc, equivalencies=units.parallax())
+            μα = self.gaia.pmra * mas_yr
+            μδ = self.gaia.pmdec * mas_yr
+            μ = μα**2 + μδ**2
+            sa = (μ * d).to(units.m / units.second / units.year,
+                            equivalencies=units.dimensionless_angles())
+
+        except AttributeError:
+            try:
+                self.simbad
+            except AttributeError:
+                if self.verbose:
+                    logger.error('no information from simbad, cannot remove secular acceleration')
+                return
+
+            if self.simbad.plx is None:
+                if self.verbose:
+                    logger.error('no parallax from simbad, cannot remove secular acceleration')
+                return
 
-        if self.simbad.plx_value is None:
             if self.verbose:
-                logger.error('no parallax from simbad, cannot remove secular acceleration')
-            return
+                logger.info('using Simbad information to remove secular acceleration')
 
-        #as_yr = units.arcsec / units.year
-        mas_yr = units.milliarcsecond / units.year
-        mas = units.milliarcsecond
+            if epoch is None:
+                epoch = 55500
+
+            π = self.simbad.plx * mas
+            d = π.to(units.pc, equivalencies=units.parallax())
+            μα = self.simbad.pmra * mas_yr
+            μδ = self.simbad.pmdec * mas_yr
+            μ = μα**2 + μδ**2
+            sa = (μ * d).to(units.m / units.second / units.year,
+                            equivalencies=units.dimensionless_angles())
+
+        if just_compute:
+            return sa
 
-        π = self.simbad.plx_value * mas
-        d = π.to(units.pc, equivalencies=units.parallax())
-        μα = self.simbad.pmra * mas_yr
-        μδ = self.simbad.pmdec * mas_yr
-        μ = μα**2 + μδ**2
-        sa = (μ * d).to(units.m / units.second / units.year,
-                        equivalencies=units.dimensionless_angles())
         sa = sa.value
 
         if self.verbose:
@@ -961,12 +1048,41 @@ class RV:
                     continue
 
                 s.vrad = s.vrad - sa * (s.time - epoch) / 365.25
-            
+
             self._build_arrays()
 
         self._did_secular_acceleration = True
+        self._did_secular_acceleration_epoch = epoch
+        self._did_secular_acceleration_simbad = force_simbad
+
         if return_self:
             return self
+    
+    def _undo_secular_acceleration(self):
+        if self._did_secular_acceleration:
+            _old_verbose = self.verbose
+            self.verbose = False
+            sa = self.secular_acceleration(just_compute=True,
+                                           force_simbad=self._did_secular_acceleration_simbad)
+            self.verbose = _old_verbose
+            sa = sa.value
+
+            if self._child:
+                self.vrad = self.vrad + sa * (self.time - self._did_secular_acceleration_epoch) / 365.25
+            else:
+                for inst in self.instruments:
+                    if 'HIRES' in inst:  # never remove it from HIRES...
+                        continue
+                    if 'NIRPS' in inst:  # never remove it from NIRPS...
+                        continue
+
+                    s = getattr(self, inst)
+
+                    s.vrad = s.vrad + sa * (s.time - self._did_secular_acceleration_epoch) / 365.25
+
+                self._build_arrays()
+
+            self._did_secular_acceleration = False
 
     def sigmaclip(self, sigma=5, instrument=None, strict=True):
         """ Sigma-clip RVs (per instrument!) """
@@ -1012,7 +1128,7 @@ class RV:
 
     def clip_maxerror(self, maxerror:float):
         """ Mask out points with RV error larger than a given value
-        
+
         Args:
             maxerror (float): Maximum error to keep.
         """
@@ -1038,10 +1154,10 @@ class RV:
 
         WARNING: This creates and returns a new object and does not modify self.
         """
-        
+
         # create copy of self to be returned
         snew = deepcopy(self)
-        
+
         all_bad_quantities = []
 
         for inst in snew.instruments:
@@ -1050,7 +1166,7 @@ class RV:
             # only one observation?
             if s.N == 1:
                 continue
-        
+
             # are all observations masked?
             if s.mtime.size == 0:
                 continue
@@ -1101,7 +1217,7 @@ class RV:
                     with warnings.catch_warnings():
                         warnings.filterwarnings('ignore', category=RuntimeWarning)
                         try:
-                            _, yb = binRV(s.mtime, Q[s.mask], 
+                            _, yb = binRV(s.mtime, Q[s.mask],
                                         stat=np.nanmean, tstat=np.nanmean)
                             setattr(s, q, yb)
                         except TypeError:
@@ -1116,7 +1232,7 @@ class RV:
 
             s.time = tb
             s.mask = np.full(tb.shape, True)
-        
+
         if snew.verbose and len(all_bad_quantities) > 0:
             logger.warning('\nnew object will not have these non-float quantities')
 
@@ -1162,6 +1278,11 @@ class RV:
         for inst in self.instruments:
             s = getattr(self, inst)
 
+            if s.mtime.size == 0:
+                if self.verbose:
+                    logger.info(f'all observations of {inst} are masked')
+                continue
+
             if s.N == 1:
                 if self.verbose:
                     msg = (f'only 1 observation for {inst}, '
@@ -1173,28 +1294,37 @@ class RV:
 
             s.rv_mean = wmean(s.mvrad, s.msvrad)
             s.vrad -= s.rv_mean
+
             if self.verbose:
                 logger.info(f'subtracted weighted average from {inst:10s}: ({s.rv_mean:.3f} {self.units})')
+
             if just_rv:
                 continue
-            # log_msg = 'same for '
+
             for i, other in enumerate(others):
                 y, ye = getattr(s, other), getattr(s, other + '_err')
                 m = wmean(y[s.mask], ye[s.mask])
                 setattr(s, f'{other}_mean', m)
                 setattr(s, other, getattr(s, other) - m)
-                # log_msg += other
-                # if i < len(others) - 1:
-                #     log_msg += ', '
-            
-            # if self.verbose:
-            #     logger.info(log_msg)
 
         self._build_arrays()
         self._did_adjust_means = True
         if return_self:
             return self
 
+    def add_to_vrad(self, values):
+        """ Add an array of values to the RVs of all instruments """
+        if values.size != self.vrad.size:
+            raise ValueError(f"incompatible sizes: len(values) must equal self.N, got {values.size} != {self.vrad.size}")
+
+        for inst in self.instruments:
+            s = getattr(self, inst)
+            mask = self.instrument_array == inst
+            s.vrad += values[mask]
+
+        self._build_arrays()
+
+
     def put_at_systemic_velocity(self):
         """
         For instruments in which mean(RV) < ptp(RV), "move" RVs to the systemic
@@ -1237,7 +1367,7 @@ class RV:
             self._build_arrays()
 
 
-    def save(self, directory=None, instrument=None, full=False,
+    def save(self, directory=None, instrument=None, full=False, postfix=None,
              save_masked=False, save_nans=True):
         """ Save the observations in .rdb files.
 
@@ -1246,9 +1376,10 @@ class RV:
                 Directory where to save the .rdb files.
             instrument (str, optional):
                 Instrument for which to save observations.
-            full (bool, optional): 
-                Whether to save just RVs and errors (False) or more indicators
-                (True).
+            full (bool, optional):
+                Save just RVs and errors (False) or more indicators (True).
+            postfix (str, optional):
+                Postfix to add to the filenames ([star]_[instrument]_[postfix].rdb).
             save_nans (bool, optional)
                 Whether to save NaN values in the indicators, if they exist. If
                 False, the full observation is not saved.
@@ -1301,8 +1432,11 @@ class RV:
                 else:
                     d = np.c_[_s.mtime, _s.mvrad, _s.msvrad]
                 header = 'bjd\tvrad\tsvrad\n---\t----\t-----'
-            
+
             file = f'{star_name}_{inst}.rdb'
+            if postfix is not None:
+                file = f'{star_name}_{inst}_{postfix}.rdb'
+
             files.append(file)
             file = os.path.join(directory, file)
 
@@ -1310,7 +1444,7 @@ class RV:
 
             if self.verbose:
                 logger.info(f'saving to {file}')
-        
+
         return files
 
     def checksum(self, write_to=None):
@@ -1325,7 +1459,7 @@ class RV:
 
 
     #
-    def run_lbl(self, instrument=None, data_dir=None, 
+    def run_lbl(self, instrument=None, data_dir=None,
                 skysub=False, tell=False, limit=None, **kwargs):
         from .lbl_wrapper import run_lbl, NIRPS_create_telluric_corrected_S2D
 
@@ -1339,7 +1473,7 @@ class RV:
                     logger.error(f"No data from instrument '{instrument}'")
                     logger.info(f'available: {self.instruments}')
                     return
-            
+
             if isinstance(instrument, str):
                 instruments = [instrument]
             else:
@@ -1394,7 +1528,7 @@ class RV:
                     logger.error(f"No data from instrument '{instrument}'")
                     logger.info(f'available: {self.instruments}')
                     return
-            
+
             if isinstance(instrument, str):
                 instruments = [instrument]
             else:
@@ -1430,22 +1564,50 @@ class RV:
         return self._planets
 
 
-def fit_sine(t, y, yerr, period='gls', fix_period=False):
+def fit_sine(t, y, yerr=None, period='gls', fix_period=False):
+    """ Fit a sine curve of the form y = A * sin(2π * t / P + φ) + c
+
+    Args:
+        t (ndarray):
+            Time array
+        y (ndarray):
+            Array of observed values
+        yerr (ndarray, optional):
+            Array of uncertainties. Defaults to None.
+        period (str or float, optional):
+            Initial guess for period or 'gls' to get it from Lomb-Scargle
+            periodogram. Defaults to 'gls'.
+        fix_period (bool, optional):
+            Whether to fix the period. Defaults to False.
+
+    Returns:
+        p (ndarray):
+            Best-fit parameters [A, P, φ, c] or [A, φ, c]
+        f (callable):
+            Function that returns the best-fit sine curve for input times
+    """
     from scipy.optimize import leastsq
     if period == 'gls':
         from astropy.timeseries import LombScargle
         gls = LombScargle(t, y, yerr)
         freq, power = gls.autopower()
         period = 1 / freq[power.argmax()]
-
-    if fix_period and period is None:
-        logger.error('period is fixed but no value provided')
-        return
-
-    def sine(t, p):
-        return p[0] * np.sin(2 * np.pi * t / p[1] + p[2]) + p[3]
-
-    p0 = [y.ptp(), period, 0.0, 0.0]
-    xbest, _ = leastsq(lambda p, t, y, ye: (sine(t, p) - y) / ye,
-                       p0, args=(t, y, yerr))
-    return xbest, partial(sine, p=xbest)
+    else:
+        period = float(period)
+
+    if yerr is None:
+        yerr = np.ones_like(y)
+
+    if fix_period:
+        def sine(t, p):
+            return p[0] * np.sin(2 * np.pi * t / period + p[1]) + p[2]
+        f = lambda p, t, y, ye: (sine(t, p) - y) / ye
+        p0 = [y.ptp(), 0.0, 0.0]
+    else:
+        def sine(t, p):
+            return p[0] * np.sin(2 * np.pi * t / p[1] + p[2]) + p[3]
+        f = lambda p, t, y, ye: (sine(t, p) - y) / ye
+        p0 = [y.ptp(), period, 0.0, 0.0]
+
+    xbest, _ = leastsq(f, p0, args=(t, y, yerr))
+    return xbest, partial(sine, p=xbest)

arvi/translations.py

@@ -1,3 +1,5 @@
+import re
+
 STARS = {
     'Barnard': 'GJ699',
     "Barnard's": 'GJ699',
@@ -5,6 +7,15 @@ STARS = {
 
 
 def translate(star):
+    # known translations
     if star in STARS:
         return STARS[star]
+
+    # regex translations
+    NGC_match = re.match(r'NGC([\s\d]+)No([\s\d]+)', star)
+    if NGC_match:
+        cluster = NGC_match.group(1).replace(' ', '')
+        target = NGC_match.group(2).replace(' ', '')
+        return f'Cl* NGC {cluster} MMU {target}'
+
     return star

arvi/utils.py

@@ -1,4 +1,5 @@
 import os
+import time
 from contextlib import contextmanager
 try:
     from unittest.mock import patch
@@ -19,6 +20,9 @@ except ImportError:
     tqdm = lambda x, *args, **kwargs: x
     trange = lambda *args, **kwargs: range(*args, **kwargs)
 
+from .setup_logger import logger
+from . import config
+
 
 def create_directory(directory):
     """ Create a directory if it does not exist """
@@ -61,6 +65,23 @@ def all_logging_disabled():
     finally:
         logging.disable(previous_level)
 
+
+@contextmanager
+def timer():
+    """ A simple context manager to time a block of code """
+    if not config.debug:
+        yield
+        return
+
+    logger.debug(f'starting timer')
+    start = time.time()
+    try:
+        yield
+    finally:
+        end = time.time()
+        logger.debug(f'elapsed time: {end - start:.2f} seconds')
+
+
 def strtobool(val):
     """Convert a string representation of truth to true (1) or false (0).