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

arvi/__init__.py

@@ -8,6 +8,7 @@ def __getattr__(name: str):
     if name in (
         '_ipython_canary_method_should_not_exist_',
         '_repr_mimebundle_',
+        '__wrapped__'
     ):
         return
 

arvi/ariadne_wrapper.py

@@ -0,0 +1,70 @@
+import os
+import sys
+from matplotlib import pyplot as plt
+
+try:
+    from astroARIADNE.star import Star
+    from astroARIADNE.fitter import Fitter
+except ImportError:
+    print('This module requires astroARIADNE. Install with `pip install astroARIADNE`')
+    sys.exit(0)
+
+
+def run_ariadne(self, fit=True, plot=True, priors={},
+                models = ('phoenix', 'btsettl', 'btnextgen', 'btcond', 'kurucz', 'ck04'),
+                nlive=300, dlogz=1, threads=6, dynamic=False, **kwargs):
+    if hasattr(self, 'gaia'):
+        s = Star(self.star, self.gaia.ra, self.gaia.dec, g_id=self.gaia.dr3_id,
+                 search_radius=1)
+    else:
+        s = Star(self.star, self.simbad.ra, self.simbad.dec, g_id=self.simbad.gaia_id,
+                 search_radius=1)
+
+    out_folder = f'{self.star}_ariadne'
+
+    setup = dict(engine='dynesty', nlive=nlive, dlogz=dlogz,
+                 bound='multi', sample='auto', threads=threads, dynamic=dynamic)
+    setup = list(setup.values())
+
+    f = Fitter()
+    f.star = s
+    f.setup = setup
+    f.av_law = 'fitzpatrick'
+    f.out_folder = out_folder
+    f.bma = True
+    f.models = models
+    f.n_samples = 10_000
+
+    f.prior_setup = {
+            'teff': priors.get('teff', ('default')),
+            'logg': ('default'),
+            'z': priors.get('feh', ('default')),
+            'dist': ('default'),
+            'rad': ('default'),
+            'Av': ('default')
+    }
+
+    if fit:
+        f.initialize()
+        f.fit_bma()
+
+    if plot:
+        from pkg_resources import resource_filename
+        from astroARIADNE.plotter import SEDPlotter
+        modelsdir = resource_filename('astroARIADNE', 'Datafiles/models')
+        artist = SEDPlotter(os.path.join(out_folder, 'BMA.pkl'), out_folder, models_dir=modelsdir)
+
+        artist.plot_SED_no_model()
+        try:
+            artist.plot_SED()
+        except FileNotFoundError as e:
+            print('No model found:', e)
+        except IndexError as e:
+            print('Error!')
+        artist.plot_bma_hist()
+        artist.plot_bma_HR(10)
+        artist.plot_corner()
+        plt.close('all')
+        return s, f, artist
+
+    return s, f

arvi/config.py

@@ -1 +1,8 @@
+# whether to return self from (some) RV methods
 return_self = False
+
+# whether to check internet connection before querying DACE
+check_internet = False
+
+# make all DACE requests without using a .dacerc file
+request_as_public = False

arvi/dace_wrapper.py

@@ -5,10 +5,15 @@ import numpy as np
 from dace_query import DaceClass
 from dace_query.spectroscopy import SpectroscopyClass, Spectroscopy as default_Spectroscopy
 from .setup_logger import logger
-from .utils import create_directory, all_logging_disabled, stdout_disabled
+from .utils import create_directory, all_logging_disabled, stdout_disabled, tqdm
 
 
 def load_spectroscopy() -> SpectroscopyClass:
+    from .config import request_as_public
+    if request_as_public:
+        with all_logging_disabled():
+            dace = DaceClass(dace_rc_config_path='none')
+        return SpectroscopyClass(dace_instance=dace)
     if 'DACERC' in os.environ:
         dace = DaceClass(dace_rc_config_path=os.environ['DACERC'])
         return SpectroscopyClass(dace_instance=dace)
@@ -156,6 +161,13 @@ def check_existing(output_directory, files, type):
         f.partition('.fits')[0] for f in os.listdir(output_directory)
         if type in f
     ]
+
+    # also check for lowercase type
+    existing += [
+        f.partition('.fits')[0] for f in os.listdir(output_directory)
+        if type.lower() in f
+    ]
+
     if os.name == 'nt':  # on Windows, be careful with ':' in filename
         import re
         existing = [re.sub(r'T(\d+)_(\d+)_(\d+)', r'T\1:\2:\3', f) for f in existing]
@@ -174,9 +186,9 @@ def check_existing(output_directory, files, type):
 def download(files, type, output_directory):
     """ Download files from DACE """
     Spectroscopy = load_spectroscopy()
-    # with stdout_disabled(), all_logging_disabled():
-    Spectroscopy.download_files(files, file_type=type,
-                                output_directory=output_directory)
+    with stdout_disabled(), all_logging_disabled():
+        Spectroscopy.download_files(files, file_type=type.lower(),
+                                    output_directory=output_directory)
 
 def extract_fits(output_directory):
     """ Extract fits files from tar.gz file """
@@ -194,86 +206,99 @@ def extract_fits(output_directory):
     return files
 
 
-def do_download_ccf(raw_files, output_directory, clobber=False, verbose=True):
-    """ Download CCFs from DACE """
+def do_download_filetype(type, raw_files, output_directory, clobber=False,
+                         verbose=True, chunk_size=20):
+    """ Download CCFs / S1Ds / S2Ds from DACE """
     raw_files = np.atleast_1d(raw_files)
 
     create_directory(output_directory)
 
     # check existing files to avoid re-downloading
     if not clobber:
-        raw_files = check_existing(output_directory, raw_files, 'CCF')
+        raw_files = check_existing(output_directory, raw_files, type)
+
+    n = raw_files.size
 
     # any file left to download?
-    if raw_files.size == 0:
+    if n == 0:
         if verbose:
             logger.info('no files to download')
         return
 
-    if verbose:
-        n = raw_files.size
-        logger.info(f"Downloading {n} CCFs into '{output_directory}'...")
-
-    download(raw_files, 'ccf', output_directory)
+    # avoid an empty chunk
+    if chunk_size > n:
+        chunk_size = n
 
     if verbose:
-        logger.info('Extracting .fits files')
+        if chunk_size < n:
+            msg = f"Downloading {n} {type}s "
+            msg += f"(in chunks of {chunk_size}) "
+            msg += f"into '{output_directory}'..."
+            logger.info(msg)
+        else:
+            msg = f"Downloading {n} {type}s into '{output_directory}'..."
+            logger.info(msg)
 
-    extract_fits(output_directory)
+    iterator = [raw_files[i:i + chunk_size] for i in range(0, n, chunk_size)]
+    for files in tqdm(iterator, total=len(iterator)):
+        download(files, type, output_directory)
+        extract_fits(output_directory)
 
+    logger.info('Extracted .fits files')
 
-def do_download_s1d(raw_files, output_directory, clobber=False, verbose=True):
-    """ Download S1Ds from DACE """
-    raw_files = np.atleast_1d(raw_files)
 
-    create_directory(output_directory)
+# def do_download_s1d(raw_files, output_directory, clobber=False, verbose=True):
+#     """ Download S1Ds from DACE """
+#     raw_files = np.atleast_1d(raw_files)
 
-    # check existing files to avoid re-downloading
-    if not clobber:
-        raw_files = check_existing(output_directory, raw_files, 'S1D')
+#     create_directory(output_directory)
 
-    # any file left to download?
-    if raw_files.size == 0:
-        if verbose:
-            logger.info('no files to download')
-        return
+#     # check existing files to avoid re-downloading
+#     if not clobber:
+#         raw_files = check_existing(output_directory, raw_files, 'S1D')
 
-    if verbose:
-        n = raw_files.size
-        logger.info(f"Downloading {n} S1Ds into '{output_directory}'...")
+#     # any file left to download?
+#     if raw_files.size == 0:
+#         if verbose:
+#             logger.info('no files to download')
+#         return
 
-    download(raw_files, 's1d', output_directory)
+#     if verbose:
+#         n = raw_files.size
+#         logger.info(f"Downloading {n} S1Ds into '{output_directory}'...")
 
-    if verbose:
-        logger.info('Extracting .fits files')
+#     download(raw_files, 's1d', output_directory)
 
-    extract_fits(output_directory)
+#     if verbose:
+#         logger.info('Extracting .fits files')
 
+#     extract_fits(output_directory)
 
-def do_download_s2d(raw_files, output_directory, clobber=False, verbose=True):
-    """ Download S2Ds from DACE """
-    raw_files = np.atleast_1d(raw_files)
 
-    create_directory(output_directory)
+# def do_download_s2d(raw_files, output_directory, clobber=False, verbose=True):
+#     """ Download S2Ds from DACE """
+#     raw_files = np.atleast_1d(raw_files)
 
-    # check existing files to avoid re-downloading
-    if not clobber:
-        raw_files = check_existing(output_directory, raw_files, 'S2D')
+#     create_directory(output_directory)
 
-    # any file left to download?
-    if raw_files.size == 0:
-        if verbose:
-            logger.info('no files to download')
-        return
+#     # check existing files to avoid re-downloading
+#     if not clobber:
+#         raw_files = check_existing(output_directory, raw_files, 'S2D')
 
-    if verbose:
-        n = raw_files.size
-        logger.info(f"Downloading {n} S2Ds into '{output_directory}'...")
+#     # any file left to download?
+#     if raw_files.size == 0:
+#         if verbose:
+#             logger.info('no files to download')
+#         return
 
-    download(raw_files, 's2d', output_directory)
+#     if verbose:
+#         n = raw_files.size
+#         logger.info(f"Downloading {n} S2Ds into '{output_directory}'...")
 
-    if verbose:
-        logger.info('Extracting .fits files')
+#     download(raw_files, 's2d', output_directory)
+
+#     if verbose:
+#         logger.info('Extracting .fits files')
 
-    extracted_files = extract_fits(output_directory)
-    return extracted_files
+#     extracted_files = extract_fits(output_directory)
+#     return extracted_files

arvi/extra_data.py

@@ -19,15 +19,24 @@ def get_extra_data(star, instrument=None, path=None, verbose=True):
     # print(metadata)
 
     files = glob(os.path.join(path, star + '*'))
+    files = [f for f in files if os.path.isfile(f)]
+    files = [f for f in files if not os.path.basename(f).endswith('.zip')]
+
     if len(files) == 0:
         raise FileNotFoundError
 
-    instruments = [os.path.basename(f).split('.')[0] for f in files]
-    instruments = [i.split('_', maxsplit=1)[1] for i in instruments]
+    def get_instruments(files):
+        instruments = [os.path.basename(f).split('.')[0] for f in files]
+        instruments = [i.split('_', maxsplit=1)[1] for i in instruments]
+        return instruments
+    
+    instruments = get_instruments(files)
 
     if instrument is not None:
         if not any([instrument in i for i in instruments]):
             raise FileNotFoundError
+        files = [f for f in files if instrument in f]
+        instruments = get_instruments(files)
 
     if verbose:
         logger.info(f'loading extra data for {star}')

arvi/gaia_wrapper.py

@@ -0,0 +1,94 @@
+import os
+from io import StringIO
+from csv import DictReader
+from dataclasses import dataclass, field
+import requests
+
+from astropy.coordinates import SkyCoord
+import pysweetcat
+
+DATA_PATH = os.path.dirname(__file__)
+DATA_PATH = os.path.join(DATA_PATH, 'data')
+
+QUERY = """
+SELECT TOP 20 gaia_source.designation,gaia_source.source_id,gaia_source.ra,gaia_source.dec,gaia_source.parallax,gaia_source.pmra,gaia_source.pmdec,gaia_source.ruwe,gaia_source.phot_g_mean_mag,gaia_source.bp_rp,gaia_source.radial_velocity,gaia_source.phot_variable_flag,gaia_source.non_single_star,gaia_source.has_xp_continuous,gaia_source.has_xp_sampled,gaia_source.has_rvs,gaia_source.has_epoch_photometry,gaia_source.has_epoch_rv,gaia_source.has_mcmc_gspphot,gaia_source.has_mcmc_msc,gaia_source.teff_gspphot,gaia_source.logg_gspphot,gaia_source.mh_gspphot,gaia_source.distance_gspphot,gaia_source.azero_gspphot,gaia_source.ag_gspphot,gaia_source.ebpminrp_gspphot
+FROM gaiadr3.gaia_source 
+WHERE 
+CONTAINS(
+	POINT('ICRS',gaiadr3.gaia_source.ra,gaiadr3.gaia_source.dec),
+	CIRCLE(
+		'ICRS',
+		COORD1(EPOCH_PROP_POS({ra},{dec},{plx},{pmra},{pmdec},{rv},2000,2016.0)),
+		COORD2(EPOCH_PROP_POS({ra},{dec},{plx},{pmra},{pmdec},{rv},2000,2016.0)),
+		0.001388888888888889)
+)=1
+"""
+
+def run_query(query):
+    url = 'https://gea.esac.esa.int/tap-server/tap/sync'
+    data = dict(query=query, request='doQuery', lang='ADQL', format='csv')
+    try:
+        response = requests.post(url, data=data, timeout=10)
+    except requests.ReadTimeout as err:
+        raise IndexError(err)
+    except requests.ConnectionError as err:
+        raise IndexError(err)
+    return response.content.decode()
+
+def parse_csv(csv):
+    reader = DictReader(StringIO(csv))
+    return list(reader)
+
+
+class gaia:
+    """
+    A very simple wrapper around a TAP query to gaia for a given target. This
+    class simply runs a few TAP queries and stores the result as attributes.
+
+    Attributes:
+        ra (float): right ascension
+        dec (float): declination
+        coords (SkyCoord): coordinates as a SkyCoord object
+        dr3_id (int): Gaia DR3 identifier
+        plx (float): parallax
+        radial_velocity (float): radial velocity
+    """
+    def __init__(self, star:str, simbad=None):
+        """
+        Args:
+            star (str): The name of the star to query simbad
+        """
+        self.star = star
+
+        if simbad is None:
+            from .simbad_wrapper import simbad as Simbad
+            simbad = Simbad(star)
+
+        ra = simbad.ra
+        dec = simbad.dec
+        plx = simbad.plx
+        pmra = simbad.pmra
+        pmdec = simbad.pmdec
+        rv = simbad.rvz_radvel
+        args = dict(ra=ra, dec=dec, plx=plx, pmra=pmra, pmdec=pmdec, rv=rv)          
+
+        try:
+            table1 = run_query(query=QUERY.format(**args))
+            results = parse_csv(table1)[0]
+        except IndexError:
+            raise ValueError(f'Gaia query for {star} failed')
+
+        self.dr3_id = int(results['source_id'])
+
+        self.ra = float(results['ra'])
+        self.dec = float(results['dec'])
+        self.pmra = float(results['pmra'])
+        self.pmdec = float(results['pmdec'])
+        self.coords = SkyCoord(self.ra, self.dec, unit='deg')
+        self.plx = float(results['parallax'])
+        self.radial_velocity = float(results['radial_velocity'])
+
+        return
+
+    def __repr__(self):
+        return f'{self.star} (DR3 id={self.dr3_id})'

arvi/plots.py

@@ -1,6 +1,8 @@
 import os
 from functools import partial, partialmethod
+from itertools import cycle
 
+import matplotlib.collections
 import numpy as np
 import matplotlib
 import matplotlib.pyplot as plt
@@ -15,7 +17,7 @@ from . import config
 
 def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
          remove_50000=False, tooltips=False, label=None, N_in_label=False,
-         versus_n=False, show_histogram=False, **kwargs):
+         versus_n=False, show_histogram=False, bw=False, **kwargs):
     """ Plot the RVs
 
     Args:
@@ -38,6 +40,8 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
         show_histogram (bool, optional)
             Whether to show a panel with the RV histograms (per intrument).
             Defaults to False.
+        bw (bool, optional):
+            Adapt plot to black and white. Defaults to False.
 
     Returns:
         Figure: the figure
@@ -58,7 +62,7 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
             ax, axh = ax
         fig = ax.figure
 
-    kwargs.setdefault('marker', 'o')
+
     kwargs.setdefault('ls', '')
     kwargs.setdefault('capsize', 0)
     kwargs.setdefault('ms', 4)
@@ -66,10 +70,22 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
     if remove_50000:
         time_offset = 50000
 
-    instruments = self._check_instrument(instrument)
+    strict = kwargs.pop('strict', False)
+    instruments = self._check_instrument(instrument, strict=strict)
+
+    if bw:
+        markers = cycle(('o', 'P', 's', '^', '*'))
+    else:
+        markers = cycle(('o',) * len(instruments))
+
+    try:
+        zorders = cycle(-np.argsort([getattr(self, i).error for i in instruments])[::-1])
+    except AttributeError:
+        zorders = cycle([1] * len(instruments))
 
     cursors = {}
-    for inst in instruments:
+    containers = {}
+    for _i, inst in enumerate(instruments):
         s = self if self._child else getattr(self, inst)
         if s.mask.sum() == 0:
             continue
@@ -81,37 +97,42 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
                 p = p.replace('_', '.')
                 _label = f'{i}-{p}'
         else:
-            _label = label
+            if isinstance(label, list):
+                _label = label[_i]
+            else:
+                _label = label
 
         if versus_n:
-            container = ax.errorbar(np.arange(1, s.mtime.size + 1),
-                                    s.mvrad, s.msvrad, label=_label, picker=True, **kwargs)
+            container = ax.errorbar(np.arange(1, s.mtime.size + 1), s.mvrad, s.msvrad,
+                                    label=_label, picker=True, marker=next(markers), zorder=next(zorders),
+                                    **kwargs)
         else:
-            container = ax.errorbar(s.mtime - time_offset,
-                                    s.mvrad, s.msvrad, label=_label, picker=True, **kwargs)
+            container = ax.errorbar(s.mtime - time_offset, s.mvrad, s.msvrad,
+                                    label=_label, picker=True, marker=next(markers), zorder=next(zorders),
+                                    **kwargs)
+
+        containers[inst] = list(container)
 
         if show_histogram:
             kw = dict(histtype='step', bins='doane', orientation='horizontal')
             hlabel = f'{s.mvrad.std():.2f} {self.units}'
             axh.hist(s.mvrad, **kw, label=hlabel)
 
-        if tooltips:
-            cursors[inst] = crsr = mplcursors.cursor(container, multiple=False)
-
-            @crsr.connect("add")
-            def _(sel):
-                inst = sel.artist.get_label()
-                _s = getattr(self, inst)
-                vrad, svrad = _s.vrad[sel.index], _s.svrad[sel.index]
-                sel.annotation.get_bbox_patch().set(fc="white")
-                text = f'{inst}\n'
-                text += f'BJD: {sel.target[0]:9.5f}\n'
-                text += f'RV: {vrad:.3f} ± {svrad:.3f}'
-                # if fig.canvas.manager.toolmanager.get_tool('infotool').toggled:
-                #     text += '\n\n'
-                #     text += f'date: {_s.date_night[sel.index]}\n'
-                #     text += f'mask: {_s.ccf_mask[sel.index]}'
-                sel.annotation.set_text(text)
+                # cursors[inst] = crsr = mplcursors.cursor(container, multiple=False)
+                # @crsr.connect("add")
+                # def _(sel):
+                #     inst = sel.artist.get_label()
+                #     _s = getattr(self, inst)
+                #     vrad, svrad = _s.vrad[sel.index], _s.svrad[sel.index]
+                #     sel.annotation.get_bbox_patch().set(fc="white")
+                #     text = f'{inst}\n'
+                #     text += f'BJD: {sel.target[0]:9.5f}\n'
+                #     text += f'RV: {vrad:.3f} ± {svrad:.3f}'
+                #     # if fig.canvas.manager.toolmanager.get_tool('infotool').toggled:
+                #     #     text += '\n\n'
+                #     #     text += f'date: {_s.date_night[sel.index]}\n'
+                #     #     text += f'mask: {_s.ccf_mask[sel.index]}'
+                #     sel.annotation.set_text(text)
 
         if show_masked:
             if versus_n:
@@ -147,9 +168,43 @@ def plot(self, ax=None, show_masked=False, instrument=None, time_offset=0,
                 fig.canvas.draw()
             except ValueError:
                 pass
-
     plt.connect('pick_event', on_pick_legend)
 
+    if tooltips:
+        annotations = []
+        def on_pick_point(event):
+            print('annotations:', annotations)
+            for text in annotations:
+                text.remove()
+                annotations.remove(text)
+
+            artist = event.artist
+            if isinstance(artist, (matplotlib.lines.Line2D, matplotlib.collections.LineCollection)):
+                print(event.ind, artist)
+                if isinstance(artist, matplotlib.lines.Line2D):
+                    matching_instrument = [k for k, v in containers.items() if artist in v]
+                    print(matching_instrument)
+                    if len(matching_instrument) == 0:
+                        return
+                    inst = matching_instrument[0]
+                    _s = getattr(self, inst)
+                    ind = event.ind[0]
+                    # print(_s.mtime[ind], _s.mvrad[ind], _s.msvrad[ind])
+
+                    text = f'{inst}\n'
+                    text += f'{_s.mtime[ind]:9.5f}\n'
+                    text += f'RV: {_s.mvrad[ind]:.1f} ± {_s.msvrad[ind]:.1f}'
+
+                    annotations.append(
+                        ax.annotate(text, (_s.mtime[ind], _s.mvrad[ind]), xycoords='data',
+                                    xytext=(5, 10), textcoords='offset points', fontsize=9,
+                                    bbox={'boxstyle': 'round', 'fc': 'w'}, arrowprops=dict(arrowstyle="-"))
+                    )
+                    # ax.annotate(f'{inst}', (0.5, 0.5), xycoords=artist, ha='center', va='center')
+                    fig.canvas.draw()
+                    # print(event.ind, artist.get_label())
+        plt.connect('pick_event', on_pick_point)
+
 
     if show_histogram:
         axh.legend()
@@ -340,6 +395,31 @@ def gls(self, ax=None, label=None, fap=True, picker=True, instrument=None, **kwa
     if label is not None:
         ax.legend()
 
+    if ax.get_legend() is not None:
+        leg = ax.get_legend()
+        for text in leg.get_texts():
+            text.set_picker(True)
+
+        def on_pick_legend(event):
+            handles, labels = ax.get_legend_handles_labels()
+            artist = event.artist
+            if isinstance(artist, matplotlib.text.Text):
+                # print('handles:', handles)
+                # print('labels:', labels)
+                # print(artist.get_text())
+                try:
+                    h = handles[labels.index(artist.get_text())]
+                    alpha_text = {None:0.2, 1.0: 0.2, 0.2:1.0}[artist.get_alpha()]
+                    alpha_point = {None: 0.0, 1.0: 0.0, 0.2: 1.0}[artist.get_alpha()]
+                    h.set_alpha(alpha_point)
+                    artist.set_alpha(alpha_text)
+                    fig.canvas.draw()
+                except ValueError:
+                    pass
+
+        if 'pick_event' not in fig.canvas.callbacks.callbacks:
+            plt.connect('pick_event', on_pick_legend)
+
 
     if config.return_self:
         return self