arvi 0.1.3__py3-none-any.whl → 0.1.6__py3-none-any.whl

arvi/timeseries.py

@@ -55,6 +55,8 @@ class RV:
     _raise_on_error: bool = field(init=True, repr=False, default=True)
 
     def __repr__(self):
+        if self.N == 0:
+            return f"RV(star='{self.star}', N=0)"
         if self.time.size == self.mtime.size:
             return f"RV(star='{self.star}', N={self.N})"
         else:
@@ -73,7 +75,7 @@ class RV:
             if self.verbose:
                 logger.info(f'querying DACE for {self.__star__}...')
             try:
-                self.dace_result = get_observations(self.__star__, self.instrument,
+                self.dace_result = get_observations(self.__star__, self.instrument, 
                                                     verbose=self.verbose)
             except ValueError as e:
                 if self._raise_on_error:
@@ -96,20 +98,18 @@ class RV:
             arrays = get_arrays(self.dace_result, verbose=self.verbose)
             for (inst, pipe, mode), data in arrays:
                 child = RV.from_dace_data(self.star, inst, pipe, mode, data, _child=True)
+                inst = inst.replace('-', '_')
                 setattr(self, inst, child)
 
         # build joint arrays
         if not self._child:
             #! sorted?
-            self.instruments = sorted(list(self.dace_result.keys()))
+            self.instruments = [
+                inst.replace('-', '_') 
+                for inst in sorted(list(self.dace_result.keys()))
+            ]
             # self.pipelines =
-            # "observatory" (or instrument id)
-            self.obs = np.concatenate(
-                [np.full(getattr(self, inst).N, i+1) for i, inst in enumerate(self.instruments)],
-                dtype=int
-            )
-            # mask
-            self.mask = np.full_like(self.obs, True, dtype=bool)
+
             # all other quantities
             self._build_arrays()
 
@@ -202,6 +202,18 @@ class RV:
 
     error = sigma  # alias!
 
+    @property
+    def _time_sorter(self):
+        return np.argsort(self.time)
+
+    @property
+    def _mtime_sorter(self):
+        return np.argsort(self.mtime)
+
+    @property
+    def _tt(self):
+        return np.linspace(self.mtime.min(), self.mtime.max(), 20*self.N)
+
     @classmethod
     def from_dace_data(cls, star, inst, pipe, mode, data, **kwargs):
         s = cls(star, **kwargs)
@@ -211,9 +223,14 @@ class RV:
         s.time = data['rjd'][ind]
         s.vrad = data['rv'][ind]
         s.svrad = data['rv_err'][ind]
+
         # mask
         s.mask = np.full_like(s.time, True, dtype=bool)
         s.mask[np.isnan(s.svrad)] = False
+        ## be careful with bogus values
+        s.mask[s.svrad < 0] = False
+
+
         # all other quantities
         s._quantities = []
         for arr in data.keys():
@@ -223,11 +240,22 @@ class RV:
                     setattr(s, 'ccf_mask', data[arr][ind])
                     s._quantities.append('ccf_mask')
                 else:
+                    # be careful with bogus values in rhk and rhk_err
+                    if arr in ('rhk', 'rhk_err'):
+                        mask99999 = (data[arr] == -99999) | (data[arr] == -99)
+                        data[arr][mask99999] = np.nan
+
                     setattr(s, arr, data[arr][ind])
                     s._quantities.append(arr)
-        #
+
         s._quantities = np.array(s._quantities)
 
+        # mask out drs_qc = False
+        if not s.drs_qc.all():
+            n = (~s.drs_qc).sum()
+            logger.warning(f'masking {n} points where DRS QC failed for {inst}')
+            s.mask &= s.drs_qc
+
         s.instruments = [inst]
         s.pipelines = [pipe]
         s.modes = [mode]
@@ -258,22 +286,119 @@ class RV:
         return s
 
     @classmethod
-    def from_snapshot(cls, file):
+    def from_snapshot(cls, file=None, star=None):
         import pickle
         from datetime import datetime
-        assert file.endswith('.pkl'), 'expected a .pkl file'
-        star, timestamp = file.replace('.pkl', '').split('_')
+        if star is None:
+            assert file.endswith('.pkl'), 'expected a .pkl file'
+            star, timestamp = file.replace('.pkl', '').split('_')
+        else:
+            try:
+                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('_')
+
         dt = datetime.fromtimestamp(float(timestamp))
         logger.info(f'Reading snapshot of {star} from {dt}')
         return pickle.load(open(file, 'rb'))
 
-    def _check_instrument(self, instrument):
+    @classmethod
+    def from_rdb(cls, files, star=None, units='ms', **kwargs):
+        if isinstance(files, str):
+            files = [files]
+
+        if star is None:
+            star_ = np.unique([os.path.splitext(f)[0].split('_')[0] for f in files])
+            if star_.size == 1:
+                logger.info(f'assuming star is {star_[0]}')
+                star = star_[0]
+        
+        instruments = np.array([os.path.splitext(f)[0].split('_')[1] for f in files])
+        logger.info(f'assuming instruments: {instruments}')
+
+        if instruments.size == 1 and len(files) > 1:
+            instruments = np.repeat(instruments, len(files))
+
+        factor = 1e3 if units == 'kms' else 1.0
+
+        s = cls(star, _child=True, **kwargs)
+
+        for i, (f, instrument) in enumerate(zip(files, instruments)):
+            data = np.loadtxt(f, skiprows=2, usecols=range(3), unpack=True)
+            _s = cls(star, _child=True, **kwargs)
+            time = data[0]
+            _s.time = time
+            _s.vrad = data[1] * factor
+            _s.svrad = data[2] * factor
+
+            _quantities = []
+            #! hack
+            data = np.genfromtxt(f, names=True, dtype=None, comments='--', encoding=None)
+
+            if 'fwhm' in data.dtype.fields:
+                _s.fwhm = data['fwhm']
+                if 'sfwhm' in data.dtype.fields:
+                    _s.fwhm_err = data['sfwhm']
+                else:
+                    _s.fwhm_err = 2 * _s.svrad
+            else:
+                _s.fwhm = np.zeros_like(time)
+                _s.fwhm_err = np.full_like(time, np.nan)
+
+            _quantities.append('fwhm')
+            _quantities.append('fwhm_err')
+
+            if 'rhk' in data.dtype.fields:
+                _s.rhk = data['rhk']
+                if 'srhk' in data.dtype.fields:
+                    _s.rhk_err = data['srhk']
+            else:
+                _s.rhk = np.zeros_like(time)
+                _s.rhk_err = np.full_like(time, np.nan)
+
+            _quantities.append('rhk')
+            _quantities.append('rhk_err')
+
+            _s.bispan = np.zeros_like(time)
+            _s.bispan_err = np.full_like(time, np.nan)
+            #! end hack
+
+            _s.mask = np.ones_like(time, dtype=bool)
+            _s.obs = np.full_like(time, i + 1)
+
+            _s.instruments = [instrument]
+            _s._quantities = np.array(_quantities)
+            setattr(s, instrument, _s)
+
+        s._child = False
+        s.instruments = list(instruments)
+        s._build_arrays()
+
+        if kwargs.get('do_adjust_means', False):
+            s.adjust_means()
+
+        return s
+
+    def _check_instrument(self, instrument, strict=False):
+        """
+        Check if there are observations from `instrument`.
+
+        Args:
+            instrument (str, None): Instrument name to check
+            strict (bool): Whether to match `instrument` exactly
+        Returns:
+            instruments (list):
+                List of instruments matching `instrument`, or None if there
+                are no matches.
+        """
         if instrument is None:
             return self.instruments
+        if not strict:
+            if any([instrument in inst for inst in self.instruments]):
+                return [inst for inst in self.instruments if instrument in inst]
         if instrument in self.instruments:
             return [instrument]
-        if any([instrument in inst for inst in self.instruments]):
-            return [inst for inst in self.instruments if instrument in inst]
         
 
     def _build_arrays(self):
@@ -298,6 +423,13 @@ class RV:
             [getattr(self, inst).mask for inst in self.instruments]
         )
 
+        # "observatory" (or instrument id)
+        self.obs = np.concatenate(
+            [np.full(getattr(self, inst).N, i+1) for i, inst in enumerate(self.instruments)],
+            dtype=int
+        )
+
+
         # all other quantities
         self._quantities = getattr(self, self.instruments[0])._quantities
         if len(self.instruments) > 1:
@@ -376,18 +508,19 @@ class RV:
         extracted_files = do_download_s2d(files[:limit], directory)
 
 
-    from .plots import plot, plot_fwhm, plot_bis, plot_rhk
+    from .plots import plot, plot_fwhm, plot_bis, plot_rhk, plot_quantity
     from .plots import gls, gls_fwhm, gls_bis, gls_rhk
     from .reports import report
 
     from .instrument_specific import known_issues
 
 
-    def remove_instrument(self, instrument):
+    def remove_instrument(self, instrument, strict=False):
         """ Remove all observations from one instrument
         
         Args:
             instrument (str): The instrument for which to remove observations.
+            strict (bool): Whether to match `instrument` exactly
         
         Note:
             A common name can be used to remove observations for several subsets
@@ -406,34 +539,37 @@ class RV:
 
             will remove observations from the specific subset.
         """
-        if instrument not in self.instruments:
+        instruments = self._check_instrument(instrument, strict)
+
+        if instruments is None:
             logger.error(f"No data from instrument '{instrument}'")
             logger.info(f'available: {self.instruments}')
             return
 
-        ind = self.instruments.index(instrument) + 1
-        remove = np.where(self.obs == ind)
-        self.obs = np.delete(self.obs, remove)
-        self.obs[self.obs > ind] -= 1
-        #
-        self.time = np.delete(self.time, remove)
-        self.vrad = np.delete(self.vrad, remove)
-        self.svrad = np.delete(self.svrad, remove)
-        #
-        self.mask = np.delete(self.mask, remove)
-        #
-        # all other quantities
-        for q in self._quantities:
-            if q not in ('rjd', 'rv', 'rv_err'):
-                new = np.delete(getattr(self, q), remove)
-                setattr(self, q, new)
-        #
-        self.instruments.remove(instrument)
-        #
-        delattr(self, instrument)
+        for instrument in instruments:
+            ind = self.instruments.index(instrument) + 1
+            remove = np.where(self.obs == ind)
+            self.obs = np.delete(self.obs, remove)
+            self.obs[self.obs > ind] -= 1
+            #
+            self.time = np.delete(self.time, remove)
+            self.vrad = np.delete(self.vrad, remove)
+            self.svrad = np.delete(self.svrad, remove)
+            #
+            self.mask = np.delete(self.mask, remove)
+            #
+            # all other quantities
+            for q in self._quantities:
+                if q not in ('rjd', 'rv', 'rv_err'):
+                    new = np.delete(getattr(self, q), remove)
+                    setattr(self, q, new)
+            #
+            self.instruments.remove(instrument)
+            #
+            delattr(self, instrument)
 
-        if self.verbose:
-            logger.info(f"Removed observations from '{instrument}'")
+            if self.verbose:
+                logger.info(f"Removed observations from '{instrument}'")
 
         if return_self:
             return self
@@ -453,6 +589,9 @@ class RV:
             logger.errors(f'index {index} is out of bounds for N={self.N}')
             return
 
+        if self.verbose:
+            logger.info(f'removing points {index}')
+
         self.mask[index] = False
         self._propagate_mask_changes()
         # for i, inst in zip(index, instrument):
@@ -461,6 +600,45 @@ class RV:
         if return_self:
             return self
 
+    def remove_non_public(self):
+        if self.verbose:
+            n = (~self.public).sum()
+            logger.info(f'masking non-public observations ({n})')
+        self.mask = self.mask & self.public
+        self._propagate_mask_changes()
+
+    def remove_single_observations(self):
+        """ Remove instruments for which there is a single observation """
+        instruments = deepcopy(self.instruments)
+        for inst in instruments:
+            if getattr(self, inst).mtime.size == 1:
+                self.remove_instrument(inst)
+
+    def remove_prog_id(self, prog_id):
+        from glob import has_magic
+        if has_magic(prog_id):
+            from fnmatch import filter
+            matching = np.unique(filter(self.prog_id, prog_id))
+            mask = np.full_like(self.time, False, dtype=bool)
+            for m in matching:
+                mask |= np.isin(self.prog_id, m)
+            ind = np.where(mask)[0]
+            self.remove_point(ind)
+        else:
+            if prog_id in self.prog_id:
+                ind = np.where(self.prog_id == prog_id)[0]
+                self.remove_point(ind)
+            else:
+                if self.verbose:
+                    logger.warning(f'no observations for prog_id "{prog_id}"')
+
+
+    def remove_after_bjd(self, bjd):
+        if (self.time > bjd).any():
+            ind = np.where(self.time > bjd)[0]
+            self.remove_point(ind)
+
+
     def _propagate_mask_changes(self):
         """ link self.mask with each self.`instrument`.mask """
         masked = np.where(~self.mask)[0]
@@ -529,31 +707,43 @@ class RV:
         if return_self:
             return self
 
-    def sigmaclip(self, sigma=3):
-        """ Sigma-clip RVs """
+    def sigmaclip(self, sigma=5):
+        """ Sigma-clip RVs (per instrument!) """
+        #from scipy.stats import sigmaclip as dosigmaclip
+        from .stats import sigmaclip_median as dosigmaclip
+
         if self._child or self._did_sigma_clip:
             return
-        from scipy.stats import sigmaclip as dosigmaclip
-        result = dosigmaclip(self.vrad, low=sigma, high=sigma)
-        n = self.vrad.size - result.clipped.size
 
-        if self.verbose and n > 0:
-            s = 's' if (n == 0 or n > 1) else ''
-            logger.warning(f'sigma-clip RVs will remove {n} point' + s)
+        for inst in self.instruments:
+            m = self.instrument_array == inst
+            result = dosigmaclip(self.vrad[m], low=sigma, high=sigma)
+            n = self.vrad[m].size - result.clipped.size
 
-        ind = (self.vrad > result.lower) & (self.vrad < result.upper)
+            ind = m & ((self.vrad < result.lower) | (self.vrad > result.upper))
 
-        # check if going to remove all observations from one instrument
-        if n in self.NN.values(): # all observations
-            insts = np.unique(self.instrument_array[~ind])
-            if insts.size == 1: # of the same instrument?
-                if self.verbose:
-                    logger.warning(f'would remove all observations from {insts[0]}, skipping')
-                if return_self:
-                    return self
+            if self.verbose and n > 0:
+                s = 's' if (n == 0 or n > 1) else ''
+                logger.warning(f'sigma-clip RVs will remove {n} point{s} for {inst}')
+
+            # # check if going to remove all observations from one instrument
+            # if n in self.NN.values(): # all observations
+            #     # insts = np.unique(self.instrument_array[~ind])
+            #     # if insts.size == 1: # of the same instrument?
+            #     if self.verbose:
+            #         logger.warning(f'would remove all observations from {insts[0]}, skipping')
+            #     if return_self:
+            #         return self
+            #     continue
+
+            self.mask[ind] = False
 
-        self.mask[~ind] = False
         self._propagate_mask_changes()
+
+        if self._did_adjust_means:
+            self._did_adjust_means = False
+            self.adjust_means()
+
         if return_self:
             return self
 
@@ -601,31 +791,38 @@ class RV:
             bad_quantities = []
 
             for q in s._quantities:
+                Q = getattr(s, q)
+
                 # treat date_night specially, basically doing a group-by
                 if q == 'date_night':
                     inds = binRV(s.mtime, None, None, binning_indices=True)
-                    setattr(s, q, getattr(s, q)[s.mask][inds])
+                    setattr(s, q, Q[s.mask][inds])
                     continue
 
-                if getattr(s, q).dtype != np.float64:
+                if Q.dtype != np.float64:
                     bad_quantities.append(q)
                     all_bad_quantities.append(q)
                     continue
 
-                if np.isnan(getattr(s, q)).all():
+                if np.isnan(Q).all():
                     yb = np.full_like(tb, np.nan)
                     setattr(s, q, yb)
+
                 elif q + '_err' in s._quantities:
-                    _, yb, eb = binRV(s.mtime, 
-                                      getattr(s, q)[s.mask], 
-                                      getattr(s, q + '_err')[s.mask])
+                    Qerr = getattr(s, q + '_err')
+                    if (Qerr == 0.0).all():
+                        _, yb = binRV(s.mtime, Q[s.mask], stat='mean', tstat='mean')
+                    else:
+                        _, yb, eb = binRV(s.mtime, Q[s.mask], Qerr[s.mask])
+                        setattr(s, q + '_err', eb)
+
                     setattr(s, q, yb)
-                    setattr(s, q + '_err', eb)
+
                 elif not q.endswith('_err'):
                     with warnings.catch_warnings():
                         warnings.filterwarnings('ignore', category=RuntimeWarning)
                         try:
-                            _, yb = binRV(s.mtime, getattr(s, q)[s.mask], 
+                            _, yb = binRV(s.mtime, Q[s.mask], 
                                         stat=np.nanmean, tstat=np.nanmean)
                             setattr(s, q, yb)
                         except TypeError:
@@ -651,7 +848,33 @@ class RV:
         snew._build_arrays()
         return snew
 
+    def nth_day_mean(self, n=1.0):
+        mask = np.abs(self.mtime[:, None] - self.mtime[None, :]) < n
+        z = np.full((self.mtime.size, self.mtime.size), np.nan)
+        z[mask] = np.repeat(self.mvrad[:, None], self.mtime.size, axis=1)[mask]
+        return np.nanmean(z, axis=0)
+
+    def subtract_mean(self):
+        """ Subtract (single) mean RV from all instruments """
+        self._meanRV = meanRV = self.mvrad.mean()
+        for inst in self.instruments:
+            s = getattr(self, inst)
+            s.vrad -= meanRV
+        self._build_arrays()
+
+    def _add_back_mean(self):
+        """ Add the (single) mean RV removed by self.subtract_mean() """
+        if not hasattr(self, '_meanRV'):
+            logger.warning("no mean RV stored, run 'self.subtract_mean()'")
+            return
+
+        for inst in self.instruments:
+            s = getattr(self, inst)
+            s.vrad += self._meanRV
+        self._build_arrays()
+
     def adjust_means(self, just_rv=False):
+        """ Subtract individual mean RVs from each instrument """
         if self._child or self._did_adjust_means:
             return
 
@@ -692,6 +915,42 @@ class RV:
         if return_self:
             return self
 
+    def put_at_systemic_velocity(self):
+        """
+        For instruments in which mean(RV) < ptp(RV), "move" RVs to the systemic
+        velocity from simbad. This is useful if some instruments are centered
+        at zero while others are not, and instead of calling `.adjust_means()`,
+        but it only works when the systemic velocity is smaller than ptp(RV).
+        """
+        changed = False
+        for inst in self.instruments:
+            s = getattr(self, inst)
+            if s.mask.any():
+                if np.abs(s.mvrad.mean()) < s.mvrad.ptp():
+                    s.vrad += self.simbad.rvz_radvel * 1e3
+                    changed = True
+            else:  # all observations are masked, use non-masked arrays
+                if np.abs(s.vrad.mean()) < s.vrad.ptp():
+                    s.vrad += self.simbad.rvz_radvel * 1e3
+                    changed = True
+        if changed:
+            self._build_arrays()
+
+    def sort_instruments(self, by_first_observation=True, by_last_observation=False):
+        if by_last_observation:
+            by_first_observation = False
+        # if by_first_observation and by_last_observation:
+        #     logger.error("'by_first_observation' and 'by_last_observation' can't both be true")
+        #     return
+        if by_first_observation:
+            fun = lambda i: getattr(self, i).time.min()
+            self.instruments = sorted(self.instruments, key=fun)
+            self._build_arrays()
+        if by_last_observation:
+            fun = lambda i: getattr(self, i).time.max()
+            self.instruments = sorted(self.instruments, key=fun)
+            self._build_arrays()
+
     #
 
     def save(self, directory=None, instrument=None, full=False):
@@ -707,6 +966,11 @@ class RV:
         """
         star_name = self.star.replace(' ', '')
 
+        if directory is None:
+            directory = '.'
+        else:
+            os.makedirs(directory, exist_ok=True)
+
         files = []
 
         for inst in self.instruments:
@@ -714,26 +978,27 @@ class RV:
                 if instrument not in inst:
                     continue
 
-            file = f'{star_name}_{inst}.rdb'
-            files.append(file)
-
-            if directory is not None:
-                os.makedirs(directory, exist_ok=True)
-                file = os.path.join(directory, file)
-
             _s = getattr(self, inst)
 
+            if not _s.mask.any():  # all observations are masked, don't save
+                continue
+
             if full:
                 d = np.c_[
                     _s.mtime, _s.mvrad, _s.msvrad,
-                    _s.fwhm[_s.mask], _s.fwhm_err[_s.mask]
+                    _s.fwhm[_s.mask], _s.fwhm_err[_s.mask],
+                    _s.rhk[_s.mask], _s.rhk_err[_s.mask],
                 ]
-                header =  'bjd\tvrad\tsvrad\tfwhm\tsfwhm\n'
-                header += '---\t----\t-----\t----\t-----'
+                header =  'bjd\tvrad\tsvrad\tfwhm\tsfwhm\trhk\tsrhk\n'
+                header += '---\t----\t-----\t----\t-----\t---\t----'
             else:
                 d = np.c_[_s.mtime, _s.mvrad, _s.msvrad]
                 header = 'bjd\tvrad\tsvrad\n---\t----\t-----'
             
+            file = f'{star_name}_{inst}.rdb'
+            files.append(file)
+            file = os.path.join(directory, file)
+
             np.savetxt(file, d, fmt='%9.5f', header=header, delimiter='\t', comments='')
 
             if self.verbose:
@@ -741,6 +1006,16 @@ class RV:
         
         return files
 
+    def checksum(self, write_to=None):
+        from hashlib import md5
+        d = np.r_[self.time, self.vrad, self.svrad]
+        H = md5(d.data.tobytes()).hexdigest()
+        if write_to is not None:
+            with open(write_to, 'w') as f:
+                f.write(H)
+        return H
+
+
     #
     def run_lbl(self, instrument=None, data_dir=None, 
                 skysub=False, tell=False, limit=None, **kwargs):
@@ -796,9 +1071,10 @@ class RV:
 
             run_lbl(self, instrument, files[:limit], **kwargs)
 
-    def load_lbl(self, instrument=None, tell=False):
+    def load_lbl(self, instrument=None, tell=False, id=None):
         if hasattr(self, '_did_load_lbl') and self._did_load_lbl: # don't do it twice
             return
+
         from .lbl_wrapper import load_lbl
 
         if instrument is None:
@@ -818,11 +1094,25 @@ class RV:
                 instruments = instrument
 
         for inst in instruments:
-            load_lbl(self, inst, tell=tell)
+            if self.verbose:
+                logger.info(f'loading LBL data for {inst}')
 
+            load_lbl(self, inst, tell=tell, id=id)
+            # self.instruments.append(f'{inst}_LBL')
+
+        # self._build_arrays()
         self._did_load_lbl = True
 
 
+    #
+    @property
+    def planets(self):
+        from .nasaexo_wrapper import Planets
+        if not hasattr(self, '_planets'):
+            self._planets = Planets(self)
+        return self._planets
+
+
 def fit_sine(t, y, yerr, period='gls', fix_period=False):
     from scipy.optimize import leastsq
     if period == 'gls':