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

arvi/__init__.py

@@ -2,18 +2,33 @@ __all__ = ['RV']
 
 from .timeseries import RV
 
-_ran_once = False
+## OLD
+# # the __getattr__ function is always called twice, so we need this
+# # to only build and return the RV object on the second time
+# _ran_once = False
 
 def __getattr__(name: str):
     if name in (
         '_ipython_canary_method_should_not_exist_',
+        '_ipython_display_',
         '_repr_mimebundle_',
         '__wrapped__'
     ):
         return
 
-    global _ran_once  # can't do it any other way :(
-    if _ran_once:
-        return RV(name)
-    else:
-        _ran_once = True
+    try:
+        globals()[name] = RV(name)
+        return globals()[name]
+    except ValueError as e:
+        raise ImportError(e) from None
+        # raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    ## OLD
+    # # can't do it any other way :(
+    # global _ran_once
+
+    # if _ran_once:
+    #     _ran_once = False
+    #     return RV(name)
+    # else:
+    #     _ran_once = True

arvi/berv.py

@@ -0,0 +1,437 @@
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+
+from arvi.headers import get_headers
+from barycorrpy import get_BC_vel
+from astropy.coordinates import SkyCoord
+from astropy.time import Time
+from astropy import units as u
+from astropy import constants as const
+from astropy.timeseries import LombScargle
+from tqdm import tqdm
+
+from .setup_logger import logger
+from . import config
+
+
+def correct_rvs(self, simple=False, H=None, save_files=False, plot=True):
+    """
+    """
+    import pickle
+
+    if hasattr(self, '_did_correct_berv') and self._did_correct_berv:
+        logger.info('Already corrected for the BERV! Not doing anything.')
+        return
+
+    path = os.path.dirname(__file__)
+    path = os.path.join(path, 'data')
+    pkl = os.path.join(path, 'berv_espresso_sine.pkl')
+    berv_espresso = pickle.load(open(pkl, 'rb'))
+    
+    if simple:
+        logger.info('Correcting RVs with a previously-fitted sinusoid function')
+        _f = berv_espresso['func'].replace('lambda t: ', '')
+        logger.info(f': {_f}')
+        f = eval(berv_espresso['func'])
+        if plot:
+            _, ax = self.plot()
+            ax.plot(self._tt, f(self._tt) + self.vrad.mean(), 'k')
+            _, axgls = self.gls(label='before')
+        
+        self.vrad = self.vrad + f(self.time)
+
+        if plot:
+            self.gls(ax=axgls, label='after')
+
+        return f(self.time)
+
+    else:
+        logger.info('Correcting RVs with actual difference between BERVs')
+        logger.info('(basically, use BERV_barycorrpy for BERV correction)')
+
+        old_vrad = self.vrad.copy()
+
+        _, berv = BERV(self, H=H, use_gaia_meassurements=True, plx=self.gaia.plx, 
+                       plot=False, ignore_mask=True)
+
+        if plot:
+            fig, axs = plt.subplots(2, 1, constrained_layout=True, height_ratios=(3, 1), sharex=True)
+            _, ax = self.plot(ax=axs[0])
+            _, axgls = self.gls(label='before')
+
+        # undo secular acceleration, if it was done
+        _did_secular_acceleration = self._did_secular_acceleration
+        self._undo_secular_acceleration()
+
+        # transform RVs: RV --> RV - BERVpipe + BERVbarycorrpy
+
+        diff = berv[self.star]['berv_barycorrpy'] - berv[self.star]['berv_pipeline']
+
+        if save_files:
+            i_inst = np.hstack([np.arange(n) for n in self.NN.values()])
+            with open(f'{self.star}_berv_correction.rdb', 'w') as rdb:
+                rdb.write('# time\n')
+                rdb.write('# vrad\n')
+                rdb.write('# svrad\n')
+                rdb.write('# berv - BERV value from header\n')
+                rdb.write('# berv_pipe - BERV from header corrected for 1.55e-8 factor\n')
+                rdb.write('# berv_barycorrpy - BERV value from barycorrpy\n')
+                rdb.write('# diff - difference between berv_barycorrpy and berv_pipe\n')
+                rdb.write('# vrad_berv_corrected = vrad + diff\n')
+                rdb.write('# instrument\n')
+                rdb.write('# i - index\n')
+                rdb.write('# i_inst - index within the instrument\n')
+                rdb.write('#\n')
+                rdb.write('# --> TO CORRECT vrad, we ** add the diff column **\n')
+                rdb.write('# --> the result of this operation is in column vrad_berv_corrected\n')
+                rdb.write('# --> vrad_berv_corrected is already corrected for the secular acceleration, vrad is not\n')
+                rdb.write('#\n')
+                # 
+                cols = [
+                    'time', 'vrad', 'svrad', 
+                    'berv', 'berv_pipe', 'berv_barycorrpy', 'diff', 'vrad_berv_corrected',
+                    'instrument', 'i', 'i_inst'
+                ]
+                rdb.write('# ' + '\t'.join(cols) + '\n')
+                for i, t in enumerate(self.time):
+                    rdb.write(f'{t:11.5f}\t')
+                    # if _did_secular_acceleration:
+                    #     rdb.write(f'{old_vrad[i]:13.5f}\t')
+                    # else:
+                    rdb.write(f'{self.vrad[i]:13.7f}\t')
+                    rdb.write(f'{self.svrad[i]:13.7f}\t')
+                    rdb.write(f'{self.berv[i]:15.7f}\t')
+                    rdb.write(f'{berv[self.star]["berv_pipeline"][i]/1e3:15.7f}\t')
+                    rdb.write(f'{berv[self.star]["berv_barycorrpy"][i]/1e3:15.7f}\t')
+                    rdb.write(f'{diff[i]:15.7f}\t')
+                    rdb.write(f'{self.vrad[i] + diff[i]:13.7f}\t')
+                    rdb.write(f'{self.instrument_array[i]}\t')
+                    rdb.write(f'{i}\t')
+                    rdb.write(f'{i_inst[i]}\t')
+                    rdb.write('\n')
+
+        self.add_to_vrad(diff)
+        self._did_correct_berv = True
+        self._did_secular_acceleration = True # "automatically", by using BERV_barycorrpy
+        self._did_secular_acceleration_simbad = False
+        self._did_secular_acceleration_epoch = Time('J2016').jd - 24e5
+
+        # the secular acceleration hadn't been done, but it was introduced by
+        # BERV_barycorrpy, so we need to undo it
+        if not _did_secular_acceleration:
+            self._undo_secular_acceleration()
+
+        if plot:
+            self.plot(ax=axs[0], marker='+', ms=5)
+            axs[1].plot(self.time, old_vrad - self.vrad, '.k', label='old RV - new RV')
+            ma = np.abs(axs[1].get_ylim()).max()
+            axs[1].set(ylim=(-ma, ma), xlabel=axs[0].get_xlabel(), ylabel='RV difference [m/s]')
+            self.gls(ax=axgls, label='after')
+
+        return diff
+
+def get_A_and_V_from_lesta(self, username=config.username):
+    try:
+        import paramiko
+    except ImportError:
+        raise ImportError("paramiko is not installed. Please install it with 'pip install paramiko'")
+
+    logs = []
+    for f in self.raw_file:
+        f = f.replace('espresso/', '/projects/astro/ESPRESSODRS/')
+        f = f.replace('nirps/', '/projects/astro/NIRPSDRS/')
+        f = f.replace('.fits', '_SCIENCE_FP.log')
+        f = f.replace('reduced', 'log')
+        f = f.replace('r.ESPRE', 'ESPRESSO')
+        logs.append(f)
+
+    A, V = [], []
+
+    try:
+        ssh = paramiko.SSHClient()
+        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        ssh.connect("lesta02.astro.unige.ch", username=username, timeout=5)
+    except Exception as e:
+        if 'getaddrinfo failed' in str(e):
+            jump = paramiko.SSHClient()
+            jump.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+            jump.connect('login01.astro.unige.ch', username=username, timeout=5)
+            jump_transport = jump.get_transport()
+            jump_channel = jump_transport.open_channel('direct-tcpip', ('10.194.64.162', 22), ('129.194.64.20', 22))
+            ssh = paramiko.SSHClient()
+            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+            ssh.connect('lesta02.astro.unige.ch', username=username, sock=jump_channel)
+        else:
+            raise e
+
+    with ssh.open_sftp() as sftp:
+        pbar = tqdm(logs, total=len(logs), unit='file', desc='Reading logs')
+        for f in pbar:
+            with sftp.open(f) as fp:
+                pattern1 = 'Sun-Earth'
+                pattern2 = "Barycentric Observer's Velocity"
+                for line in fp:
+                    if pattern1 in line:
+                        value = line.strip().split(':')[-1].replace('\x1b[32m', '').replace('\x1b[0m', '').replace(' ', '')
+                        A.append(float(value))
+                    if pattern2 in line:
+                        value = line.strip().split(':')[-1].replace('\x1b[32m', '').replace('\x1b[0m', '').replace(' ', '')
+                        V.append(float(value))
+
+    ssh.close()
+
+    return np.array(A), np.array(V)
+
+
+def BERV(self, H=None, use_gaia_meassurements=False, plx=None,
+         A=None, V=None, plot=True, ignore_mask=False, verbose=False, dpi=None):
+    """ Calculate Barycentric Radial Velocity with barycorr and compare with pipeline
+
+    Args:
+        H (list, optional):
+            List of (CCF/S1D/etc) headers for the target. If None, try to
+            download the CCF files to get the headers.
+        use_gaia_meassurements (bool, optional):
+            Use Gaia coordinates and proper motions instead of those in the headers.
+        plx (float, optional):
+            Value of stellar parallax [mas] to use in barycorr.
+        A (array, optional):
+            Earth-Sun distance [AU] for each BJD (found in the pipeline logs).
+        V (array, optional):
+            Earth's orbital velocity [km/s] for each BJD (found in the pipeline logs).
+        plot (bool, optional):
+            Plot the results.
+    """
+    if H is None:
+        H = get_headers(self, check_lesta=False, check_exo2=False, instrument='ESPRE')
+
+    if len(H) != self.N:
+        raise ValueError(f'Expected {self.N} headers (in `H`), got {len(H)}')
+
+    if 'HARPS' in H[0]['INSTRUME'] or 'NIRPS' in H[0]['INSTRUME']:
+        obsname = 'lasilla'
+    elif 'ESPRESSO' in H[0]['INSTRUME']:
+        obsname = 'paranal'
+    else:
+        raise ValueError('unknown instrument')
+
+    bjd = np.array([h['HIERARCH ESO QC BJD'] for h in H])
+    bjd -= 24e5
+    berv_pipeline = np.array([h['HIERARCH ESO QC BERV'] for h in H])
+
+    # in the pipeline, the BERV is used to shift wavelenghts with this formula
+    # berv_factor = (1 + 1.55e-8) * (1 + BERV/c)
+    # The 1.55e-8 factor is an average of some relativistic effects, which are
+    # probably already included in the BERV calculated from barycorrpy.
+    # Therefore, we compute an "effective" BERV from the pipeline doing
+    # (1 + 1.55e-8) * (1 + BERV/c) = 1 + effBERV/c 
+    # => effBERV = ((1 + 1.55e-8) * (1 + BERV/c) - 1) * c
+
+    if A is None and V is None:
+        if verbose:
+            logger.info("Using mean value for Earth-Sun distance and Earth's orbital velocity")
+
+    if A is None:
+        Φobs = const.G * const.M_sun / const.au + const.G * const.M_earth / const.R_earth
+    else:
+        A = np.atleast_1d(A) * u.km
+        Φobs = const.G * const.M_sun / A + const.G * const.M_earth / const.R_earth
+
+    if V is None:
+        V = 29785 *u.m / u.second
+    else:
+        V = np.atleast_1d(V) * u.km / u.second
+
+    f = 1 / (1 - Φobs / const.c**2 - V**2 / (2*const.c**2))
+    c = const.c.to(u.km / u.second).value
+    berv_pipeline = (f * (1 + berv_pipeline/c) - 1) * c
+
+
+    tmmean = np.array([h['HIERARCH ESO QC TMMEAN USED'] for h in H])
+    mjdobs = np.array([h['MJD-OBS'] for h in H])
+    texp = np.array([h['EXPTIME'] for h in H])
+    jd = mjdobs + 24e5 + 0.5 + (texp * tmmean)/60/60/24
+
+    if verbose:
+        logger.info(f"Unique exposure times: {np.unique(texp)}")
+
+    berv = []
+    if verbose:
+        pbar = enumerate(jd)
+    else:
+        pbar = tqdm(enumerate(jd), total=len(jd),
+                    unit='observation', desc='Computing BERV')
+
+    for i, _jd in pbar:
+        if use_gaia_meassurements:
+            if not hasattr(self, 'gaia'):
+                raise ValueError('No Gaia data available')
+            
+            target = self.gaia.coords
+            pmra = self.gaia.pmra
+            pmdec = self.gaia.pmdec
+            epoch = Time('J2016').jd
+        else:
+            ra = H[i]['* TARG ALPHA'][0]
+            ra = f'{ra:09.2f}'
+            ra = ra[:2] + 'h' + ra[2:4] + 'm' + ra[4:] + 's'
+
+            dec = H[i]['* TARG DELTA'][0]
+            if dec < 0:
+                dec = f'{dec:010.2f}'
+            else:
+                dec = f'{dec:09.2f}'
+            if dec.startswith('-'):
+                dec = dec[:3] + 'd' + dec[3:5] + 'm' + dec[5:] + 's'
+            else:
+                dec = dec[:2] + 'd' + dec[2:4] + 'm' + dec[4:] + 's'
+
+            target = SkyCoord(ra, dec)
+            pmra = H[i]['* TARG PMA'][0] * 1e3
+            pmdec = H[i]['* TARG PMD'][0] * 1e3
+            epoch = Time('J2000').jd
+
+        if verbose:
+            logger.info(f'jd: {_jd}')
+            logger.info(f'\t ra: {target.ra}')
+            logger.info(f'\t dec: {target.dec}')
+            logger.info(f'\t pmra: {pmra}')
+            logger.info(f'\t pmdec: {pmdec}')
+                        
+
+        px = plx or 0.0
+        out = get_BC_vel(_jd, obsname=obsname, rv=0.0, px=px, zmeas=0.0, epoch=epoch,
+                         ra=target.ra.value, dec=target.dec.value, pmra=pmra, pmdec=pmdec)
+        # print(out[1][3])
+        berv.append(out[0])
+    
+    berv = np.array(berv).flatten()
+
+    if ignore_mask: # ignore the system's masked points
+        pass
+    else: # mask points in the BERV output as well
+        bjd = bjd[self.mask]
+        berv = berv[self.mask]
+        berv_pipeline = berv_pipeline[self.mask]
+
+    fig = None
+    if plot:
+        fig, axs = plt.subplots(2, 1, figsize=(8, 6), dpi=dpi, sharex=True,
+                                constrained_layout=True)
+    
+        axs[0].set_title(f'{self.star}', loc='right')
+        axs[0].plot(bjd, berv_pipeline*1e3, '.', label='pipeline', alpha=0.5)
+        axs[0].plot(bjd, berv, '.', label='barycorrpy', alpha=0.5)
+        axs[0].legend(bbox_to_anchor=(0.0, 1.15), loc=2, borderaxespad=0., ncol=2)
+        axs[0].set(xlabel='BJD - 2450000', ylabel='BERV [m/s]')
+
+
+        if plx is not None:
+            epoch = 55500
+            sa = self.secular_acceleration(just_compute=True)
+            print('sa:', sa)
+            sec_acc = sa.value * (bjd - epoch) / 365.25
+
+            axs[0].plot(bjd, sec_acc)
+
+            # fitp = np.polyfit(bjd - epoch, diff, 1)
+            # axs[1].plot(bjd, np.polyval(fitp, bjd - epoch))
+            # axs[1].plot(bjd, np.mean(diff) + diff - np.polyval(fitp, bjd - epoch), '.')
+
+        if plx is None:
+            diff = berv - berv_pipeline*1e3
+            label=r'BERV$_{\rm barycorrpy}$ - BERV$_{\rm pipeline}$'
+        else:
+            diff = berv + sec_acc - berv_pipeline*1e3
+            label=r'BERV$_{\rm barycorrpy}$ (+SA) - BERV$_{\rm pipeline}$'
+
+        axs[1].plot(bjd, diff, 'k.', label=label)
+        axs[1].axhline(np.mean(diff), ls='--', c='k', alpha=0.1)
+
+        from adjustText import adjust_text
+        text = axs[1].text(bjd.max(), diff.min() + 0.1*diff.ptp(), 
+                           f'ptp: {diff.ptp()*1e2:.2f} cm/s',
+                           ha='right', va='bottom', color='g', alpha=0.8)
+        axs[1].plot([bjd[np.argmax(diff)], bjd.max() + 0.05 * bjd.ptp()], 
+                    [np.max(diff), np.max(diff)], 'g--', alpha=0.3)
+        axs[1].plot([bjd[np.argmin(diff)], bjd.max() + 0.05 * bjd.ptp()], 
+                    [np.min(diff), np.min(diff)], 'g--', alpha=0.3)
+
+        ax = axs[1].twinx()
+        diff_cms = 1e2*(diff - np.mean(diff))
+        ax.plot(bjd, diff_cms, alpha=0)
+        ma = np.max(np.abs(ax.get_ylim()))
+        ax.set_ylim(-1 - 5*round(ma/5), 1 + 5*round(ma/5))
+        ax.set(ylabel='diff - mean(diff) [cm/s]')
+        axs[1].set_ylim(np.mean(diff)-ma/100, np.mean(diff)+ma/100)
+
+        axs[1].legend(bbox_to_anchor=(0.0, 1.15), loc=2, borderaxespad=0.)
+        axs[1].set(xlabel='BJD - 2450000', ylabel='diff [m/s]')
+    
+        # adjust_text([text], va='bottom')
+
+    return fig, {
+        self.star: {
+            'bjd': bjd,
+            'berv_pipeline': berv_pipeline*1e3,
+            'berv_barycorrpy': berv
+        }
+    }
+
+
+def plot_BERV_correction(self, H, A, V, berv2=None, berv6=None,
+                         inset=False, inset_range=(3, 5)):
+    fig, axs = plt.subplot_mosaic('ab\ncc\ndd\nee', constrained_layout=True, figsize=(2*3.57, 10))
+
+    if berv2 is None:
+        _, berv2 = BERV(self, H, plot=False)
+    if berv6 is None:
+        _, berv6 = BERV(self, H, A=A, V=V, plot=False)
+
+    self.plot(ax=axs['a'], ms=2)
+    axs['a'].set_title('original', loc='right')
+    self.gls(ax=axs['e'], label='original', color='r', alpha=0.5,
+             fill_between=True, samples_per_peak=20)
+
+    temp_vrad = self.vrad.copy()
+    self.vrad[self.mask] = self.vrad[self.mask] - berv2[self.star]['berv_pipeline'].value + berv6[self.star]['berv_pipeline'].value
+
+    self.plot(ax=axs['b'], ms=2)
+    axs['b'].set_title('after correction', loc='right')
+
+    diff = temp_vrad[self.mask] - self.vrad[self.mask]
+
+    axs['c'].plot(self.mtime, diff, 'k.')
+    axs['c'].set_title('RV difference', loc='right')
+    axs['c'].set(xlabel='BJD - 2450000', ylabel='RV diff [m/s]')
+
+    text = axs['c'].text(self.mtime.max(), diff.min() + 0.1*diff.ptp(), 
+                         f'ptp: {diff.ptp()*1e2:.2f} cm/s',
+                         ha='right', va='bottom', color='g', alpha=0.8)
+    axs['c'].plot([self.mtime[np.argmax(diff)], self.mtime.max() + 0.05 * self.mtime.ptp()], 
+                [np.max(diff), np.max(diff)], 'g--', alpha=0.3)
+    axs['c'].plot([self.mtime[np.argmin(diff)], self.mtime.max() + 0.05 * self.mtime.ptp()], 
+                [np.min(diff), np.min(diff)], 'g--', alpha=0.3)
+
+
+    f, p = LombScargle(self.mtime, diff).autopower(maximum_frequency=1.0, samples_per_peak=10)
+    axs['d'].semilogx(1/f, p, color='k', alpha=0.8)
+    axs['d'].vlines([365.25, 365.25/2], 0, 1, color='k', ls='--', alpha=0.3)
+    axs['d'].set(xlabel='Period [days]', ylabel='normalized power', ylim=(0, 1))
+    axs['d'].set_title('GLS of RV difference', loc='right')
+
+    if inset:
+        inset = axs['d'].inset_axes(bounds=[0.15, 0.3, 0.3, 0.6])
+        m = (1/f > inset_range[0]) & (1/f < inset_range[1])
+        inset.plot(1/f[m], p[m], color='k', alpha=0.8)
+        inset.set(xlim=inset_range, yticks=[])
+        inset.minorticks_on()
+
+    self.gls(ax=axs['e'], label='after correction', color='g', alpha=1,
+             lw=0.8, samples_per_peak=20)
+    axs['e'].set(xlabel='Period [days]', ylabel='normalized power')
+    axs['e'].sharex(axs['d'])
+
+    self.vrad = temp_vrad
+    return fig

arvi/dace_wrapper.py

@@ -193,7 +193,7 @@ def get_observations(star, instrument=None, main_id=None, verbose=True):
             raise ValueError(msg) from None
     else:
         try:
-            result = get_observations_from_instrument(star, instrument, main_id, verbose)
+            result = get_observations_from_instrument(star, instrument, main_id)
         except ValueError:
             msg = f'no {instrument} observations for {star}'
             raise ValueError(msg) from None

arvi/headers.py

@@ -0,0 +1,47 @@
+from tqdm import tqdm
+from astropy.io import fits
+import iCCF
+
+from . import config
+
+def get_headers(self, check_lesta=False, lesta_username=config.username,
+                check_exo2=False, instrument=None):
+    try:
+        import paramiko
+    except ImportError:
+        raise ImportError("paramiko is not installed. Please install it with 'pip install paramiko'")
+
+    H = []
+
+    if check_lesta:
+        with paramiko.SSHClient() as ssh:
+            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+            ssh.connect("lesta02.astro.unige.ch", username=lesta_username)
+            sftp = ssh.open_sftp()
+
+            pbar = tqdm(self.raw_file, total=len(self.raw_file), unit='file', desc='Reading headers')
+            for f in pbar:
+                f = f.replace('espresso/', '/projects/astro/ESPRESSODRS/')
+                f = f.replace('nirps/', '/projects/astro/NIRPSDRS/')
+                f = f.replace('.fits', '_CCF_A.fits')#.replace(':', r'\:')
+                with sftp.open(f) as fp:
+                    header = fits.getheader(fp)
+                    H.append(header)
+
+    if len(H) == 0 and check_exo2:
+        raise NotImplementedError('getting headers from exo2 not yet implemented')
+
+    if len(H) == 0:
+        self.download_ccf()
+        if instrument is None:
+            I = iCCF.from_file(f'{self.star}_downloads/*CCF_A.fits')
+        else:
+            I = iCCF.from_file(f'{self.star}_downloads/r.{instrument}.*CCF_A.fits',
+                               guess_instrument='HARPS' not in instrument)
+        H = [i.HDU[0].header for i in I]
+
+    # sort by BJD
+    H = sorted(H, key=lambda x: x['HIERARCH ESO QC BJD'])
+
+    return H
+

arvi/plots.py

@@ -407,6 +407,7 @@ def plot_quantity(self, quantity, ax=None, show_masked=False, instrument=None,
         'bispan': f'BIS [{self.units}]',
         'rhk': r"$\log$ R'$_{HK}$",
         'berv': f'BERV [km/s]',
+        quantity: quantity,
     }
     ax.set_ylabel(ylabel[quantity])
 

arvi/programs.py

@@ -20,12 +20,14 @@ def get_star(star, instrument=None):
 
 
 class LazyRV:
-    def __init__(self, stars: list, instrument: str = None):
+    def __init__(self, stars: list, instrument: str = None,
+                 _parallel_limit=10):
         self.stars = stars
         if isinstance(self.stars, str):
             self.stars = [self.stars]
         self.instrument = instrument
         self._saved = None
+        self._parallel_limit = _parallel_limit
 
     @property
     def N(self):
@@ -35,7 +37,7 @@ class LazyRV:
         return f"RV({self.N} stars)"
 
     def _get(self):
-        if self.N > 10:
+        if self.N > self._parallel_limit:
             # logger.info('Querying DACE...')
             _get_star = partial(get_star, instrument=self.instrument)
             with multiprocessing.Pool() as pool:

arvi/stellar.py

@@ -0,0 +1,89 @@
+import numpy as np
+
+class prot_age_result:
+    prot_n84: float
+    prot_n84_err: float
+    prot_m08: float
+    prot_m08_err: float
+    age_m08: float
+    age_m08_err: float
+    def __init__(self):
+        pass
+    def __repr__(self):
+        s = f'{self.prot_n84=:.2f} ± {self.prot_n84_err:.2f}, '
+        s += f'{self.prot_m08=:.2f} ± {self.prot_m08_err:.2f}, '
+        s += f'{self.age_m08=:.2f} ± {self.age_m08_err:.2f}'
+        return s.replace('self.', '')
+
+
+def calc_prot_age(self, bv=None):
+    """
+    Calculate rotation period and age from logR'HK activity level, based on the
+    empirical relations of Noyes et al. (1984) and Mamajek & Hillenbrand (2008).
+
+    Args:
+        self (`arvi.RV`):
+            RV object
+        bv (float, optional):
+            B-V colour. If None, use value from Simbad
+
+    Returns:
+        An object with the following attributes:
+
+        prot_n84 (float, array):
+            Chromospheric rotational period via Noyes et al. (1984)
+        prot_n84_err (float, array):
+            Error on 'prot_n84'
+        prot_m08 (float, array):
+            Chromospheric rotational period via Mamajek & Hillenbrand (2008)
+        prot_m08_err (float, array):
+            Error on 'prot_m08'
+        age_m08 (float, array):
+            Gyrochronology age via Mamajek & Hillenbrand (2008)
+        age_m08_err (float, array):
+            Error on 'age_m08'
+
+    Range of logR'HK-Prot relation: -5.5 < logR'HK < -4.3
+    Range of Mamajek & Hillenbrand (2008) relation for ages: 0.5 < B-V < 0.9
+    """
+    log_rhk = np.nanmean(self.rhk[self.mask])
+    if bv is None:
+        bv = self.simbad.B - self.simbad.V
+
+    # Calculate chromospheric Prot:
+    if np.any(log_rhk < -4.3) & np.any(log_rhk > -5.5):
+        if bv < 1:
+            tau = 1.362 - 0.166*(1-bv) + 0.025*(1-bv)**2 - 5.323*(1-bv)**3
+        else:
+            tau = 1.362 - 0.14*(1-bv)
+
+        prot_n84 = 0.324 - 0.400*(5 + log_rhk) - 0.283*(5 + log_rhk)**2 - 1.325*(5 + log_rhk)**3 + tau
+        prot_n84 = 10**prot_n84
+        prot_n84_err = np.log(10)*0.08*prot_n84
+
+        prot_m08 = (0.808 - 2.966*(log_rhk + 4.52))*10**tau
+        prot_m08_err = 4.4*bv*1.7 - 1.7
+    else:
+        prot_n84 = np.nan
+        prot_n84_err = np.nan
+        prot_m08 = np.nan
+        prot_m08_err = np.nan
+
+    # Calculate gyrochronology age:
+    if np.any(prot_m08 > 0.0) & (bv > 0.50) & (bv < 0.9):
+        age_m08 = 1e-3*(prot_m08/0.407/(bv - 0.495)**0.325)**(1./0.566)
+        #age_m08_err = 0.05*np.log(10)*age_m08
+        age_m08_err  = 0.2 * age_m08 * np.log(10) # using 0.2 dex typical error from paper
+    else:
+        age_m08 = np.nan
+        age_m08_err = np.nan
+
+    r = prot_age_result()
+    r.prot_n84 = prot_n84
+    r.prot_n84_err = prot_n84_err
+    r.prot_m08 = prot_m08
+    r.prot_m08_err = prot_m08_err
+    r.age_m08 = age_m08
+    r.age_m08_err = age_m08_err
+    return r
+    # return prot_n84, prot_n84_err, prot_m08, prot_m08_err, age_m08, age_m08_err

arvi/timeseries.py

@@ -11,7 +11,7 @@ import numpy as np
 from astropy import units
 
 from .setup_logger import logger
-from .config import return_self, check_internet, debug
+from . import config
 from .translations import translate
 from .dace_wrapper import do_download_filetype, do_symlink_filetype, get_observations, get_arrays
 from .simbad_wrapper import simbad
@@ -23,6 +23,9 @@ from .HZ import getHZ_period
 from .utils import strtobool, there_is_internet, timer
 
 
+class ExtraFields:
+    pass
+
 @dataclass
 class RV:
     """
@@ -71,7 +74,7 @@ class RV:
         self.__star__ = translate(self.star)
 
         if not self._child:
-            if check_internet and not there_is_internet():
+            if config.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
@@ -103,9 +106,9 @@ class RV:
                 logger.info(f'querying DACE for {self.__star__}...')
             try:
                 with timer():
+                    mid = self.simbad.main_id if hasattr(self, 'simbad') else None
                     self.dace_result = get_observations(self.__star__, self.instrument,
-                                                        main_id=self.simbad.main_id,
-                                                        verbose=self.verbose)
+                                                        main_id=mid, verbose=self.verbose)
             except ValueError as e:
                 # querying DACE failed, should we raise an error?
                 if self._raise_on_error:
@@ -218,6 +221,7 @@ class RV:
         self._did_secular_acceleration = False
         self._did_sigma_clip = False
         self._did_adjust_means = False
+        self._did_correct_berv = False
         self.__post_init__()
 
     def snapshot(self):
@@ -382,6 +386,7 @@ class RV:
         s.mask = kwargs.get('mask', np.full_like(s.time, True, dtype=bool))
 
         s.instruments = [inst]
+        s._quantities = np.array([])
 
         return s
 
@@ -443,6 +448,13 @@ class RV:
 
         s = cls(star, _child=True, **kwargs)
 
+        def find_column(data, names):
+            has_col = np.array([name in data.dtype.fields for name in names])
+            if any(has_col):
+                col = np.where(has_col)[0][0]
+                return data[names[col]]
+            return False
+
         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)
@@ -471,10 +483,11 @@ class RV:
             else:
                 data = np.array([], dtype=np.dtype([]))
 
-            if 'fwhm' in data.dtype.fields:
-                _s.fwhm = data['fwhm']
-                if 'sfwhm' in data.dtype.fields:
-                    _s.fwhm_err = data['sfwhm']
+            # try to find FWHM and uncertainty
+            if (v := find_column(data, ['fwhm'])) is not False:  # walrus !!
+                _s.fwhm = v
+                if (sv := find_column(data, ['sfwhm', 'fwhm_err', 'sig_fwhm'])) is not False:
+                    _s.fwhm_err = sv
                 else:
                     _s.fwhm_err = 2 * _s.svrad
             else:
@@ -484,12 +497,11 @@ class RV:
             _quantities.append('fwhm')
             _quantities.append('fwhm_err')
 
-            if 'rhk' in data.dtype.fields:
-                _s.rhk = data['rhk']
+            if (v := find_column(data, ['rhk'])) is not False:
+                _s.rhk = v
                 _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]
+                if (sv := find_column(data, ['srhk', 'rhk_err', 'sig_rhk'])) is not False:
+                    _s.rhk_err = sv
             else:
                 _s.rhk = np.zeros_like(time)
                 _s.rhk_err = np.full_like(time, np.nan)
@@ -516,6 +528,12 @@ class RV:
                 setattr(_s, q, np.full(time.size, True))
                 _quantities.append(q)
 
+            _s.extra_fields = ExtraFields()
+            for field in data.dtype.names:
+                if field not in _quantities:
+                    setattr(_s.extra_fields, field, data[field])
+                    # _quantities.append(field)
+
             #! end hack
 
             _s.mask = np.ones_like(time, dtype=bool)
@@ -543,31 +561,75 @@ class RV:
             logger.error('iCCF is not installed. Please install it with `pip install iCCF`')
             return
 
+        verbose = kwargs.get('verbose', True)
+
         if isinstance(files, str):
             files = [files]
 
-        I = iCCF.from_file(files)
+        CCFs = iCCF.from_file(files)
+
+        if not isinstance(CCFs, list):
+            CCFs = [CCFs]
 
-        objects = np.unique([i.HDU[0].header['OBJECT'].replace(' ', '') for i in I])
+        objects = np.unique([i.HDU[0].header['OBJECT'].replace(' ', '') for i in CCFs])
         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)
+        instruments = list(np.unique([i.instrument for i in CCFs]))
 
-        # 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])
+        for instrument in instruments:
+            # time, RVs, uncertainties
+            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.fwhm = np.array([i.FWHM*1e3 for i in I])
-        s.fwhm_err = np.array([i.FWHMerror*1e3 for i in I])
+            _quantities = []
 
-        # mask
-        s.mask = np.full_like(s.time, True, dtype=bool)
+            _s.fwhm = np.array([i.FWHM*1e3 for i in CCFs])
+            _s.fwhm_err = np.array([i.FWHMerror*1e3 for i in CCFs])
+
+            _quantities.append('fwhm')
+            _quantities.append('fwhm_err')
+
+            _s.contrast = np.array([i.contrast for i in CCFs])
+            _s.contrast_err = np.array([i.contrast_error for i in CCFs])
+
+            _quantities.append('contrast')
+            _quantities.append('contrast_err')
+
+            _s.texp = np.array([i.HDU[0].header['EXPTIME'] for i in CCFs])
+            _quantities.append('texp')
+
+            _s.date_night = np.array([
+                i.HDU[0].header['DATE-OBS'].split('T')[0] for i in CCFs
+            ])
+            _quantities.append('date_night')
+
+            _s.mask = np.full_like(_s.time, True, dtype=bool)
 
-        s.instruments = list(np.unique([i.instrument for i in I]))
+            _s.drs_qc = np.array([i.HDU[0].header['HIERARCH ESO QC SCIRED CHECK'] for i in CCFs], dtype=bool)
+            # mask out drs_qc = False
+            if not _s.drs_qc.all():
+                n = (~ _s.drs_qc).sum()
+                if verbose:
+                    logger.warning(f'masking {n} points where DRS QC failed for {instrument}')
+                _s.mask &= _s.drs_qc
+            print(_s.mask)
+
+            _s._quantities = np.array(_quantities)
+            setattr(s, instrument, _s)
+
+        s._child = False
+        s.instruments = instruments
+        s._build_arrays()
+
+        if instruments == ['ESPRESSO']:
+            from .instrument_specific import divide_ESPRESSO
+            divide_ESPRESSO(s)
 
         return s
 
@@ -610,7 +672,6 @@ class RV:
             logger.info(f'available: {self.instruments}')
             return
 
-
     def _build_arrays(self):
         """ build all concatenated arrays of `self` from each of the `.inst`s """
         if self._child:
@@ -653,7 +714,6 @@ class RV:
                 )
                 setattr(self, q, arr)
 
-
     def download_ccf(self, instrument=None, index=None, limit=None,
                      directory=None, symlink=False, **kwargs):
         """ Download CCFs from DACE
@@ -828,7 +888,7 @@ class RV:
             if self.verbose:
                 logger.info(f"Removed observations from '{instrument}'")
 
-        if return_self:
+        if config.return_self:
             return self
 
     def remove_condition(self, condition):
@@ -856,7 +916,7 @@ class RV:
         index = np.atleast_1d(index)
         try:
             instrument_index = self.obs[index]
-            instrument = np.array(self.instruments)[instrument_index - 1]
+            np.array(self.instruments)[instrument_index - 1]
         except IndexError:
             logger.errors(f'index {index} is out of bounds for N={self.N}')
             return
@@ -869,7 +929,7 @@ class RV:
         # for i, inst in zip(index, instrument):
         #     index_in_instrument = i - (self.obs < instrument_index).sum()
         #     getattr(self, inst).mask[index_in_instrument] = False
-        if return_self:
+        if config.return_self:
             return self
 
     def remove_non_public(self):
@@ -932,7 +992,7 @@ class RV:
         instruments = self._check_instrument(instrument)
         rng = np.random.default_rng(seed=seed)
         for inst in instruments:
-            s = getattr(self, inst)
+            # 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:
@@ -1037,13 +1097,18 @@ class RV:
             self.vrad = self.vrad - sa * (self.time - epoch) / 365.25
         else:
             for inst in self.instruments:
+                s = getattr(self, inst)
+
+                # if RVs come from a publication, don't remove the secular
+                # acceleration
+                if np.all(s.pub_reference != ''):
+                    continue
+
                 if 'HIRES' in inst:  # never remove it from HIRES...
                     continue
                 if 'NIRPS' in inst:  # never remove it from NIRPS...
                     continue
 
-                s = getattr(self, inst)
-
                 if hasattr(s, '_did_secular_acceleration') and s._did_secular_acceleration:
                     continue
 
@@ -1055,7 +1120,7 @@ class RV:
         self._did_secular_acceleration_epoch = epoch
         self._did_secular_acceleration_simbad = force_simbad
 
-        if return_self:
+        if config.return_self:
             return self
     
     def _undo_secular_acceleration(self):
@@ -1111,7 +1176,7 @@ class RV:
             #     # 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:
+            #     if config.return_self:
             #         return self
             #     continue
 
@@ -1123,7 +1188,7 @@ class RV:
             self._did_adjust_means = False
             self.adjust_means()
 
-        if return_self:
+        if config.return_self:
             return self
 
     def clip_maxerror(self, maxerror:float):
@@ -1145,7 +1210,7 @@ class RV:
             logger.warning(f'clip_maxerror ({maxerror} {self.units}) removed {n} point' + s)
 
         self._propagate_mask_changes()
-        if return_self:
+        if config.return_self:
             return self
 
     def bin(self):
@@ -1309,7 +1374,7 @@ class RV:
 
         self._build_arrays()
         self._did_adjust_means = True
-        if return_self:
+        if config.return_self:
             return self
 
     def add_to_vrad(self, values):
@@ -1546,6 +1611,8 @@ class RV:
 
 
     #
+    from .stellar import calc_prot_age
+
     @property
     def HZ(self):
         if not hasattr(self, 'star_mass'):

arvi/translations.py

@@ -1,8 +1,10 @@
 import re
 
 STARS = {
+    '_51Peg': '51 Peg',
     'Barnard': 'GJ699',
     "Barnard's": 'GJ699',
+    'Ross128': 'Ross  128',
 }
 
 

{arvi-0.1.13.dist-info → arvi-0.1.14.dist-info}/METADATA

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

{arvi-0.1.13.dist-info → arvi-0.1.14.dist-info}/RECORD

@@ -1,23 +1,26 @@
 arvi/HZ.py,sha256=u7rguhlILRBW-LOczlY3dkIB4LM8p8W7Xfg4FnNaYG0,2850
-arvi/__init__.py,sha256=cECMjLbPiolmft2lsnNoszU6_1a-dKOVY8cb2D1CBfg,366
+arvi/__init__.py,sha256=stLL3UtXmR1jpNEoMyVohjvkCatunz1w5ZK0ApmWuYQ,814
 arvi/ariadne_wrapper.py,sha256=jv8Wl35LfHl1UH1EklbvxHcQHqaEDhRGNogSYjFt7Y4,2141
+arvi/berv.py,sha256=5avwcmc2nkYH1KDo-z4eMPsS1ElNvold2DWkziV13gE,17633
 arvi/binning.py,sha256=jJ_resqSvfI-2BT0cnGPFckTIei_k2Mk95oZyE5b5zQ,15250
 arvi/config.py,sha256=wyj6FTxN7QOZj8LaHAe_ZdPKVfrNfobNNOHRNLH-S7Q,339
-arvi/dace_wrapper.py,sha256=VjC1YjZ7URkHkgRO7rUKjkxBtdMMjoEqcqTtZQ8M7NY,17034
+arvi/dace_wrapper.py,sha256=HmsyThieljcjorSSZgmGTZPGBqNrm9QoPwjoR_GpVKo,17025
 arvi/extra_data.py,sha256=WEEaYeLh52Zdv0uyHO72Ys5MWS3naTAP4wJV2BJ1mbk,2551
 arvi/gaia_wrapper.py,sha256=YxqqlWkLAKBY7AGHfDpNbkOLYWK0lt0L5GZenn0FfxM,3555
+arvi/headers.py,sha256=uvdJebw1M5YkGjE3vJJwYBOnLikib75uuZE9FXB5JJM,1673
 arvi/instrument_specific.py,sha256=fhkvR5jJzpJH7XbT8Fbzkaz6tGeoFYugkJCIAJgSR7o,4746
 arvi/lbl_wrapper.py,sha256=_ViGVkpakvuBR_xhu9XJRV5EKHpj5Go6jBZGJZMIS2Y,11850
 arvi/nasaexo_wrapper.py,sha256=mWt7eHgSZe4MBKCmUvMPTyUPGuiwGTqKugNBvmjOg9s,7306
-arvi/plots.py,sha256=2W1IwUeMdNq5--l6nQXn3EFft0UQgEyPyudRxTCRM40,25566
-arvi/programs.py,sha256=WpFE2cXYG-4ax2xmih0fFhvQbcVhnOEofVcwjePNmKQ,4505
+arvi/plots.py,sha256=53_rHVh7SeyBylM91PP6mNr7WTGmIoRnR8JfKJPdu7E,25594
+arvi/programs.py,sha256=C0Fbldjf-QEZYYJp5wBKP3h7zraD0O2mJC7Su967STg,4607
 arvi/reports.py,sha256=yrdajC-zz5_kH1Ucc6cU70DK-5dpG0Xyeru-EITKpNo,3355
 arvi/setup_logger.py,sha256=pBzaRTn0hntozjbaRVx0JIbWGuENkvYUApa6uB-FsRo,279
 arvi/simbad_wrapper.py,sha256=fwO0NbMvxAW3RyC58FPkLjXqLs9OkRyK0sYcb0ZVBQA,5571
 arvi/spectra.py,sha256=pTAWSW4vk96DWRQ-6l5mNJHUhiAyaPR-QDjZdOT6Ak0,7489
 arvi/stats.py,sha256=MQiyLvdiAFxIC29uajTy8kuxD-b2Y6mraL4AfWkRJkM,2576
-arvi/timeseries.py,sha256=PCQntiNK8MnlgKeyGOTd4BrQXm8NSjZRtGfLa4hfaE0,58257
-arvi/translations.py,sha256=AljjMsHRqzfcirdeavHsyVJFLvwKrM1wDLMv2RApG58,446
+arvi/stellar.py,sha256=MdO1_-ZXHnMHbkG-a5LvrnYeULy6MCAnWMKPvR4NZxQ,2954
+arvi/timeseries.py,sha256=sdm9l_LHItZHDAVI6r_wsQwkrfbuMX46q3Wqs0y-KGA,60787
+arvi/translations.py,sha256=SUIrJHt3JZdL_GQh3OJyg2Gm3X5ut86w5zW8hZpxHe0,498
 arvi/utils.py,sha256=nuIOuUdysMesWF9ute4v-kbeo6DylWjWW-SJhCsHn9I,4078
 arvi/data/info.svg,sha256=0IMI6W-eFoTD8acnury79WJJakpBwLa4qKS4JWpsXiI,489
 arvi/data/obs_affected_ADC_issues.dat,sha256=tn93uOL0eCTYhireqp1wG-_c3CbxPA7C-Rf-pejVY8M,10853
@@ -25,8 +28,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.1.13.dist-info/LICENSE,sha256=6JfQgl7SpM55t0EHMFNMnNh-AdkpGW25MwMiTnhdWQg,1068
-arvi-0.1.13.dist-info/METADATA,sha256=9wDBsoyGzFt-Eo186zVKLVJT_VdUlr_2C2uzrqQ5Fxs,1306
-arvi-0.1.13.dist-info/WHEEL,sha256=cpQTJ5IWu9CdaPViMhC9YzF8gZuS5-vlfoFihTBC86A,91
-arvi-0.1.13.dist-info/top_level.txt,sha256=4EeiKDVLD45ztuflTGfQ3TU8GVjJg5Y95xS5XjI-utU,5
-arvi-0.1.13.dist-info/RECORD,,
+arvi-0.1.14.dist-info/LICENSE,sha256=6JfQgl7SpM55t0EHMFNMnNh-AdkpGW25MwMiTnhdWQg,1068
+arvi-0.1.14.dist-info/METADATA,sha256=duTaG6amS-M9iVUNvfo2RFwnA_fz1gY1ysUB2N_yCaE,1306
+arvi-0.1.14.dist-info/WHEEL,sha256=Wyh-_nZ0DJYolHNn1_hMa4lM7uDedD_RGVwbmTjyItk,91
+arvi-0.1.14.dist-info/top_level.txt,sha256=4EeiKDVLD45ztuflTGfQ3TU8GVjJg5Y95xS5XjI-utU,5
+arvi-0.1.14.dist-info/RECORD,,

{arvi-0.1.13.dist-info → arvi-0.1.14.dist-info}/WHEEL

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