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

arvi/plots.py

@@ -6,13 +6,101 @@ import matplotlib.collections
 import numpy as np
 import matplotlib
 import matplotlib.pyplot as plt
-from matplotlib.collections import LineCollection
 import mplcursors
 
 from astropy.timeseries import LombScargle
 
 from .setup_logger import logger
 from . import config
+from .stats import wmean
+
+
+class BlittedCursor:
+    """ A cross-hair cursor using blitting for faster redraw. """
+    def __init__(self, axes, vertical=True, horizontal=True, show_text=None,
+                 transforms_x=None, transforms_y=None):
+        if isinstance(axes, matplotlib.axes.Axes):
+            axes = [axes]
+        self.axes = axes
+        self.background = None
+        self.vertical = vertical
+        self.horizontal = horizontal
+        
+        self.transforms_x = [lambda x:x for _ in axes] if transforms_x is None else transforms_x
+        self.transforms_y = [lambda x:x for _ in axes] if transforms_y is None else transforms_y
+
+        if horizontal:
+            self.horizontal_line = [ax.axhline(color='k', lw=0.8, ls='--') for ax in axes]
+        if vertical:
+            self.vertical_line = [ax.axvline(color='k', lw=0.8, ls='--') for ax in axes]
+
+        self.show_text = show_text
+        if show_text is not None: # text location in axes coordinates
+            self.text = [ax.text(0.72, 0.9, '', transform=ax.transAxes) for ax in axes]
+
+        self._creating_background = False
+        for ax in axes:
+            ax.figure.canvas.mpl_connect('draw_event', self.on_draw)
+
+    def on_draw(self, event):
+        self.create_new_background()
+
+    def set_cross_hair_visible(self, visible):
+        if self.horizontal:
+            need_redraw = [line.get_visible() != visible for line in self.horizontal_line]
+        else:
+            need_redraw = [line.get_visible() != visible for line in self.vertical_line]
+        if self.horizontal:
+            [line.set_visible(visible) for line in self.horizontal_line]
+        if self.vertical:
+            [line.set_visible(visible) for line in self.vertical_line]
+        if self.show_text:
+            self.text.set_visible(visible)
+        return need_redraw
+
+    def create_new_background(self):
+        if self._creating_background:
+            # discard calls triggered from within this function
+            return
+        self._creating_background = True
+        self.set_cross_hair_visible(False)
+        for ax in self.axes:
+            ax.figure.canvas.draw()
+        self.backgrounds = [ax.figure.canvas.copy_from_bbox(ax.bbox) for ax in self.axes]
+        self.set_cross_hair_visible(True)
+        self._creating_background = False
+
+    def on_mouse_move(self, event):
+        if self.background is None:
+            self.create_new_background()
+        if not event.inaxes:
+            need_redraw = self.set_cross_hair_visible(False)
+            if any(need_redraw):
+                for ax, bkgd in zip(self.axes, self.backgrounds):
+                    ax.figure.canvas.restore_region(bkgd)
+                    ax.figure.canvas.blit(ax.bbox)
+        else:
+            self.set_cross_hair_visible(True)
+            # update the line positions
+            x, y = event.xdata, event.ydata
+            X = [trans(x) for trans in self.transforms_x]
+            Y = [trans(y) for trans in self.transforms_y]
+            if self.horizontal:
+                [line.set_ydata([y]) for line, y in zip(self.horizontal_line, Y)]
+            if self.vertical:
+                [line.set_xdata([x]) for line, x in zip(self.vertical_line, X)]
+            if self.show_text:
+                self.text.set_text(f'x={x:1.2f}, y={y:1.2f}')
+
+            for ax, bkgd in zip(self.axes, self.backgrounds):
+                ax.figure.canvas.restore_region(bkgd)
+                if self.horizontal:
+                    [ax.draw_artist(line) for line in self.horizontal_line]
+                if self.vertical:
+                    [ax.draw_artist(line) for line in self.vertical_line]
+                if self.show_text:
+                    ax.draw_artist(self.text)
+                ax.figure.canvas.blit(ax.bbox)
 
 
 def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
@@ -72,11 +160,12 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
 
     strict = kwargs.pop('strict', False)
     instruments = self._check_instrument(instrument, strict=strict)
+    marker = kwargs.pop('marker', 'o')
 
     if bw:
-        markers = cycle(('o', 'P', 's', '^', '*'))
+        markers = cycle((marker, 'P', 's', '^', '*'))
     else:
-        markers = cycle(('o',) * len(instruments))
+        markers = cycle((marker,) * len(instruments))
 
     try:
         zorders = cycle(-np.argsort([getattr(self, i).error for i in instruments])[::-1])
@@ -313,12 +402,13 @@ def plot_quantity(self, quantity, ax=None, show_masked=False, instrument=None,
     ax.legend()
     ax.minorticks_on()
 
-    if quantity == 'fwhm':
-        ax.set_ylabel(f'FWHM [{self.units}]')
-    elif quantity == 'bispan':
-        ax.set_ylabel(f'BIS [{self.units}]')
-    elif quantity == 'rhk':
-        ax.set_ylabel(r"$\log$ R'$_{HK}$")
+    ylabel = {
+        'fwhm': f'FWHM [{self.units}]',
+        'bispan': f'BIS [{self.units}]',
+        'rhk': r"$\log$ R'$_{HK}$",
+        'berv': f'BERV [km/s]',
+    }
+    ax.set_ylabel(ylabel[quantity])
 
     if remove_50000:
         ax.set_xlabel('BJD - 2450000 [days]')
@@ -334,9 +424,11 @@ def plot_quantity(self, quantity, ax=None, show_masked=False, instrument=None,
 plot_fwhm = partialmethod(plot_quantity, quantity='fwhm')
 plot_bis = partialmethod(plot_quantity, quantity='bispan')
 plot_rhk = partialmethod(plot_quantity, quantity='rhk')
+plot_berv = partialmethod(plot_quantity, quantity='berv')
 
 
-def gls(self, ax=None, label=None, fap=True, picker=True, instrument=None, **kwargs):
+def gls(self, ax=None, label=None, fap=True, instrument=None, adjust_means=config.adjust_means_gls,
+        picker=True, **kwargs):
     """
     Calculate and plot the Generalised Lomb-Scargle periodogram of the radial
     velocities.
@@ -348,42 +440,82 @@ def gls(self, ax=None, label=None, fap=True, picker=True, instrument=None, **kwa
         label (str):
             The label to use for the plot.
         fap (bool):
-            Whether to show the false alarm probability.
+            Whether to show the false alarm probability. Default is True.
         instrument (str or list):
-            Which instruments' data to include in the periodogram.
+            Which instruments' data to include in the periodogram. Default is
+            all instruments.
+        adjust_means (bool):
+            Whether to adjust (subtract) the weighted means of each instrument.
+            Default is `config.adjust_means_gls`.
     """
     if self.N == 0:
         if self.verbose:
             logger.error('no data to compute gls')
         return
 
-    if ax is None:
-        fig, ax = plt.subplots(1, 1, constrained_layout=True)
-    else:
-        fig = ax.figure
+    if not self._did_adjust_means and not adjust_means:
+        logger.warning('gls() called before adjusting instrument means, '
+                       'consider using `adjust_means` argument')
 
     if instrument is not None:
         strict = kwargs.pop('strict', False)
-        instrument = self._check_instrument(instrument, strict=strict)
-        if instrument is not None:
-            instrument_mask = np.isin(self.instrument_array, instrument)
-            t = self.time[instrument_mask & self.mask]
-            y = self.vrad[instrument_mask & self.mask]
-            e = self.svrad[instrument_mask & self.mask]
+        instrument = self._check_instrument(instrument, strict=strict, log=True)
+        if instrument is None:
+            return
+
+        instrument_mask = np.isin(self.instrument_array, instrument)
+        t = self.time[instrument_mask & self.mask].copy()
+        y = self.vrad[instrument_mask & self.mask].copy()
+        e = self.svrad[instrument_mask & self.mask].copy()
+        if self.verbose:
+            logger.info(f'calculating periodogram for instrument {instrument}')
+
+        if adjust_means:
             if self.verbose:
-                logger.info(f'calculating periodogram for instrument {instrument}')
+                logger.info('adjusting instrument means before gls')
+            means = np.empty_like(y)
+            for i in instrument:
+                mask = self.instrument_array[instrument_mask & self.mask] == i
+                if len(y[mask]) > 0:
+                    means += wmean(y[mask], e[mask]) * mask
+            y = y - means
+
     else:
-        t = self.time[self.mask]
-        y = self.vrad[self.mask]
-        e = self.svrad[self.mask]
+        t = self.time[self.mask].copy()
+        y = self.vrad[self.mask].copy()
+        e = self.svrad[self.mask].copy()
+
+        if adjust_means:
+            if self.verbose:
+                logger.info('adjusting instrument means before gls')
+            means = np.empty_like(y)
+            for i in self.instruments:
+                mask = self.instrument_array[self.mask] == i
+                if len(y[mask]) > 0:
+                    means += wmean(y[mask], e[mask]) * mask
+            y = y - means
 
     self._gls = gls = LombScargle(t, y, e)
+
     maximum_frequency = kwargs.pop('maximum_frequency', 1.0)
     minimum_frequency = kwargs.pop('minimum_frequency', None)
+    samples_per_peak = kwargs.pop('samples_per_peak', 10)
+
     freq, power = gls.autopower(maximum_frequency=maximum_frequency,
                                 minimum_frequency=minimum_frequency,
-                                samples_per_peak=10)
-    ax.semilogx(1/freq, power, picker=picker, label=label, **kwargs)
+                                samples_per_peak=samples_per_peak)
+
+    if ax is None:
+        fig, ax = plt.subplots(1, 1, constrained_layout=True)
+    else:
+        fig = ax.figure
+
+    if kwargs.pop('fill_between', False):
+        kwargs.pop('lw', None)
+        ax.fill_between(1/freq, 0, power, label=label, lw=0, **kwargs)
+        ax.set_xscale('log')
+    else:
+        ax.semilogx(1/freq, power, picker=picker, label=label, **kwargs)
 
     if fap:
         ax.axhline(gls.false_alarm_level(0.01),
@@ -471,6 +603,73 @@ gls_bis = partialmethod(gls_quantity, quantity='bispan')
 gls_rhk = partialmethod(gls_quantity, quantity='rhk')
 
 
+
+def window_function(self, ax1=None, ax2=None, instrument=None, crosshair=False, **kwargs):
+    """
+    Calculate and plot the window function of the observed times.
+
+    Args:
+        ax1 (matplotlib.axes.Axes):
+            An axes to plot the window function vs period. If None, a new figure
+            will be created.
+        ax2 (matplotlib.axes.Axes):
+            An axes to plot the periodogram vs frequency. If None, a new figure
+            will be created.
+        instrument (str or list):
+            Which instruments' data to include in the window function.
+        crosshair (bool):
+            If True, a crosshair will be drawn on the plot.
+    """
+    if self.N == 0:
+        if self.verbose:
+            logger.error('no data to compute window function')
+        return
+
+    if ax1 is None:
+        fig, (ax1, ax2) = plt.subplots(2, 1, constrained_layout=True)
+    else:
+        fig = ax1.figure
+
+    if instrument is not None:
+        strict = kwargs.pop('strict', False)
+        instrument = self._check_instrument(instrument, strict=strict)
+        if instrument is not None:
+            instrument_mask = np.isin(self.instrument_array, instrument)
+            t = self.time[instrument_mask & self.mask]
+            ye = self.svrad[instrument_mask & self.mask]
+            if self.verbose:
+                logger.info(f'calculating window function for instrument {instrument}')
+    else:
+        t = self.time[self.mask]
+        ye = self.svrad[self.mask]
+
+    wf = LombScargle(t, np.ones_like(t), ye / np.std(ye),
+                     fit_mean=False, center_data=False)
+
+    freq, power = wf.autopower(maximum_frequency=1.1, 
+                               samples_per_peak=20, method='cython')
+    ax1.semilogx(1/freq, power, **kwargs)
+    ax1.set(xlabel='Period [days]', ylabel='Window function')
+
+    ax2.plot(freq, power, **kwargs)
+    ax2.set(xlabel='Frequency [1/day]', ylabel='Window function')
+
+    for x in (365.25, 1.0, 1 - 1.0/365.25):
+        ax1.axvline(x, color='k', alpha=0.2, zorder=-1)
+        ax2.axvline(1/x, color='k', alpha=0.2, zorder=-1)
+    
+    if crosshair:
+        blitted_cursor = BlittedCursor((ax1, ax2), horizontal=False,
+                                    transforms_x=(lambda x:x, lambda x:1/x))
+        fig.canvas.mpl_connect('motion_notify_event', blitted_cursor.on_mouse_move)
+        return fig, (ax1, ax2), blitted_cursor
+        # from matplotlib.widgets import MultiCursor
+        # cursor = MultiCursor(fig.canvas, (ax1, ax2), color='r',
+        #                      lw=0.5, horizOn=False, vertOn=True)
+        # return fig, (ax1, ax2), (cursor)
+    else:
+        return fig, (ax1, ax2)
+
 def histogram_svrad(self, ax=None, instrument=None, label=None):
     """ Plot an histogram of the radial velocity uncertainties. 
     

arvi/simbad_wrapper.py

@@ -1,5 +1,4 @@
 import os
-from dataclasses import dataclass, field
 import requests
 
 from astropy.coordinates import SkyCoord
@@ -126,7 +125,11 @@ class simbad:
         except IndexError:
             raise ValueError(f'simbad query for {star} failed')
 
-        self.gaia_id = int([i for i in self.ids if 'Gaia DR3' in i][0].split('Gaia DR3')[-1])
+        try:
+            self.gaia_id = int([i for i in self.ids if 'Gaia DR3' in i][0]
+                               .split('Gaia DR3')[-1])
+        except IndexError:
+            self.gaia_id = None
 
         for col, val in zip(cols, values):
             if col == 'oid':

arvi/spectra.py

@@ -0,0 +1,208 @@
+import os
+from glob import glob
+import pickle
+import numpy as np
+import matplotlib.pyplot as plt
+
+from .setup_logger import logger
+
+from tqdm import tqdm
+import astropy.units as u, astropy.constants as const
+from astropy.io import fits
+
+def doppler_shift(wave: np.ndarray, flux: np.ndarray, velocity: float):
+    """ Doppler shift a spectrum by a given velocity
+
+    Args:
+        wave (np.ndarray): wavelength array
+        flux (np.ndarray): flux array
+        velocity (float): velocity in km/s
+    """
+    c = const.c.to(u.km/u.second).value
+    doppler_factor = np.sqrt((1 + velocity/c) / (1 - velocity/c))
+    new_wavelength = wave * doppler_factor
+    new_flux = np.interp(new_wavelength, wave, flux)
+    return new_wavelength, new_flux
+
+def fit_gaussian_to_line(wave, flux, center_wavelength, around=0.15 * u.angstrom,
+                         careful_continuum=False, plot=True, ax=None):
+    from scipy.optimize import curve_fit
+    if center_wavelength < wave.min() or center_wavelength > wave.max():
+        raise ValueError('`center_wavelength` is outside the wavelength range')
+    if center_wavelength < wave[np.nonzero(flux)].min() or center_wavelength > wave[np.nonzero(flux)].max():
+        raise ValueError('`center_wavelength` is outside the wavelength range where flux is not zero')
+
+    try:
+        wave <<= u.angstrom
+    except u.UnitConversionError as e:
+        raise ValueError(f'could not convert `wave` to Angstroms: {e}') from None
+
+    try:
+        center_wavelength <<= u.angstrom
+    except u.UnitConversionError as e:
+        raise ValueError(f'could not convert `center_wavelength` to Angstroms: {e}') from None
+
+    try:
+        around <<= u.angstrom
+    except u.UnitConversionError as e:
+        raise ValueError(f'could not convert `around` to Angstroms: {e}') from None
+
+
+    def gaussian(x, amp, cen, wid, off):
+        return amp * np.exp(-(x-cen)**2 / (2*wid**2)) + off
+
+    wave_around = (wave > center_wavelength - around) & (wave < center_wavelength + around)
+    w, f = wave[wave_around].value, flux[wave_around]
+
+    if careful_continuum:
+        wave_around_continuum = (wave > center_wavelength - 10*around) & (wave < center_wavelength + 10*around)
+        wc, fc = wave[wave_around_continuum].value, flux[wave_around_continuum]
+        lim = np.percentile(fc, 80)
+        wc = wc[fc > lim]
+        fc = fc[fc > lim]
+        w, f = np.r_[wc, w], np.r_[fc, f]
+        f = f[np.argsort(w)]
+        w = np.sort(w)
+
+    lower, upper = np.array([
+        [-np.inf, 0], 
+        [-np.inf, np.inf], 
+        [0.0, 0.11], 
+        [0.9*f.max(), 1.1*f.max()]
+    ]).T
+
+    try:
+        popt, pcov = curve_fit(gaussian, w, f, p0=[-f.ptp(), center_wavelength.value, 0.1, f.max()],
+                               bounds=(lower, upper))
+    except RuntimeError as e:
+        logger.warning(f'fit_gaussian_to_line: {e}')
+        return None, np.nan, np.nan
+
+    EW = A = (np.sqrt(2) * np.abs(popt[0]) * np.abs(popt[2]) * np.sqrt(np.pi)) / popt[3]
+    perr = np.sqrt(np.diag(pcov))
+    EW_err = (np.sqrt(2) * np.abs(perr[0]) * np.abs(perr[2]) * np.sqrt(np.pi)) / perr[3]
+
+    if plot:
+        if ax is None:
+            fig, ax = plt.subplots(figsize=(8, 4))
+        if careful_continuum:
+            ax.plot(w, f, 'ko', ms=4, zorder=1)
+            wave_around_plot = wave_around_continuum
+        else:
+            ax.plot(wave[wave_around], flux[wave_around], 'ko', ms=4, zorder=1)
+            ax.plot(wave[wave_around], flux[wave_around] - gaussian(w, *popt), 'o', ms=2)
+            wave_around_plot = (wave > center_wavelength - 2*around) & (wave < center_wavelength + 2*around)
+        # ax.plot(wave[wave_around_plot], flux[wave_around_plot], 'o', ms=2)
+        w = wave[wave_around_plot].value
+        ax.plot(w, gaussian(w, *popt), 'r-')
+        ax.fill_between([popt[1]-A, popt[1]+A], popt[3]+popt[0], popt[3],
+                         color='C2', alpha=0.1, lw=0)
+
+    return popt, EW*1e3, EW_err*1e3
+
+def detrend(w, f):
+    if w.shape[0] > w.shape[1]:
+        w = np.copy(w).T
+        f = np.copy(f).T
+
+    f_detrended = np.zeros_like(f)
+    for i, (ww, ff) in enumerate(zip(w, f)):
+        m = np.nonzero(ff)
+        fit = np.polyval(np.polyfit(ww[m] - np.median(ww[m]), ff[m], 1), ww - np.median(ww[m]))
+        f_detrended[i] = ff - fit + np.median(ff[m])
+    return w, f_detrended
+
+def build_master(self, limit=None, plot=True):
+    files = sorted(glob(f'{self.star}_downloads/*S1D_A.fits'))
+    if self.verbose:
+        logger.info(f'Found {len(files)} S1D files')
+
+    files = files[:limit]
+
+    if len(files) == 0:
+        if self.verbose:
+            logger.warning('Should probably run `download_s1d` first')
+        return
+
+    if plot:
+        fig, axs = plt.subplots(2, 1, figsize=(9, 6), sharex=True, constrained_layout=True)
+        for ax in axs:
+            ax.set(xlabel=r'wavelength air [$\AA$]', ylabel='flux')
+        axs[0].set_title(self.star, loc='right', fontsize=10)
+
+    w0 = fits.getdata(files[0])['wavelength_air']
+    master_flux = np.zeros_like(w0)
+    for file in files:
+        rv = fits.getval(file, '*CCF RV')[0]
+        flux = fits.getdata(file)['flux']
+        _, new_flux = doppler_shift(w0, flux, rv)
+        master_flux += new_flux
+        if plot:
+            axs[0].plot(w0, new_flux, alpha=0.5)
+
+    master_flux /= len(files)
+    if plot:
+        axs[1].plot(w0, master_flux, 'k', label='master')
+        axs[1].legend()
+        axs[0].legend([], [], title=f'{len(files)} S1D spectra')
+
+    return w0, master_flux
+
+
+def determine_stellar_parameters(self, linelist: str, plot=True, **kwargs):
+    try:
+        from juliacall import Main as jl
+        jl.seval("using Korg")
+        Korg = jl.Korg
+    except ModuleNotFoundError:
+        msg = 'this function requires juliacall and Korg.jl, please `pip install juliacall`'
+        logger.error(msg)
+        return
+
+    w, f = build_master(self, plot=plot)
+
+    linelist = np.genfromtxt(linelist, dtype=None, encoding=None, names=True)
+    lines = [
+        Korg.Line(line['wl'], line['loggf'], Korg.Species(line['elem'].replace('Fe', 'Fe ')), line['EP'])
+        for line in linelist
+    ]
+
+    if self.verbose:
+        logger.info(f'Found {len(lines)} lines in linelist')
+        logger.info('Measuring EWs...')
+
+    EW = []
+    pbar = tqdm(linelist)
+    for line in pbar:
+        pbar.set_description(f'{line["elem"]} {line["wl"]}')
+        _, ew, _ = fit_gaussian_to_line(w, f, line['wl'], plot=plot,
+                                        careful_continuum=kwargs.pop('careful_continuum', False))
+        EW.append(ew)
+
+    lines = list(np.array(lines)[~np.isnan(EW)])
+    EW = np.array(EW)[~np.isnan(EW)]
+
+    if self.verbose:
+        logger.info('Determining stellar parameters (can take a few minutes)...')
+
+    callback = lambda p, r, A: print('current parameters:', p)
+    result = Korg.Fit.ews_to_stellar_parameters(lines, EW, callback=callback)
+    par, stat_err, sys_err = result
+
+    if self.verbose:
+        logger.info(f'Best fit stellar parameters:')
+        logger.info(f'  Teff: {par[0]:.0f} ± {sys_err[0]:.0f} K')
+        logger.info(f'  logg: {par[1]:.2f} ± {sys_err[1]:.2f} dex')
+        logger.info(f'  m/H :  {par[3]:.2f} ± {sys_err[3]:.2f} dex')
+
+    r = {
+        'teff': (par[0], sys_err[0]),
+        'logg': (par[1], sys_err[1]),
+        'vmic': (par[2], sys_err[2]),
+        'moh': (par[3], sys_err[3]),
+    }
+
+    with open(f'{self.star}_stellar_parameters.pkl', 'wb') as f:
+        pickle.dump(r, f)
+    
+    return r