astro-nest 0.5.5__tar.gz → 0.5.6__tar.gz

{astro_nest-0.5.5 → astro_nest-0.5.6}/NEST/NEST.py

@@ -1,5 +1,6 @@
 import numpy as np
 import pickle,json,os,warnings,shutil,zipfile,urllib.request,json
+
 try:
     from tqdm import tqdm
     has_tqdm = True
@@ -15,7 +16,8 @@ try:
     has_matplotlib = True
 except ImportError:
     has_matplotlib = False
-NNSA_DIR = os.path.dirname(os.path.abspath(__file__))
+
+NEST_DIR = os.path.dirname(os.path.abspath(__file__))
 
 loaded_isochrones = {}
 
@@ -24,16 +26,25 @@ def custom_warning(message, category, filename, lineno, file=None, line=None):
 
 warnings.showwarning = custom_warning
 
-def get_mode(arr,min_age=0,max_age=14):
-    hist, bins = np.histogram(arr,bins=100,range=(min_age,max_age))
+def sanitize_input(input):
+    if input is not None and type(input) is not list:
+        if hasattr(input,'tolist'):
+            input = input.tolist()
+        else:
+            input = [input]
+    return input
+
+def get_mode(arr,min_age=0,max_age=14,nbins=280):
+    #TODO: choose number of bins appropriately
+    hist, bins = np.histogram(arr,bins=nbins,range=(min_age,max_age))
     return bins[np.argmax(hist)] + (bins[1]-bins[0])/2
 
 def download_isochrones():
     if input('Isochrone curves for plots do not exist. Download them ? (15.6Mb) [Y/n] (default: Y)') in ['n', 'N']:
             return None
     iso_url = "https://github.com/star-age/star-age.github.io/archive/refs/heads/main.zip"
-    iso_dir = os.path.join(NNSA_DIR, 'isochrones')
-    tmp_zip = os.path.join(NNSA_DIR, 'isochrones_tmp.zip')
+    iso_dir = os.path.join(NEST_DIR, 'isochrones')
+    tmp_zip = os.path.join(NEST_DIR, 'isochrones_tmp.zip')
 
     print("Downloading isochrones from GitHub...")
     if has_requests and has_tqdm:
@@ -54,49 +65,52 @@ def download_isochrones():
 
     with zipfile.ZipFile(tmp_zip, 'r') as zip_ref:
         members = [m for m in zip_ref.namelist() if m.startswith('star-age.github.io-main/isochrones/')]
-        zip_ref.extractall(NNSA_DIR, members)
+        zip_ref.extractall(NEST_DIR, members)
 
-    src = os.path.join(NNSA_DIR, 'star-age.github.io-main', 'isochrones')
+    src = os.path.join(NEST_DIR, 'star-age.github.io-main', 'isochrones')
     if os.path.exists(iso_dir):
         shutil.rmtree(iso_dir)
     shutil.move(src, iso_dir)
-    shutil.rmtree(os.path.join(NNSA_DIR, 'star-age.github.io-main'))
+    shutil.rmtree(os.path.join(NEST_DIR, 'star-age.github.io-main'))
     os.remove(tmp_zip)
     print("Isochrones downloaded and extracted.")
 
 def get_isochrones(model):
     if model.model_name in loaded_isochrones:
         return loaded_isochrones[model.model_name]
-    if os.path.exists(os.path.join(NNSA_DIR, 'isochrones')) == False:
+    if os.path.exists(os.path.join(NEST_DIR, 'isochrones')) == False:
         download_isochrones()
-    isochrone_path = os.path.join(NNSA_DIR, 'isochrones', 'isochrones_' + model.model_name + '.json')
+    isochrone_path = os.path.join(NEST_DIR, 'isochrones', 'isochrones_' + model.model_name + '.json')
     if os.path.exists(isochrone_path):
         loaded_isochrones[model.model_name] = json.load(open(isochrone_path, 'r'))
         return loaded_isochrones[model.model_name]
     return None
 
 def available_models():
-    model_list = ['BaSTI','PARSEC','MIST','Geneva','Dartmouth','YaPSI']
+    model_list = ['BaSTI','PARSEC','MIST','Geneva','Dartmouth','YaPSI','BaSTI_HST']
     model_sources = [
         'http://basti-iac.oa-abruzzo.inaf.it/',
         'https://stev.oapd.inaf.it/PARSEC/',
         'https://waps.cfa.harvard.edu/MIST/',
-        'https://www.unige.ch/sciences/astro/evolution/en/database/syclist',
+        'https://www.unige.ch/sciences/astro/evolution/en/database/syclist/',
         'https://rcweb.dartmouth.edu/stellar/',
-        'http://www.astro.yale.edu/yapsi/'
+        'http://www.astro.yale.edu/yapsi/',
+        'http://basti-iac.oa-abruzzo.inaf.it/'
     ]
     for model,source in zip(model_list,model_sources):
         print(model + 'Model (' + source + ')')
 
 class AgeModel:
-    def __init__(self,model_name,cut=False,use_sklearn=True,use_tqdm=True):
+    def __init__(self,model_name,use_sklearn=True,use_tqdm=True,photometric_type=None):
         self.model_name = model_name
         self.use_sklearn = use_sklearn
         self.use_tqdm = use_tqdm
+        if photometric_type == None:
+            photometric_type = 'Gaia'
+        self.photometric_type = photometric_type
         if not has_tqdm and self.use_tqdm:
             self.use_tqdm = False
-        self.cut = cut
-        domain_path = os.path.join(NNSA_DIR, 'domain.pkl')
+        domain_path = os.path.join(NEST_DIR, 'domain.pkl')
         domain = pickle.load(open(domain_path, 'rb'))
         if model_name in domain:
             self.domain = domain[model_name]
@@ -109,8 +123,6 @@ class AgeModel:
             self.space_col = None
             self.space_mag = None
             self.space_met = None
-        if self.cut:
-            self.model_name = self.model_name + '_cut'
         self.neural_networks = {}
         self.scalers = {}
         self.samples = None
@@ -121,13 +133,21 @@ class AgeModel:
         self.stds = None
         self.load_neural_network(self.model_name)
 
+    def __repr__(self):
+        return self.__str__()
+
     def __str__(self):
-        return self.model_name + ' Age Model'
+        _str = self.model_name + ' Age Model'
+        if self.photometric_type == 'Gaia':
+            _str += ', Gaia photometry: mag=MG, col=(GBP-GRP).'
+        elif self.photometric_type == 'HST':
+            _str += ', HST photometry: mag=F814W, col=(F606W-F814W).'
+        return _str
 
     def load_neural_network(self, model_name):
         if self.use_sklearn:
-            model_path_full = os.path.join(NNSA_DIR, 'models', f'{model_name}.sav')
-            model_path_reduced = os.path.join(NNSA_DIR, 'models', f'{model_name}_BPRP.sav')
+            model_path_full = os.path.join(NEST_DIR, 'models', f'{model_name}.sav')
+            model_path_reduced = os.path.join(NEST_DIR, 'models', f'{model_name}_BPRP.sav')
             if os.path.exists(model_path_full):
                 nn = pickle.load(open(model_path_full, 'rb'))
                 self.neural_networks['full'] = nn['NN']
@@ -137,8 +157,8 @@ class AgeModel:
                 self.neural_networks['reduced'] = nn['NN']
                 self.scalers['reduced'] = nn['Scaler']
         else:
-            model_path_full = os.path.join(NNSA_DIR, 'models', f'NN_{model_name}.json')
-            model_path_reduced = os.path.join(NNSA_DIR, 'models', f'NN_{model_name}_BPRP.json')
+            model_path_full = os.path.join(NEST_DIR, 'models', f'NN_{model_name}.json')
+            model_path_reduced = os.path.join(NEST_DIR, 'models', f'NN_{model_name}_BPRP.json')
             if os.path.exists(model_path_full):
                 json_nn = json.load(open(model_path_full, 'r'))
                 self.neural_networks['full'] = {
@@ -169,56 +189,17 @@ class AgeModel:
                         store_samples=True,
                         min_age=0,max_age=14):
         
-        if met is not None and type(met) is not list:
-            if hasattr(met,'tolist'):
-                met = met.tolist()
-            else:
-                met = [met]
-        if mag is not None and type(mag) is not list:
-            if hasattr(mag,'tolist'):
-                mag = mag.tolist()
-            else:
-                mag = [mag]
-        if col is not None and type(col) is not list:
-            if hasattr(col,'tolist'):
-                col = col.tolist()
-            else:
-                col = [col]
-        if emet is not None and type(emet) is not list:
-            if hasattr(emet,'tolist'):
-                emet = emet.tolist()
-            else:
-                emet = [emet]
-        if emag is not None and type(emag) is not list:
-            if hasattr(emag,'tolist'):
-                emag = emag.tolist()
-            else:
-                emag = [emag]
-        if ecol is not None and type(ecol) is not list:
-            if hasattr(ecol,'tolist'):
-                ecol = ecol.tolist()
-            else:
-                ecol = [ecol]
-        if GBP is not None and type(GBP) is not list:
-            if hasattr(GBP,'tolist'):
-                GBP = GBP.tolist()
-            else:
-                GBP = [GBP]
-        if GRP is not None and type(GRP) is not list:
-            if hasattr(GRP,'tolist'):
-                GRP = GRP.tolist()
-            else:
-                GRP = [GRP]
-        if eGBP is not None and type(eGBP) is not list:
-            if hasattr(eGBP,'tolist'):
-                eGBP = eGBP.tolist()
-            else:
-                eGBP = [eGBP]
-        if eGRP is not None and type(eGRP) is not list:
-            if hasattr(eGRP,'tolist'):
-                eGRP = eGRP.tolist()
-            else:
-                eGRP = [eGRP]
+        
+        met = sanitize_input(met)
+        mag = sanitize_input(mag)
+        col = sanitize_input(col)
+        emet = sanitize_input(emet)
+        emag = sanitize_input(emag)
+        ecol = sanitize_input(ecol)
+        GBP = sanitize_input(GBP)
+        GRP = sanitize_input(GRP)
+        eGBP = sanitize_input(eGBP)
+        eGRP = sanitize_input(eGRP)
 
         if store_samples and n*len(met) > 1e6:
             warnings.warn('Storing samples for {} stars with {} samples for each will take a lot of memory. Consider setting store_samples=False to only store mean,median,mode and std of individual age distributions.'.format(len(met),n))
@@ -306,36 +287,12 @@ class AgeModel:
     def check_domain(self,met,mag,col,emet=None,emag=None,ecol=None):
         if self.domain is None:
             raise ValueError('No domain defined for this model')
-        if met is not None and type(met) is not list:
-            if hasattr(met,'tolist'):
-                met = met.tolist()
-            else:
-                met = [met]
-        if mag is not None and type(mag) is not list:
-            if hasattr(mag,'tolist'):
-                mag = mag.tolist()
-            else:
-                mag = [mag]
-        if col is not None and type(col) is not list:
-            if hasattr(col,'tolist'):
-                col = col.tolist()
-            else:
-                col = [col]
-        if emet is not None and type(emet) is not list:
-            if hasattr(emet,'tolist'):
-                emet = emet.tolist()
-            else:
-                emet = [emet]
-        if emag is not None and type(emag) is not list:
-            if hasattr(emag,'tolist'):
-                emag = emag.tolist()
-            else:
-                emag = [emag]
-        if ecol is not None and type(ecol) is not list:
-            if hasattr(ecol,'tolist'):
-                ecol = ecol.tolist()
-            else:
-                ecol = [ecol]
+        met = sanitize_input(met)
+        mag = sanitize_input(mag)
+        col = sanitize_input(col)
+        emet = sanitize_input(emet)
+        emag = sanitize_input(emag)
+        ecol = sanitize_input(ecol)
         
         has_errors = emet != None and emag != None and ecol != None
         
@@ -357,9 +314,6 @@ class AgeModel:
             max_i_mag = np.minimum(np.digitize(mag[i] + errors[1],self.space_mag) - 1,self.space_mag.size-2)
             min_i_met = np.maximum(np.digitize(met[i] - errors[2],self.space_met) - 1,0)
             max_i_met = np.minimum(np.digitize(met[i] + errors[2],self.space_met) - 1,self.space_met.size-2)
-            if self.cut and (self.space_col[min_i_col] > 1.25 or self.space_mag[min_i_mag] > 4):
-                in_domain[i] = False
-                continue
             in_domain[i] = bool(np.any(self.domain[min_i_col:max_i_col+1,min_i_mag:max_i_mag+1,min_i_met:max_i_met+1]) == 1)
         
         return in_domain
@@ -413,7 +367,6 @@ class AgeModel:
     def mode_ages(self):
         if self.ages is None:
             raise ValueError('No age predictions have been made yet')
-        #TODO: choose number of bins appropriately
         modes = []
         min_age = max(0,self.ages.min())
         max_age = max(14,self.ages.max())
@@ -458,6 +411,9 @@ class AgeModel:
         if check_domain and met is not None:
             in_domain = self.check_domain(met,mag,col)
             age = age[in_domain]
+            met = np.array(sanitize_input(met))
+            mag = np.array(sanitize_input(mag))
+            col = np.array(sanitize_input(col))
             met = met[in_domain]
             mag = mag[in_domain]
             col = col[in_domain]
@@ -475,7 +431,7 @@ class AgeModel:
             print('Metallicity not available for this model (-2,-1,0), using [M/H]=0 instead')
             isochrone_met = '0'
         isochrones = isochrones[isochrone_met]
-        # Store line objects and their ages for interactivity
+        
         lines = []
         ages = []
         for isochrone in isochrones:
@@ -503,7 +459,6 @@ class AgeModel:
             )
             plt.colorbar(scatter,label='Age [Gyr]')
 
-        # Add hover info box for isochrone lines
         annot = ax.annotate("", xy=(0,0), xytext=(0,15), textcoords="offset points",
                             ha='center',
                             bbox=dict(boxstyle="round", fc="w"),
@@ -535,46 +490,40 @@ class AgeModel:
 
         ax.set_xlim(-1,3)
         ax.set_ylim(10,-5)
-        ax.set_xlabel(r'$(G_{BP}-G_{RP})_0$ [mag]')
-        ax.set_ylabel(r'$M_G$ [mag]')
+        if self.photometric_type == 'Gaia':
+            ax.set_xlabel(r'$(G_{BP}-G_{RP})_0$ [mag]')
+            ax.set_ylabel(r'$M_G$ [mag]')
+        elif self.photometric_type == 'HST':
+            ax.set_xlabel(r'$(F606W-F814W)$ [mag]')
+            ax.set_ylabel(r'$F606W$ [mag]')
 
         plt.tight_layout()
         plt.show()
 
 class BaSTIModel(AgeModel):
-    def __init__(self,cut=False,use_sklearn=True,use_tqdm=True):
-        super().__init__('BaSTI',cut,use_sklearn,use_tqdm)
-
-'''
-class BaSTI2Model(AgeModel):
-    def __init__(self,cut=False,use_sklearn=True,use_tqdm=True):
-        super().__init__('BaSTI2',cut,use_sklearn,use_tqdm)
-
-class BaSTI_HSTModel(AgeModel):
-    def __init__(self,cut=False,use_sklearn=True,use_tqdm=True):
-        super().__init__('BaSTI_HST',cut,use_sklearn,use_tqdm)
-
-class BaSTI_HST_alpha_zeroModel(AgeModel):
-    def __init__(self,cut=False,use_sklearn=True,use_tqdm=True):
-        super().__init__('BaSTI_HST_alpha_zero',cut,use_sklearn,use_tqdm)
-'''
+    def __init__(self,use_sklearn=True,use_tqdm=True):
+        super().__init__('BaSTI',use_sklearn,use_tqdm)
 
 class PARSECModel(AgeModel):
-    def __init__(self,cut=False,use_sklearn=True,use_tqdm=True):
-        super().__init__('PARSEC',cut,use_sklearn,use_tqdm)
+    def __init__(self,use_sklearn=True,use_tqdm=True):
+        super().__init__('PARSEC',use_sklearn,use_tqdm)
 
 class MISTModel(AgeModel):
-    def __init__(self,cut=False,use_sklearn=True,use_tqdm=True):
-        super().__init__('MIST',cut,use_sklearn,use_tqdm)
+    def __init__(self,use_sklearn=True,use_tqdm=True):
+        super().__init__('MIST',use_sklearn,use_tqdm)
 
 class GenevaModel(AgeModel):
-    def __init__(self,cut=False,use_sklearn=True,use_tqdm=True):
-        super().__init__('Geneva',cut,use_sklearn,use_tqdm)
+    def __init__(self,use_sklearn=True,use_tqdm=True):
+        super().__init__('Geneva',use_sklearn,use_tqdm)
 
 class DartmouthModel(AgeModel):
-    def __init__(self,cut=False,use_sklearn=True,use_tqdm=True):
-        super().__init__('Dartmouth',cut,use_sklearn,use_tqdm)
+    def __init__(self,use_sklearn=True,use_tqdm=True):
+        super().__init__('Dartmouth',use_sklearn,use_tqdm)
 
 class YaPSIModel(AgeModel):
-    def __init__(self,cut=False,use_sklearn=True,use_tqdm=True):
-        super().__init__('YaPSI',cut,use_sklearn,use_tqdm)
+    def __init__(self,use_sklearn=True,use_tqdm=True):
+        super().__init__('YaPSI',use_sklearn,use_tqdm)
+
+class BaSTI_HSTModel(AgeModel):
+    def __init__(self,use_sklearn=True,use_tqdm=True,photometric_type='HST'):
+        super().__init__('BaSTI_HST',use_sklearn,use_tqdm,photometric_type)