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

arvi/dace_wrapper.py

@@ -193,12 +193,13 @@ def get_observations_from_instrument(star, instrument, user=None, main_id=None,
         except TypeError:
             msg = f'no {instrument} observations for {star}'
             raise ValueError(msg) from None
+
     if (isinstance(instrument, str)):
         filters = {
             "ins_name": {"contains": [instrument]},
             "obj_id_daceid": {"contains": [dace_id]}
         }
-    elif (isinstance(instrument, list)):
+    elif (isinstance(instrument, (list, tuple, np.ndarray))):
         filters = {
             "ins_name": {"contains": instrument},
             "obj_id_daceid": {"contains": [dace_id]}
@@ -240,7 +241,9 @@ def get_observations_from_instrument(star, instrument, user=None, main_id=None,
                     'rv': result['spectro_ccf_rv'][mask3],
                     'rv_err': result['spectro_ccf_rv_err'][mask3],
                     'berv': result['spectro_cal_berv'][mask3],
-                    'ccf_noise': _nan,
+                    'ccf_noise': np.sqrt(
+                        np.square(result['spectro_ccf_rv_err'][mask3]) - np.square(result['spectro_cal_drift_noise'][mask3])
+                    ),
                     'rhk': result['spectro_analysis_rhk'][mask3],
                     'rhk_err': result['spectro_analysis_rhk_err'][mask3],
                     'contrast': result['spectro_ccf_contrast'][mask3],

arvi/instrument_specific.py

@@ -123,7 +123,10 @@ def HARPS_commissioning(self, mask=True, plot=True):
     if check(self, 'HARPS') is None:
         return
 
-    affected = self.time < HARPS_start
+    affected = np.logical_and(
+        self.instrument_array == 'HARPS03', 
+        self.time < HARPS_start
+    )
     total_affected = affected.sum()
 
     if self.verbose:
@@ -133,7 +136,7 @@ def HARPS_commissioning(self, mask=True, plot=True):
 
     if mask:
         self.mask[affected] = False
-        self._propagate_mask_changes()
+        self._propagate_mask_changes(_remove_instrument=False)
 
         if plot:
             self.plot(show_masked=True)
@@ -155,7 +158,14 @@ def HARPS_fiber_commissioning(self, mask=True, plot=True):
     if check(self, 'HARPS') is None:
         return
 
-    affected = (self.time >= HARPS_technical_intervention_range[0]) & (self.time <= HARPS_technical_intervention_range[1])
+    affected = np.logical_and(
+        self.time >= HARPS_technical_intervention_range[0],
+        self.time <= HARPS_technical_intervention_range[1]
+    )
+    affected = np.logical_and(
+        affected,
+        np.char.find(self.instrument_array, 'HARPS') == 0
+    )
     total_affected = affected.sum()
 
     if self.verbose:
@@ -165,7 +175,7 @@ def HARPS_fiber_commissioning(self, mask=True, plot=True):
 
     if mask:
         self.mask[affected] = False
-        self._propagate_mask_changes()
+        self._propagate_mask_changes(_remove_instrument=False)
 
         if plot:
             self.plot(show_masked=True)
@@ -187,7 +197,10 @@ def ESPRESSO_commissioning(self, mask=True, plot=True):
     if check(self, 'ESPRESSO') is None:
         return
 
-    affected = self.time < ESPRESSO_start
+    affected = np.logical_and(
+        self.instrument_array == 'ESPRESSO18',
+        self.time < ESPRESSO_start
+    )
     total_affected = affected.sum()
 
     if self.verbose:
@@ -197,7 +210,7 @@ def ESPRESSO_commissioning(self, mask=True, plot=True):
 
     if mask:
         self.mask[affected] = False
-        self._propagate_mask_changes()
+        self._propagate_mask_changes(_remove_instrument=False)
 
         if plot and total_affected > 0:
             self.plot(show_masked=True)
@@ -246,7 +259,7 @@ def ADC_issues(self, mask=True, plot=True, check_headers=False):
 
     if mask:
         self.mask[intersect] = False
-        self._propagate_mask_changes()
+        self._propagate_mask_changes(_remove_instrument=False)
 
         if plot:
             self.plot(show_masked=True)
@@ -282,7 +295,7 @@ def blue_cryostat_issues(self, mask=True, plot=True):
 
     if mask:
         self.mask[intersect] = False
-        self._propagate_mask_changes()
+        self._propagate_mask_changes(_remove_instrument=False)
 
         if plot:
             self.plot(show_masked=True)
@@ -330,7 +343,7 @@ def qc_scired_issues(self, plot=False, **kwargs):
         return
 
     self.mask[affected] = False
-    self._propagate_mask_changes()
+    self._propagate_mask_changes(_remove_instrument=False)
 
     if plot:
         self.plot(show_masked=True)
@@ -364,6 +377,7 @@ class ISSUES:
                 logger.error('are the data binned? cannot proceed to mask these points...')
         
         results = list(filter(lambda x: x is not None, results))
+        self._propagate_mask_changes()
 
         try:
             return np.logical_or.reduce(results)

arvi/kepmodel_wrapper.py

@@ -0,0 +1,296 @@
+import io
+from contextlib import redirect_stdout, contextmanager
+from string import ascii_lowercase
+from matplotlib import pyplot as plt
+import numpy as np
+
+from kepmodel.rv import RvModel
+from spleaf.cov import merge_series
+from spleaf.term import Error, InstrumentJitter
+
+from .setup_logger import setup_logger
+from .utils import timer, adjust_lightness
+
+
+class model:
+    logger = setup_logger()
+    # periodogram settings
+    Pmin, Pmax, nfreq = 1.5, 10_000, 100_000
+
+    @contextmanager
+    def ew(self):
+        for name in self.model.keplerian:
+            self.model.set_keplerian_param(name, param=['P', 'la0', 'K', 'e', 'w'])
+        try:
+            yield
+        finally:
+            for name in self.model.keplerian:
+                self.model.set_keplerian_param(name, param=['P', 'la0', 'K', 'esinw', 'ecosw'])
+
+    @property
+    def nu0(self):
+        return 2 * np.pi / self.Pmax
+    
+    @property
+    def dnu(self):
+        return (2 * np.pi / self.Pmin - self.nu0) / (self.nfreq - 1)
+
+    def __init__(self, s):
+        self.s = s
+        self.instruments = s.instruments
+        self.Pmax = 2 * np.ptp(s.time)
+        ts = self.ts = self.s._mtime_sorter
+
+        # t, y_ ye, series_index = merge_series(
+        #     [_s.mtime for _s in s],
+        #     [_s.mvrad for _s in s],
+        #     [_s.msvrad for _s in s],
+        # )
+
+        inst_jit = self._get_jitters()
+
+        self.model = RvModel(self.s.mtime[ts], self.s.mvrad[ts],
+                             err=Error(self.s.msvrad[ts]), **inst_jit)
+        self.np = 0
+        self._add_means()
+
+    def _add_means(self):
+        for inst in self.s.instruments:
+            # if inst not in self.s.instrument_array[self.s.mask]:
+            #     continue
+            mask = self.s.instrument_array[self.s.mask][self.ts] == inst
+            self.model.add_lin(
+                derivative=1.0 * mask,
+                name=f"offset_{inst}",
+                value=getattr(self.s, inst).mvrad.mean(),
+            )
+        self.model.fit_lin()
+    
+    def _get_jitters(self):
+        inst_jit = {}
+        for inst in self.s.instruments:
+            inst_jit[f'jit_{inst}'] = InstrumentJitter(
+                indices=self.s.instrument_array[self.s.mask] == inst,
+                sig=self.s.svrad[self.s.mask].min()
+            )
+        return inst_jit
+
+    def _set_jitters(self, value=0.0):
+        for par in self.model.cov.param:
+            if 'jit' in par:
+                self.model.set_param(value, f'cov.{par}')
+        # self.model.fit()
+    
+    def _equal_coralie_offsets(self):
+        if self.s._check_instrument('CORALIE') is None:
+            return
+        mask = np.char.find(self.s.instrument_array, 'CORALIE') == 0
+        mean = self.s.vrad[mask].mean()
+        for inst in self.instruments:
+            if 'CORALIE' in inst:
+                self.model.set_param(mean, f'lin.offset_{inst}')
+
+
+    def __repr__(self):
+        with self.ew():
+            with io.StringIO() as buf, redirect_stdout(buf):
+                self.model.show_param()
+                output = buf.getvalue()
+        return output
+
+    def to_table(self, **kwargs):
+        from .utils import pretty_print_table
+        lines = repr(self)
+        lines = lines.replace(' [deg]', '_[deg]')
+        lines = lines.encode().replace(b'\xc2\xb1', b'').decode()
+        lines = lines.split('\n')
+        lines = [line.split() for line in lines]
+        lines = [[col.replace('_[deg]', ' [deg]') for col in line] for line in lines]
+        pretty_print_table(lines[:-2], **kwargs)
+
+    @property
+    def fit_param(self):
+        return self.model.fit_param
+
+    def fit(self):
+        # fit offsets
+        self.model.fit_param = [f'lin.offset_{inst}' for inst in self.instruments]
+        # fit jitters
+        self.model.fit_param += [f'cov.{par}' for par in self.model.cov.param]
+        # fit keplerian(s)
+        self.model.fit_param += [
+            f'kep.{k}.{p}' 
+            for k, v in self.model.keplerian.items() 
+            for p in v._param
+        ]
+        # if self.np == 0:
+        #     self._set_jitters(0.1 * np.std(self.model.y - self.model.model()))
+        try:
+            self.model.fit()
+        except Exception as e:
+            print(e)
+            
+
+    def plot(self, **kwargs):
+        fig, ax = self.s.plot(**kwargs)
+        tt = self.s._tt()
+        time_offset = 50000 if 'remove_50000' in kwargs else 0
+
+        for i, inst in enumerate(self.s):
+            inst_name = inst.instruments[0].replace('-', '_')
+            val = self.model.get_param(f'lin.offset_{inst_name}')
+            x = np.array([inst.mtime.min(), inst.mtime.max()]) - time_offset
+            y = [val, val]
+            ax.plot(x, y, ls='--', color=f'C{i}')
+            mask = (tt > inst.mtime.min()) & (tt < inst.mtime.max())
+            color = adjust_lightness(f'C{i}', 1.2)
+            ax.plot(tt[mask] - time_offset,
+                    val + self.model.keplerian_model(tt)[mask],
+                    color=color)
+
+        return fig, ax
+
+    def plot_phasefolding(self, planets=None, ax=None):
+        t = self.model.t
+        res = self.model.residuals()
+        sig = np.sqrt(self.model.cov.A)
+
+        Msmooth = np.linspace(0, 360, 1000)
+
+        if planets is None:
+            planets = list(self.model.keplerian.keys())
+
+        if ax is None:
+            fig, axs = plt.subplots(
+                1, len(planets), sharex=True, sharey=True, constrained_layout=True,
+                squeeze=False
+            )
+        else:
+            axs = np.atleast_1d(ax)
+            fig = axs[0].figure
+
+        for p, ax in zip(planets, axs.flat):
+            self.model.set_keplerian_param(p, param=['P', 'la0', 'K', 'e', 'w'])
+            kep = self.model.keplerian[p]
+            P = self.model.get_param(f'kep.{p}.P')
+            M0 = (180 / np.pi * (self.model.get_param(f'kep.{p}.la0') - self.model.get_param(f'kep.{p}.w')))
+            M = (M0 + 360 / P * t) % 360
+            reskep = res + kep.rv(t)
+            tsmooth = (Msmooth - M0) * P / 360
+            mod = kep.rv(tsmooth)
+
+            ax.plot([0, 360], [0, 0], '--', c='gray', lw=1)
+            ax.plot(Msmooth, mod, 'k-', lw=3, rasterized=True)
+            for inst in self.instruments:
+                sel = self.s.instrument_array[self.s.mask] == inst
+                ax.errorbar(M[sel], reskep[sel], sig[sel], fmt='.',
+                            rasterized=True, alpha=0.7)
+
+            ax.set(ylabel='RV [m/s]', xlabel='Mean anomaly [deg]',
+                   xticks=np.arange(0, 361, 90))
+            ax.minorticks_on()
+            self.model.set_keplerian_param(p, param=['P', 'la0', 'K', 'esinw', 'ecosw'])
+
+        # handles, labels = ax.get_legend_handles_labels()
+        # fig.legend(handles, labels, loc='center left', bbox_to_anchor=(0.9, 0.5))
+        return fig, axs
+
+    def add_planet_from_period(self, period):
+        self.model.add_keplerian_from_period(period, fit=True)
+        self.model.fit()
+        self.np += 1
+
+    def _plot_periodogram(self, P=None, power=None, kmax=None, faplvl=None,
+                          **kwargs):
+        if P is None and power is None:
+            with timer('periodogram'):
+                nu, power = self.model.periodogram(self.nu0, self.dnu, self.nfreq)
+                P = 2 * np.pi / nu
+
+        if 'ax' in kwargs:
+            ax = kwargs.pop('ax')
+            fig = ax.figure
+        else:
+            fig, ax = plt.subplots(1, 1, constrained_layout=True)
+        ax.semilogx(P, power, 'k', lw=1, rasterized=True)
+        ax.set_ylim(0, 1.2 * power.max())
+        ax.set(xlabel='Period [days]', ylabel='Normalized power')
+
+        if kmax is None:
+            kmax = np.argmax(power)
+        ax.plot(P[kmax], power[kmax], 'or', ms=4)
+        ax.text(P[kmax], power[kmax] * 1.1, f'{P[kmax]:.3f} d',
+                ha='right', va='center', color='r')
+
+        if faplvl is None:
+            faplvl = self.model.fap(power[kmax], nu.max())
+        ax.text(0.99, 0.95, f'FAP = {faplvl:.2g}', transform=ax.transAxes,
+                ha='right', va='top')
+
+        return fig, ax
+
+    def add_keplerian_from_periodogram(self, fap_max=0.001, plot=False,
+                                       fit_first=True):
+        if fit_first and self.np == 0:
+            self.fit()
+
+        self._equal_coralie_offsets()
+
+        with timer('periodogram'):
+            nu, power = self.model.periodogram(self.nu0, self.dnu, self.nfreq)
+
+        P = 2 * np.pi / nu
+        # Compute FAP
+        kmax = np.argmax(power)
+        faplvl = self.model.fap(power[kmax], nu.max())
+        self.logger.info('highest periodogram peak:')
+        self.logger.info(f'P={P[kmax]:.4f} d, power={power[kmax]:.3f}, FAP={faplvl:.2e}')
+        if plot:
+            self._plot_periodogram(P, power, kmax, faplvl)
+
+        if faplvl > fap_max:
+            print('non-significant peak')
+            self.fit()
+            return False
+
+        # add new planet
+        letter = ascii_lowercase[1:][self.np]
+        self.model.add_keplerian_from_period(P[kmax], name=letter,
+                                             guess_kwargs={'emax': 0.8})
+        # self.model.set_keplerian_param(letter, param=['P', 'la0', 'K', 'e', 'w'])
+        self.model.set_keplerian_param(letter, param=['P', 'la0', 'K', 'esinw', 'ecosw'])
+        self.np += 1
+        self.fit()
+        
+        if plot:
+            self.plot()
+
+        return True
+    
+    @property
+    def offsets(self):
+        names = [f'lin.offset_{inst}' for inst in self.instruments]
+        return {
+            name.replace('lin.', ''): self.model.get_param(name)
+            for name in names
+        }
+
+    @property
+    def jitters(self):
+        names = [f'cov.{par}' for par in self.model.cov.param]
+        return {
+            name.replace('cov.', '').replace('.sig', ''): self.model.get_param(name)
+            for name in names
+        }
+    
+    @property
+    def keplerians(self):
+        keps = {name: {} for name in self.model.keplerian.keys()}
+        for name in keps:
+            params = self.model.keplerian[name]._param
+            pars = [f'kep.{name}.{p}' for p in params]
+            keps[name] = {
+                par.replace(f'kep.{name}.', ''): self.model.get_param(par)
+                for par in pars
+            }
+        return keps

arvi/nasaexo_wrapper.py

@@ -6,7 +6,7 @@ from io import StringIO
 import numpy as np
 from astropy.timeseries import LombScargle
 
-from .setup_logger import logger
+from .setup_logger import setup_logger
 from kepmodel.rv import RvModel
 from spleaf.term import Error
 
@@ -32,6 +32,7 @@ def run_query(query):
 
 class Planets:
     def __init__(self, system):
+        logger = setup_logger()
         self.s = system
         self.verbose = system.verbose
 
@@ -163,6 +164,7 @@ class Planets:
         self.model.show_param()
 
     def fit_all(self, adjust_data=False):
+        logger = setup_logger()
         self.model.fit()
 
         newP = np.array([self.model.get_param(f'kep.{i}.P') for i in range(self.np)])
@@ -187,5 +189,7 @@ class Planets:
             self.s._build_arrays()
 
     def __repr__(self):
-        return f'{self.star}({self.np} planets, '\
-               f'P={list(self.P)}, K={list(self.K)}, e={list(self.e)})'
+        P = list(map(float, self.P))
+        K = list(map(float, self.K))
+        e = list(map(float, self.e))
+        return f'{self.star}({self.np} planets, {P=}, {K=}, {e=})'

arvi/reports.py

@@ -40,6 +40,7 @@ class REPORTS:
             rows.append([self.star] + [''] * len(self.instruments) + [''])
             rows.append([''] + self.instruments + ['full'])
             rows.append(['N'] + list(self.NN.values()) + [self.N])
+            rows.append(['T span'] + [np.ptp(s.mtime).round(1) for s in self] + [np.ptp(self.mtime).round(1)])
             rows.append(['RV span'] + [np.ptp(s.mvrad).round(3) for s in self] + [np.ptp(self.mvrad).round(3)])
             rows.append(['RV std'] + [s.mvrad.std().round(3) for s in self] + [self.mvrad.std().round(3)])
             rows.append(['eRV mean'] + [s.msvrad.mean().round(3) for s in self] + [self.msvrad.mean().round(3)])
@@ -201,4 +202,110 @@ class REPORTS:
                     pdf.savefig(fig)
                 # os.system(f'evince {save} &')
 
-        return fig
+        return fig
+
+
+def kepmodel_report(self, fit_keplerians=3, save=None, nasaexo_title=False):
+    import matplotlib.pyplot as plt
+    import matplotlib.gridspec as gridspec
+    from matplotlib.backends.backend_pdf import PdfPages
+    logger = setup_logger()
+
+    def set_align_for_column(table, col, align="left"):
+        cells = [key for key in table._cells if key[1] == col]
+        for cell in cells:
+            table._cells[cell]._loc = align
+            table._cells[cell]._text.set_horizontalalignment(align) 
+
+    from .kepmodel_wrapper import model
+    m = model(self)
+
+    while fit_keplerians > 0:
+        if m.add_keplerian_from_periodogram():
+            fit_keplerians -= 1
+        else:
+            break
+
+    m.fit()
+
+
+    # size = A4
+    size = 8.27, 11.69
+    fig = plt.figure(figsize=size, constrained_layout=True)
+    gs = gridspec.GridSpec(5, 3, figure=fig, height_ratios=[2, 1, 1, 1, 1])
+
+    # first row, all columns
+    ax1 = plt.subplot(gs[0, :])
+
+    if nasaexo_title:
+        title = str(self.planets).replace('(', '\n').replace(')', '')
+        star, planets = title.split('\n')
+        planets = planets.replace('planets,', 'known planets\n')
+        ax1.set_title(star, loc='left', fontsize=14)
+        ax1.set_title(planets, loc='right', fontsize=10)
+    else:
+        title = f'{self.star}'
+        ax1.set_title(title, loc='left', fontsize=14)
+    # ax1.set_title(r"\href{http://www.google.com}{link}", color='blue',
+    #               loc='center')
+
+    m.plot(ax=ax1, N_in_label=True, tooltips=False, remove_50000=True)
+
+    ax1.legend().remove()
+    legend_ax = plt.subplot(gs[1, -1])
+    legend_ax.axis('off')
+    leg = plt.legend(*ax1.get_legend_handles_labels(),
+                    prop={'family': 'monospace'})
+    legend_ax.add_artist(leg)
+    
+    ax2 = plt.subplot(gs[1, :-1])
+    m._plot_periodogram(ax=ax2)
+
+    ax3 = plt.subplot(gs[2, 0])
+    ax3.axis('off')
+    items = list(m.offsets.items())
+    items = [[item[0].replace('offset_', 'offet '), item[1].round(3)] for item in items]
+    table = ax3.table(items, loc='center', edges='open')
+    table.auto_set_font_size(False)
+    table.set_fontsize(9)
+    set_align_for_column(table, 1, align="left")
+
+    ax4 = plt.subplot(gs[2, 1])
+    ax4.axis('off')
+    items = list(m.jitters.items())
+    items = [[item[0].replace('jit_', 'jitter '), item[1].round(3)] for item in items]
+    table = ax4.table(items, loc='center', edges='open')
+    table.auto_set_font_size(False)
+    table.set_fontsize(9)
+    set_align_for_column(table, 1, align="left")
+
+    ax5 = plt.subplot(gs[2, 2])
+    ax5.axis('off')
+    items = [
+        ['N', m.model.n],
+        [r'N$_{\rm free}$', len(m.model.fit_param)],
+        [r'$\chi^2$', round(m.model.chi2(), 2)],
+        [r'$\chi^2_r$', round(m.model.chi2() / (m.model.n - len(m.model.fit_param)), 2)],
+        [r'$\log L$', round(m.model.loglike(), 2)],
+    ]
+    table = ax5.table(items, loc='center', edges='open')
+    table.auto_set_font_size(False)
+    table.set_fontsize(9)
+    set_align_for_column(table, 1, align="left")
+
+    for i, name in enumerate(m.keplerians):
+        ax = plt.subplot(gs[3, i])
+        m.plot_phasefolding(planets=name, ax=ax)
+
+        ax = plt.subplot(gs[4, i])
+        ax.axis('off')
+        with m.ew():
+            items = list(m.keplerians[name].items())
+        items = [[item[0], item[1].round(3)] for item in items]
+        table = ax.table(items, loc='center', edges='open')
+        table.auto_set_font_size(False)
+        table.set_fontsize(9)
+        set_align_for_column(table, 1, align="left")
+
+
+    return fig, m

arvi/stats.py

@@ -1,4 +1,6 @@
+from functools import partial
 import numpy as np
+from scipy.stats import norm
 
 def wmean(a, e):
     """Weighted mean of array `a`, with uncertainties given by `e`.
@@ -50,16 +52,39 @@ def wrms(a, e, ignore_nans=False):
     w = 1 / e**2
     return np.sqrt(np.sum(w * (a - np.average(a, weights=w))**2) / sum(w))
 
+# from https://stackoverflow.com/questions/20601872/numpy-or-scipy-to-calculate-weighted-median
+def weighted_quantiles_interpolate(values, weights, quantiles):
+    i = np.argsort(values)
+    c = np.cumsum(weights[i])
+    q = np.searchsorted(c, quantiles * c[-1])
+    # Ensure right-end isn't out of bounds. Thanks @Jeromino!
+    q_plus1 = np.clip(q + 1, a_min=None, a_max=values.shape[0] - 1)
+    return np.where(
+        c[q] / c[-1] == quantiles,
+        0.5 * (values[i[q]] + values[i[q_plus1]]),
+        values[i[q]],
+    )
 
-def sigmaclip_median(a, low=4.0, high=4.0):
+weighted_median = partial(weighted_quantiles_interpolate, quantiles=0.5)
+
+
+
+def sigmaclip_median(a, low=4.0, high=4.0, k=1/norm.ppf(3/4)):
     """
     Same as scipy.stats.sigmaclip but using the median and median absolute
     deviation instead of the mean and standard deviation.
 
     Args:
-        a (array): Array containing data
-        low (float): Number of MAD to use for the lower clipping limit
-        high (float): Number of MAD to use for the upper clipping limit
+        a (array):
+            Array containing data
+        low (float):
+            Number of MAD to use for the lower clipping limit
+        high (float):
+            Number of MAD to use for the upper clipping limit
+        k (float):
+            Scale factor for the MAD to be an estimator of the standard
+            deviation. Depends on the (assumed) distribution of the data.
+            Default value is for the normal distribution (=1/norm.ppf(3/4)).
     Returns:
         SigmaclipResult: Object with the following attributes:
             - `clipped`: Masked array of data
@@ -71,7 +96,7 @@ def sigmaclip_median(a, low=4.0, high=4.0):
     c = np.asarray(a).ravel()
     delta = 1
     while delta:
-        c_mad = median_abs_deviation(c)
+        c_mad = median_abs_deviation(c) * k
         c_median = np.median(c)
         size = c.size
         critlower = c_median - c_mad * low

arvi/timeseries.py

@@ -4,10 +4,9 @@ from typing import Union
 from functools import partial, partialmethod
 from glob import glob
 import warnings
-from copy import deepcopy
+from copy import copy, deepcopy
 from datetime import datetime, timezone
 
-# import lazy_loader as lazy
 import numpy as np
 
 from .setup_logger import setup_logger
@@ -26,7 +25,8 @@ from .HZ import getHZ_period
 from .instrument_specific import ISSUES
 from .reports import REPORTS
 from .utils import sanitize_path, strtobool, there_is_internet, timer, chdir
-# from .utils import lazy_import
+from .setup_logger import setup_logger
+logger = setup_logger()
 
 # units = lazy_import('astropy.units')
 # units = lazy.load('astropy.units')
@@ -411,35 +411,71 @@ class RV(ISSUES, REPORTS):
         self._did_correct_berv = False
         self.__post_init__()
 
-    def snapshot(self, directory=None, delete_others=False):
-        import pickle
+    def snapshot(self, directory=None, delete_others=False, compress=False):
+        if compress:
+            try:
+                import compress_pickle as pickle
+            except ImportError:
+                logger.warning('compress_pickle not installed, not compressing')
+                import pickle
+                compress = False
+        else:
+            import pickle
+        import re
         from datetime import datetime
+
         ts = datetime.now().timestamp()
         star_name = self.star.replace(' ', '')
         file = f'{star_name}_{ts}.pkl'
 
+        server = None
         if directory is None:
             directory = '.'
         else:
-            os.makedirs(directory, exist_ok=True)
-
-        file = os.path.join(directory, file)
-
-        if delete_others:
-            import re
-            other_pkls = [
-                f for f in os.listdir(directory)
-                if re.search(fr'{star_name}_\d+.\d+.pkl', f)
-            ]
-            for pkl in other_pkls:
-                os.remove(os.path.join(directory, pkl))
+            if ':' in directory:
+                server, directory = directory.split(':')
+                delete_others = False
+            else:
+                os.makedirs(directory, exist_ok=True)
 
         metadata = {
             'star': self.star,
             'timestamp': ts,
             'description': 'arvi snapshot'
         }
-        pickle.dump((self, metadata), open(file, 'wb'), protocol=0)
+
+
+        if server:
+            import posixpath
+            from .utils import server_sftp, server_file
+            with server_sftp(server=server) as sftp:
+                try:
+                    sftp.chdir(directory)
+                except FileNotFoundError:
+                    sftp.mkdir(directory)
+                finally:
+                    sftp.chdir(directory)
+                with sftp.open(file, 'wb') as f:
+                    print('saving snapshot to server...', end='', flush=True)
+                    pickle.dump((self, metadata), f, protocol=0)
+                    print('done')
+            file = posixpath.join(directory, file)
+        else:
+            if delete_others:
+                other_pkls = [
+                    f for f in os.listdir(directory)
+                    if re.search(fr'{star_name}_\d+.\d+.pkl', f)
+                ]
+                for pkl in other_pkls:
+                    os.remove(os.path.join(directory, pkl))
+
+            file = os.path.join(directory, file)
+
+            if compress:
+                file += '.gz'
+
+            with open(file, 'wb') as f:
+                pickle.dump((self, metadata), f)
 
         if self.verbose:
             logger.info(f'saved snapshot to {file}')
@@ -514,6 +550,15 @@ class RV(ISSUES, REPORTS):
     def instrument_array(self):
         return np.concatenate([[i] * n for i, n in self.NN.items()])
 
+    def _instrument_mask(self, instrument):
+        if isinstance(instrument, str):
+            return np.char.find(self.instrument_array, instrument) == 0
+        elif isinstance(instrument, (list, tuple, np.ndarray)):
+            m = np.full_like(self.time, False, dtype=bool)
+            for i in instrument:
+                m |= np.char.find(self.instrument_array, i) == 0
+            return m
+
     @property
     def rms(self) -> float:
         """ Weighted rms of the (masked) radial velocities """
@@ -540,6 +585,11 @@ class RV(ISSUES, REPORTS):
     def _mtime_sorter(self):
         return np.argsort(self.mtime)
 
+    @property
+    def timespan(self):
+        """ Total time span of the (masked) observations """
+        return np.ptp(self.mtime)
+
     def _index_from_instrument_index(self, index, instrument):
         ind = np.where(self.instrument_array == instrument)[0]
         return ind[getattr(self, instrument).mask][index]
@@ -633,22 +683,28 @@ class RV(ISSUES, REPORTS):
         import pickle
         from datetime import datetime
         if star is None:
-            assert file.endswith('.pkl'), 'expected a .pkl file'
-            star, timestamp = file.replace('.pkl', '').split('_')
+            assert file.endswith(('.pkl', '.pkl.gz')), 'expected a .pkl file'
+            basefile = os.path.basename(file)
+            star, timestamp = basefile.replace('.pkl.gz', '').replace('.pkl', '').split('_')
         else:
             try:
-                file = sorted(glob(f'{star}_*.*.pkl'))[-1]
+                file = sorted(glob(f'{star}_*.*.pkl*'))[-1]
             except IndexError:
                 raise ValueError(f'cannot find any file matching {star}_*.pkl')
-            star, timestamp = file.replace('.pkl', '').split('_')
+            star, timestamp = file.replace('.pkl.gz', '').replace('.pkl', '').split('_')
 
         dt = datetime.fromtimestamp(float(timestamp))
         if verbose:
             logger.info(f'reading snapshot of {star} from {dt}')
 
-        s = pickle.load(open(file, 'rb'))
+        with open(file, 'rb') as f:
+            if file.endswith('.gz'):
+                import compress_pickle as pickle
+            s = pickle.load(f)
+
         if isinstance(s, tuple) and len(s) == 2:
             s, _metadata = s
+
         s._snapshot = file
         return s
 
@@ -1508,7 +1564,7 @@ class RV(ISSUES, REPORTS):
         """ Remove all observations that satisfy a condition
 
         Args:
-            condition (np.ndarray):
+            condition (ndarray):
                 Boolean array of the same length as the observations
         """
         if self.verbose:
@@ -1668,16 +1724,18 @@ class RV(ISSUES, REPORTS):
         self._propagate_mask_changes()
 
 
-    def _propagate_mask_changes(self):
+    def _propagate_mask_changes(self, _remove_instrument=True):
         """ link self.mask with each self.`instrument`.mask """
         masked = np.where(~self.mask)[0]
         for m in masked:
             inst = self.instruments[self.obs[m] - 1]
             n_before = (self.obs < self.obs[m]).sum()
             getattr(self, inst).mask[m - n_before] = False
-        for inst in self.instruments:
-            if getattr(self, inst).mtime.size == 0:
-                self.remove_instrument(inst, strict=True)
+        if _remove_instrument:
+            instruments = copy(self.instruments)
+            for inst in instruments:
+                if getattr(self, inst).mtime.size == 0:
+                    self.remove_instrument(inst, strict=True)
 
     def secular_acceleration(self, epoch=None, just_compute=False, force_simbad=False):
         """ 
@@ -1695,9 +1753,12 @@ class RV(ISSUES, REPORTS):
             force_simbad (bool, optional):
                 Use Simbad proper motions even if Gaia is available
         """
-        if self._did_secular_acceleration and not just_compute:  # don't do it twice
+        # don't do it twice
+        if self._did_secular_acceleration and not just_compute:
             return
 
+        from astropy import units
+
         #as_yr = units.arcsec / units.year
         mas_yr = units.milliarcsecond / units.year
         mas = units.milliarcsecond
@@ -1829,15 +1890,21 @@ class RV(ISSUES, REPORTS):
 
             self._did_secular_acceleration = False
 
-    def sigmaclip(self, sigma=5, instrument=None, strict=True):
+    def sigmaclip(self, sigma=5, quantity='vrad', instrument=None,
+                  strict=True):
         """
-        Sigma-clip RVs (per instrument!), by MAD away from the median.
+        Sigma-clip RVs or other quantities (per instrument!), by MAD away from
+        the median.
 
         Args:
             sigma (float):
-                Number of MADs to clip
+                Number of MADs away from the median
+            quantity (str):
+                Quantity to sigma-clip (by default the RVs)
             instrument (str, list):
                 Instrument(s) to sigma-clip
+            strict (bool):
+                Passed directly to self._check_instrument
         """
         #from scipy.stats import sigmaclip as dosigmaclip
         from .stats import sigmaclip_median as dosigmaclip
@@ -1846,20 +1913,26 @@ class RV(ISSUES, REPORTS):
             return
 
         instruments = self._check_instrument(instrument, strict)
+        if instruments is None:
+            return
         changed_instruments = []
 
         for inst in instruments:
             m = self.instrument_array == inst
-            result = dosigmaclip(self.vrad[m], low=sigma, high=sigma)
+            d = getattr(self, quantity)
+
+            if np.isnan(d[m]).all():
+                continue
+
+            result = dosigmaclip(d[m], low=sigma, high=sigma)
             # n = self.vrad[m].size - result.clipped.size
 
-            ind = m & self.mask & \
-                  ((self.vrad < result.lower) | (self.vrad > result.upper))
+            ind = m & self.mask & ((d < result.lower) | (d > result.upper))
             n = ind.sum()
 
             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}')
+                logger.warning(f'sigma-clip {quantity} will remove {n} point{s} for {inst}')
 
             if n > 0:
                 self.mask[ind] = False
@@ -1884,21 +1957,32 @@ class RV(ISSUES, REPORTS):
         if config.return_self:
             return self
 
-    def clip_maxerror(self, maxerror:float):
-        """ Mask out points with RV error larger than a given value
+    def clip_maxerror(self, maxerror:float, instrument=None):
+        """
+        Mask out points with RV error larger than a given value. If `instrument`
+        is given, mask only observations from that instrument.
 
         Args:
             maxerror (float): Maximum error to keep.
+            instrument (str, list, tuple, ndarray): Instrument(s) to clip
         """
         if self._child:
             return
 
         self.maxerror = maxerror
+
+        if instrument is None:
+            inst_mask = np.ones_like(self.svrad, dtype=bool)
+        else:
+            inst_mask = self._instrument_mask(instrument)
+        
         above = self.svrad > maxerror
-        n = above.sum()
-        self.mask[above] = False
+        old_mask = self.mask.copy()
+
+        self.mask[inst_mask & above] = False
 
         if self.verbose and above.sum() > 0:
+            n = (above[inst_mask] & old_mask[inst_mask]).sum()
             s = 's' if (n == 0 or n > 1) else ''
             logger.warning(f'clip_maxerror ({maxerror} {self.units}) removed {n} point' + s)
 
@@ -1906,6 +1990,36 @@ class RV(ISSUES, REPORTS):
         if config.return_self:
             return self
 
+    def sigmaclip_ew(self, sigma=5):
+        """ Sigma-clip EW (FWHM x contrast), by MAD away from the median """
+        from .stats import sigmaclip_median as dosigmaclip, weighted_median
+
+        S = deepcopy(self)
+        for _s in S:
+            m = _s.mask
+            _s.fwhm -= weighted_median(_s.fwhm[m], 1 / _s.fwhm_err[m])
+            _s.contrast -= weighted_median(_s.contrast[m], 1 / _s.contrast_err[m])
+        S._build_arrays()
+        ew = S.fwhm * S.contrast
+        ew_err = np.hypot(S.fwhm_err * S.contrast, S.fwhm * S.contrast_err)
+
+        wmed = weighted_median(ew[S.mask], 1 / ew_err[S.mask])
+        data = (ew - wmed) / ew_err
+        result = dosigmaclip(data, low=sigma, high=sigma)
+        ind = (data < result.lower) | (data > result.upper)
+        self.mask[ind] = False
+
+        if self.verbose and ind.sum() > 0:
+            n = ind.sum()
+            s = 's' if (n == 0 or n > 1) else ''
+            logger.warning(f'sigmaclip_ew removed {n} point' + s)
+
+        self._propagate_mask_changes()
+        if config.return_self:
+            return self
+
+
+
     def bin(self):
         """
         Nightly bin the observations.
@@ -1949,7 +2063,8 @@ class RV(ISSUES, REPORTS):
 
                 # treat ccf_mask specially, doing a 'unique' bin
                 if q == 'ccf_mask':
-                    setattr(s, q, bin_ccf_mask(s.mtime, getattr(s, q)))
+                    ccf_mask = getattr(s, q)[s.mask]
+                    setattr(s, q, bin_ccf_mask(s.mtime, ccf_mask))
                     continue
 
                 if Q.dtype != np.float64:
@@ -2317,7 +2432,7 @@ class RV(ISSUES, REPORTS):
         self.units = new_units
 
 
-    def put_at_systemic_velocity(self, factor=1.0):
+    def put_at_systemic_velocity(self, factor=1.0, ignore=None):
         """
         For instruments in which mean(RV) < `factor` * ptp(RV), "move" RVs to
         the systemic velocity from simbad. This is useful if some instruments
@@ -2327,6 +2442,9 @@ class RV(ISSUES, REPORTS):
         """
         changed = False
         for inst in self.instruments:
+            if ignore is not None:
+                if inst in ignore or any([i in inst for i in ignore]):
+                    continue
             changed_inst = False
             s = getattr(self, inst)
             if s.mask.any():
@@ -2360,6 +2478,11 @@ class RV(ISSUES, REPORTS):
             self.instruments = sorted(self.instruments, key=lambda i: getattr(self, i).time.max())
             self._build_arrays()
 
+    def put_instrument_last(self, instrument):
+        if not self._check_instrument(instrument, strict=True, log=True):
+            return
+        self.instruments = [i for i in self.instruments if i != instrument] + [instrument]
+        self._build_arrays()
 
     def save(self, directory=None, instrument=None, format='rdb',
              indicators=False, join_instruments=False, postfix=None,
@@ -2383,7 +2506,7 @@ class RV(ISSUES, REPORTS):
                 Postfix to add to the filenames ([star]_[instrument]_[postfix].rdb).
             save_masked (bool, optional)
                 If True, also save masked observations (those for which
-                self.mask == True)
+                self.mask == False)
             save_nans (bool, optional)
                 Whether to save NaN values in the indicators, if they exist. If
                 False, the full observation which contains NaN values is not saved.

arvi/utils.py

@@ -2,6 +2,8 @@ import os
 import sys
 import time
 from contextlib import contextmanager
+from functools import partial
+from collections import defaultdict
 
 try:
     from unittest.mock import patch
@@ -68,6 +70,50 @@ def all_logging_disabled():
         logging.disable(previous_level)
 
 
+class record_removals:
+    def __init__(self, s, storage=None):
+        """
+        A simple context manager to record removed files
+
+        Args:
+            s (RV):
+                An `RV` object
+            storage (dict): 
+                A dictionary to store the removed files, with keys 'raw_file'
+                and 'reason' as lists.
+        
+        Examples:
+            >>> with record_removals(s) as rec:
+              :    s.remove_instrument('HARPS')
+              :    rec.store('removed HARPS')
+            >>> rec.storage
+        """
+        self.s = s
+        if storage is None:
+            self.storage = defaultdict(list)
+        else:
+            if 'raw_file' not in storage:
+                storage['raw_file'] = []
+            if 'reason' not in storage:
+                storage['reason'] = []
+            self.storage = storage
+        self.raw_file_start = self.s.raw_file.copy()
+
+    def store(self, reason):
+        missing = ~ np.isin(self.raw_file_start, self.s.raw_file[self.s.mask])
+        if missing.any():
+            lost = self.raw_file_start[missing]
+            self.storage['raw_file'].extend(lost)
+            self.storage['reason'].extend(len(lost) * [reason])
+            self.raw_file_start = self.s.raw_file[self.s.mask].copy()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        pass
+
+
 @contextmanager
 def timer(name=None):
     """ A simple context manager to time a block of code """

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

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

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

@@ -4,35 +4,36 @@ arvi/ariadne_wrapper.py,sha256=YvilopJa9T4NwPcj3Nah_U8smSeSAU5-HYZMb_GJ-BQ,2232
 arvi/berv.py,sha256=eKnpuPC1w45UrUEyFRbs9F9j3bXz3kxYzNXbnRgvFQM,17596
 arvi/binning.py,sha256=NK9y9bUrdyWCbh79LkcRABHG-n5MtlETMHMvLj1z-OM,15437
 arvi/config.py,sha256=JkHSwF-EEqwwbcc8thGgbFc9udDZPjQH-9XFjqDepBY,2337
-arvi/dace_wrapper.py,sha256=08hNZMCt2Kd9x8BvJHZHml4cDCcbBsY1lBheuXk9yGk,25952
+arvi/dace_wrapper.py,sha256=i3aIRKSc2MylEzM5WbbU-BnuLbz8JlG1ez6sK_vvQH0,26125
 arvi/exofop_wrapper.py,sha256=8S7UEcrBAgANIweMV0-CvaWaVTPgGVo8vQQk_KRa0nU,2414
 arvi/extra_data.py,sha256=Xi65pI5kkzqlMmHGl9xFoumtH699611pJJ5PV-a_IfU,3397
 arvi/gaia_wrapper.py,sha256=HTuigIduin3raWfSC7QYuQxDk2dEXYH_4egRkzzg7Xw,4379
 arvi/headers.py,sha256=uvdJebw1M5YkGjE3vJJwYBOnLikib75uuZE9FXB5JJM,1673
-arvi/instrument_specific.py,sha256=ycLhtT3oeNtSREm9bmWICaT8uureYcl3NFzbdDYRMVY,11624
+arvi/instrument_specific.py,sha256=94oMb6UeH6tp7H8YXnXHpxEhIz2evz0iYsT_HrNOCTo,12105
+arvi/kepmodel_wrapper.py,sha256=mmHudetAZ4cBxKDwzQzgUydzkjhomCWw5VVuyiKfXq8,10288
 arvi/kima_wrapper.py,sha256=GrAZWkDCg8ukhW41M1VTadSbab0GBa6BIzjtAtvjk58,3891
 arvi/lbl_wrapper.py,sha256=_ViGVkpakvuBR_xhu9XJRV5EKHpj5Go6jBZGJZMIS2Y,11850
-arvi/nasaexo_wrapper.py,sha256=mWt7eHgSZe4MBKCmUvMPTyUPGuiwGTqKugNBvmjOg9s,7306
+arvi/nasaexo_wrapper.py,sha256=ZKY3IUClqsJuysxDv0Gu51EnzMX7101zQb7UQy_urhI,7431
 arvi/plots.py,sha256=U4VUNyIx4h_rEFd7ZWgBcawUcIGcURES0A4VXIBKp3U,35240
 arvi/programs.py,sha256=M8o8hXr6W22dMiIX3Nxz4pgb8lsJXApDlq7HStyTfqs,9047
-arvi/reports.py,sha256=CKmtg5rewMyT26gbWeoZDYrL0z5Sbb6cTJry0HWk_rs,7445
+arvi/reports.py,sha256=a38EZNhyGoSSzJh63wBQCAt3_xhqbpVGcDOXaZWTLXs,11127
 arvi/setup_logger.py,sha256=dHzO2gPjw6CaKWpYZd2f83z09tmxgi--qpp7k1jROjI,615
 arvi/simbad_wrapper.py,sha256=uZc8mcfNijXsQi29LReRTabZb2hRPhYdLsDLMgq1OEI,9927
 arvi/sophie_wrapper.py,sha256=KUeWccXud5_Lrx72S1HSemHIZRdjd2oLvqyofwsL0QQ,3440
 arvi/spectra.py,sha256=ebF1ocodTastLx0CyqLSpE8EZNDXBF8riyfxMr3L6H0,7491
-arvi/stats.py,sha256=ilzzGL9ew-SyVa9eEdrYCpD3DliOAwhoNUg9LIlHjzU,2583
+arvi/stats.py,sha256=gvMkKzP83AV8_Oi71JHmA8QH8Y1z1viYykV9ELVDqZI,3547
 arvi/stellar.py,sha256=GQ7yweuBRnfkJ0M5eWjvLd8uvGq_by81PbXfidBvWis,4918
-arvi/timeseries.py,sha256=9-EjmhMFMo8IJ_Erqf5SnEeKi4J2x0Pt518V30jQ8bw,100162
+arvi/timeseries.py,sha256=suXPLmFlNbnFbfixXufNMwqwZWz-FwKdckJ1x2PtrXQ,104664
 arvi/translations.py,sha256=PUSrn4zvYO2MqGzUxlFGwev_tBkgJaJrIYs6NKHzbWo,951
-arvi/utils.py,sha256=x_zDTW1vp672CZe-m9-KXo5IVk-JKOs2wme_ta4t8MU,9402
+arvi/utils.py,sha256=MuAgjyXr297Sm_T6QmB1riVUktyT9ud1qngGMgKlXMc,10863
 arvi/data/info.svg,sha256=0IMI6W-eFoTD8acnury79WJJakpBwLa4qKS4JWpsXiI,489
 arvi/data/obs_affected_ADC_issues.dat,sha256=tn93uOL0eCTYhireqp1wG-_c3CbxPA7C-Rf-pejVY8M,10853
 arvi/data/obs_affected_blue_cryostat_issues.dat,sha256=z4AK17xfz8tGTDv1FjRvQFnio4XA6PNNfDXuicewHk4,1771
 arvi/data/extra/HD86226_PFS1.rdb,sha256=vfAozbrKHM_j8dYkCBJsuHyD01KEM1asghe2KInwVao,3475
 arvi/data/extra/HD86226_PFS2.rdb,sha256=F2P7dB6gVyzCglUjNheB0hIHVClC5RmARrGwbrY1cfo,4114
 arvi/data/extra/metadata.json,sha256=C69hIw6CohyES6BI9vDWjxwSz7N4VOYX0PCgjXtYFmU,178
-arvi-0.2.9.dist-info/licenses/LICENSE,sha256=6JfQgl7SpM55t0EHMFNMnNh-AdkpGW25MwMiTnhdWQg,1068
-arvi-0.2.9.dist-info/METADATA,sha256=KLtdToTvrCF-TT0P04dctgqKRHH4eTOL47hspjj8qaU,1932
-arvi-0.2.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-arvi-0.2.9.dist-info/top_level.txt,sha256=4EeiKDVLD45ztuflTGfQ3TU8GVjJg5Y95xS5XjI-utU,5
-arvi-0.2.9.dist-info/RECORD,,
+arvi-0.2.11.dist-info/licenses/LICENSE,sha256=6JfQgl7SpM55t0EHMFNMnNh-AdkpGW25MwMiTnhdWQg,1068
+arvi-0.2.11.dist-info/METADATA,sha256=G2yojj2vhNVuEPtoY81O3aWPvgMMxNw5ax4_40zVhb8,1933
+arvi-0.2.11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+arvi-0.2.11.dist-info/top_level.txt,sha256=4EeiKDVLD45ztuflTGfQ3TU8GVjJg5Y95xS5XjI-utU,5
+arvi-0.2.11.dist-info/RECORD,,