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

arvi/programs.py

@@ -1,5 +1,7 @@
 import os
-import concurrent.futures
+import multiprocessing
+from functools import partial
+from itertools import chain
 from collections import namedtuple
 from tqdm import tqdm
 # import numpy as np
@@ -13,7 +15,8 @@ path = os.path.join(os.path.dirname(__file__), 'data')
 
 
 def get_star(star, instrument=None):
-    return RV(star, verbose=False, instrument=instrument, _raise_on_error=False)
+    return RV(star, instrument=instrument,
+              _raise_on_error=False, verbose=False, load_extra_data=False)
 
 
 class LazyRV:
@@ -22,6 +25,7 @@ class LazyRV:
         if isinstance(self.stars, str):
             self.stars = [self.stars]
         self.instrument = instrument
+        self._saved = None
 
     @property
     def N(self):
@@ -31,30 +35,49 @@ class LazyRV:
         return f"RV({self.N} stars)"
 
     def _get(self):
-        result = []
-        # use a with statement to ensure threads are cleaned up promptly
-        with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool:
-            star_to_RV = {
-                pool.submit(get_star, star, self.instrument): star 
-                for star in self.stars
-            }
+        if self.N > 10:
+            # logger.info('Querying DACE...')
+            _get_star = partial(get_star, instrument=self.instrument)
+            with multiprocessing.Pool() as pool:
+                result = list(tqdm(pool.imap(_get_star, self.stars), 
+                                   total=self.N, unit='star', desc='Querying DACE'))
+                # result = pool.map(get_star, self.stars)
+        else:
+            result = []
             logger.info('Querying DACE...')
-            pbar = tqdm(concurrent.futures.as_completed(star_to_RV),
-                        total=self.N, unit='star')
-            for future in pbar:
-                star = star_to_RV[future]
+            pbar = tqdm(self.stars, total=self.N, unit='star')
+            for star in pbar:
                 pbar.set_description(star)
-                try:
-                    result.append(future.result())
-                except Exception:
-                    print(f'{star} generated an exception')
+                result.append(get_star(star, self.instrument))
+
         return result
 
+        # # use a with statement to ensure threads are cleaned up promptly
+        # with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool:
+        #     star_to_RV = {
+        #         pool.submit(get_star, star, self.instrument): star 
+        #         for star in self.stars
+        #     }
+        #     logger.info('Querying DACE...')
+        #     pbar = tqdm(concurrent.futures.as_completed(star_to_RV),
+        #                 total=self.N, unit='star')
+        #     for future in pbar:
+        #         star = star_to_RV[future]
+        #         pbar.set_description(star)
+        #         try:
+        #             result.append(future.result())
+        #         except ValueError:
+        #             print(f'{star} generated an exception')
+        #             result.append(None)
+        # return result
+
     def __iter__(self):
         return self._get()
 
     def __call__(self):
-        return self._get()
+        if not self._saved:
+            self._saved = self._get()
+        return self._saved
 
 
 # sorted by spectral type

arvi/simbad_wrapper.py

@@ -1,9 +1,12 @@
+import os
 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 basic.OID,
@@ -45,6 +48,8 @@ def run_query(query):
         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_table(table, cols=None, values=None):
@@ -80,7 +85,8 @@ class simbad:
         dec (float): declination
         coords (SkyCoord): coordinates as a SkyCoord object
         main_id (str): main identifier
-        plx_value (float): parallax
+        gaia_id (int): Gaia DR3 identifier
+        plx (float): parallax
         rvz_radvel (float): radial velocity
         sp_type (str): spectral type
         B (float): B magnitude 
@@ -94,6 +100,16 @@ class simbad:
         """
         self.star = star
 
+        if 'kobe' in self.star.lower():
+            fname = os.path.join(DATA_PATH, 'KOBE-translate.csv')
+            kobe_translate = {}
+            if os.path.exists(fname):
+                with open(fname) as f:
+                    for line in f.readlines():
+                        kobe_id, catname = line.strip().split(',')
+                        kobe_translate[kobe_id] = catname
+                self.star = star = kobe_translate[self.star]
+
         # oid = run_query(query=OID_QUERY.format(star=star))
         # self.oid = str(oid.split()[-1])
 
@@ -110,6 +126,8 @@ 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])
+
         for col, val in zip(cols, values):
             if col == 'oid':
                 setattr(self, col, str(val))
@@ -123,6 +141,9 @@ class simbad:
 
         if self.plx_value == '':
             self.plx_value = None
+        
+        self.plx = self._plx_value = self.plx_value
+        del self.plx_value
 
         try:
             swc_data = pysweetcat.get_data()

arvi/stats.py

@@ -19,15 +19,19 @@ def wmean(a, e):
         raise ValueError
     return np.average(a, weights=1 / e**2)
 
-def rms(a):
+def rms(a, ignore_nans=False):
     """ Root mean square of array `a`
 
     Args:
         a (array): Array containing data
     """
+    if ignore_nans:
+        a = a[~np.isnan(a)]
+    if len(a) == 0:
+        return np.nan
     return np.sqrt((a**2).mean())
 
-def wrms(a, e):
+def wrms(a, e, ignore_nans=False):
     """ Weighted root mean square of array `a`, with uncertanty given by `e`.
     The weighted rms is calculated using the weighted mean, where the weights
     are equal to 1/e**2.
@@ -36,6 +40,16 @@ def wrms(a, e):
         a (array): Array containing data 
         e (array): Uncertainties on `a`
     """
+    if ignore_nans:
+        nans = np.logical_or(np.isnan(a), np.isnan(e))
+        a = a[~nans]
+        e = e[~nans]
+    if (e == 0).any():
+        raise ZeroDivisionError('uncertainty cannot be zero')
+    if (e < 0).any():
+        raise ValueError('uncertainty cannot be negative')
+    if (a.shape != e.shape):
+        raise ValueError('arrays must have the same shape')
     w = 1 / e**2
     return np.sqrt(np.sum(w * (a - np.average(a, weights=w))**2) / sum(w))