astro-nest 0.5.0__py3-none-any.whl → 0.5.2__py3-none-any.whl

NEST/NEST.py

@@ -1,12 +1,24 @@
 import numpy as np
-import pickle,json,os,warnings
+import pickle,json,os,warnings,shutil,zipfile,urllib.request,json
 try:
     from tqdm import tqdm
     has_tqdm = True
 except ImportError:
     has_tqdm = False
+try:
+    import requests
+    has_requests = True
+except ImportError:
+    has_requests = False
+try:
+    import matplotlib.pyplot as plt
+    has_matplotlib = True
+except ImportError:
+    has_matplotlib = False
 NNSA_DIR = os.path.dirname(os.path.abspath(__file__))
 
+loaded_isochrones = {}
+
 def custom_warning(message, category, filename, lineno, file=None, line=None):
     print(f"{category.__name__}: {message}")
 
@@ -16,6 +28,53 @@ def get_mode(arr,min_age=0,max_age=14):
     hist, bins = np.histogram(arr,bins=100,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')
+
+    print("Downloading isochrones from GitHub...")
+    if has_requests and has_tqdm:
+        response = requests.get(iso_url, stream=True)
+        total = int(response.headers.get('content-length', 0))
+        with open(tmp_zip, 'wb') as file, tqdm(
+            desc="Downloading isochrones",
+            total=total,
+            unit='B',
+            unit_scale=True,
+            unit_divisor=1024,
+        ) as bar:
+            for data in response.iter_content(chunk_size=1024):
+                size = file.write(data)
+                bar.update(size)
+    else:
+        urllib.request.urlretrieve(iso_url, tmp_zip)
+
+    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)
+
+    src = os.path.join(NNSA_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'))
+    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:
+        download_isochrones()
+    isochrone_path = os.path.join(NNSA_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_sources = [
@@ -298,13 +357,12 @@ 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)
-            #cells[i] = np.array([[min_i_col,max_i_col],[min_i_mag,max_i_mag],[min_i_met,max_i_met]])
             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#,cells
+        return in_domain
     
     def propagate(self,X,neural_network,scaler):
         if self.use_sklearn:
@@ -371,6 +429,117 @@ class AgeModel:
         self.stds = np.std(self.ages,axis=1)
         return self.stds
     
+    def HR_diagram(self,met=None,mag=None,col=None,age=None,isochrone_met=0,age_type='median',check_domain=True,**kwargs):
+        if has_matplotlib == False:
+            raise ImportError('matplotlib is required for HR diagram plotting')
+        
+        if age_type not in ('median','mean','mode'):
+            print('Age type not available, using median instead')
+            age_type = 'median'
+        
+        if met is not None and mag is not None and col is not None:
+            if age is None:
+                result = self.ages_prediction(met,mag,col,store_samples=False)
+                age = result[age_type]    
+        elif met is None and mag is None and col is None and age is None:
+            if self.samples is not None and len(self.samples) > 0 and self.ages is not None and len(self.ages) == len(self.samples):
+                if age_type == 'median':
+                    age = self.median_ages()
+                elif age_type == 'mean':
+                    age = self.mean_ages()
+                elif age_type == 'mode':
+                    age = self.mode_ages()
+                met = np.median(self.samples,axis=1)[:,0]
+                mag = np.median(self.samples,axis=1)[:,1]
+                col = np.median(self.samples,axis=1)[:,2]
+        else:
+            raise ValueError('Not a valid combination of arguments.')
+
+        if check_domain and met is not None:
+            in_domain = self.check_domain(met,mag,col)
+            age = age[in_domain]
+            met = met[in_domain]
+            mag = mag[in_domain]
+            col = col[in_domain]
+        
+        fig,ax = plt.subplots(figsize=(10,7))
+
+        isochrones = get_isochrones(self)
+        if isochrones is None:
+            raise ValueError('Isochrones not available for this model')
+        if type(isochrone_met) is not str:
+            if type (isochrone_met) is not float and type(isochrone_met) is not int:
+                raise ValueError('isochrone_met must be a float or int')
+            isochrone_met = str(round(isochrone_met))
+        if isochrone_met not in isochrones.keys():
+            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:
+            iso_age = isochrone['age']
+            iso_mag = isochrone['MG']
+            iso_col = isochrone['BP-RP']
+            if 'c' in kwargs:
+                kwargs.pop('c')
+            kwargs['color'] = 'k'
+            if 'lw' in kwargs:
+                kwargs.pop('lw')
+            kwargs['linewidth'] = 0.5
+            line, = ax.plot(iso_col, iso_mag, **kwargs)
+            lines.append(line)
+            ages.append(iso_age)
+
+        if mag is not None and col is not None:
+            scatter = ax.scatter(
+                col,
+                mag,
+                c=age,
+                cmap='viridis',
+                zorder=4
+            )
+            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"),
+                            arrowprops=dict(arrowstyle="->"),zorder=5)
+        annot.set_visible(False)
+
+        def update_annot(line, event, age):
+            x, y = event.xdata, event.ydata
+            annot.xy = (x, y)
+            annot.set_text(f"Age: {age:.2f} Gyr")
+            annot.get_bbox_patch().set_facecolor('white')
+            annot.get_bbox_patch().set_alpha(0.8)
+
+        def hover(event):
+            vis = annot.get_visible()
+            if event.inaxes == ax:
+                for line, age in zip(lines, ages):
+                    cont, ind = line.contains(event)
+                    if cont:
+                        update_annot(line, event, age)
+                        annot.set_visible(True)
+                        fig.canvas.draw_idle()
+                        return
+            if vis:
+                annot.set_visible(False)
+                fig.canvas.draw_idle()
+
+        fig.canvas.mpl_connect("motion_notify_event", hover)
+
+        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]')
+
+        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)
@@ -407,6 +576,4 @@ class DartmouthModel(AgeModel):
 
 class YaPSIModel(AgeModel):
     def __init__(self,cut=False,use_sklearn=True,use_tqdm=True):
-        super().__init__('YaPSI',cut,use_sklearn,use_tqdm)
-
-#TODO: add flavors to models (e.g. trained on cut CMD for optimal performance)
+        super().__init__('YaPSI',cut,use_sklearn,use_tqdm)

{astro_nest-0.5.0.dist-info → astro_nest-0.5.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: astro-nest
-Version: 0.5.0
+Version: 0.5.2
 Summary: A python package that provides neural networks to estimate ages of stars/star populations.
 Author-email: Tristan Boin <tristan.boin@gmail.com>
 License-Expression: MIT
@@ -17,6 +17,7 @@ Based on an upcoming paper by Boin et al. 2025.
 - numpy
 - scikit (optional but recommended for speed)
 - tqdm (optional)
+- matplotlib (optional)
 
 The [documentation website](https://star-age.github.io/NEST-docs/) guides you through the package usage.
 

{astro_nest-0.5.0.dist-info → astro_nest-0.5.2.dist-info}/RECORD

@@ -1,4 +1,4 @@
-NEST/NEST.py,sha256=Xdjp0p6N3ct2fwfP7RA8qqLc9COtrAYxzzY62aaYFdY,15865
+NEST/NEST.py,sha256=3F0ejOiHD03-PXmjDtTAIRHPPLkL42X6mshYuLUQgdk,22452
 NEST/__init__.py,sha256=OKJbgoTPkmXisdQCOj7rSIZf-kxym0077bfFPa0zxb0,19
 NEST/domain.pkl,sha256=Z0zv80P6Kc49Bw932gXe4q2ScANEjTvCetlaY9NMNCo,1353728
 NEST/tutorial.ipynb,sha256=_YcbYokFZcWgxO0kZZjJc4v5sBjd0t8td00flc1Mxh4,67643
@@ -66,7 +66,7 @@ NEST/models/scaler_SYCLIST.sav,sha256=mpr_hzBcB3jso80j2PNfMHD2IqRdHupp4LrAggBjrP
 NEST/models/scaler_SYCLIST_BPRP.sav,sha256=Eb9F6BAx8MCFFSRLgZUonGLAY1sII6X4xnyn8Qil23w,522
 NEST/models/scaler_YaPSI.sav,sha256=MzhMbTcjE3o9U_i1IKhZvw0pOjQe040X_qJJPDCJV4A,570
 NEST/models/scaler_YaPSI_BPRP.sav,sha256=j9gCxIQndfTPd84fhSNGyCWX3jcyBeLJEZZCjvgFrro,522
-astro_nest-0.5.0.dist-info/METADATA,sha256=EPPShbUBYUO-5ZYpEjfsrlaP1HrHqeFNPE14pBZYp3w,943
-astro_nest-0.5.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-astro_nest-0.5.0.dist-info/top_level.txt,sha256=TpsNVEvaG60SlJ4tjyQu4aOsE5ov9RW9lQtLpgBAx4E,5
-astro_nest-0.5.0.dist-info/RECORD,,
+astro_nest-0.5.2.dist-info/METADATA,sha256=SpX-YgOKtgPOIRDh2ysojH9ykqQAPSEHok1JEp9g1zQ,967
+astro_nest-0.5.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+astro_nest-0.5.2.dist-info/top_level.txt,sha256=TpsNVEvaG60SlJ4tjyQu4aOsE5ov9RW9lQtLpgBAx4E,5
+astro_nest-0.5.2.dist-info/RECORD,,