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

arvi/timeseries.py

@@ -11,13 +11,15 @@ import numpy as np
 from astropy import units
 
 from .setup_logger import logger
-from .config import return_self
+from .config import return_self, check_internet
 from .translations import translate
-from .dace_wrapper import get_observations, get_arrays
-from .dace_wrapper import do_download_ccf, do_download_s1d, do_download_s2d
+from .dace_wrapper import do_download_filetype, get_observations, get_arrays
 from .simbad_wrapper import simbad
+from .extra_data import get_extra_data
 from .stats import wmean, wrms
-from .binning import binRV
+from .binning import bin_ccf_mask, binRV
+from .HZ import getHZ_period
+from .utils import strtobool, there_is_internet
 
 
 @dataclass
@@ -41,13 +43,13 @@ class RV:
     """
     star: str
     instrument: str = field(init=True, repr=False, default=None)
-    N: int = field(init=False, repr=True)
     verbose: bool = field(init=True, repr=False, default=True)
     do_maxerror: Union[bool, float] = field(init=True, repr=False, default=False)
     do_secular_acceleration: bool = field(init=True, repr=False, default=True)
     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)
     #
     _child: bool = field(init=True, repr=False, default=False)
     _did_secular_acceleration: bool = field(init=False, repr=False, default=False)
@@ -68,17 +70,28 @@ class RV:
         self.__star__ = translate(self.star)
 
         if not self._child:
-            try:
-                self.simbad = simbad(self.__star__)
-            except ValueError as e:
-                logger.error(e)
+            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', '')):
+                try:
+                    self.simbad = simbad(target)
+                    break
+                except ValueError:
+                    continue
+            else:
+                if self.verbose:
+                    logger.error(f'simbad 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)
             except ValueError as e:
+                # querying DACE failed, should we raise an error?
                 if self._raise_on_error:
                     raise e
                 else:
@@ -87,7 +100,6 @@ class RV:
                     self.units = ''
                     return
 
-
             # store the date of the last DACE query
             time_stamp = datetime.now(timezone.utc)  #.isoformat().split('.')[0]
             self._last_dace_query = time_stamp
@@ -101,7 +113,8 @@ class RV:
                                 verbose=self.verbose)
 
             for (inst, pipe, mode), data in arrays:
-                child = RV.from_dace_data(self.star, inst, pipe, mode, data, _child=True)
+                child = RV.from_dace_data(self.star, inst, pipe, mode, data, _child=True,
+                                          verbose=self.verbose)
                 inst = inst.replace('-', '_')
                 pipe = pipe.replace('.', '_').replace('__', '_')
                 if self.only_latest_pipeline:
@@ -129,6 +142,22 @@ class RV:
             # all other quantities
             self._build_arrays()
 
+
+            if self.load_extra_data:
+                if isinstance(self.load_extra_data, str):
+                    path = self.load_extra_data
+                else:
+                    path = None
+                try:
+                    self.__add__(get_extra_data(self.star, instrument=self.instrument, path=path),
+                                 inplace=True)
+
+                except FileNotFoundError:
+                    pass
+
+                # all other quantities
+                self._build_arrays()
+
         # do clip_maxerror, secular_acceleration, sigmaclip, adjust_means
         if not self._child:
             if self.do_maxerror:
@@ -143,6 +172,31 @@ class RV:
             if self.do_adjust_means:
                 self.adjust_means()
 
+    def __add__(self, other, inplace=False):
+        # if not isinstance(other, self.__class__):
+        #     raise TypeError('unsupported operand type(s) for +: '
+        #                     f"'{self.__class__.__name__}' and '{other.__class__.__name__}'")
+
+        if np.isin(self.instruments, other.instruments).any():
+            logger.error('the two objects share instrument(s), cannot add them')
+            return
+
+        if inplace:
+            #? could it be as simple as this?
+            for i in other.instruments:
+                self.instruments.append(i)
+                setattr(self, i, getattr(other, i))
+            self._build_arrays()
+        else:
+            # make a copy of ourselves
+            new_self = deepcopy(self)
+            #? could it be as simple as this?
+            for i in other.instruments:
+                new_self.instruments.append(i)
+                setattr(new_self, i, getattr(other, i))
+            new_self._build_arrays()
+            return new_self
+
 
     def reload(self):
         self._did_secular_acceleration = False
@@ -161,23 +215,20 @@ class RV:
             logger.info(f'Saved snapshot to {file}')
 
     @property
-    def N(self):
+    def N(self) -> int:
         """Total number of observations"""
         return self.time.size
 
-    @N.setter
-    def N(self, value):
-        if not isinstance(value, property):
-            logger.error('Cannot set N directly')
-
     @property
     def NN(self):
         """ Total number of observations per instrument """
         return {inst: getattr(self, inst).N for inst in self.instruments}
 
     @property
-    def N_nights(self):
+    def N_nights(self) -> int:
         """ Number of individual nights """
+        if self.mtime.size == 0:
+            return 0
         return binRV(self.mtime, None, None, binning_bins=True).size - 1
 
     @property
@@ -185,15 +236,26 @@ class RV:
         return {inst: getattr(self, inst).N_nights for inst in self.instruments}
 
     @property
-    def mtime(self):
-        return self.time[self.mask]
+    def _NN_as_table(self) -> str:
+        table = ''
+        table += ' | '.join(self.instruments) + '\n'
+        table += ' | '.join([i*'-' for i in map(len, self.instruments)]) + '\n'
+        table += ' | '.join(map(str, self.NN.values())) + '\n'
+        return table
 
     @property
-    def mvrad(self):
+    def mtime(self) -> np.ndarray:
+        """ Masked array of times """
+        return self.time[self.mask]
+
+    @property  
+    def mvrad(self) -> np.ndarray:
+        """ Masked array of radial velocities """
         return self.vrad[self.mask]
 
     @property
-    def msvrad(self):
+    def msvrad(self) -> np.ndarray:
+        """ Masked array of radial velocity uncertainties """
         return self.svrad[self.mask]
 
     @property
@@ -201,7 +263,7 @@ class RV:
         return np.concatenate([[i] * n for i, n in self.NN.items()])
 
     @property
-    def rms(self):
+    def rms(self) -> float:
         """ Weighted rms of the (masked) radial velocities """
         if self.mask.sum() == 0:  # only one point
             return np.nan
@@ -210,7 +272,7 @@ class RV:
 
     @property
     def sigma(self):
-        """ Average error bar """
+        """ Average radial velocity uncertainty """
         if self.mask.sum() == 0:  # only one point
             return np.nan
         else:
@@ -227,11 +289,12 @@ class RV:
         return np.argsort(self.mtime)
 
     @property
-    def _tt(self):
+    def _tt(self) -> np.ndarray:
         return np.linspace(self.mtime.min(), self.mtime.max(), 20*self.N)
 
     @classmethod
     def from_dace_data(cls, star, inst, pipe, mode, data, **kwargs):
+        verbose = kwargs.pop('verbose', False)
         s = cls(star, **kwargs)
         #
         ind = np.argsort(data['rjd'])
@@ -269,7 +332,8 @@ class RV:
         # 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}')
+            if verbose:
+                logger.warning(f'masking {n} points where DRS QC failed for {inst}')
             s.mask &= s.drs_qc
 
         s.instruments = [inst]
@@ -302,7 +366,7 @@ class RV:
         return s
 
     @classmethod
-    def from_snapshot(cls, file=None, star=None):
+    def from_snapshot(cls, file=None, star=None, verbose=True):
         import pickle
         from datetime import datetime
         if star is None:
@@ -316,21 +380,36 @@ class RV:
             star, timestamp = file.replace('.pkl', '').split('_')
 
         dt = datetime.fromtimestamp(float(timestamp))
-        logger.info(f'Reading snapshot of {star} from {dt}')
+        if verbose:
+            logger.info(f'Reading snapshot of {star} from {dt}')
         return pickle.load(open(file, 'rb'))
 
     @classmethod
     def from_rdb(cls, files, star=None, instrument=None, units='ms', **kwargs):
+        """ Create an RV object from an rdb file or a list of rdb files
+
+        Args:
+            files (str, list):
+                File name or list of file names
+            star (str, optional):
+                Name of the star. If None, try to infer it from file name
+            instrument (str, list, optional):
+                Name of the instrument(s). If None, try to infer it from file name
+            units (str, optional):
+                Units of the radial velocities. Defaults to 'ms'.
+
+        Examples:
+            s = RV.from_rdb('star_HARPS.rdb')
+        """
         if isinstance(files, str):
             files = [files]
 
         if star is None:
-            star_ = np.unique([os.path.splitext(f)[0].split('_')[0] for f in files])
+            star_ = np.unique([os.path.splitext(os.path.basename(f))[0].split('_')[0] for f in files])
             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}')
@@ -353,8 +432,24 @@ class RV:
             _s.svrad = data[2] * factor
 
             _quantities = []
+
             #! hack
-            data = np.genfromtxt(f, names=True, dtype=None, comments='--', encoding=None)
+            with open(f) as ff:
+                header = ff.readline().strip()
+                if '\t' in header:
+                    names = header.split('\t')
+                else:
+                    names = header.split()
+
+            if len(names) > 3:
+                kw = dict(skip_header=0, comments='--', names=True, dtype=None, encoding=None)
+                if '\t' in header:
+                    data = np.genfromtxt(f, **kw, delimiter='\t')
+                else:
+                    data = np.genfromtxt(f, **kw)
+                # data.dtype.names = names
+            else:
+                data = np.array([], dtype=np.dtype([]))
 
             if 'fwhm' in data.dtype.fields:
                 _s.fwhm = data['fwhm']
@@ -371,8 +466,10 @@ class RV:
 
             if 'rhk' in data.dtype.fields:
                 _s.rhk = data['rhk']
-                if 'srhk' in data.dtype.fields:
-                    _s.rhk_err = data['srhk']
+                _s.rhk_err = np.full_like(time, np.nan)
+                for possible_name in ['srhk', 'rhk_err']:
+                    if possible_name in data.dtype.fields:
+                        _s.rhk_err = data[possible_name]
             else:
                 _s.rhk = np.zeros_like(time)
                 _s.rhk_err = np.full_like(time, np.nan)
@@ -382,6 +479,23 @@ class RV:
 
             _s.bispan = np.zeros_like(time)
             _s.bispan_err = np.full_like(time, np.nan)
+
+            # other quantities, but all NaNs
+            for q in ['bispan', 'caindex', 'ccf_asym', 'contrast', 'haindex', 'naindex', 'sindex']:
+                setattr(_s, q, np.full_like(time, np.nan))
+                setattr(_s, q + '_err', np.full_like(time, np.nan))
+                _quantities.append(q)
+                _quantities.append(q + '_err')
+            for q in ['berv', 'texp']:
+                setattr(_s, q, np.full_like(time, np.nan))
+                _quantities.append(q)
+            for q in ['ccf_mask', 'date_night', 'prog_id', 'raw_file', 'pub_reference']:
+                setattr(_s, q, np.full(time.size, ''))
+                _quantities.append(q)
+            for q in ['drs_qc']:
+                setattr(_s, q, np.full(time.size, True))
+                _quantities.append(q)
+            
             #! end hack
 
             _s.mask = np.ones_like(time, dtype=bool)
@@ -400,7 +514,45 @@ class RV:
 
         return s
 
-    def _check_instrument(self, instrument, strict=False):
+    @classmethod
+    def from_ccf(cls, files, star=None, instrument=None, **kwargs):
+        """ Create an RV object from a CCF file or a list of CCF files """
+        try:
+            import iCCF
+        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])
+        if objects.size != 1:
+            logger.warning(f'found {objects.size} different stars in the CCF files, '
+                           'choosing the first one')
+        star = objects[0]
+
+        s = cls(star, _child=True)
+
+        # time, RVs, uncertainties
+        s.time = np.array([i.bjd for i in I])
+        s.vrad = np.array([i.RV*1e3 for i in I])
+        s.svrad = np.array([i.RVerror*1e3 for i in I])
+
+        s.fwhm = np.array([i.FWHM*1e3 for i in I])
+        s.fwhm_err = np.array([i.FWHMerror*1e3 for i in I])
+
+        # mask
+        s.mask = np.full_like(s.time, True, dtype=bool)
+
+        s.instruments = list(np.unique([i.instrument for i in I]))
+
+        return s
+
+
+    def _check_instrument(self, instrument, strict=False):# -> list | None:
         """
         Check if there are observations from `instrument`.
 
@@ -414,12 +566,25 @@ class RV:
         """
         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 isinstance(instrument, list):
+            if strict:
+                return [inst for inst in instrument if inst in self.instruments]
+            else:
+                r = []
+                for i in instrument:
+                    if any([i in inst for inst in self.instruments]):
+                        r += [inst for inst in self.instruments if i in inst]
+                return r
+
+        else:
+            if strict:
+                if instrument in self.instruments:
+                    return [instrument]
+            else:
+                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):
         """ build all concatenated arrays of `self` from each of the `.inst`s """
@@ -464,11 +629,12 @@ class RV:
                 setattr(self, q, arr)
 
 
-    def download_ccf(self, instrument=None, limit=None, directory=None):
+    def download_ccf(self, instrument=None, index=None, limit=None, directory=None, **kwargs):
         """ Download CCFs from DACE
 
         Args:
             instrument (str): Specific instrument for which to download data
+            index (int): Specific index of point for which to download data (0-based)
             limit (int): Maximum number of files to download.
             directory (str): Directory where to store data.
         """
@@ -478,18 +644,27 @@ class RV:
         if instrument is None:
             files = [file for file in self.raw_file if file.endswith('.fits')]
         else:
-            instrument = self._check_instrument(instrument)
+            strict = kwargs.pop('strict', False)
+            instrument = self._check_instrument(instrument, strict=strict)
             files = []
             for inst in instrument:
                 files += list(getattr(self, inst).raw_file)
 
-        do_download_ccf(files[:limit], directory)
+        if index is not None:
+            index = np.atleast_1d(index)
+            files = list(np.array(files)[index])
+
+        # remove empty strings
+        files = list(filter(None, files))
 
-    def download_s1d(self, instrument=None, limit=None, directory=None):
+        do_download_filetype('CCF', files[:limit], directory, **kwargs)
+
+    def download_s1d(self, instrument=None, index=None, limit=None, directory=None, **kwargs):
         """ Download S1Ds from DACE
 
         Args:
             instrument (str): Specific instrument for which to download data
+            index (int): Specific index of point for which to download data (0-based)
             limit (int): Maximum number of files to download.
             directory (str): Directory where to store data.
         """
@@ -499,18 +674,27 @@ class RV:
         if instrument is None:
             files = [file for file in self.raw_file if file.endswith('.fits')]
         else:
-            instrument = self._check_instrument(instrument)
+            strict = kwargs.pop('strict', False)
+            instrument = self._check_instrument(instrument, strict=strict)
             files = []
             for inst in instrument:
                 files += list(getattr(self, inst).raw_file)
 
-        do_download_s1d(files[:limit], directory)
+        if index is not None:
+            index = np.atleast_1d(index)
+            files = list(np.array(files)[index])
+
+        # remove empty strings
+        files = list(filter(None, files))
+
+        do_download_filetype('S1D', files[:limit], directory, **kwargs)
 
-    def download_s2d(self, instrument=None, limit=None, directory=None):
+    def download_s2d(self, instrument=None, index=None, limit=None, directory=None, **kwargs):
         """ Download S2Ds from DACE
 
         Args:
             instrument (str): Specific instrument for which to download data
+            index (int): Specific index of point for which to download data (0-based)
             limit (int): Maximum number of files to download.
             directory (str): Directory where to store data.
         """
@@ -520,12 +704,20 @@ class RV:
         if instrument is None:
             files = [file for file in self.raw_file if file.endswith('.fits')]
         else:
-            instrument = self._check_instrument(instrument)
+            strict = kwargs.pop('strict', False)
+            instrument = self._check_instrument(instrument, strict=strict)
             files = []
             for inst in instrument:
                 files += list(getattr(self, inst).raw_file)
 
-        extracted_files = do_download_s2d(files[:limit], directory)
+        if index is not None:
+            index = np.atleast_1d(index)
+            files = list(np.array(files)[index])
+
+        # remove empty strings
+        files = list(filter(None, files))
+
+        do_download_filetype('S2D', files[:limit], directory, **kwargs)
 
 
     from .plots import plot, plot_fwhm, plot_bis, plot_rhk, plot_quantity
@@ -539,8 +731,10 @@ class RV:
         """ Remove all observations from one instrument
         
         Args:
-            instrument (str): The instrument for which to remove observations.
-            strict (bool): Whether to match `instrument` exactly
+            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
@@ -557,7 +751,7 @@ class RV:
             s.remove_instrument('HARPS03')
             ```
 
-            will remove observations from the specific subset.
+            will only remove observations from the specific subset.
         """
         instruments = self._check_instrument(instrument, strict)
 
@@ -595,7 +789,9 @@ class RV:
             return self
 
     def remove_point(self, index):
-        """ Remove individual observations at a given index (or indices)
+        """
+        Remove individual observations at a given index (or indices).
+        NOTE: Like Python, the index is 0-based.
         
         Args:
             index (int, list, ndarray):
@@ -621,6 +817,7 @@ class RV:
             return self
 
     def remove_non_public(self):
+        """ Remove non-public observations """
         if self.verbose:
             n = (~self.public).sum()
             logger.info(f'masking non-public observations ({n})')
@@ -635,6 +832,7 @@ class RV:
                 self.remove_instrument(inst)
 
     def remove_prog_id(self, prog_id):
+        """ Remove observations from a given program ID """
         from glob import has_magic
         if has_magic(prog_id):
             from fnmatch import filter
@@ -652,12 +850,46 @@ class RV:
                 if self.verbose:
                     logger.warning(f'no observations for prog_id "{prog_id}"')
 
-
     def remove_after_bjd(self, bjd):
+        """ Remove observations after a given BJD """
         if (self.time > bjd).any():
             ind = np.where(self.time > bjd)[0]
             self.remove_point(ind)
 
+    def remove_before_bjd(self, bjd):
+        """ Remove observations before a given BJD """
+        if (self.time < bjd).any():
+            ind = np.where(self.time < bjd)[0]
+            self.remove_point(ind)
+
+    def choose_n_points(self, n, seed=None, instrument=None):
+        """ Randomly choose `n` observations and mask out the remaining ones
+
+        Args:
+            n (int):
+                Number of observations to keep.
+            seed (int, optional):
+                Random seed for reproducibility.
+            instrument (str or list, optional):
+                For which instrument to choose points (default is all).
+        """
+        instruments = self._check_instrument(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
+        self._propagate_mask_changes()
+
 
     def _propagate_mask_changes(self):
         """ link self.mask with each self.`instrument`.mask """
@@ -690,6 +922,11 @@ class RV:
                 logger.error('no information 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
+
         #as_yr = units.arcsec / units.year
         mas_yr = units.milliarcsecond / units.year
         mas = units.milliarcsecond
@@ -719,6 +956,10 @@ class RV:
                     continue
 
                 s = getattr(self, inst)
+
+                if hasattr(s, '_did_secular_acceleration') and s._did_secular_acceleration:
+                    continue
+
                 s.vrad = s.vrad - sa * (s.time - epoch) / 365.25
             
             self._build_arrays()
@@ -727,7 +968,7 @@ class RV:
         if return_self:
             return self
 
-    def sigmaclip(self, sigma=5):
+    def sigmaclip(self, sigma=5, instrument=None, strict=True):
         """ Sigma-clip RVs (per instrument!) """
         #from scipy.stats import sigmaclip as dosigmaclip
         from .stats import sigmaclip_median as dosigmaclip
@@ -735,7 +976,9 @@ class RV:
         if self._child or self._did_sigma_clip:
             return
 
-        for inst in self.instruments:
+        instruments = self._check_instrument(instrument, strict)
+
+        for inst in instruments:
             m = self.instrument_array == inst
             result = dosigmaclip(self.vrad[m], low=sigma, high=sigma)
             n = self.vrad[m].size - result.clipped.size
@@ -767,12 +1010,11 @@ class RV:
         if return_self:
             return self
 
-    def clip_maxerror(self, maxerror:float, plot=False):
+    def clip_maxerror(self, maxerror:float):
         """ Mask out points with RV error larger than a given value
         
         Args:
             maxerror (float): Maximum error to keep.
-            plot (bool): Whether to plot the masked points.
         """
         if self._child:
             return
@@ -828,6 +1070,11 @@ class RV:
                     setattr(s, q, Q[s.mask][inds])
                     continue
 
+                # treat ccf_mask specially, doing a 'unique' bin
+                if q == 'ccf_mask':
+                    setattr(s, q, bin_ccf_mask(s.mtime, getattr(s, q)))
+                    continue
+
                 if Q.dtype != np.float64:
                     bad_quantities.append(q)
                     all_bad_quantities.append(q)
@@ -881,6 +1128,7 @@ class RV:
         return snew
 
     def nth_day_mean(self, n=1.0):
+        """ Calculate the n-th day rolling mean of the radial velocities """
         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]
@@ -969,11 +1217,16 @@ class RV:
             self._build_arrays()
 
     def sort_instruments(self, by_first_observation=True, by_last_observation=False):
+        """ Sort instruments by first or last observation date.
+
+        Args:
+            by_first_observation (bool, optional):
+                Sort by first observation date.
+            by_last_observation (bool, optional):
+                Sort by last observation data.
+        """
         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)
@@ -983,10 +1236,9 @@ class RV:
             self.instruments = sorted(self.instruments, key=fun)
             self._build_arrays()
 
-    #
 
     def save(self, directory=None, instrument=None, full=False,
-             save_nans=True):
+             save_masked=False, save_nans=True):
         """ Save the observations in .rdb files.
 
         Args:
@@ -1021,11 +1273,18 @@ class RV:
                 continue
 
             if full:
-                d = np.c_[
-                    _s.mtime, _s.mvrad, _s.msvrad,
-                    _s.fwhm[_s.mask], _s.fwhm_err[_s.mask],
-                    _s.rhk[_s.mask], _s.rhk_err[_s.mask],
-                ]
+                if save_masked:
+                    d = np.c_[
+                        _s.time, _s.vrad, _s.svrad,
+                        _s.fwhm, _s.fwhm_err,
+                        _s.rhk, _s.rhk_err,
+                    ]
+                else:
+                    d = np.c_[
+                        _s.mtime, _s.mvrad, _s.msvrad,
+                        _s.fwhm[_s.mask], _s.fwhm_err[_s.mask],
+                        _s.rhk[_s.mask], _s.rhk_err[_s.mask],
+                    ]
                 if not save_nans:
                     if np.isnan(d).any():
                         # remove observations where any of the indicators are # NaN
@@ -1037,7 +1296,10 @@ class RV:
                 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]
+                if save_masked:
+                    d = np.c_[_s.time, _s.vrad, _s.svrad]
+                else:
+                    d = np.c_[_s.mtime, _s.mvrad, _s.msvrad]
                 header = 'bjd\tvrad\tsvrad\n---\t----\t-----'
             
             file = f'{star_name}_{inst}.rdb'
@@ -1052,6 +1314,7 @@ class RV:
         return files
 
     def checksum(self, write_to=None):
+        """ Calculate a hash based on the data """
         from hashlib import md5
         d = np.r_[self.time, self.vrad, self.svrad]
         H = md5(d.data.tobytes()).hexdigest()
@@ -1101,7 +1364,6 @@ class RV:
                 logger.error(f"not all required files exist in {data_dir}")
                 logger.error(f"missing {np.logical_not(exist).sum()} / {len(files)}")
 
-                from distutils.util import strtobool
                 go_on = input('continue? (y/N) ')
                 if go_on == '' or not bool(strtobool(go_on)):
                     return
@@ -1150,8 +1412,18 @@ class RV:
 
 
     #
+    @property
+    def HZ(self):
+        if not hasattr(self, 'star_mass'):
+            self.star_mass = float(input('stellar mass (Msun): '))
+        if not hasattr(self, 'lum'):
+            self.lum = float(input('luminosity (Lsun): '))
+        return getHZ_period(self.simbad.teff, self.star_mass, 1.0, self.lum)
+
+
     @property
     def planets(self):
+        """ Query the NASA Exoplanet Archive for any known planets """
         from .nasaexo_wrapper import Planets
         if not hasattr(self, '_planets'):
             self._planets = Planets(self)